From 81b8d1e3ca302ee591535eaf7d49b067a9dfd30d Mon Sep 17 00:00:00 2001
From: Benjamin Saunders <ben.e.saunders@gmail.com>
Date: Fri, 19 Mar 2021 17:07:12 -0700
Subject: [PATCH] Fix handshake deadlock

When the server's first flight is interrupted by anti-amplification
and the client's second flight is lost, the client has no
ack-eliciting packets in flight but will not be able to proceed until
an anti-amplification deadlock prevention probe is sent.
---
 quinn-proto/src/connection/mod.rs |  5 +++-
 quinn-proto/src/tests/mod.rs      | 46 +++++++++++++++++++++++++++++++
 2 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/quinn-proto/src/connection/mod.rs b/quinn-proto/src/connection/mod.rs
index d48bd6da0..0aef8e39a 100644
--- a/quinn-proto/src/connection/mod.rs
+++ b/quinn-proto/src/connection/mod.rs
@@ -1292,7 +1292,10 @@ where
             }
             let last_ack_eliciting = match self.spaces[space].time_of_last_ack_eliciting_packet {
                 Some(time) => time,
-                None => continue,
+                None => {
+                    debug_assert!(!self.peer_completed_address_validation());
+                    now
+                }
             };
             let pto = last_ack_eliciting + duration;
             if result.map_or(true, |(earliest_pto, _)| pto < earliest_pto) {
diff --git a/quinn-proto/src/tests/mod.rs b/quinn-proto/src/tests/mod.rs
index 9feb05a0e..df89cedc1 100644
--- a/quinn-proto/src/tests/mod.rs
+++ b/quinn-proto/src/tests/mod.rs
@@ -15,6 +15,7 @@ use tracing::info;
 use super::*;
 use crate::cid_generator::{ConnectionIdGenerator, RandomConnectionIdGenerator};
 use crate::crypto::Session as _;
+use crate::{Certificate, CertificateChain, PrivateKey};
 mod util;
 use util::*;
 
@@ -1716,3 +1717,48 @@ fn repeated_request_response() {
         let _ = chunks.finalize();
     }
 }
+
+/// Ensures that the client sends an anti-deadlock probe after an incomplete server's first flight
+#[test]
+fn handshake_anti_deadlock_probe() {
+    let _guard = subscribe();
+
+    // Configure a big fat certificate that can't fit inside the initial anti-amplification limit
+    let cert = rcgen::generate_simple_self_signed(
+        Some("localhost".into())
+            .into_iter()
+            .chain((0..1000).map(|x| format!("foo_{}", x)))
+            .collect::<Vec<_>>(),
+    )
+    .unwrap();
+    let key = PrivateKey::from_der(&cert.serialize_private_key_der()).unwrap();
+    let cert = Certificate::from_der(&cert.serialize_der().unwrap()).unwrap();
+
+    let mut server = server_config();
+    server
+        .certificate(CertificateChain::from_certs(Some(cert.clone())), key)
+        .unwrap();
+    let mut client = client_config();
+    client.add_certificate_authority(cert).unwrap();
+    let mut pair = Pair::new(Default::default(), server);
+
+    let client_ch = pair.begin_connect(client);
+    // Client sends initial
+    pair.drive_client();
+    // Server sends first flight, gets blocked on anti-amplification
+    pair.drive_server();
+    // Client acks...
+    pair.drive_client();
+    // ...but it's lost, so the server doesn't get anti-amplification credit from it
+    pair.server.inbound.clear();
+    // Client sends an anti-deadlock probe, and the handshake completes as usual.
+    pair.drive();
+    assert_matches!(
+        pair.client_conn_mut(client_ch).poll(),
+        Some(Event::HandshakeDataReady)
+    );
+    assert_matches!(
+        pair.client_conn_mut(client_ch).poll(),
+        Some(Event::Connected { .. })
+    );
+}