From afc7d7f8ae3ef690e7da4db7beadd6c1b07eae03 Mon Sep 17 00:00:00 2001 From: Phoenix Kahlo Date: Mon, 16 Dec 2024 10:51:24 -0600 Subject: [PATCH] proto: Factor out IncomingToken::from_header Moves more logic from Endpoint::handle_first_packet to the proto::token module. Due to improved decoupling between modules, Token::from_bytes and ValidationError can now be made private. --- quinn-proto/src/endpoint.rs | 46 +++++++++------------------ quinn-proto/src/token.rs | 62 +++++++++++++++++++++++++++++++++++-- 2 files changed, 73 insertions(+), 35 deletions(-) diff --git a/quinn-proto/src/endpoint.rs b/quinn-proto/src/endpoint.rs index 4c3d70ae6..95974c436 100644 --- a/quinn-proto/src/endpoint.rs +++ b/quinn-proto/src/endpoint.rs @@ -29,7 +29,7 @@ use crate::{ ConnectionEvent, ConnectionEventInner, ConnectionId, DatagramConnectionEvent, EcnCodepoint, EndpointEvent, EndpointEventInner, IssuedCid, }, - token::{self, IncomingToken}, + 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,10 +524,7 @@ impl Endpoint { }, rest, crypto, - token: IncomingToken { - retry_src_cid, - orig_dst_cid, - }, + token, incoming_idx, improper_drop_warner: IncomingImproperDropWarner, })) diff --git a/quinn-proto/src/token.rs b/quinn-proto/src/token.rs index 4ff145ab7..59dcc5575 100644 --- a/quinn-proto/src/token.rs +++ b/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 { + 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, address: &SocketAddr) { @@ -99,7 +116,7 @@ fn decode_addr(buf: &mut B) -> Option { /// 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 /// @@ -185,6 +202,45 @@ pub(crate) struct IncomingToken { 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 { + 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")))]