Set retry token on every outgoing Initial, in any state

quinn-proto/src/connection/mod.rs

@@ -178,6 +178,8 @@ pub struct Connection {
     authentication_failures: u64,
     /// Why the connection was lost, if it has been
     error: Option<ConnectionError>,
+    /// Sent in every outgoing Initial packet. Always empty for servers.
+    retry_token: Bytes,
 
     //
     // Queued non-retransmittable 1-RTT data
@@ -288,6 +290,7 @@ impl Connection {
             timers: TimerTable::default(),
             authentication_failures: 0,
             error: None,
+            retry_token: Bytes::new(),
 
             path_response: None,
             close: false,
@@ -2100,8 +2103,9 @@ 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.state = State::Handshake(state::Handshake {
-                    token: packet.payload.freeze().split_to(token_len),
+                    token: Bytes::new(),
                     rem_cid_set: false,
                     client_hello: None,
                 });

quinn-proto/src/connection/packet_builder.rs

@@ -4,7 +4,7 @@ use bytes::Bytes;
 use rand::Rng;
 use tracing::{trace, trace_span};
 
-use super::{spaces::SentPacket, Connection, SentFrames, State};
+use super::{spaces::SentPacket, Connection, SentFrames};
 use crate::{
     frame::{self, Close},
     packet::{Header, LongType, PacketNumber, PartialEncode, SpaceId, FIXED_BIT},
@@ -39,7 +39,6 @@ impl PacketBuilder {
         conn: &mut Connection,
         version: u32,
     ) -> Option<PacketBuilder> {
-        let is_client = conn.side().is_client();
         // Initiate key update if we're approaching the confidentiality limit
         let confidentiality_limit = conn.spaces[space_id]
             .crypto
@@ -107,10 +106,7 @@ impl PacketBuilder {
             SpaceId::Initial => Header::Initial {
                 src_cid: conn.handshake_cid,
                 dst_cid: conn.rem_cids.active(),
-                token: match conn.state {
-                    State::Handshake(ref state) if is_client => state.token.clone(),
-                    _ => Bytes::new(),
-                },
+                token: conn.retry_token.clone(),
                 number,
                 version,
             },