proto: Token encode refactors

quinn-proto/src/endpoint.rs

@@ -29,10 +29,10 @@ use crate::{
         ConnectionEvent, ConnectionEventInner, ConnectionId, DatagramConnectionEvent, EcnCodepoint,
         EndpointEvent, EndpointEventInner, IssuedCid,
     },
-    token::{IncomingToken, InvalidRetryTokenError},
+    token::{IncomingToken, InvalidRetryTokenError, RetryToken},
     transport_parameters::{PreferredAddress, TransportParameters},
-    Duration, Instant, ResetToken, RetryToken, Side, Transmit, TransportConfig, TransportError,
-    INITIAL_MTU, MAX_CID_SIZE, MIN_INITIAL_SIZE, RESET_TOKEN_SIZE,
+    Duration, Instant, ResetToken, Side, Transmit, TransportConfig, TransportError, INITIAL_MTU,
+    MAX_CID_SIZE, MIN_INITIAL_SIZE, RESET_TOKEN_SIZE,
 };
 
 /// The main entry point to the library
@@ -743,14 +743,11 @@ impl Endpoint {
         let loc_cid = self.local_cid_generator.generate_cid();
 
         let token = RetryToken {
+            address: incoming.addresses.remote,
             orig_dst_cid: incoming.packet.header.dst_cid,
             issued: server_config.time_source.now(),
         }
-        .encode(
-            &*server_config.token_key,
-            incoming.addresses.remote,
-            loc_cid,
-        );
+        .encode(&*server_config.token_key, loc_cid);
 
         let header = Header::Retry {
             src_cid: loc_cid,

quinn-proto/src/lib.rs

@@ -85,7 +85,7 @@ pub use crate::cid_generator::{
 };
 
 mod token;
-use token::{ResetToken, RetryToken};
+use token::ResetToken;
 
 #[cfg(feature = "arbitrary")]
 use arbitrary::Arbitrary;

quinn-proto/src/token.rs

@@ -1,13 +1,13 @@
 use std::{
-    fmt, io,
+    fmt,
     net::{IpAddr, SocketAddr},
 };
 
 use bytes::{Buf, BufMut};
 
 use crate::{
     coding::{BufExt, BufMutExt},
-    crypto::{CryptoError, HandshakeTokenKey, HmacKey},
+    crypto::{HandshakeTokenKey, HmacKey},
     packet::InitialHeader,
     shared::ConnectionId,
     Duration, ServerConfig, SystemTime, RESET_TOKEN_SIZE, UNIX_EPOCH,
@@ -33,27 +33,35 @@ impl IncomingToken {
             orig_dst_cid: header.dst_cid,
         };
 
+        // Decode token or short-circuit
         if header.token.is_empty() {
             return Ok(unvalidated);
         }
 
-        let result = RetryToken::from_bytes(
-            &*server_config.token_key,
-            remote_address,
-            header.dst_cid,
-            &header.token,
-        );
-
-        let retry = match result {
-            Ok(retry) => retry,
-            Err(ValidationError::Unusable) => return Ok(unvalidated),
-            Err(ValidationError::InvalidRetry) => return Err(InvalidRetryTokenError),
+        // In cases where a token cannot be decrypted/decoded, we must allow for the possibility
+        // that this is caused not by client malfeasance, but by the token having been generated by
+        // an incompatible endpoint, e.g. a different version or a neighbor behind the same load
+        // balancer. In such cases we proceed as if there was no token.
+        //
+        // [_RFC 9000 § 8.1.3:_](https://www.rfc-editor.org/rfc/rfc9000.html#section-8.1.3-10)
+        //
+        // > If the token is invalid, then the server SHOULD proceed as if the client did not have
+        // > a validated address, including potentially sending a Retry packet.
+        let Some(retry) =
+            RetryToken::decode(&*server_config.token_key, header.dst_cid, &header.token)
+        else {
+            return Ok(unvalidated);
         };
 
+        // Validate token
+        if retry.address != remote_address {
+            return Err(InvalidRetryTokenError);
+        }
         if retry.issued + server_config.retry_token_lifetime < server_config.time_source.now() {
             return Err(InvalidRetryTokenError);
         }
 
+        // Convert token into Self
         Ok(Self {
             retry_src_cid: Some(header.dst_cid),
             orig_dst_cid: retry.orig_dst_cid,
@@ -67,6 +75,8 @@ impl IncomingToken {
 pub(crate) struct InvalidRetryTokenError;
 
 pub(crate) struct RetryToken {
+    /// The client's address
+    pub(crate) address: SocketAddr,
     /// The destination connection ID set in the very first packet from the client
     pub(crate) orig_dst_cid: ConnectionId,
     /// The time at which this token was issued
@@ -77,58 +87,64 @@ impl RetryToken {
     pub(crate) fn encode(
         &self,
         key: &dyn HandshakeTokenKey,
-        address: SocketAddr,
         retry_src_cid: ConnectionId,
     ) -> Vec<u8> {
-        let aead_key = key.aead_from_hkdf(&retry_src_cid);
-
         let mut buf = Vec::new();
-        encode_addr(&mut buf, address);
+
+        // Encode payload
+        encode_addr(&mut buf, self.address);
         self.orig_dst_cid.encode_long(&mut buf);
-        buf.write::<u64>(
-            self.issued
-                .duration_since(UNIX_EPOCH)
-                .map(|x| x.as_secs())
-                .unwrap_or(0),
-        );
+        encode_unix_secs(&mut buf, self.issued);
 
+        // Encrypt
+        let aead_key = key.aead_from_hkdf(&retry_src_cid);
         aead_key.seal(&mut buf, &[]).unwrap();
 
         buf
     }
 
-    fn from_bytes(
+    fn decode(
         key: &dyn HandshakeTokenKey,
-        address: SocketAddr,
         retry_src_cid: ConnectionId,
         raw_token_bytes: &[u8],
-    ) -> Result<Self, ValidationError> {
+    ) -> Option<Self> {
+        // Decrypt
         let aead_key = key.aead_from_hkdf(&retry_src_cid);
         let mut sealed_token = raw_token_bytes.to_vec();
+        let data = aead_key.open(&mut sealed_token, &[]).ok()?;
+
+        // Decode payload
+        let mut reader = &data[..];
+        let address = decode_addr(&mut reader)?;
+        let orig_dst_cid = ConnectionId::decode_long(&mut reader)?;
+        let issued = decode_unix_secs(&mut reader)?;
 
-        let data = aead_key.open(&mut sealed_token, &[])?;
-        let mut reader = io::Cursor::new(data);
-        let token_addr = decode_addr(&mut reader).ok_or(ValidationError::Unusable)?;
-        if token_addr != address {
-            return Err(ValidationError::InvalidRetry);
+        if !reader.is_empty() {
+            // Consider extra bytes a decoding error (it may be from an incompatible endpoint)
+            return None;
         }
-        let orig_dst_cid =
-            ConnectionId::decode_long(&mut reader).ok_or(ValidationError::Unusable)?;
-        let issued = UNIX_EPOCH
-            + Duration::new(
-                reader.get::<u64>().map_err(|_| ValidationError::Unusable)?,
-                0,
-            );
 
-        Ok(Self {
+        Some(Self {
+            address,
             orig_dst_cid,
             issued,
         })
     }
 }
 
 fn encode_addr(buf: &mut Vec<u8>, address: SocketAddr) {
-    match address.ip() {
+    encode_ip(buf, address.ip());
+    buf.put_u16(address.port());
+}
+
+fn decode_addr<B: Buf>(buf: &mut B) -> Option<SocketAddr> {
+    let ip = decode_ip(buf)?;
+    let port = buf.get().ok()?;
+    Some(SocketAddr::new(ip, port))
+}
+
+fn encode_ip(buf: &mut Vec<u8>, ip: IpAddr) {
+    match ip {
         IpAddr::V4(x) => {
             buf.put_u8(0);
             buf.put_slice(&x.octets());
@@ -138,50 +154,26 @@ fn encode_addr(buf: &mut Vec<u8>, address: SocketAddr) {
             buf.put_slice(&x.octets());
         }
     }
-    buf.put_u16(address.port());
 }
 
-fn decode_addr<B: Buf>(buf: &mut B) -> Option<SocketAddr> {
-    let ip = match buf.get_u8() {
-        0 => IpAddr::V4(buf.get().ok()?),
-        1 => IpAddr::V6(buf.get().ok()?),
-        _ => return None,
-    };
-    let port = buf.get_u16();
-    Some(SocketAddr::new(ip, port))
+fn decode_ip<B: Buf>(buf: &mut B) -> Option<IpAddr> {
+    match buf.get::<u8>().ok()? {
+        0 => buf.get().ok().map(IpAddr::V4),
+        1 => buf.get().ok().map(IpAddr::V6),
+        _ => None,
+    }
 }
 
-/// Error for a token failing to validate a client's address
-#[derive(Debug, Copy, Clone)]
-enum ValidationError {
-    /// Token may have come from a NEW_TOKEN frame (including from a different server or a previous
-    /// run of this server with different keys), and was not valid
-    ///
-    /// It should be silently ignored.
-    ///
-    /// In cases where a token cannot be decrypted/decoded, we must allow for the possibility that
-    /// this is caused not by client malfeasance, but by the token having been generated by an
-    /// incompatible endpoint, e.g. a different version or a neighbor behind the same load
-    /// balancer. In such cases we proceed as if there was no token.
-    ///
-    /// [_RFC 9000 § 8.1.3:_](https://www.rfc-editor.org/rfc/rfc9000.html#section-8.1.3-10)
-    ///
-    /// > If the token is invalid, then the server SHOULD proceed as if the client did not have a
-    /// > validated address, including potentially sending a Retry packet.
-    ///
-    /// That said, this may also be used when a token _can_ be unambiguously decrypted/decoded as a
-    /// token from a NEW_TOKEN frame, but is simply not valid.
-    Unusable,
-    /// Token was unambiguously from a Retry packet, and was not valid
-    ///
-    /// The connection cannot be established.
-    InvalidRetry,
+fn encode_unix_secs(buf: &mut Vec<u8>, time: SystemTime) {
+    buf.write::<u64>(
+        time.duration_since(UNIX_EPOCH)
+            .unwrap_or_default()
+            .as_secs(),
+    );
 }
 
-impl From<CryptoError> for ValidationError {
-    fn from(CryptoError: CryptoError) -> Self {
-        Self::Unusable
-    }
+fn decode_unix_secs<B: Buf>(buf: &mut B) -> Option<SystemTime> {
+    Some(UNIX_EPOCH + Duration::new(buf.get::<u64>().ok()?, 0))
 }
 
 /// Stateless reset token
@@ -256,16 +248,18 @@ mod test {
 
         let prk = hkdf::Salt::new(hkdf::HKDF_SHA256, &[]).extract(&master_key);
 
-        let addr = SocketAddr::new(Ipv6Addr::LOCALHOST.into(), 4433);
+        let address = SocketAddr::new(Ipv6Addr::LOCALHOST.into(), 4433);
         let retry_src_cid = RandomConnectionIdGenerator::new(MAX_CID_SIZE).generate_cid();
         let token = RetryToken {
+            address,
             orig_dst_cid: RandomConnectionIdGenerator::new(MAX_CID_SIZE).generate_cid(),
             issued: UNIX_EPOCH + Duration::new(42, 0), // Fractional seconds would be lost
         };
-        let encoded = token.encode(&prk, addr, retry_src_cid);
+        let encoded = token.encode(&prk, retry_src_cid);
 
-        let decoded = RetryToken::from_bytes(&prk, addr, retry_src_cid, &encoded)
-            .expect("token didn't validate");
+        let decoded =
+            RetryToken::decode(&prk, retry_src_cid, &encoded).expect("token didn't validate");
+        assert_eq!(token.address, decoded.address);
         assert_eq!(token.orig_dst_cid, decoded.orig_dst_cid);
         assert_eq!(token.issued, decoded.issued);
     }
@@ -276,7 +270,6 @@ mod test {
         use crate::cid_generator::{ConnectionIdGenerator, RandomConnectionIdGenerator};
         use crate::MAX_CID_SIZE;
         use rand::RngCore;
-        use std::net::Ipv6Addr;
 
         let rng = &mut rand::thread_rng();
 
@@ -285,7 +278,6 @@ mod test {
 
         let prk = hkdf::Salt::new(hkdf::HKDF_SHA256, &[]).extract(&master_key);
 
-        let addr = SocketAddr::new(Ipv6Addr::LOCALHOST.into(), 4433);
         let retry_src_cid = RandomConnectionIdGenerator::new(MAX_CID_SIZE).generate_cid();
 
         let mut invalid_token = Vec::new();
@@ -295,6 +287,6 @@ mod test {
         invalid_token.put_slice(&random_data);
 
         // Assert: garbage sealed data returns err
-        assert!(RetryToken::from_bytes(&prk, addr, retry_src_cid, &invalid_token).is_err());
+        assert!(RetryToken::decode(&prk, retry_src_cid, &invalid_token).is_none());
     }
 }