feat(proto)!: Don't require a HKDF construction in HandshakeTokenKey (#480)

* Implement HandshakeTokenKeys directly using AES-GCM

* Revert to old HDKF-based construction, but with changed trait

* Cleanup

* Clippy fix

Author: matheus23

noq-proto/src/config/mod.rs

@@ -400,18 +400,10 @@ impl ServerConfig {
     ///
     /// Uses a randomized handshake token key.
     pub fn with_crypto(crypto: Arc<dyn crypto::ServerConfig>) -> Self {
-        #[cfg(all(feature = "aws-lc-rs", not(feature = "ring")))]
-        use aws_lc_rs::hkdf;
-        use rand::RngCore;
-        #[cfg(feature = "ring")]
-        use ring::hkdf;
-
-        let rng = &mut rand::rng();
-        let mut master_key = [0u8; 64];
-        rng.fill_bytes(&mut master_key);
-        let master_key = hkdf::Salt::new(hkdf::HKDF_SHA256, &[]).extract(&master_key);
+        use crate::crypto::ring_like::RetryTokenKey;
 
-        Self::new(crypto, Arc::new(master_key))
+        let retry_token_key = RetryTokenKey::new(&mut rand::rng());
+        Self::new(crypto, Arc::new(retry_token_key))
     }
 }
 

noq-proto/src/crypto.rs

@@ -204,22 +204,17 @@ pub trait HmacKey: Send + Sync {
 #[derive(Debug, PartialEq, Eq)]
 pub struct ExportKeyingMaterialError;
 
-/// A pseudo random key for HKDF
+/// The trait for encrypting tokens that power retry packets and NEW_TOKEN frames.
 pub trait HandshakeTokenKey: Send + Sync {
-    /// Derive AEAD using hkdf
-    fn aead_from_hkdf(&self, random_bytes: &[u8]) -> Box<dyn AeadKey>;
-}
-
-/// A key for sealing data with AEAD-based algorithms
-pub trait AeadKey {
-    /// Method for sealing message `data`
-    fn seal(&self, data: &mut Vec<u8>, additional_data: &[u8]) -> Result<(), CryptoError>;
-    /// Method for opening a sealed message `data`
-    fn open<'a>(
-        &self,
-        data: &'a mut [u8],
-        additional_data: &[u8],
-    ) -> Result<&'a mut [u8], CryptoError>;
+    /// Method for sealing a token in-place.
+    ///
+    /// The nonce doesn't need to be attached to the ciphertext,
+    /// but the authentication tag is expected to be appended to `data`.
+    fn seal(&self, token_nonce: u128, data: &mut Vec<u8>) -> Result<(), CryptoError>;
+    /// Method for opening a sealed message `data` in-place.
+    ///
+    /// Returns the portion of `data` that contains the decrypted plaintext.
+    fn open<'a>(&self, token_nonce: u128, data: &'a mut [u8]) -> Result<&'a [u8], CryptoError>;
 }
 
 /// Generic crypto errors

noq-proto/src/crypto/ring_like.rs

@@ -19,34 +19,54 @@ impl crypto::HmacKey for hmac::Key {
     }
 }
 
-impl crypto::HandshakeTokenKey for hkdf::Prk {
-    fn aead_from_hkdf(&self, random_bytes: &[u8]) -> Box<dyn crypto::AeadKey> {
-        let mut key_buffer = [0u8; 32];
-        let info = [random_bytes];
-        let okm = self.expand(&info, hkdf::HKDF_SHA256).unwrap();
+/// Implements retry token encryption using HKDF followed by AES-256-GCM.
+///
+/// This construction was originally defined in https://github.com/quinn-rs/quinn/issues/783
+/// The goal is to produce tokens that look random to clients, but contain decryptable
+/// information for the server.
+///
+/// The problem is the 12-byte nonce in AES-GCM: I suspect the original authors of this
+/// code didn't like the fact that it limits you to ~2^32 safe encryptions before you're getting
+/// into "dangerous" nonce-reuse territory with randomly-generated nonces.
+/// You can't use sequence numbers either, though, because that doesn't look random to
+/// clients and leaks information.
+pub(crate) struct RetryTokenKey(hkdf::Prk);
+
+impl RetryTokenKey {
+    pub(crate) fn new(rng: &mut impl rand::Rng) -> Self {
+        let mut master_key = [0u8; 64];
+        rng.fill_bytes(&mut master_key);
+        let master_key = hkdf::Salt::new(hkdf::HKDF_SHA256, &[]).extract(&master_key);
+        Self(master_key)
+    }
+
+    fn derive_aead(&self, token_nonce: u128) -> aead::LessSafeKey {
+        let nonce_bytes = token_nonce.to_le_bytes();
+        let info = &[&nonce_bytes[..]];
+        let okm = self.0.expand(info, hkdf::HKDF_SHA256).unwrap();
 
+        let mut key_buffer = [0u8; 32];
         okm.fill(&mut key_buffer).unwrap();
 
         let key = aead::UnboundKey::new(&aead::AES_256_GCM, &key_buffer).unwrap();
-        Box::new(aead::LessSafeKey::new(key))
+        aead::LessSafeKey::new(key)
     }
 }
 
-impl crypto::AeadKey for aead::LessSafeKey {
-    fn seal(&self, data: &mut Vec<u8>, additional_data: &[u8]) -> Result<(), CryptoError> {
-        let aad = aead::Aad::from(additional_data);
-        let zero_nonce = aead::Nonce::assume_unique_for_key([0u8; 12]);
-        Ok(self.seal_in_place_append_tag(zero_nonce, aad, data)?)
-    }
-
-    fn open<'a>(
-        &self,
-        data: &'a mut [u8],
-        additional_data: &[u8],
-    ) -> Result<&'a mut [u8], CryptoError> {
-        let aad = aead::Aad::from(additional_data);
-        let zero_nonce = aead::Nonce::assume_unique_for_key([0u8; 12]);
-        Ok(self.open_in_place(zero_nonce, aad, data)?)
+impl crypto::HandshakeTokenKey for RetryTokenKey {
+    fn seal(&self, token_nonce: u128, data: &mut Vec<u8>) -> Result<(), CryptoError> {
+        let aead_key = self.derive_aead(token_nonce);
+        let nonce = aead::Nonce::assume_unique_for_key([0u8; 12]); // See docs for RetryTokenKey
+        let aad = aead::Aad::empty();
+        aead_key.seal_in_place_append_tag(nonce, aad, data)?;
+        Ok(())
+    }
+
+    fn open<'a>(&self, token_nonce: u128, data: &'a mut [u8]) -> Result<&'a [u8], CryptoError> {
+        let aead_key = self.derive_aead(token_nonce);
+        let aad = aead::Aad::empty();
+        let nonce = aead::Nonce::assume_unique_for_key([0u8; 12]); // See docs for RetryTokenKey
+        Ok(aead_key.open_in_place(nonce, aad, data)?)
     }
 }
 

noq-proto/src/crypto/rustls.rs

@@ -3,8 +3,6 @@ use std::{any::Any, io, str, sync::Arc};
 use aes_gcm::{KeyInit, aead::AeadMutInPlace};
 use bytes::BytesMut;
 pub use rustls::Error;
-#[cfg(feature = "__rustls-post-quantum-test")]
-use rustls::NamedGroup;
 use rustls::{
     self, CipherSuite,
     pki_types::{CertificateDer, ServerName},
@@ -268,7 +266,7 @@ pub struct HandshakeData {
     pub server_name: Option<String>,
     /// The key exchange group negotiated with the peer
     #[cfg(feature = "__rustls-post-quantum-test")]
-    pub negotiated_key_exchange_group: NamedGroup,
+    pub negotiated_key_exchange_group: rustls::NamedGroup,
 }
 
 /// A QUIC-compatible TLS client configuration

noq-proto/src/token.rs

@@ -1,6 +1,5 @@
 use std::{
     fmt,
-    mem::size_of,
     net::{IpAddr, SocketAddr},
 };
 
@@ -236,8 +235,7 @@ impl Token {
         }
 
         // Encrypt
-        let aead_key = key.aead_from_hkdf(&self.nonce.to_le_bytes());
-        aead_key.seal(&mut buf, &[]).unwrap();
+        key.seal(self.nonce, &mut buf).unwrap();
         buf.extend(&self.nonce.to_le_bytes());
 
         buf
@@ -247,30 +245,27 @@ impl Token {
     fn decode(key: &dyn HandshakeTokenKey, raw_token_bytes: &[u8]) -> Option<Self> {
         // Decrypt
 
-        let nonce_slice_start = raw_token_bytes.len().checked_sub(size_of::<u128>())?;
-        let (sealed_token, nonce_bytes) = raw_token_bytes.split_at_checked(nonce_slice_start)?;
+        let (sealed_token, nonce_bytes) = raw_token_bytes.split_last_chunk()?;
 
-        let nonce = u128::from_le_bytes(nonce_bytes.try_into().unwrap());
+        let nonce = u128::from_le_bytes(*nonce_bytes);
 
-        let aead_key = key.aead_from_hkdf(nonce_bytes);
         let mut sealed_token = sealed_token.to_vec();
-        let data = aead_key.open(&mut sealed_token, &[]).ok()?;
+        let mut data = key.open(nonce, &mut sealed_token).ok()?;
 
         // Decode payload
-        let mut reader = &data[..];
-        let payload = match TokenType::from_byte((&mut reader).get::<u8>().ok()?)? {
+        let payload = match TokenType::from_byte((&mut data).get::<u8>().ok()?)? {
             TokenType::Retry => TokenPayload::Retry {
-                address: decode_addr(&mut reader)?,
-                orig_dst_cid: ConnectionId::decode_long(&mut reader)?,
-                issued: decode_unix_secs(&mut reader)?,
+                address: decode_addr(&mut data)?,
+                orig_dst_cid: ConnectionId::decode_long(&mut data)?,
+                issued: decode_unix_secs(&mut data)?,
             },
             TokenType::Validation => TokenPayload::Validation {
-                ip: decode_ip(&mut reader)?,
-                issued: decode_unix_secs(&mut reader)?,
+                ip: decode_ip(&mut data)?,
+                issued: decode_unix_secs(&mut data)?,
             },
         };
 
-        if !reader.is_empty() {
+        if !data.is_empty() {
             // Consider extra bytes a decoding error (it may be from an incompatible endpoint)
             return None;
         }
@@ -408,21 +403,17 @@ impl fmt::Display for ResetToken {
 
 #[cfg(all(test, any(feature = "aws-lc-rs", feature = "ring")))]
 mod test {
+    use crate::crypto::ring_like::RetryTokenKey;
+
     use super::*;
-    #[cfg(all(feature = "aws-lc-rs", not(feature = "ring")))]
-    use aws_lc_rs::hkdf;
     use rand::prelude::*;
-    #[cfg(feature = "ring")]
-    use ring::hkdf;
 
     fn token_round_trip(payload: TokenPayload) -> TokenPayload {
         let rng = &mut rand::rng();
         let token = Token::new(payload, rng);
-        let mut master_key = [0; 64];
-        rng.fill_bytes(&mut master_key);
-        let prk = hkdf::Salt::new(hkdf::HKDF_SHA256, &[]).extract(&master_key);
-        let encoded = token.encode(&prk);
-        let decoded = Token::decode(&prk, &encoded).expect("token didn't decrypt / decode");
+        let master_key = RetryTokenKey::new(rng);
+        let encoded = token.encode(&master_key);
+        let decoded = Token::decode(&master_key, &encoded).expect("token didn't decrypt / decode");
         assert_eq!(token.nonce, decoded.nonce);
         decoded.payload
     }
@@ -485,14 +476,8 @@ mod test {
     #[test]
     fn invalid_token_returns_err() {
         use super::*;
-        use rand::RngCore;
-
-        let rng = &mut rand::rng();
-
-        let mut master_key = [0; 64];
-        rng.fill_bytes(&mut master_key);
 
-        let prk = hkdf::Salt::new(hkdf::HKDF_SHA256, &[]).extract(&master_key);
+        let master_key = RetryTokenKey::new(&mut rand::rng());
 
         let mut invalid_token = Vec::new();
 
@@ -501,6 +486,6 @@ mod test {
         invalid_token.put_slice(&random_data);
 
         // Assert: garbage sealed data returns err
-        assert!(Token::decode(&prk, &invalid_token).is_none());
+        assert!(Token::decode(&master_key, &invalid_token).is_none());
     }
 }