Allow client to use NEW_TOKEN frames

When a client receives a token from a NEW_TOKEN frame, it submits it to
a ValidationTokenStore object for storage. When an endpoint connects to
a server, it queries the ValidationTokenStore object for a token
applicable to the server name, and uses it if one is retrieved.

quinn-proto/src/config.rs

@@ -20,8 +20,9 @@ use crate::{
     congestion,
     crypto::{self, HandshakeTokenKey, HmacKey},
     shared::ConnectionId,
-    Duration, RandomConnectionIdGenerator, TokenLog, VarInt, VarIntBoundsExceeded,
-    DEFAULT_SUPPORTED_VERSIONS, INITIAL_MTU, MAX_CID_SIZE, MAX_UDP_PAYLOAD,
+    Duration, RandomConnectionIdGenerator, TokenLog, ValidationTokenMemoryCache,
+    ValidationTokenStore, VarInt, VarIntBoundsExceeded, DEFAULT_SUPPORTED_VERSIONS, INITIAL_MTU,
+    MAX_CID_SIZE, MAX_UDP_PAYLOAD,
 };
 
 /// Parameters governing the core QUIC state machine
@@ -1061,6 +1062,9 @@ pub struct ClientConfig {
     /// Cryptographic configuration to use
     pub(crate) crypto: Arc<dyn crypto::ClientConfig>,
 
+    /// Validation token store to use
+    pub(crate) validation_token_store: Option<Arc<dyn ValidationTokenStore>>,
+
     /// Provider that populates the destination connection ID of Initial Packets
     pub(crate) initial_dst_cid_provider: Arc<dyn Fn() -> ConnectionId + Send + Sync>,
 
@@ -1074,6 +1078,7 @@ impl ClientConfig {
         Self {
             transport: Default::default(),
             crypto,
+            validation_token_store: Some(Arc::new(ValidationTokenMemoryCache::<2>::default())),
             initial_dst_cid_provider: Arc::new(|| {
                 RandomConnectionIdGenerator::new(MAX_CID_SIZE).generate_cid()
             }),
@@ -1103,6 +1108,21 @@ impl ClientConfig {
         self
     }
 
+    /// Set a custom [`ValidationTokenStore`]
+    ///
+    /// Defaults to an in-memory store limited to 256 servers and 2 tokens per server. This default
+    /// is chosen to complement `rustls`'s default
+    /// [`ClientSessionStore`][rustls::client::ClientSessionStore].
+    ///
+    /// Setting to `None` disables the use of tokens from NEW_TOKEN frames as a client.
+    pub fn validation_token_store(
+        &mut self,
+        validation_token_store: Option<Arc<dyn ValidationTokenStore>>,
+    ) -> &mut Self {
+        self.validation_token_store = validation_token_store;
+        self
+    }
+
     /// Set the QUIC version to use
     pub fn version(&mut self, version: u32) -> &mut Self {
         self.version = version;
@@ -1134,6 +1154,7 @@ impl fmt::Debug for ClientConfig {
     fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
         fmt.debug_struct("ClientConfig<T>")
             .field("transport", &self.transport)
+            // validation_token_store not debug
             // crypto not debug
             .field("version", &self.version)
             .finish_non_exhaustive()

quinn-proto/src/connection/mod.rs

@@ -33,8 +33,8 @@ use crate::{
     token::{ResetToken, Token, TokenInner},
     transport_parameters::TransportParameters,
     Dir, Duration, EndpointConfig, Frame, Instant, Side, StreamId, Transmit, TransportError,
-    TransportErrorCode, VarInt, INITIAL_MTU, MAX_CID_SIZE, MAX_STREAM_COUNT, MIN_INITIAL_SIZE,
-    TIMER_GRANULARITY,
+    TransportErrorCode, ValidationTokenStore, VarInt, INITIAL_MTU, MAX_CID_SIZE, MAX_STREAM_COUNT,
+    MIN_INITIAL_SIZE, TIMER_GRANULARITY,
 };
 
 mod ack_frequency;
@@ -194,7 +194,7 @@ pub struct Connection {
     error: Option<ConnectionError>,
     /// Sent in every outgoing Initial packet. Always empty for servers and after Initial keys are
     /// discarded.
-    retry_token: Bytes,
+    token: Bytes,
     /// Identifies Data-space packet numbers to skip. Not used in earlier spaces.
     packet_number_filter: PacketNumberFilter,
 
@@ -227,6 +227,9 @@ pub struct Connection {
     /// no outgoing application data.
     app_limited: bool,
 
+    validation_token_store: Option<Arc<dyn ValidationTokenStore>>,
+    server_name: Option<String>,
+
     streams: StreamsState,
     /// Surplus remote CIDs for future use on new paths
     rem_cids: CidQueue,
@@ -258,6 +261,8 @@ impl Connection {
         allow_mtud: bool,
         rng_seed: [u8; 32],
         path_validated: bool,
+        validation_token_store: Option<Arc<dyn ValidationTokenStore>>,
+        server_name: Option<String>,
     ) -> Self {
         let side = if server_config.is_some() {
             Side::Server
@@ -274,6 +279,10 @@ impl Connection {
             client_hello: None,
         });
         let mut rng = StdRng::from_seed(rng_seed);
+        let token = validation_token_store
+            .as_ref()
+            .and_then(|store| store.take(server_name.as_ref().unwrap()))
+            .unwrap_or_default();
         let mut this = Self {
             endpoint_config,
             server_config,
@@ -324,7 +333,7 @@ impl Connection {
             timers: TimerTable::default(),
             authentication_failures: 0,
             error: None,
-            retry_token: Bytes::new(),
+            token,
             #[cfg(test)]
             packet_number_filter: match config.deterministic_packet_numbers {
                 false => PacketNumberFilter::new(&mut rng),
@@ -346,6 +355,9 @@ impl Connection {
             receiving_ecn: false,
             total_authed_packets: 0,
 
+            validation_token_store,
+            server_name,
+
             streams: StreamsState::new(
                 side,
                 config.max_concurrent_uni_streams,
@@ -2105,7 +2117,7 @@ impl Connection {
         trace!("discarding {:?} keys", space_id);
         if space_id == SpaceId::Initial {
             // No longer needed
-            self.retry_token = Bytes::new();
+            self.token = Bytes::new();
         }
         let space = &mut self.spaces[space_id];
         space.crypto = None;
@@ -2424,7 +2436,7 @@ impl Connection {
                 self.streams.retransmit_all_for_0rtt();
 
                 let token_len = packet.payload.len() - 16;
-                self.retry_token = packet.payload.freeze().split_to(token_len);
+                self.token = packet.payload.freeze().split_to(token_len);
                 self.state = State::Handshake(state::Handshake {
                     expected_token: Bytes::new(),
                     rem_cid_set: false,
@@ -2866,7 +2878,9 @@ impl Connection {
                         return Err(TransportError::FRAME_ENCODING_ERROR("empty token"));
                     }
                     trace!("got new token");
-                    // TODO: Cache, or perhaps forward to user?
+                    if let Some(store) = self.validation_token_store.as_ref() {
+                        store.store(self.server_name.as_ref().unwrap(), token);
+                    }
                 }
                 Frame::Datagram(datagram) => {
                     if self

quinn-proto/src/connection/packet_builder.rs

@@ -113,7 +113,7 @@ impl PacketBuilder {
             SpaceId::Initial => Header::Initial(InitialHeader {
                 src_cid: conn.handshake_cid,
                 dst_cid,
-                token: conn.retry_token.clone(),
+                token: conn.token.clone(),
                 number,
                 version,
             }),

quinn-proto/src/endpoint.rs

@@ -32,7 +32,8 @@ use crate::{
     token::{TokenDecodeError, TokenInner},
     transport_parameters::{PreferredAddress, TransportParameters},
     Duration, Instant, ResetToken, Side, SystemTime, Token, Transmit, TransportConfig,
-    TransportError, INITIAL_MTU, MAX_CID_SIZE, MIN_INITIAL_SIZE, RESET_TOKEN_SIZE,
+    TransportError, ValidationTokenStore, INITIAL_MTU, MAX_CID_SIZE, MIN_INITIAL_SIZE,
+    RESET_TOKEN_SIZE,
 };
 
 /// The main entry point to the library
@@ -432,6 +433,8 @@ impl Endpoint {
             None,
             config.transport,
             true,
+            config.validation_token_store,
+            Some(server_name.into()),
         );
         Ok((ch, conn))
     }
@@ -685,6 +688,8 @@ impl Endpoint {
             Some(server_config),
             transport_config,
             remote_address_validated,
+            None,
+            None,
         );
         self.index.insert_initial(dst_cid, ch);
 
@@ -851,6 +856,8 @@ impl Endpoint {
         server_config: Option<Arc<ServerConfig>>,
         transport_config: Arc<TransportConfig>,
         path_validated: bool,
+        new_token_store: Option<Arc<dyn ValidationTokenStore>>,
+        server_name: Option<String>,
     ) -> Connection {
         let mut rng_seed = [0; 32];
         self.rng.fill_bytes(&mut rng_seed);
@@ -875,6 +882,8 @@ impl Endpoint {
             self.allow_mtud,
             rng_seed,
             path_validated,
+            new_token_store,
+            server_name,
         );
 
         let mut cids_issued = 0;

quinn-proto/src/lib.rs

@@ -96,6 +96,9 @@ mod bloom_token_log;
 #[cfg(feature = "fastbloom")]
 pub use bloom_token_log::BloomTokenLog;
 
+mod validation_token_store;
+pub use validation_token_store::{ValidationTokenMemoryCache, ValidationTokenStore};
+
 #[cfg(feature = "arbitrary")]
 use arbitrary::Arbitrary;