From 37bc65830655ed97ce4e36f8f34d057942fb9c4e Mon Sep 17 00:00:00 2001
From: Tony Arcieri <bascule@gmail.com>
Date: Sat, 25 Feb 2023 19:00:09 -0700
Subject: [PATCH] ascon: `Key` type alias; cache key in `Ascon` state (#42)

- Adds a type alias for Ascon keys (16-byte)
- Stores key internally in `struct Ascon` so it doesn't need to be
  passed separately to `Ascon::finalize`.
---
 Cargo.lock        | 22 ----------------------
 ascon/src/aead.rs | 14 +++++++-------
 ascon/src/lib.rs  | 31 +++++++++++++++++++------------
 3 files changed, 26 insertions(+), 41 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 7f62d02..680a29b 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -5,25 +5,3 @@ version = 3
 [[package]]
 name = "ascon"
 version = "0.2.0-pre"
-
-[[package]]
-name = "cpufeatures"
-version = "0.2.5"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "28d997bd5e24a5928dd43e46dc529867e207907fe0b239c3477d924f7f2ca320"
-dependencies = [
- "libc",
-]
-
-[[package]]
-name = "keccak"
-version = "0.1.3"
-dependencies = [
- "cpufeatures",
-]
-
-[[package]]
-name = "libc"
-version = "0.2.137"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "fc7fcc620a3bff7cdd7a365be3376c97191aeaccc2a603e600951e452615bf89"
diff --git a/ascon/src/aead.rs b/ascon/src/aead.rs
index 3f1ce58..81c6dd6 100644
--- a/ascon/src/aead.rs
+++ b/ascon/src/aead.rs
@@ -1,6 +1,6 @@
 //! Authenticated Encryption with Associated Data.
 
-use crate::{Ascon, KEY_LEN, RATE, S_SIZE};
+use crate::{Ascon, Key, KEY_SIZE, RATE, S_SIZE};
 use alloc::{vec, vec::Vec};
 
 /// Ascon(a,b) `b`-parameter.
@@ -23,7 +23,7 @@ pub enum DecryptFail {
 }
 
 /// AEAD encryption.
-pub fn encrypt(key: &[u8], iv: &[u8], message: &[u8], aad: &[u8]) -> (Vec<u8>, Tag) {
+pub fn encrypt(key: &Key, iv: &[u8], message: &[u8], aad: &[u8]) -> (Vec<u8>, Tag) {
     let s = aad.len() / RATE + 1;
     let t = message.len() / RATE + 1;
     let l = message.len() % RATE;
@@ -70,20 +70,20 @@ pub fn encrypt(key: &[u8], iv: &[u8], message: &[u8], aad: &[u8]) -> (Vec<u8>, T
 
     // tag
     let mut tag = Tag::default();
-    tag.copy_from_slice(&ss.finalize(key)[S_SIZE - KEY_LEN..]);
+    tag.copy_from_slice(&ss.finalize()[S_SIZE - KEY_SIZE..]);
 
     (output, tag)
 }
 
 /// AEAD decryption.
 pub fn decrypt(
-    key: &[u8],
+    key: &Key,
     iv: &[u8],
     ciphertext: &[u8],
     aad: &[u8],
     tag: &[u8],
 ) -> Result<Vec<u8>, DecryptFail> {
-    if tag.len() != KEY_LEN {
+    if tag.len() != KEY_SIZE {
         Err(DecryptFail::TagLengthError)?
     };
 
@@ -128,9 +128,9 @@ pub fn decrypt(
     ss.state[l] ^= 0x80;
 
     // finalization
-    let expected_tag = ss.finalize(key);
+    let expected_tag = ss.finalize();
 
-    if ct_eq(&expected_tag[S_SIZE - KEY_LEN..], tag) {
+    if ct_eq(&expected_tag[S_SIZE - KEY_SIZE..], tag) {
         Ok(mm[..ciphertext.len()].into())
     } else {
         Err(DecryptFail::AuthenticationFail)
diff --git a/ascon/src/lib.rs b/ascon/src/lib.rs
index d9e25db..3840a1f 100644
--- a/ascon/src/lib.rs
+++ b/ascon/src/lib.rs
@@ -22,8 +22,8 @@ pub mod aead;
 
 use core::convert::TryInto;
 
-/// Key length.
-const KEY_LEN: usize = 16;
+/// Size of an Ascon key in bytes.
+const KEY_SIZE: usize = 16;
 
 /// State size in bits.
 const S_BITS: usize = 320;
@@ -34,24 +34,31 @@ const S_SIZE: usize = S_BITS / 8;
 /// Rate: Sᵣ.
 const RATE: usize = 128 / 8;
 
+/// Ascon permutation key.
+pub type Key = [u8; KEY_SIZE];
+
+/// Ascon permutation state.
+type State = [u8; S_SIZE];
+
 /// Ascon(a,b) permutation.
 pub struct Ascon<const A: usize = 12, const B: usize = 8> {
-    state: [u8; S_SIZE],
+    key: Key,
+    state: State,
 }
 
 impl<const A: usize, const B: usize> Ascon<A, B> {
     /// Initialize Ascon permutation.
     // TODO(tarcieri): validate length of key and nonce
-    pub fn new(key: &[u8], nonce: &[u8]) -> Self {
+    pub fn new(key: &Key, nonce: &[u8]) -> Self {
         let mut state = [0; S_SIZE];
-        state[0] = KEY_LEN as u8 * 8;
+        state[0] = KEY_SIZE as u8 * 8;
         state[1] = RATE as u8 * 8;
         state[2] = A as u8;
         state[3] = B as u8;
 
-        let mut pos = S_SIZE - 2 * KEY_LEN;
+        let mut pos = S_SIZE - 2 * KEY_SIZE;
         state[pos..pos + key.len()].copy_from_slice(key);
-        pos += KEY_LEN;
+        pos += KEY_SIZE;
         state[pos..pos + nonce.len()].copy_from_slice(nonce);
 
         permutation(&mut state, 12 - A, A);
@@ -60,7 +67,7 @@ impl<const A: usize, const B: usize> Ascon<A, B> {
             state[pos + i] ^= b;
         }
 
-        Self { state }
+        Self { key: *key, state }
     }
 
     /// Perform Ascon permutation on internal state.
@@ -71,17 +78,17 @@ impl<const A: usize, const B: usize> Ascon<A, B> {
     /// Finalize Ascon permutation.
     #[inline(always)]
     #[must_use]
-    pub fn finalize(self, key: &[u8]) -> [u8; S_SIZE] {
+    pub fn finalize(self) -> [u8; S_SIZE] {
         let mut s = self.state;
 
-        for (i, &b) in key.iter().enumerate() {
+        for (i, &b) in self.key.iter().enumerate() {
             s[RATE + i] ^= b;
         }
 
         permutation(&mut s, 12 - A, A);
 
-        for (i, &b) in key.iter().enumerate() {
-            s[S_SIZE - KEY_LEN + i] ^= b;
+        for (i, &b) in self.key.iter().enumerate() {
+            s[S_SIZE - KEY_SIZE + i] ^= b;
         }
 
         s