fix(transport): harden QUIC admission boundaries

core/src/tls.rs

@@ -1,6 +1,6 @@
 use std::fmt;
 use std::io::Cursor;
-use std::sync::Arc;
+use std::sync::{Arc, Mutex};
 
 use rustls::client::ResolvesClientCert;
 use rustls::pki_types::CertificateDer;
@@ -13,7 +13,7 @@ use crate::wire::NativeTunnelCredential;
 pub(crate) type SignCallback = dyn Fn(&[u8]) -> Result<[u8; 64]> + Send + Sync;
 
 pub(crate) struct CallbackSigningKey {
-    callback: Arc<SignCallback>,
+    callback: Arc<Mutex<Option<Arc<SignCallback>>>>,
 }
 
 impl fmt::Debug for CallbackSigningKey {
@@ -24,7 +24,9 @@ impl fmt::Debug for CallbackSigningKey {
 
 impl CallbackSigningKey {
     pub(crate) fn new(callback: Arc<SignCallback>) -> Self {
-        Self { callback }
+        Self {
+            callback: Arc::new(Mutex::new(Some(callback))),
+        }
     }
 }
 
@@ -40,7 +42,7 @@ impl SigningKey for CallbackSigningKey {
     }
 }
 
-struct CallbackSigner(Arc<SignCallback>);
+struct CallbackSigner(Arc<Mutex<Option<Arc<SignCallback>>>>);
 
 impl fmt::Debug for CallbackSigner {
     fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
@@ -50,7 +52,13 @@ impl fmt::Debug for CallbackSigner {
 
 impl Signer for CallbackSigner {
     fn sign(&self, message: &[u8]) -> std::result::Result<Vec<u8>, rustls::Error> {
-        (self.0)(message)
+        let callback = self
+            .0
+            .lock()
+            .map_err(|_| rustls::Error::General("client signing state failed".to_owned()))?
+            .take()
+            .ok_or_else(|| rustls::Error::General("client signer already used".to_owned()))?;
+        callback(message)
             .map(|signature| signature.to_vec())
             .map_err(|_| rustls::Error::General("client signing failed".to_owned()))
     }
@@ -112,3 +120,46 @@ pub(crate) fn client_config(
     config.enable_early_data = false;
     Ok(config)
 }
+
+#[cfg(test)]
+mod tests {
+    use std::sync::atomic::{AtomicBool, Ordering};
+
+    use super::*;
+
+    struct DropMarker(Arc<AtomicBool>);
+
+    impl Drop for DropMarker {
+        fn drop(&mut self) {
+            self.0.store(true, Ordering::SeqCst);
+        }
+    }
+
+    #[test]
+    fn tls_callback_is_consumed_and_released_after_one_signature() {
+        let dropped = Arc::new(AtomicBool::new(false));
+        let marker = DropMarker(Arc::clone(&dropped));
+        let callback: Arc<SignCallback> = Arc::new(move |_| {
+            let _ = &marker;
+            Ok([7; 64])
+        });
+        let key = CallbackSigningKey::new(callback);
+        let signer = key
+            .choose_scheme(&[SignatureScheme::ED25519])
+            .expect("ED25519 signer");
+        let second_signer = key
+            .choose_scheme(&[SignatureScheme::ED25519])
+            .expect("second ED25519 signer");
+        drop(key);
+
+        assert_eq!(
+            signer.sign(b"handshake").expect("first signature"),
+            vec![7; 64]
+        );
+        assert!(dropped.load(Ordering::SeqCst), "callback remained retained");
+        assert!(
+            second_signer.sign(b"second request").is_err(),
+            "signer was reusable"
+        );
+    }
+}