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.

As of this commit, the only provided implementation of TokenStore is
NoneTokenStore, which is equivalent to the lack of a token store, and is
the default.

quinn-proto/src/config/mod.rs

@@ -17,8 +17,8 @@ use crate::{
     cid_generator::{ConnectionIdGenerator, HashedConnectionIdGenerator},
     crypto::{self, HandshakeTokenKey, HmacKey},
     shared::ConnectionId,
-    Duration, NoneTokenLog, RandomConnectionIdGenerator, SystemTime, TokenLog, VarInt,
-    VarIntBoundsExceeded, DEFAULT_SUPPORTED_VERSIONS, MAX_CID_SIZE,
+    Duration, NoneTokenLog, NoneTokenStore, RandomConnectionIdGenerator, SystemTime, TokenLog,
+    TokenStore, VarInt, VarIntBoundsExceeded, DEFAULT_SUPPORTED_VERSIONS, MAX_CID_SIZE,
 };
 
 mod transport;
@@ -458,6 +458,9 @@ pub struct ClientConfig {
     /// Cryptographic configuration to use
     pub(crate) crypto: Arc<dyn crypto::ClientConfig>,
 
+    /// Validation token store to use
+    pub(crate) token_store: 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>,
 
@@ -471,6 +474,7 @@ impl ClientConfig {
         Self {
             transport: Default::default(),
             crypto,
+            token_store: Arc::new(NoneTokenStore),
             initial_dst_cid_provider: Arc::new(|| {
                 RandomConnectionIdGenerator::new(MAX_CID_SIZE).generate_cid()
             }),
@@ -500,6 +504,15 @@ impl ClientConfig {
         self
     }
 
+    /// Set a custom [`TokenStore`]
+    ///
+    /// Defaults to [`NoneTokenStore`], which disables the use of tokens from NEW_TOKEN frames as a
+    /// client.
+    pub fn token_store(&mut self, store: 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;
@@ -532,6 +545,7 @@ impl fmt::Debug for ClientConfig {
         fmt.debug_struct("ClientConfig")
             .field("transport", &self.transport)
             // crypto not debug
+            // token_store not debug
             .field("version", &self.version)
             .finish_non_exhaustive()
     }

quinn-proto/src/connection/mod.rs

@@ -31,9 +31,9 @@ use crate::{
     },
     token::{ResetToken, Token, TokenPayload},
     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,
+    Dir, Duration, EndpointConfig, Frame, Instant, Side, StreamId, TokenStore, Transmit,
+    TransportError, TransportErrorCode, VarInt, INITIAL_MTU, MAX_CID_SIZE, MAX_STREAM_COUNT,
+    MIN_INITIAL_SIZE, TIMER_GRANULARITY,
 };
 
 mod ack_frequency;
@@ -2853,14 +2853,19 @@ impl Connection {
                     }
                 }
                 Frame::NewToken(NewToken { token }) => {
-                    if self.side.is_server() {
+                    let ConnectionSide::Client {
+                        token_store,
+                        server_name,
+                        ..
+                    } = &self.side
+                    else {
                         return Err(TransportError::PROTOCOL_VIOLATION("client sent NEW_TOKEN"));
-                    }
+                    };
                     if token.is_empty() {
                         return Err(TransportError::FRAME_ENCODING_ERROR("empty token"));
                     }
                     trace!("got new token");
-                    // TODO: Cache, or perhaps forward to user?
+                    token_store.insert(server_name, token);
                 }
                 Frame::Datagram(datagram) => {
                     if self
@@ -3669,6 +3674,8 @@ enum ConnectionSide {
     Client {
         /// Sent in every outgoing Initial packet. Always empty after Initial keys are discarded
         token: Bytes,
+        token_store: Arc<dyn TokenStore>,
+        server_name: String,
     },
     Server {
         server_config: Arc<ServerConfig>,
@@ -3702,8 +3709,13 @@ impl ConnectionSide {
 impl From<SideArgs> for ConnectionSide {
     fn from(side: SideArgs) -> Self {
         match side {
-            SideArgs::Client => Self::Client {
-                token: Bytes::new(),
+            SideArgs::Client {
+                token_store,
+                server_name,
+            } => Self::Client {
+                token: token_store.take(&server_name).unwrap_or_default(),
+                token_store,
+                server_name,
             },
             SideArgs::Server {
                 server_config,
@@ -3716,7 +3728,10 @@ impl From<SideArgs> for ConnectionSide {
 
 /// Parameters to `Connection::new` specific to it being client-side or server-side
 pub(crate) enum SideArgs {
-    Client,
+    Client {
+        token_store: Arc<dyn TokenStore>,
+        server_name: String,
+    },
     Server {
         server_config: Arc<ServerConfig>,
         pref_addr_cid: Option<ConnectionId>,

quinn-proto/src/endpoint.rs

@@ -364,7 +364,10 @@ impl Endpoint {
             now,
             tls,
             config.transport,
-            SideArgs::Client,
+            SideArgs::Client {
+                token_store: config.token_store,
+                server_name: server_name.into(),
+            },
         );
         Ok((ch, conn))
     }

quinn-proto/src/lib.rs

@@ -86,7 +86,7 @@ pub use crate::cid_generator::{
 
 mod token;
 use token::ResetToken;
-pub use token::{NoneTokenLog, TokenLog, TokenReuseError};
+pub use token::{NoneTokenLog, NoneTokenStore, TokenLog, TokenReuseError, TokenStore};
 
 #[cfg(feature = "arbitrary")]
 use arbitrary::Arbitrary;

quinn-proto/src/token.rs

@@ -4,7 +4,7 @@ use std::{
     net::{IpAddr, SocketAddr},
 };
 
-use bytes::{Buf, BufMut};
+use bytes::{Buf, BufMut, Bytes};
 use rand::Rng;
 
 use crate::{
@@ -75,6 +75,33 @@ impl TokenLog for NoneTokenLog {
     }
 }
 
+/// Responsible for storing validation tokens received from servers and retrieving them for use in
+/// subsequent connections
+pub trait TokenStore: Send + Sync {
+    /// Potentially store a token for later one-time use
+    ///
+    /// Called when a NEW_TOKEN frame is received from the server.
+    fn insert(&self, server_name: &str, token: Bytes);
+
+    /// Try to find and take a token that was stored with the given server name
+    ///
+    /// The same token must never be returned from `take` twice, as doing so can be used to
+    /// de-anonymize a client's traffic.
+    ///
+    /// Called when trying to connect to a server. It is always ok for this to return `None`.
+    fn take(&self, server_name: &str) -> Option<Bytes>;
+}
+
+/// Null implementation of [`TokenStore`], which does not store any tokens
+pub struct NoneTokenStore;
+
+impl TokenStore for NoneTokenStore {
+    fn insert(&self, _: &str, _: Bytes) {}
+    fn take(&self, _: &str) -> Option<Bytes> {
+        None
+    }
+}
+
 /// State in an `Incoming` determined by a token or lack thereof
 #[derive(Debug)]
 pub(crate) struct IncomingToken {

quinn/src/lib.rs

@@ -65,9 +65,9 @@ pub use proto::{
     congestion, crypto, AckFrequencyConfig, ApplicationClose, Chunk, ClientConfig, ClosedStream,
     ConfigError, ConnectError, ConnectionClose, ConnectionError, ConnectionId,
     ConnectionIdGenerator, ConnectionStats, Dir, EcnCodepoint, EndpointConfig, FrameStats,
-    FrameType, IdleTimeout, MtuDiscoveryConfig, NoneTokenLog, PathStats, ServerConfig, Side,
-    StdSystemTime, StreamId, TimeSource, TokenLog, TokenReuseError, Transmit, TransportConfig,
-    TransportErrorCode, UdpStats, VarInt, VarIntBoundsExceeded, Written,
+    FrameType, IdleTimeout, MtuDiscoveryConfig, NoneTokenLog, NoneTokenStore, PathStats,
+    ServerConfig, Side, StdSystemTime, StreamId, TimeSource, TokenLog, TokenReuseError, TokenStore,
+    Transmit, TransportConfig, TransportErrorCode, UdpStats, VarInt, VarIntBoundsExceeded, Written,
 };
 #[cfg(any(feature = "rustls-aws-lc-rs", feature = "rustls-ring"))]
 pub use rustls;