Factor out IncomingToken

quinn-proto/src/endpoint.rs

@@ -29,7 +29,7 @@ use crate::{
         ConnectionEvent, ConnectionEventInner, ConnectionId, DatagramConnectionEvent, EcnCodepoint,
         EndpointEvent, EndpointEventInner, IssuedCid,
     },
-    token,
+    token::{IncomingToken, InvalidRetryTokenError},
     transport_parameters::{PreferredAddress, TransportParameters},
     Duration, Instant, ResetToken, RetryToken, Side, Transmit, TransportConfig, TransportError,
     INITIAL_MTU, MAX_CID_SIZE, MIN_INITIAL_SIZE, RESET_TOKEN_SIZE,
@@ -494,33 +494,18 @@ impl Endpoint {
 
         let server_config = self.server_config.as_ref().unwrap().clone();
 
-        let (retry_src_cid, orig_dst_cid) = if header.token.is_empty() {
-            (None, header.dst_cid)
-        } else {
-            match RetryToken::from_bytes(
-                &*server_config.token_key,
-                &addresses.remote,
-                &header.dst_cid,
-                &header.token,
-            ) {
-                Ok(token)
-                    if token.issued + server_config.retry_token_lifetime
-                        > server_config.time_source.now() =>
-                {
-                    (Some(header.dst_cid), token.orig_dst_cid)
-                }
-                Err(token::ValidationError::Unusable) => (None, header.dst_cid),
-                _ => {
-                    debug!("rejecting invalid stateless retry token");
-                    return Some(DatagramEvent::Response(self.initial_close(
-                        header.version,
-                        addresses,
-                        &crypto,
-                        &header.src_cid,
-                        TransportError::INVALID_TOKEN(""),
-                        buf,
-                    )));
-                }
+        let token = match IncomingToken::from_header(&header, &server_config, addresses.remote) {
+            Ok(token) => token,
+            Err(InvalidRetryTokenError) => {
+                debug!("rejecting invalid retry token");
+                return Some(DatagramEvent::Response(self.initial_close(
+                    header.version,
+                    addresses,
+                    &crypto,
+                    &header.src_cid,
+                    TransportError::INVALID_TOKEN(""),
+                    buf,
+                )));
             }
         };
 
@@ -539,8 +524,7 @@ impl Endpoint {
             },
             rest,
             crypto,
-            retry_src_cid,
-            orig_dst_cid,
+            token,
             incoming_idx,
             improper_drop_warner: IncomingImproperDropWarner,
         }))
@@ -630,8 +614,8 @@ impl Endpoint {
             &mut self.rng,
         );
         params.stateless_reset_token = Some(ResetToken::new(&*self.config.reset_key, &loc_cid));
-        params.original_dst_cid = Some(incoming.orig_dst_cid);
-        params.retry_src_cid = incoming.retry_src_cid;
+        params.original_dst_cid = Some(incoming.token.orig_dst_cid);
+        params.retry_src_cid = incoming.token.retry_src_cid;
         let mut pref_addr_cid = None;
         if server_config.preferred_address_v4.is_some()
             || server_config.preferred_address_v6.is_some()
@@ -1192,8 +1176,7 @@ pub struct Incoming {
     packet: InitialPacket,
     rest: Option<BytesMut>,
     crypto: Keys,
-    retry_src_cid: Option<ConnectionId>,
-    orig_dst_cid: ConnectionId,
+    token: IncomingToken,
     incoming_idx: usize,
     improper_drop_warner: IncomingImproperDropWarner,
 }
@@ -1216,12 +1199,12 @@ impl Incoming {
     /// This means that the sender of the initial packet has proved that they can receive traffic
     /// sent to `self.remote_address()`.
     pub fn remote_address_validated(&self) -> bool {
-        self.retry_src_cid.is_some()
+        self.token.retry_src_cid.is_some()
     }
 
     /// The original destination connection ID sent by the client
     pub fn orig_dst_cid(&self) -> &ConnectionId {
-        &self.orig_dst_cid
+        &self.token.orig_dst_cid
     }
 }
 
@@ -1232,8 +1215,7 @@ impl fmt::Debug for Incoming {
             .field("ecn", &self.ecn)
             // packet doesn't implement debug
             // rest is too big and not meaningful enough
-            .field("retry_src_cid", &self.retry_src_cid)
-            .field("orig_dst_cid", &self.orig_dst_cid)
+            .field("token", &self.token)
             .field("incoming_idx", &self.incoming_idx)
             // improper drop warner contains no information
             .finish_non_exhaustive()

quinn-proto/src/token.rs

@@ -8,8 +8,9 @@ use bytes::{Buf, BufMut};
 use crate::{
     coding::{BufExt, BufMutExt},
     crypto::{CryptoError, HandshakeTokenKey, HmacKey},
+    packet::InitialHeader,
     shared::ConnectionId,
-    Duration, SystemTime, RESET_TOKEN_SIZE, UNIX_EPOCH,
+    Duration, ServerConfig, SystemTime, RESET_TOKEN_SIZE, UNIX_EPOCH,
 };
 
 pub(crate) struct RetryToken {
@@ -43,7 +44,7 @@ impl RetryToken {
         buf
     }
 
-    pub(crate) fn from_bytes(
+    fn from_bytes(
         key: &dyn HandshakeTokenKey,
         address: &SocketAddr,
         retry_src_cid: &ConnectionId,
@@ -71,6 +72,22 @@ impl RetryToken {
             issued,
         })
     }
+
+    /// Ensure that this token validates an `Incoming`, and construct its token state
+    fn validate(
+        &self,
+        header: &InitialHeader,
+        server_config: &ServerConfig,
+    ) -> Result<IncomingToken, ValidationError> {
+        if self.issued + server_config.retry_token_lifetime < server_config.time_source.now() {
+            return Err(ValidationError::InvalidRetry);
+        }
+
+        Ok(IncomingToken {
+            retry_src_cid: Some(header.dst_cid),
+            orig_dst_cid: self.orig_dst_cid,
+        })
+    }
 }
 
 fn encode_addr(buf: &mut Vec<u8>, address: &SocketAddr) {
@@ -99,7 +116,7 @@ fn decode_addr<B: Buf>(buf: &mut B) -> Option<SocketAddr> {
 
 /// Error for a token failing to validate a client's address
 #[derive(Debug, Copy, Clone)]
-pub(crate) enum ValidationError {
+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
     ///
@@ -178,6 +195,52 @@ impl fmt::Display for ResetToken {
     }
 }
 
+/// State in an `Incoming` determined by a token or lack thereof
+#[derive(Debug)]
+pub(crate) struct IncomingToken {
+    pub(crate) retry_src_cid: Option<ConnectionId>,
+    pub(crate) orig_dst_cid: ConnectionId,
+}
+
+impl IncomingToken {
+    /// Construct for an `Incoming` which is not validated by a token
+    fn unvalidated(header: &InitialHeader) -> Self {
+        Self {
+            retry_src_cid: None,
+            orig_dst_cid: header.dst_cid,
+        }
+    }
+
+    /// Construct for an `Incoming` given the first packet header, or error if the connection
+    /// cannot be established
+    pub(crate) fn from_header(
+        header: &InitialHeader,
+        server_config: &ServerConfig,
+        remote_address: SocketAddr,
+    ) -> Result<Self, InvalidRetryTokenError> {
+        if header.token.is_empty() {
+            return Ok(Self::unvalidated(header));
+        }
+
+        RetryToken::from_bytes(
+            &*server_config.token_key,
+            &remote_address,
+            &header.dst_cid,
+            &header.token,
+        )
+        .and_then(|token| token.validate(header, server_config))
+        .or_else(|e| match e {
+            ValidationError::Unusable => Ok(Self::unvalidated(header)),
+            ValidationError::InvalidRetry => Err(InvalidRetryTokenError),
+        })
+    }
+}
+
+/// Error for a token being unambiguously from a Retry packet, and not valid
+///
+/// The connection cannot be established.
+pub(crate) struct InvalidRetryTokenError;
+
 #[cfg(all(test, any(feature = "aws-lc-rs", feature = "ring")))]
 mod test {
     #[cfg(all(feature = "aws-lc-rs", not(feature = "ring")))]