edhoc: fixing mutable references. gstar.

openwsn-berkeley · Nov 2, 2024 · bb2d03b · bb2d03b
1 parent 7615093
commit bb2d03b
Showing 1 changed file with 34 additions and 32 deletions.
diff --git a/lib/src/edhoc.rs b/lib/src/edhoc.rs
@@ -1,56 +1,64 @@
 use lakers_shared::{Crypto as CryptoTrait, *};
 
-// Implementation of edhoc_exporter for 3 or 4 messages
 pub trait Done {
-    fn get_prk_exporter_mut(&mut self) -> &mut [u8; SHA256_DIGEST_LEN];
     fn get_prk_exporter(&self) -> &[u8; SHA256_DIGEST_LEN];
-    fn get_prk_out_mut(&mut self) -> &mut [u8; SHA256_DIGEST_LEN];
     fn get_prk_out(&self) -> &[u8; SHA256_DIGEST_LEN];
+    fn update_keys(
+        &mut self,
+        new_prk_out: &[u8; SHA256_DIGEST_LEN],
+        new_prk_exporter: &[u8; SHA256_DIGEST_LEN],
+    );
 }
 
 impl Done for WaitM4 {
-    fn get_prk_exporter_mut(&mut self) -> &mut [u8; SHA256_DIGEST_LEN] {
-        &mut self.prk_exporter
-    }
     fn get_prk_exporter(&self) -> &[u8; SHA256_DIGEST_LEN] {
         &self.prk_exporter
     }
-    fn get_prk_out_mut(&mut self) -> &mut [u8; SHA256_DIGEST_LEN] {
-        &mut self.prk_out
-    }
     fn get_prk_out(&self) -> &[u8; SHA256_DIGEST_LEN] {
         &self.prk_out
     }
+    fn update_keys(
+        &mut self,
+        new_prk_out: &[u8; SHA256_DIGEST_LEN],
+        new_prk_exporter: &[u8; SHA256_DIGEST_LEN],
+    ) {
+        self.prk_out.copy_from_slice(new_prk_out);
+        self.prk_exporter.copy_from_slice(new_prk_exporter);
+    }
 }
 
 impl Done for Completed {
-    fn get_prk_exporter_mut(&mut self) -> &mut [u8; SHA256_DIGEST_LEN] {
-        &mut self.prk_exporter
-    }
     fn get_prk_exporter(&self) -> &[u8; SHA256_DIGEST_LEN] {
         &self.prk_exporter
     }
-    fn get_prk_out_mut(&mut self) -> &mut [u8; SHA256_DIGEST_LEN] {
-        &mut self.prk_out
-    }
     fn get_prk_out(&self) -> &[u8; SHA256_DIGEST_LEN] {
         &self.prk_out
     }
+    fn update_keys(
+        &mut self,
+        new_prk_out: &[u8; SHA256_DIGEST_LEN],
+        new_prk_exporter: &[u8; SHA256_DIGEST_LEN],
+    ) {
+        self.prk_out.copy_from_slice(new_prk_out);
+        self.prk_exporter.copy_from_slice(new_prk_exporter);
+    }
 }
 
 impl Done for ProcessedM3 {
-    fn get_prk_exporter_mut(&mut self) -> &mut [u8; SHA256_DIGEST_LEN] {
-        &mut self.prk_exporter
-    }
     fn get_prk_exporter(&self) -> &[u8; SHA256_DIGEST_LEN] {
         &self.prk_exporter
     }
-    fn get_prk_out_mut(&mut self) -> &mut [u8; SHA256_DIGEST_LEN] {
-        &mut self.prk_out
-    }
     fn get_prk_out(&self) -> &[u8; SHA256_DIGEST_LEN] {
         &self.prk_out
     }
+    fn update_keys(
+        &mut self,
+        new_prk_out: &[u8; SHA256_DIGEST_LEN],
+        new_prk_exporter: &[u8; SHA256_DIGEST_LEN],
+    ) {
+        self.prk_out.copy_from_slice(new_prk_out);
+        self.prk_exporter.copy_from_slice(new_prk_exporter);
+    }
 }
 
 pub fn edhoc_exporter(
@@ -77,12 +85,10 @@ pub fn edhoc_key_update(
     context: &BytesMaxContextBuffer,
     context_len: usize,
 ) -> BytesHashLen {
-    // new PRK_out
-    let mut prk_out = *state.get_prk_out();
-    // let mut prk_exporter = *state.get_prk_exporter();
+    // Calculate new PRK_out
     let prk_new_buf = edhoc_kdf(
         crypto,
-        &prk_out,
+        state.get_prk_out(),
         11u8,
         context,
         context_len,
@@ -91,7 +97,7 @@ pub fn edhoc_key_update(
     let mut new_prk_out = [0u8; SHA256_DIGEST_LEN];
     new_prk_out.copy_from_slice(&prk_new_buf[..SHA256_DIGEST_LEN]);
 
-    // new PRK_exporter
+    // Calculate new PRK_exporter
     let prk_exporter_buf = edhoc_kdf(
         crypto,
         &new_prk_out,
@@ -103,12 +109,8 @@ pub fn edhoc_key_update(
     let mut new_prk_exporter = [0u8; SHA256_DIGEST_LEN];
     new_prk_exporter.copy_from_slice(&prk_exporter_buf[..SHA256_DIGEST_LEN]);
 
-    // Update state
-    let prk_out_mut = state.get_prk_out_mut();
-    prk_out_mut.copy_from_slice(&new_prk_out);
-
-    let prk_exporter_mut = state.get_prk_exporter_mut();
-    prk_exporter_mut.copy_from_slice(&new_prk_exporter);
+    // Update state with new keys
+    state.update_keys(&new_prk_out, &new_prk_exporter);
 
     new_prk_out
 }