Add internal method additional_data

quinn-rs · Oct 10, 2020 · 74a6296 · 74a6296
1 parent 01134c5
commit 74a6296
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 21 deletions.
diff --git 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
@@ -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(())?;