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, no implementation of TokenStore is provided, and it
defaults to None.

quinn-proto/src/config/mod.rs

@@ -17,8 +17,8 @@ use crate::{
     cid_generator::{ConnectionIdGenerator, HashedConnectionIdGenerator},
     crypto::{self, HandshakeTokenKey, HmacKey},
     shared::ConnectionId,
-    Duration, RandomConnectionIdGenerator, SystemTime, TokenLog, VarInt, VarIntBoundsExceeded,
-    DEFAULT_SUPPORTED_VERSIONS, MAX_CID_SIZE,
+    Duration, RandomConnectionIdGenerator, SystemTime, TokenLog, TokenStore, VarInt,
+    VarIntBoundsExceeded, DEFAULT_SUPPORTED_VERSIONS, MAX_CID_SIZE,
 };
 
 mod transport;
@@ -460,6 +460,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>,
 
@@ -473,6 +476,7 @@ impl ClientConfig {
         Self {
             transport: Default::default(),
             crypto,
+            token_store: None,
             initial_dst_cid_provider: Arc::new(|| {
                 RandomConnectionIdGenerator::new(MAX_CID_SIZE).generate_cid()
             }),
@@ -502,6 +506,16 @@ impl ClientConfig {
         self
     }
 
+    /// Set a custom [`TokenStore`]
+    ///
+    /// Defaults to `None`.
+    ///
+    /// 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;
@@ -534,6 +548,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;
@@ -2860,7 +2860,14 @@ impl Connection {
                         return Err(TransportError::FRAME_ENCODING_ERROR("empty token"));
                     }
                     trace!("got new token");
-                    // TODO: Cache, or perhaps forward to user?
+                    if let ConnectionSide::Client {
+                        token_store: Some(store),
+                        server_name,
+                        ..
+                    } = &self.side
+                    {
+                        store.insert(server_name, token);
+                    }
                 }
                 Frame::Datagram(datagram) => {
                     if self
@@ -3670,6 +3677,8 @@ enum ConnectionSide {
     Client {
         /// Sent in every outgoing Initial packet. Always empty after Initial keys are discarded
         token: Bytes,
+        token_store: Option<Arc<dyn TokenStore>>,
+        server_name: String,
     },
     Server {
         server_config: Arc<ServerConfig>,
@@ -3703,8 +3712,16 @@ 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
+                    .as_ref()
+                    .and_then(|store| store.take(&server_name))
+                    .unwrap_or_default(),
+                token_store,
+                server_name,
             },
             SideArgs::Server {
                 server_config,
@@ -3717,7 +3734,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: Option<Arc<dyn TokenStore>>,
+        server_name: String,
+    },
     Server {
         server_config: Arc<ServerConfig>,
         pref_addr_cid: Option<ConnectionId>,

quinn-proto/src/endpoint.rs

@@ -420,7 +420,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::{TokenLog, TokenReuseError};
+pub use token::{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::{
@@ -66,6 +66,23 @@ pub trait TokenLog: Send + Sync {
 /// Error for when a validation token may have been reused
 pub struct TokenReuseError;
 
+/// 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>;
+}
+
 /// State in an `Incoming` determined by a token or lack thereof
 #[derive(Debug)]
 pub(crate) struct IncomingToken {

quinn/src/lib.rs

@@ -66,8 +66,8 @@ pub use proto::{
     ConfigError, ConnectError, ConnectionClose, ConnectionError, ConnectionId,
     ConnectionIdGenerator, ConnectionStats, Dir, EcnCodepoint, EndpointConfig, FrameStats,
     FrameType, IdleTimeout, MtuDiscoveryConfig, PathStats, ServerConfig, Side, StdSystemTime,
-    StreamId, TimeSource, TokenLog, TokenReuseError, Transmit, TransportConfig, TransportErrorCode,
-    UdpStats, VarInt, VarIntBoundsExceeded, Written,
+    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;