Allow client to use NEW_TOKEN frames

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

quinn-proto/src/config.rs

@@ -20,8 +20,8 @@ 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, TokenMemoryCache, TokenStore, VarInt,
+    VarIntBoundsExceeded, DEFAULT_SUPPORTED_VERSIONS, INITIAL_MTU, MAX_CID_SIZE, MAX_UDP_PAYLOAD,
 };
 
 /// Parameters governing the core QUIC state machine
@@ -1066,6 +1066,9 @@ pub struct ClientConfig {
     /// Cryptographic configuration to use
     pub(crate) crypto: Arc<dyn crypto::ClientConfig>,
 
+    /// Validation token store to use
+    pub(crate) token_store: Option<Arc<dyn TokenStore>>,
+
     /// Provider that populates the destination connection ID of Initial Packets
     pub(crate) initial_dst_cid_provider: Arc<dyn Fn() -> ConnectionId + Send + Sync>,
 
@@ -1079,6 +1082,7 @@ impl ClientConfig {
         Self {
             transport: Default::default(),
             crypto,
+            token_store: Some(Arc::new(TokenMemoryCache::<2>::default())),
             initial_dst_cid_provider: Arc::new(|| {
                 RandomConnectionIdGenerator::new(MAX_CID_SIZE).generate_cid()
             }),
@@ -1108,6 +1112,19 @@ impl ClientConfig {
         self
     }
 
+    /// Set a custom [`TokenStore`]
+    ///
+    /// Defaults to a [`TokenMemoryCache`] limited to 256 servers and 2 tokens per server. This
+    /// default is chosen to complement `rustls`'s default [`ClientSessionStore`].
+    ///
+    /// [`ClientSessionStore`]: rustls::client::ClientSessionStore
+    ///
+    /// Setting to `None` disables the use of tokens from NEW_TOKEN frames as a client.
+    pub fn token_store(&mut self, store: Option<Arc<dyn TokenStore>>) -> &mut Self {
+        self.token_store = store;
+        self
+    }
+
     /// Set the QUIC version to use
     pub fn version(&mut self, version: u32) -> &mut Self {
         self.version = version;
@@ -1139,6 +1156,7 @@ impl fmt::Debug for ClientConfig {
     fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
         fmt.debug_struct("ClientConfig<T>")
             .field("transport", &self.transport)
+            // token_store not debug
             // crypto not debug
             .field("version", &self.version)
             .finish_non_exhaustive()

quinn-proto/src/connection/mod.rs

@@ -31,9 +31,9 @@ use crate::{
     },
     token::{ResetToken, Token, TokenInner},
     transport_parameters::TransportParameters,
-    Dir, Duration, EndpointConfig, Frame, Instant, Side, StreamId, SystemTime, Transmit,
-    TransportError, TransportErrorCode, VarInt, INITIAL_MTU, MAX_CID_SIZE, MAX_STREAM_COUNT,
-    MIN_INITIAL_SIZE, TIMER_GRANULARITY,
+    Dir, Duration, EndpointConfig, Frame, Instant, Side, StreamId, SystemTime, TokenStore,
+    Transmit, TransportError, TransportErrorCode, VarInt, INITIAL_MTU, MAX_CID_SIZE,
+    MAX_STREAM_COUNT, MIN_INITIAL_SIZE, TIMER_GRANULARITY,
 };
 
 mod ack_frequency;
@@ -193,7 +193,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,
 
@@ -226,6 +226,9 @@ pub struct Connection {
     /// no outgoing application data.
     app_limited: bool,
 
+    token_store: Option<Arc<dyn TokenStore>>,
+    server_name: Option<String>,
+
     streams: StreamsState,
     /// Surplus remote CIDs for future use on new paths
     rem_cids: CidQueue,
@@ -257,6 +260,8 @@ impl Connection {
         allow_mtud: bool,
         rng_seed: [u8; 32],
         path_validated: bool,
+        token_store: Option<Arc<dyn TokenStore>>,
+        server_name: Option<String>,
     ) -> Self {
         let side = if server_config.is_some() {
             Side::Server
@@ -273,6 +278,10 @@ impl Connection {
             client_hello: None,
         });
         let mut rng = StdRng::from_seed(rng_seed);
+        let token = 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,
@@ -323,7 +332,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),
@@ -345,6 +354,9 @@ impl Connection {
             receiving_ecn: false,
             total_authed_packets: 0,
 
+            token_store,
+            server_name,
+
             streams: StreamsState::new(
                 side,
                 config.max_concurrent_uni_streams,
@@ -2104,7 +2116,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;
@@ -2423,7 +2435,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,
@@ -2865,7 +2877,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.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

@@ -31,7 +31,7 @@ use crate::{
     },
     token::{TokenDecodeError, TokenInner},
     transport_parameters::{PreferredAddress, TransportParameters},
-    Duration, Instant, ResetToken, Side, SystemTime, Token, Transmit, TransportConfig,
+    Duration, Instant, ResetToken, Side, SystemTime, Token, TokenStore, Transmit, TransportConfig,
     TransportError, INITIAL_MTU, MAX_CID_SIZE, MIN_INITIAL_SIZE, RESET_TOKEN_SIZE,
 };
 
@@ -433,6 +433,8 @@ impl Endpoint {
             None,
             config.transport,
             true,
+            config.token_store,
+            Some(server_name.into()),
         );
         Ok((ch, conn))
     }
@@ -687,6 +689,8 @@ impl Endpoint {
             Some(server_config),
             transport_config,
             remote_address_validated,
+            None,
+            None,
         );
         self.index.insert_initial(dst_cid, ch);
 
@@ -853,6 +857,8 @@ impl Endpoint {
         server_config: Option<Arc<ServerConfig>>,
         transport_config: Arc<TransportConfig>,
         path_validated: bool,
+        token_store: Option<Arc<dyn TokenStore>>,
+        server_name: Option<String>,
     ) -> Connection {
         let mut rng_seed = [0; 32];
         self.rng.fill_bytes(&mut rng_seed);
@@ -877,6 +883,8 @@ impl Endpoint {
             self.allow_mtud,
             rng_seed,
             path_validated,
+            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 token_store;
+pub use token_store::{TokenMemoryCache, TokenStore};
+
 #[cfg(feature = "arbitrary")]
 use arbitrary::Arbitrary;