From 42891b9d568849c576ee983933b5ab1b97413e93 Mon Sep 17 00:00:00 2001
From: Arthur Gautier <arthur.gautier@arista.com>
Date: Wed, 6 Nov 2024 12:36:56 -0800
Subject: [PATCH] implement `KeyInit` for aes-kw::Kek and BeltKwp

---
 aes-kw/src/lib.rs   | 41 ++++++++++++++++++++++++++---------------
 belt-kwp/src/lib.rs | 33 ++++++++++++++++++++++-----------
 2 files changed, 48 insertions(+), 26 deletions(-)

diff --git a/aes-kw/src/lib.rs b/aes-kw/src/lib.rs
index f2f03a5..262b274 100644
--- a/aes-kw/src/lib.rs
+++ b/aes-kw/src/lib.rs
@@ -23,11 +23,13 @@ pub use error::{Error, Result};
 
 use aes::cipher::{
     array::Array,
-    typenum::{Unsigned, U16, U24, U32},
+    typenum::{U16, U24, U32},
     Block, BlockCipherDecBackend, BlockCipherDecClosure, BlockCipherDecrypt, BlockCipherEncBackend,
-    BlockCipherEncClosure, BlockCipherEncrypt, BlockSizeUser, KeyInit,
+    BlockCipherEncClosure, BlockCipherEncrypt, BlockSizeUser,
 };
 
+pub use aes::cipher::{self, Key, KeyInit, KeySizeUser};
+
 #[cfg(feature = "alloc")]
 use alloc::vec::Vec;
 
@@ -138,13 +140,7 @@ where
     type Error = Error;
 
     fn try_from(value: &[u8]) -> Result<Self> {
-        if value.len() == Aes::KeySize::to_usize() {
-            Ok(Kek::new(
-                &Array::try_from(value).expect("size invariant violated"),
-            ))
-        } else {
-            Err(Error::InvalidKekSize { size: value.len() })
-        }
+        Self::new_from_slice(value).map_err(|_| Error::InvalidKekSize { size: value.len() })
     }
 }
 
@@ -152,12 +148,6 @@ impl<Aes> Kek<Aes>
 where
     Aes: KeyInit + BlockSizeUser<BlockSize = U16> + BlockCipherEncrypt + BlockCipherDecrypt,
 {
-    /// Constructs a new Kek based on the appropriate raw key material.
-    pub fn new(key: &Array<u8, Aes::KeySize>) -> Self {
-        let cipher = Aes::new(key);
-        Kek { cipher }
-    }
-
     /// AES Key Wrap, as defined in RFC 3394.
     ///
     /// The `out` buffer will be overwritten, and must be exactly [`IV_LEN`]
@@ -424,6 +414,27 @@ where
     }
 }
 
+impl<Aes> KeyInit for Kek<Aes>
+where
+    Aes: KeyInit + BlockSizeUser<BlockSize = U16> + BlockCipherEncrypt + BlockCipherDecrypt,
+{
+    fn new(key: &Key<Self>) -> Self {
+        let cipher = Aes::new(key);
+        Kek { cipher }
+    }
+}
+
+impl<Aes> KeySizeUser for Kek<Aes>
+where
+    Aes: KeyInit + BlockSizeUser<BlockSize = U16> + BlockCipherEncrypt + BlockCipherDecrypt,
+{
+    type KeySize = Aes::KeySize;
+
+    fn key_size() -> usize {
+        Aes::key_size()
+    }
+}
+
 struct WCtx<'a> {
     n: usize,
     block: &'a mut Block<Self>,
diff --git a/belt-kwp/src/lib.rs b/belt-kwp/src/lib.rs
index 5383af1..952c51e 100644
--- a/belt-kwp/src/lib.rs
+++ b/belt-kwp/src/lib.rs
@@ -8,9 +8,11 @@
 #![forbid(unsafe_code)]
 #![warn(missing_docs, rust_2018_idioms)]
 
-use belt_block::{belt_wblock_dec, belt_wblock_enc};
+use belt_block::{belt_wblock_dec, belt_wblock_enc, BeltBlock};
 use core::fmt;
 
+pub use belt_block::cipher::{self, Key, KeyInit, KeySizeUser};
+
 /// Size of wrapping "header".
 pub const IV_LEN: usize = 16;
 
@@ -27,16 +29,6 @@ impl fmt::Debug for BeltKwp {
 }
 
 impl BeltKwp {
-    /// Create new [`BeltKwp`] instance.
-    #[inline]
-    pub fn new(key: &[u8; 32]) -> Self {
-        let mut res = [0u32; 8];
-        res.iter_mut()
-            .zip(key.chunks_exact(4))
-            .for_each(|(dst, src)| *dst = u32::from_le_bytes(src.try_into().unwrap()));
-        Self { key: res }
-    }
-
     /// Wrap key `x` with given `iv` and write result to `out`.
     ///
     /// Size of `x` must be bigger than 16 bytes.
@@ -107,6 +99,25 @@ impl BeltKwp {
     }
 }
 
+impl KeyInit for BeltKwp {
+    fn new(key: &Key<Self>) -> Self {
+        let mut res = [0u32; 8];
+        res.iter_mut()
+            .zip(key.chunks_exact(4))
+            .for_each(|(dst, src)| *dst = u32::from_le_bytes(src.try_into().unwrap()));
+
+        Self { key: res }
+    }
+}
+
+impl KeySizeUser for BeltKwp {
+    type KeySize = <BeltBlock as KeySizeUser>::KeySize;
+
+    fn key_size() -> usize {
+        BeltBlock::key_size()
+    }
+}
+
 /// Errors emitted from the wrap and unwrap operations.
 #[derive(Debug)]
 pub enum Error {