From 9e1f77cf9fdb57d20522b4b76d1011fe5791006c Mon Sep 17 00:00:00 2001
From: Phoenix Kahlo <kahlo.phoenix@gmail.com>
Date: Tue, 24 Dec 2024 12:37:55 -0600
Subject: [PATCH] Create Incoming.may_retry API

The ability for the server to process tokens from NEW_TOKEN frames will
create the possibility of Incoming which are validated, but may still be
retried. This commit creates an API for that. This means that rather
than Incoming.remote_address_validated being tied to retry_src_cid, it
is tied to a new `validated: bool` of `IncomingToken`.

Currently, this field is initialized to true iff retry_src_cid is some.
However, subsequent commits will introduce the possibility for
divergence.
---
 quinn-proto/src/endpoint.rs | 17 ++++++++++++++---
 quinn-proto/src/token.rs    |  3 +++
 quinn/src/incoming.rs       | 13 ++++++++++++-
 3 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/quinn-proto/src/endpoint.rs b/quinn-proto/src/endpoint.rs
index 4fc8f0725..a5135fea6 100644
--- a/quinn-proto/src/endpoint.rs
+++ b/quinn-proto/src/endpoint.rs
@@ -711,9 +711,9 @@ impl Endpoint {
 
     /// Respond with a retry packet, requiring the client to retry with address validation
     ///
-    /// Errors if `incoming.remote_address_validated()` is true.
+    /// Errors if `incoming.may_retry()` is false.
     pub fn retry(&mut self, incoming: Incoming, buf: &mut Vec<u8>) -> Result<Transmit, RetryError> {
-        if incoming.remote_address_validated() {
+        if !incoming.may_retry() {
             return Err(RetryError(incoming));
         }
 
@@ -1172,8 +1172,19 @@ impl Incoming {
     ///
     /// This means that the sender of the initial packet has proved that they can receive traffic
     /// sent to `self.remote_address()`.
+    ///
+    /// If `self.remote_address_validated()` is false, `self.may_retry()` is guaranteed to be true.
+    /// The inverse is not guaranteed.
     pub fn remote_address_validated(&self) -> bool {
-        self.token.retry_src_cid.is_some()
+        self.token.validated
+    }
+
+    /// Whether it is legal to respond with a retry packet
+    ///
+    /// If `self.remote_address_validated()` is false, `self.may_retry()` is guaranteed to be true.
+    /// The inverse is not guaranteed.
+    pub fn may_retry(&self) -> bool {
+        self.token.retry_src_cid.is_none()
     }
 
     /// The original destination connection ID sent by the client
diff --git a/quinn-proto/src/token.rs b/quinn-proto/src/token.rs
index 578f282b8..36fea1cfb 100644
--- a/quinn-proto/src/token.rs
+++ b/quinn-proto/src/token.rs
@@ -20,6 +20,7 @@ use crate::{
 pub(crate) struct IncomingToken {
     pub(crate) retry_src_cid: Option<ConnectionId>,
     pub(crate) orig_dst_cid: ConnectionId,
+    pub(crate) validated: bool,
 }
 
 impl IncomingToken {
@@ -33,6 +34,7 @@ impl IncomingToken {
         let unvalidated = Self {
             retry_src_cid: None,
             orig_dst_cid: header.dst_cid,
+            validated: false,
         };
 
         // Decode token or short-circuit
@@ -67,6 +69,7 @@ impl IncomingToken {
         Ok(Self {
             retry_src_cid: Some(header.dst_cid),
             orig_dst_cid: retry.payload.orig_dst_cid,
+            validated: true,
         })
     }
 }
diff --git a/quinn/src/incoming.rs b/quinn/src/incoming.rs
index 43dca653b..6b41d25bf 100644
--- a/quinn/src/incoming.rs
+++ b/quinn/src/incoming.rs
@@ -48,7 +48,7 @@ impl Incoming {
 
     /// Respond with a retry packet, requiring the client to retry with address validation
     ///
-    /// Errors if `remote_address_validated()` is true.
+    /// Errors if `may_retry()` is false.
     pub fn retry(mut self) -> Result<(), RetryError> {
         let state = self.0.take().unwrap();
         state.endpoint.retry(state.inner).map_err(|e| {
@@ -79,10 +79,21 @@ impl Incoming {
     ///
     /// This means that the sender of the initial packet has proved that they can receive traffic
     /// sent to `self.remote_address()`.
+    ///
+    /// If `self.remote_address_validated()` is false, `self.may_retry()` is guaranteed to be true.
+    /// The inverse is not guaranteed.
     pub fn remote_address_validated(&self) -> bool {
         self.0.as_ref().unwrap().inner.remote_address_validated()
     }
 
+    /// Whether it is legal to respond with a retry packet
+    ///
+    /// If `self.remote_address_validated()` is false, `self.may_retry()` is guaranteed to be true.
+    /// The inverse is not guaranteed.
+    pub fn may_retry(&self) -> bool {
+        self.0.as_ref().unwrap().inner.may_retry()
+    }
+
     /// The original destination CID when initiating the connection
     pub fn orig_dst_cid(&self) -> ConnectionId {
         *self.0.as_ref().unwrap().inner.orig_dst_cid()