Propagate reset token updates to the endpoint synchronously

quinn-rs · May 29, 2023 · d8a3da9 · d8a3da9
1 parent b17c6e6
commit d8a3da9
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 27 deletions.
diff --git a/quinn-proto/src/connection/mod.rs b/quinn-proto/src/connection/mod.rs
@@ -2304,8 +2304,7 @@ impl Connection {
                         }
                     }
                     if let Some(token) = params.stateless_reset_token {
-                        self.endpoint_events
-                            .push_back(EndpointEventInner::ResetToken(self.path.remote, token));
+                        endpoint.set_reset_token(self.handle, self.path.remote, token);
                     }
                     self.handle_peer_params(params)?;
                     self.issue_first_cids(now, endpoint);
@@ -2652,7 +2651,7 @@ impl Connection {
                                 .pending
                                 .retire_cids
                                 .extend(retired);
-                            self.set_reset_token(reset_token);
+                            self.set_reset_token(reset_token, endpoint);
                         }
                         Err(InsertError::ExceedsLimit) => {
                             return Err(TransportError::CONNECTION_ID_LIMIT_ERROR(""));
@@ -2673,7 +2672,7 @@ impl Connection {
                     if self.side.is_server() && self.rem_cids.active_seq() == 0 {
                         // We're a server still using the initial remote CID for the client, so
                         // let's switch immediately to enable clientside stateless resets.
-                        self.update_rem_cid();
+                        self.update_rem_cid(endpoint);
                     }
                 }
                 Frame::NewToken { token } => {
@@ -2741,7 +2740,7 @@ impl Connection {
             );
             self.migrate(now, remote);
             // Break linkability, if possible
-            self.update_rem_cid();
+            self.update_rem_cid(endpoint);
             self.spin = false;
         }
 
@@ -2792,7 +2791,7 @@ impl Connection {
     }
 
     /// Switch to a previously unused remote connection ID, if possible
-    fn update_rem_cid(&mut self) {
+    fn update_rem_cid(&mut self, endpoint: &Endpoint) {
         let (reset_token, retired) = match self.rem_cids.next() {
             Some(x) => x,
             None => return,
@@ -2803,15 +2802,11 @@ impl Connection {
             .pending
             .retire_cids
             .extend(retired);
-        self.set_reset_token(reset_token);
+        self.set_reset_token(reset_token, endpoint);
     }
 
-    fn set_reset_token(&mut self, reset_token: ResetToken) {
-        self.endpoint_events
-            .push_back(EndpointEventInner::ResetToken(
-                self.path.remote,
-                reset_token,
-            ));
+    fn set_reset_token(&mut self, reset_token: ResetToken, endpoint: &Endpoint) {
+        endpoint.set_reset_token(self.handle, self.path.remote, reset_token);
         self.peer_params.stateless_reset_token = Some(reset_token);
     }
 

diff --git a/quinn-proto/src/endpoint.rs b/quinn-proto/src/endpoint.rs
@@ -86,18 +86,6 @@ impl Endpoint {
     ) -> Option<ConnectionEvent> {
         use EndpointEventInner::*;
         match event.0 {
-            ResetToken(remote, token) => {
-                let old = self.connections.lock().unwrap()[ch]
-                    .reset_token
-                    .replace((remote, token));
-                let mut index = self.index.write().unwrap();
-                if let Some(old) = old {
-                    index.connection_reset_tokens.remove(old.0, old.1);
-                }
-                if index.connection_reset_tokens.insert(remote, token, ch) {
-                    warn!("duplicate reset token");
-                }
-            }
             Drained => {
                 let conn = self.connections.lock().unwrap().remove(ch.0);
                 self.index.write().unwrap().remove(&conn);
@@ -106,6 +94,24 @@ impl Endpoint {
         None
     }
 
+    pub(crate) fn set_reset_token(
+        &self,
+        ch: ConnectionHandle,
+        remote: SocketAddr,
+        token: ResetToken,
+    ) {
+        let old = self.connections.lock().unwrap()[ch]
+            .reset_token
+            .replace((remote, token));
+        let mut index = self.index.write().unwrap();
+        if let Some(old) = old {
+            index.connection_reset_tokens.remove(old.0, old.1);
+        }
+        if index.connection_reset_tokens.insert(remote, token, ch) {
+            warn!("duplicate reset token");
+        }
+    }
+
     pub(crate) fn retire_cid(&self, ch: ConnectionHandle, seq: u64) {
         let cid = self.connections.lock().unwrap()[ch].loc_cids.remove(&seq);
         if let Some(cid) = cid {

diff --git a/quinn-proto/src/shared.rs b/quinn-proto/src/shared.rs
@@ -45,8 +45,6 @@ impl EndpointEvent {
 pub(crate) enum EndpointEventInner {
     /// The connection has been drained
     Drained,
-    /// The reset token and/or address eligible for generating resets has been updated
-    ResetToken(SocketAddr, ResetToken),
 }
 
 /// Protocol-level identifier for a connection.