feat(s2n-quic-dc): path_secrets_ready callback can cancel handshake (#2648)

Change return type of make_application_data callback registered by
application to Result<...>, and canceling the handshake, propagating an
error message when an Error is returned.

Co-authored-by: Benjamin Barzen <[email protected]>

Author: benlcb

dc/s2n-quic-dc/src/path/secret/map.rs

@@ -215,7 +215,7 @@ impl Map {
                 receiver::State::new(),
                 dc::testing::TEST_APPLICATION_PARAMS,
                 dc::testing::TEST_REHANDSHAKE_PERIOD,
-                Arc::new(()),
+                None,
             );
             let entry = Arc::new(entry);
             provider.store.test_insert(entry);
@@ -277,7 +277,7 @@ impl Map {
                 super::receiver::State::new(),
                 params,
                 dc::testing::TEST_REHANDSHAKE_PERIOD,
-                Arc::new(()),
+                None,
             );
             let entry = Arc::new(entry);
             map.store.test_insert(entry);
@@ -297,7 +297,11 @@ impl Map {
     pub fn register_make_application_data(
         &self,
         cb: Box<
-            dyn Fn(&dyn s2n_quic_core::crypto::tls::TlsSession) -> ApplicationData + Send + Sync,
+            dyn Fn(
+                    &dyn s2n_quic_core::crypto::tls::TlsSession,
+                ) -> Result<Option<ApplicationData>, &'static str>
+                + Send
+                + Sync,
         >,
     ) {
         self.store.register_make_application_data(cb);

dc/s2n-quic-dc/src/path/secret/map/entry.rs

@@ -45,7 +45,7 @@ pub struct Entry {
     // we store this as a u8 to allow the cleaner to separately "take" accessed for id and addr
     // maps while not having two writes and wasting an extra byte of space.
     accessed: AtomicU8,
-    application_data: ApplicationData,
+    application_data: Option<ApplicationData>,
 }
 
 impl SizeOf for Entry {
@@ -73,6 +73,12 @@ impl SizeOf for Entry {
     }
 }
 
+impl SizeOf for Option<ApplicationData> {
+    fn size(&self) -> usize {
+        std::mem::size_of::<ApplicationData>()
+    }
+}
+
 impl SizeOf for ApplicationData {
     fn size(&self) -> usize {
         std::mem::size_of_val(self)
@@ -91,7 +97,7 @@ impl Entry {
         parameters: dc::ApplicationParams,
         // FIXME: remove unused parameter
         _: Duration,
-        application_data: ApplicationData,
+        application_data: Option<ApplicationData>,
     ) -> Self {
         // clamp max datagram size to a well-known value
         parameters
@@ -130,7 +136,7 @@ impl Entry {
             receiver,
             dc::testing::TEST_APPLICATION_PARAMS,
             dc::testing::TEST_REHANDSHAKE_PERIOD,
-            Arc::new(()),
+            None,
         ))
     }
 
@@ -273,7 +279,7 @@ impl Entry {
         self.secret.control_sealer()
     }
 
-    pub fn application_data(&self) -> &ApplicationData {
+    pub fn application_data(&self) -> &Option<ApplicationData> {
         &self.application_data
     }
 }

dc/s2n-quic-dc/src/path/secret/map/handshake.rs

@@ -84,7 +84,14 @@ impl dc::Path for HandshakingPath {
         &mut self,
         session: &impl s2n_quic_core::crypto::tls::TlsSession,
     ) -> Result<Vec<s2n_quic_core::stateless_reset::Token>, s2n_quic_core::transport::Error> {
-        self.application_data = Some(self.map.store.application_data(session));
+        match self.map.store.application_data(session) {
+            Ok(application_data) => {
+                self.application_data = application_data;
+            }
+            Err(msg) => {
+                return Err(s2n_quic_core::transport::Error::APPLICATION_ERROR.with_reason(msg));
+            }
+        };
 
         let mut material = Zeroizing::new([0; TLS_EXPORTER_LENGTH]);
         session
@@ -138,7 +145,7 @@ impl dc::Path for HandshakingPath {
             receiver,
             self.parameters.clone(),
             self.map.store.rehandshake_period(),
-            self.application_data.take().unwrap(),
+            self.application_data.take(),
         );
         let entry = Arc::new(entry);
         self.entry = Some(entry.clone());

dc/s2n-quic-dc/src/path/secret/map/state.rs

@@ -219,14 +219,14 @@ where
     mk_application_data: RwLock<
         Option<
             Box<
-                dyn Fn(&dyn s2n_quic_core::crypto::tls::TlsSession) -> ApplicationData
+                dyn Fn(
+                        &dyn s2n_quic_core::crypto::tls::TlsSession,
+                    ) -> Result<Option<ApplicationData>, &'static str>
                     + Send
                     + Sync,
             >,
         >,
     >,
-
-    dummy_application_data: ApplicationData,
 }
 
 // Share control sockets -- we only send on these so it doesn't really matter if there's only one
@@ -288,7 +288,6 @@ where
             subscriber,
             request_handshake: RwLock::new(None),
             mk_application_data: RwLock::new(None),
-            dummy_application_data: Arc::new(()),
         };
 
         // Growing to double our maximum inserted entries should ensure that we never grow again, see:
@@ -677,7 +676,11 @@ where
     fn register_make_application_data(
         &self,
         cb: Box<
-            dyn Fn(&dyn s2n_quic_core::crypto::tls::TlsSession) -> ApplicationData + Send + Sync,
+            dyn Fn(
+                    &dyn s2n_quic_core::crypto::tls::TlsSession,
+                ) -> Result<Option<ApplicationData>, &'static str>
+                + Send
+                + Sync,
         >,
     ) {
         // FIXME: Maybe panic if already initialized?
@@ -891,15 +894,15 @@ where
     fn application_data(
         &self,
         session: &dyn s2n_quic_core::crypto::tls::TlsSession,
-    ) -> ApplicationData {
+    ) -> Result<Option<ApplicationData>, &'static str> {
         if let Some(ctxt) = &*self
             .mk_application_data
             .read()
             .unwrap_or_else(|e| e.into_inner())
         {
             (ctxt)(session)
         } else {
-            self.dummy_application_data.clone()
+            Ok(None)
         }
     }
 }

dc/s2n-quic-dc/src/path/secret/map/state/tests.rs

@@ -139,7 +139,7 @@ impl Model {
                     receiver::State::new(),
                     dc::testing::TEST_APPLICATION_PARAMS,
                     dc::testing::TEST_REHANDSHAKE_PERIOD,
-                    Arc::new(()),
+                    None,
                 )));
 
                 self.invariants.insert(Invariant::ContainsIp(ip));

dc/s2n-quic-dc/src/path/secret/map/store.rs

@@ -106,12 +106,16 @@ pub trait Store: 'static + Send + Sync {
     fn register_make_application_data(
         &self,
         cb: Box<
-            dyn Fn(&dyn s2n_quic_core::crypto::tls::TlsSession) -> ApplicationData + Send + Sync,
+            dyn Fn(
+                    &dyn s2n_quic_core::crypto::tls::TlsSession,
+                ) -> Result<Option<ApplicationData>, &'static str>
+                + Send
+                + Sync,
         >,
     );
 
     fn application_data(
         &self,
         session: &dyn s2n_quic_core::crypto::tls::TlsSession,
-    ) -> ApplicationData;
+    ) -> Result<Option<ApplicationData>, &'static str>;
 }