From b5d929c2dae597e995575b9453e0dfd645b1685c Mon Sep 17 00:00:00 2001
From: bing <eightfilms@keemail.me>
Date: Sat, 29 Jun 2024 19:39:29 +0800
Subject: [PATCH] Key cleanup

---
 src/encryption/symmetric/aes/mod.rs   | 21 +++++++++++++------
 src/encryption/symmetric/aes/tests.rs | 30 ++++++++++-----------------
 2 files changed, 26 insertions(+), 25 deletions(-)

diff --git a/src/encryption/symmetric/aes/mod.rs b/src/encryption/symmetric/aes/mod.rs
index 0c9340b9..3b1a9ec5 100644
--- a/src/encryption/symmetric/aes/mod.rs
+++ b/src/encryption/symmetric/aes/mod.rs
@@ -1,6 +1,6 @@
 //! This module contains the implementation for the Advanced Encryption Standard (AES) encryption
 //! and decryption.
-#![doc = include_str!("./README.md")]
+#![cfg_attr(not(doctest), doc = include_str!("./README.md"))]
 
 use itertools::Itertools;
 
@@ -23,6 +23,13 @@ where [(); N / 8]: {
   inner: [u8; N / 8],
 }
 
+impl<const N: usize> Key<N>
+where [(); N / 8]:
+{
+  /// Creates a new `Key` of size `N` bits.
+  pub fn new(key_bytes: [u8; N / 8]) -> Self { Self { inner: key_bytes } }
+}
+
 impl<const N: usize> std::ops::Deref for Key<N>
 where [(); N / 8]:
 {
@@ -41,11 +48,16 @@ where [(); K / 8]:
   ///
   /// ## Example
   /// ```rust
+  /// #![feature(generic_const_exprs)]
+  ///
   /// use rand::{thread_rng, Rng};
-  /// use ronkathon::encryption::symmetric::{aes::AES, SymmetricEncryption};
+  /// use ronkathon::encryption::symmetric::{
+  ///   aes::{Key, AES},
+  ///   SymmetricEncryption,
+  /// };
   ///
   /// let mut rng = thread_rng();
-  /// let key = Key::<128> { inner: rng.gen() };
+  /// let key = Key::<128>::new(rng.gen());
   /// let plaintext = rng.gen();
   /// let encrypted = AES::encrypt(&key, &plaintext);
   /// ```
@@ -229,9 +241,6 @@ where [(); K / 8]:
     let key_words: Vec<Word> = key.chunks(4).map(|c| c.try_into().unwrap()).collect();
     expanded_key.extend(key_words);
 
-    // key len (Nk words)
-    // block size (Nb words)
-    // num rounds (Nr)
     for i in key_len..(block_num_words * (num_rounds + 1)) {
       let mut last = *expanded_key.last().unwrap();
 
diff --git a/src/encryption/symmetric/aes/tests.rs b/src/encryption/symmetric/aes/tests.rs
index 38ff1bbc..44203e09 100644
--- a/src/encryption/symmetric/aes/tests.rs
+++ b/src/encryption/symmetric/aes/tests.rs
@@ -4,12 +4,9 @@ use super::*;
 #[test]
 fn test_aes_128() {
   const KEY_LEN: usize = 128;
-  let key = Key::<KEY_LEN> {
-    inner: [
-      0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e,
-      0x0f,
-    ],
-  };
+  let key = Key::<KEY_LEN>::new([
+    0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
+  ]);
 
   let plaintext = [
     0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99, 0xaa, 0xbb, 0xcc, 0xdd, 0xee, 0xff,
@@ -27,12 +24,10 @@ fn test_aes_128() {
 #[test]
 fn test_aes_192() {
   const KEY_LEN: usize = 192;
-  let key = Key::<KEY_LEN> {
-    inner: [
-      0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e,
-      0x0f, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17,
-    ],
-  };
+  let key = Key::<KEY_LEN>::new([
+    0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
+    0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17,
+  ]);
 
   let plaintext = [
     0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99, 0xaa, 0xbb, 0xcc, 0xdd, 0xee, 0xff,
@@ -50,13 +45,10 @@ fn test_aes_192() {
 #[test]
 fn test_aes_256() {
   const KEY_LEN: usize = 256;
-  let key = Key::<KEY_LEN> {
-    inner: [
-      0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e,
-      0x0f, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d,
-      0x1e, 0x1f,
-    ],
-  };
+  let key = Key::<KEY_LEN>::new([
+    0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
+    0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f,
+  ]);
 
   let plaintext = [
     0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99, 0xaa, 0xbb, 0xcc, 0xdd, 0xee, 0xff,