diff --git a/quinn-proto/src/endpoint.rs b/quinn-proto/src/endpoint.rs index 8bb4c755ce..2ebc9d0428 100644 --- a/quinn-proto/src/endpoint.rs +++ b/quinn-proto/src/endpoint.rs @@ -557,12 +557,13 @@ where random_bytes: &random_bytes, } .encode(&*server_config.token_key, &remote, &temp_loc_cid); + let mut buf = Vec::new(); let header = Header::Retry { src_cid: temp_loc_cid, dst_cid: src_cid, }; - let mut buf = Vec::new(); + let encode = header.encode(&mut buf); buf.put_slice(&token); buf.extend_from_slice(&S::retry_tag(&dst_cid, &buf)); diff --git a/quinn-proto/src/token.rs b/quinn-proto/src/token.rs index bc1f3574c4..bf39ae80f6 100644 --- a/quinn-proto/src/token.rs +++ b/quinn-proto/src/token.rs @@ -32,6 +32,23 @@ impl<'a> RetryToken<'a> { const MAX_ADDITIONAL_DATA_SIZE: usize = 39; // max(ipv4, ipv6) + port + retry_src_cid pub const RANDOM_BYTES_LEN: usize = 32; + fn put_additional_data<'b>( + address: &SocketAddr, + retry_src_cid: &ConnectionId, + additional_data: &'b mut [u8], + ) -> &'b [u8] { + let mut cursor = &mut additional_data[..]; + match address.ip() { + IpAddr::V4(x) => cursor.put_slice(&x.octets()), + IpAddr::V6(x) => cursor.put_slice(&x.octets()), + } + cursor.write(address.port()); + retry_src_cid.encode_long(&mut cursor); + + let size = Self::MAX_ADDITIONAL_DATA_SIZE - cursor.len(); + &additional_data[..size] + } + pub fn encode( &self, key: &impl HandshakeTokenKey, @@ -51,16 +68,9 @@ impl<'a> RetryToken<'a> { ); let mut additional_data = [0u8; Self::MAX_ADDITIONAL_DATA_SIZE]; - let mut cursor = &mut additional_data[..]; - match address.ip() { - IpAddr::V4(x) => cursor.put_slice(&x.octets()), - IpAddr::V6(x) => cursor.put_slice(&x.octets()), - } - cursor.write(address.port()); - retry_src_cid.encode_long(&mut cursor); - - let size = Self::MAX_ADDITIONAL_DATA_SIZE - cursor.len(); - aead_key.seal(&mut buf, &additional_data[..size]).unwrap(); + let additional_data = + Self::put_additional_data(address, retry_src_cid, &mut additional_data); + aead_key.seal(&mut buf, additional_data).unwrap(); let mut token = Vec::new(); token.put_slice(self.random_bytes); @@ -84,16 +94,9 @@ impl<'a> RetryToken<'a> { let mut sealed_token = raw_token_bytes[Self::RANDOM_BYTES_LEN..].to_vec(); let mut additional_data = [0u8; Self::MAX_ADDITIONAL_DATA_SIZE]; - let mut cursor = &mut additional_data[..]; - match address.ip() { - IpAddr::V4(x) => cursor.put_slice(&x.octets()), - IpAddr::V6(x) => cursor.put_slice(&x.octets()), - } - cursor.write(address.port()); - retry_src_cid.encode_long(&mut cursor); - - let size = Self::MAX_ADDITIONAL_DATA_SIZE - cursor.len(); - let data = aead_key.open(&mut sealed_token, &additional_data[..size])?; + let additional_data = + Self::put_additional_data(address, retry_src_cid, &mut additional_data); + let data = aead_key.open(&mut sealed_token, additional_data)?; let mut reader = io::Cursor::new(data); let orig_dst_cid = ConnectionId::decode_long(&mut reader).ok_or(())?;