Implement caching of expanded keys as a permanent feature

Author: dmitry-tarnyagin

Cargo.toml

@@ -59,4 +59,3 @@ defmt = ["dep:defmt", "embedded-io/defmt-03", "heapless/defmt-03"]
 std = ["embedded-io/std", "embedded-io-async/std"]
 tokio = ["embedded-io-adapters/tokio-1"]
 alloc = []
-key-cache = []

src/connection.rs

@@ -44,7 +44,7 @@ where
         let server_key = key_schedule.get_key()?;
         let nonce = key_schedule.get_nonce()?;
 
-        let crypto = <CipherSuite::Cipher as KeyInit>::new(&server_key);
+        let crypto = <CipherSuite::Cipher as KeyInit>::new(server_key);
         crypto
             .decrypt_in_place(&nonce, header.data(), &mut app_data)
             .map_err(|_| TlsError::CryptoError)?;
@@ -106,7 +106,7 @@ where
     // trace!("encrypt nonce {:02x?}", nonce);
     // trace!("plaintext {} {:02x?}", buf.len(), buf.as_slice(),);
     //let crypto = Aes128Gcm::new_varkey(&self.key_schedule.get_client_key()).unwrap();
-    let crypto = <CipherSuite::Cipher as KeyInit>::new(&client_key);
+    let crypto = <CipherSuite::Cipher as KeyInit>::new(client_key);
     let len = buf.len() + <CipherSuite::Cipher as AeadCore>::TagSize::to_usize();
 
     if len > buf.capacity() {

src/key_schedule.rs

@@ -1,8 +1,6 @@
 use crate::handshake::binder::PskBinder;
 use crate::handshake::finished::Finished;
 use crate::{TlsError, config::TlsCipherSuite};
-#[cfg(feature = "key-cache")]
-use core::cell::OnceCell;
 use digest::OutputSizeUser;
 use digest::generic_array::ArrayLength;
 use hmac::{Mac, SimpleHmac};
@@ -137,11 +135,8 @@ where
 {
     traffic_secret: Secret<CipherSuite>,
     counter: u64,
-
-    #[cfg(feature = "key-cache")]
-    key: OnceCell<KeyArray<CipherSuite>>,
-    #[cfg(feature = "key-cache")]
-    iv: OnceCell<IvArray<CipherSuite>>,
+    key: KeyArray<CipherSuite>,
+    iv: IvArray<CipherSuite>,
 }
 
 impl<CipherSuite> KeyScheduleState<CipherSuite>
@@ -152,57 +147,24 @@ where
         Self {
             traffic_secret: Secret::Uninitialized,
             counter: 0,
-            #[cfg(feature = "key-cache")]
-            key: OnceCell::new(),
-            #[cfg(feature = "key-cache")]
-            iv: OnceCell::new(),
+            key: KeyArray::<CipherSuite>::default(),
+            iv: IvArray::<CipherSuite>::default(),
         }
     }
 
     #[inline]
-    pub fn get_key(&self) -> Result<KeyArray<CipherSuite>, TlsError> {
-        #[cfg(feature = "key-cache")]
-        if let Some(k) = self.key.get() {
-            Ok(k.clone())
-        } else {
-            let k = self.get_key_impl()?;
-            let _ = self.key.set(k.clone());
-            Ok(k)
-        }
-
-        #[cfg(not(feature = "key-cache"))]
-        self.get_key_impl()
+    pub fn get_key(&self) -> Result<&KeyArray<CipherSuite>, TlsError> {
+        Ok(&self.key)
     }
 
     #[inline]
-    pub fn get_iv(&self) -> Result<IvArray<CipherSuite>, TlsError> {
-        #[cfg(feature = "key-cache")]
-        if let Some(k) = self.iv.get() {
-            Ok(k.clone())
-        } else {
-            let k = self.get_iv_impl()?;
-            let _ = self.iv.set(k.clone());
-            Ok(k)
-        }
-
-        #[cfg(not(feature = "key-cache"))]
-        self.get_iv_impl()
-    }
-
-    fn get_key_impl(&self) -> Result<KeyArray<CipherSuite>, TlsError> {
-        self.traffic_secret
-            .make_expanded_hkdf_label(b"key", ContextType::None)
+    pub fn get_iv(&self) -> Result<&IvArray<CipherSuite>, TlsError> {
+        Ok(&self.iv)
     }
 
-    fn get_iv_impl(&self) -> Result<IvArray<CipherSuite>, TlsError> {
-        self.traffic_secret
-            .make_expanded_hkdf_label(b"iv", ContextType::None)
-    }
-
-
     pub fn get_nonce(&self) -> Result<IvArray<CipherSuite>, TlsError> {
         let iv = self.get_iv()?;
-        Ok(KeySchedule::<CipherSuite>::get_nonce(self.counter, &iv))
+        Ok(KeySchedule::<CipherSuite>::get_nonce(self.counter, iv))
     }
 
     fn calculate_traffic_secret(
@@ -216,11 +178,12 @@ where
             Hkdf::<CipherSuite>::from_prk(&secret).map_err(|_| TlsError::InternalError)?;
 
         self.traffic_secret.replace(traffic_secret);
-        #[cfg(feature = "key-cache")]
-        {
-            self.key = OnceCell::new();
-            self.iv = OnceCell::new();
-        }
+        self.key = self
+            .traffic_secret
+            .make_expanded_hkdf_label(b"key", ContextType::None)?;
+        self.iv = self
+            .traffic_secret
+            .make_expanded_hkdf_label(b"iv", ContextType::None)?;
         self.counter = 0;
         Ok(())
     }
@@ -449,7 +412,7 @@ where
         self.state.increment_counter();
     }
 
-    pub(crate) fn get_key(&self) -> Result<KeyArray<CipherSuite>, TlsError> {
+    pub(crate) fn get_key(&self) -> Result<&KeyArray<CipherSuite>, TlsError> {
         self.state.get_key()
     }
 
@@ -496,7 +459,7 @@ where
         &mut self.transcript_hash
     }
 
-    pub(crate) fn get_key(&self) -> Result<KeyArray<CipherSuite>, TlsError> {
+    pub(crate) fn get_key(&self) -> Result<&KeyArray<CipherSuite>, TlsError> {
         self.state.get_key()
     }