From da12556732bf34ded5420636cc2561ecc6cf6f15 Mon Sep 17 00:00:00 2001
From: "James O. D. Hunt" <james.o.hunt@intel.com>
Date: Thu, 1 Aug 2024 13:54:09 +0100
Subject: [PATCH] kbs: Refactor nonce handling

Create a common function to generate a nonce, and add a unit test for it.

Signed-off-by: James O. D. Hunt <james.o.hunt@intel.com>
---
 kbs/src/attestation/coco/builtin.rs | 10 +----
 kbs/src/attestation/coco/grpc.rs    | 10 +----
 kbs/src/attestation/mod.rs          | 63 ++++++++++++++++++++++++++---
 3 files changed, 59 insertions(+), 24 deletions(-)

diff --git a/kbs/src/attestation/coco/builtin.rs b/kbs/src/attestation/coco/builtin.rs
index 1433b5f702..fdeaf6600c 100644
--- a/kbs/src/attestation/coco/builtin.rs
+++ b/kbs/src/attestation/coco/builtin.rs
@@ -56,15 +56,7 @@ impl Attest for BuiltInCoCoAs {
                     .generate_supplemental_challenge(tee, tee_parameters)
                     .await?
             }
-            _ => {
-                let mut nonce: Vec<u8> = vec![0; 32];
-
-                thread_rng()
-                    .try_fill(&mut nonce[..])
-                    .map_err(anyhow::Error::from)?;
-
-                STANDARD.encode(&nonce)
-            }
+            _ => make_nonce().await?,
         };
 
         let challenge = Challenge {
diff --git a/kbs/src/attestation/coco/grpc.rs b/kbs/src/attestation/coco/grpc.rs
index 93fefe3f2f..1af19136eb 100644
--- a/kbs/src/attestation/coco/grpc.rs
+++ b/kbs/src/attestation/coco/grpc.rs
@@ -140,15 +140,7 @@ impl Attest for GrpcClientPool {
                     .into_inner()
                     .attestation_challenge
             }
-            _ => {
-                let mut nonce: Vec<u8> = vec![0; 32];
-
-                thread_rng()
-                    .try_fill(&mut nonce[..])
-                    .map_err(anyhow::Error::from)?;
-
-                STANDARD.encode(&nonce)
-            }
+            _ => make_nonce().await?,
         };
 
         let challenge = Challenge {
diff --git a/kbs/src/attestation/mod.rs b/kbs/src/attestation/mod.rs
index 87982340a2..bacdd56405 100644
--- a/kbs/src/attestation/mod.rs
+++ b/kbs/src/attestation/mod.rs
@@ -26,6 +26,20 @@ pub mod coco;
 #[cfg(feature = "intel-trust-authority-as")]
 pub mod intel_trust_authority;
 
+/// Number of bytes in a nonce.
+const NONCE_SIZE_BYTES: usize = 32;
+
+/// Create a nonce and return as a base-64 encoded string.
+pub async fn make_nonce() -> Result<String> {
+    let mut nonce: Vec<u8> = vec![0; NONCE_SIZE_BYTES];
+
+    thread_rng()
+        .try_fill(&mut nonce[..])
+        .map_err(anyhow::Error::from)?;
+
+    Ok(STANDARD.encode(&nonce))
+}
+
 /// Interface for Attestation Services.
 ///
 /// Attestation Service implementations should implement this interface.
@@ -42,13 +56,8 @@ pub trait Attest: Send + Sync {
 
     /// generate the Challenge to pass to attester based on Tee and nonce
     async fn generate_challenge(&self, _tee: Tee, _tee_parameters: String) -> Result<Challenge> {
-        let mut nonce: Vec<u8> = vec![0; 32];
-
-        thread_rng()
-            .try_fill(&mut nonce[..])
-            .map_err(anyhow::Error::from)?;
+        let nonce = make_nonce().await?;
 
-        let nonce = STANDARD.encode(&nonce);
         Ok(Challenge {
             nonce,
             extra_params: String::new(),
@@ -129,3 +138,45 @@ impl AttestationService {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn test_make_nonce() {
+        const BITS_PER_BYTE: usize = 8;
+
+        /// A base-64 encoded value is this many bits in length.
+        const BASE64_BITS_CHUNK: usize = 6;
+
+        /// Number of bytes that base64 encoding requires the result to align on.
+        const BASE64_ROUNDING_MULTIPLE: usize = 4;
+
+        /// The nominal base64 encoded length.
+        const BASE64_NONCE_LENGTH_UNROUNDED_BYTES: usize =
+            (NONCE_SIZE_BYTES * BITS_PER_BYTE) / BASE64_BITS_CHUNK;
+
+        /// The actual base64 encoded length is rounded up to the specified multiple.
+        const EXPECTED_LENGTH_BYTES: usize =
+            BASE64_NONCE_LENGTH_UNROUNDED_BYTES.next_multiple_of(BASE64_ROUNDING_MULTIPLE);
+
+        // Number of nonce tests to run (arbitrary)
+        let nonce_count = 13;
+
+        let mut nonces = vec![];
+
+        for _ in 0..nonce_count {
+            let nonce = make_nonce().await.unwrap();
+
+            assert_eq!(nonce.len(), EXPECTED_LENGTH_BYTES);
+
+            let found = nonces.contains(&nonce);
+
+            // The nonces should be unique
+            assert_eq!(found, false);
+
+            nonces.push(nonce);
+        }
+    }
+}