From e5cb4e3abcf2d9ae39845c4b39bcf4ef91dfe8a5 Mon Sep 17 00:00:00 2001
From: Benjamin Saunders <ben.e.saunders@gmail.com>
Date: Wed, 3 Oct 2018 19:04:49 -0700
Subject: [PATCH] Fix stateless resets

---
 quinn-proto/src/connection.rs | 8 ++++++--
 quinn-proto/src/endpoint.rs   | 8 ++++----
 quinn-proto/src/tests.rs      | 8 +-------
 3 files changed, 11 insertions(+), 13 deletions(-)

diff --git a/quinn-proto/src/connection.rs b/quinn-proto/src/connection.rs
index b564ab7c1..1f995ba6f 100644
--- a/quinn-proto/src/connection.rs
+++ b/quinn-proto/src/connection.rs
@@ -8,7 +8,7 @@ use rand::distributions::Distribution;
 use slog::Logger;
 
 use coding::{BufExt, BufMutExt};
-use crypto::{ConnectError, Crypto, TLSError, TlsSession, ACK_DELAY_EXPONENT};
+use crypto::{ConnectError, Crypto, TLSError, TlsSession, ACK_DELAY_EXPONENT, reset_token_for};
 use endpoint::{Config, Context, Event, Io, Timer};
 use packet::{
     set_payload_length, types, ConnectionId, Header, Packet, PacketNumber, AEAD_TAG_SIZE,
@@ -902,9 +902,13 @@ impl Connection {
         }; // TODO: Send close?
 
         trace!(ctx.log, "got initial");
+        let server_params = TransportParameters {
+            stateless_reset_token: Some(reset_token_for(&ctx.listen_keys.as_ref().unwrap().reset, &self.local_id)),
+            ..TransportParameters::new(&ctx.config)
+        };
         let mut tls = TlsSession::new_server(
             &ctx.config.tls_server_config,
-            &TransportParameters::new(&ctx.config),
+            &server_params,
         );
         self.read_tls(&mut tls, &frame);
         tls.process_new_packets()?;
diff --git a/quinn-proto/src/endpoint.rs b/quinn-proto/src/endpoint.rs
index bc1575010..0ecb129b0 100644
--- a/quinn-proto/src/endpoint.rs
+++ b/quinn-proto/src/endpoint.rs
@@ -125,7 +125,6 @@ pub struct Endpoint {
     connection_ids: FnvHashMap<ConnectionId, ConnectionHandle>,
     connection_remotes: FnvHashMap<SocketAddrV6, ConnectionHandle>,
     pub(crate) connections: Slab<Connection>,
-    listen_keys: Option<ListenKeys>,
 }
 
 pub struct Context {
@@ -140,6 +139,7 @@ pub struct Context {
     pub dirty_conns: FnvHashSet<ConnectionHandle>,
     pub readable_conns: FnvHashSet<ConnectionHandle>,
     pub initial_packet_number: distributions::Uniform<u64>,
+    pub listen_keys: Option<ListenKeys>,
 }
 
 impl Context {
@@ -219,8 +219,8 @@ impl Endpoint {
                 readable_conns: FnvHashSet::default(),
                 incoming: VecDeque::new(),
                 incoming_handshakes: 0,
+                listen_keys: listen,
             },
-            listen_keys: listen,
             connection_ids_initial: FnvHashMap::default(),
             connection_ids: FnvHashMap::default(),
             connection_remotes: FnvHashMap::default(),
@@ -229,7 +229,7 @@ impl Endpoint {
     }
 
     fn listen(&self) -> bool {
-        self.listen_keys.is_some()
+        self.ctx.listen_keys.is_some()
     }
 
     /// Get an application-facing event
@@ -447,7 +447,7 @@ impl Endpoint {
                 self.ctx.rng.fill_bytes(&mut buf[start..start + padding]);
             }
             buf.extend(&reset_token_for(
-                &self.listen_keys.as_ref().unwrap().reset,
+                &self.ctx.listen_keys.as_ref().unwrap().reset,
                 &dest_id,
             ));
             self.ctx.io.push_back(Io::Transmit {
diff --git a/quinn-proto/src/tests.rs b/quinn-proto/src/tests.rs
index 028b11948..2cd72f623 100644
--- a/quinn-proto/src/tests.rs
+++ b/quinn-proto/src/tests.rs
@@ -434,26 +434,20 @@ fn stateless_retry() {
 }
 */
 
-/*
 #[test]
 fn stateless_reset() {
     let mut pair = Pair::default();
     let (client_conn, _) = pair.connect();
-    assert_matches!(pair.client.poll(), Some((conn, Event::NewSessionTicket { .. })) if conn == client_conn);
-    let mut config = Config::default();
-    set_server_certificate(&mut config);
     pair.server.endpoint = Endpoint::new(
         pair.log.new(o!("peer" => "server")),
-        config,
+        Config::default(),
         Some(*LISTEN_KEYS),
     ).unwrap();
     pair.client.ping(client_conn);
     info!(pair.log, "resetting");
     pair.drive();
-    assert_matches!(pair.client.poll(), Some((conn, Event::NewSessionTicket { .. })) if conn == client_conn);
     assert_matches!(pair.client.poll(), Some((conn, Event::ConnectionLost { reason: ConnectionError::Reset })) if conn == client_conn);
 }
-*/
 
 #[test]
 fn finish_stream() {