From ce1025bef47b1a823ca35ea84fb030b47c704e89 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 | 14 +------
 kbs/src/attestation/coco/grpc.rs    | 12 +-----
 kbs/src/attestation/mod.rs          | 63 ++++++++++++++++++++++++++---
 3 files changed, 61 insertions(+), 28 deletions(-)

diff --git a/kbs/src/attestation/coco/builtin.rs b/kbs/src/attestation/coco/builtin.rs
index 1433b5f70..c5a194e8f 100644
--- a/kbs/src/attestation/coco/builtin.rs
+++ b/kbs/src/attestation/coco/builtin.rs
@@ -2,13 +2,11 @@
 // Licensed under the Apache License, Version 2.0, see LICENSE for details.
 // SPDX-License-Identifier: Apache-2.0
 
-use crate::attestation::Attest;
+use crate::attestation::{make_nonce, Attest};
 use anyhow::*;
 use async_trait::async_trait;
 use attestation_service::{config::Config as AsConfig, AttestationService, Data, HashAlgorithm};
-use base64::{engine::general_purpose::STANDARD, Engine};
 use kbs_types::{Attestation, Challenge, Tee};
-use rand::{thread_rng, Rng};
 use serde_json::json;
 use tokio::sync::RwLock;
 
@@ -56,15 +54,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 93fefe3f2..0d61f271b 100644
--- a/kbs/src/attestation/coco/grpc.rs
+++ b/kbs/src/attestation/coco/grpc.rs
@@ -2,7 +2,7 @@
 // Licensed under the Apache License, Version 2.0, see LICENSE for details.
 // SPDX-License-Identifier: Apache-2.0
 
-use crate::attestation::Attest;
+use crate::attestation::{make_nonce, Attest};
 use anyhow::*;
 use async_trait::async_trait;
 use base64::{
@@ -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 87982340a..bacdd5640 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);
+        }
+    }
+}