Add request/response IDs

eclipse-heimlig · Sep 28, 2023 · 84eec55 · 84eec55
1 parent 05c4b47
commit 84eec55
Show file tree

Hide file tree

Showing 9 changed files with 251 additions and 109 deletions.
diff --git a/examples/linux/src/main.rs b/examples/linux/src/main.rs
@@ -111,20 +111,22 @@ async fn client_task(
         // Send request
         Timer::after(Duration::from_millis(1000)).await;
         let random_output = Box::leak(Box::new([0u8; 16]));
-        info!(target: "CLIENT", "Sending request: random data (size={})", random_output.len());
-        api.get_random(random_output.as_mut_slice())
+        let request_size = random_output.len();
+        let request_id = api
+            .get_random(random_output.as_mut_slice())
             .await
             .expect("failed to call randomness API");
+        info!(target: "CLIENT", "--> request:  random data (id={request_id}) (size={request_size})");
 
         // Receive response
         loop {
             match api.recv_response().await {
                 None => Timer::after(Duration::from_millis(10)).await, // Continue waiting for response
                 Some(response) => {
                     match response {
-                        Response::GetRandom { data } => {
+                        Response::GetRandom { request_id, data } => {
                             info!(target: "CLIENT",
-                                "Received response: random data (size={}): {}",
+                                "<-- response: random data (id={request_id}) (size={}): {}",
                                 data.len(),
                                 hex::encode(&data)
                             );

diff --git a/examples/stm32h745i/cm7/src/bin/rng_single_core.rs b/examples/stm32h745i/cm7/src/bin/rng_single_core.rs
@@ -141,21 +141,24 @@ async fn client_task(
         // cannot unleak this. heapless::pool::Box would need to implement an interface similar to
         // std::Box::from_raw.
         core::mem::forget(random_buffer_alloc);
-        info!(
-            "Sending request: random data (size={})",
-            random_buffer.len()
-        );
-        api.get_random(random_buffer)
+        let request_size = random_buffer.len();
+        let request_id = api
+            .get_random(random_buffer)
             .await
             .expect("failed to call randomness API");
+        info!(
+            "--> request:  random data (id={}) (size={})",
+            request_id, request_size
+        );
 
         // Receive response
         loop {
             if let Some(response) = api.recv_response().await {
                 match response {
-                    Response::GetRandom { data } => {
+                    Response::GetRandom { request_id, data } => {
                         info!(
-                            "Received response: random data (size={}): {:02x}",
+                            "<-- response: random data (id={}) (size={}): {}",
+                            request_id,
                             data.len(),
                             data
                         );

diff --git a/heimlig/src/client/api.rs b/heimlig/src/client/api.rs
@@ -1,11 +1,12 @@
-use crate::common::jobs::{Request, Response};
+use crate::common::jobs::{Request, RequestId, Response};
 use crate::hsm::keystore::KeyId;
 use futures::{Sink, SinkExt, Stream, StreamExt};
 
 /// An interface to send [Request]s to the HSM core and receive [Response]es from it.
 pub struct Api<'data, Req: Sink<Request<'data>>, Resp: Stream<Item = Response<'data>>> {
     requests: Req,
     responses: Resp,
+    request_counter: RequestId,
 }
 
 #[derive(Copy, Clone, Eq, PartialEq, Debug)]
@@ -27,6 +28,7 @@ impl<
     /// Create a new instance of the HSM API.
     pub fn new(requests: Req, responses: Resp) -> Self {
         Api {
+            request_counter: 0,
             requests,
             responses,
         }
@@ -38,18 +40,30 @@ impl<
     }
 
     /// Request random bytes and write to provided buffer.
-    pub async fn get_random(&mut self, output: &'data mut [u8]) -> Result<(), Error> {
+    pub async fn get_random(&mut self, output: &'data mut [u8]) -> Result<RequestId, Error> {
+        let request_id = self.next_request_id();
         self.requests
-            .send(Request::GetRandom { output })
+            .send(Request::GetRandom { request_id, output })
             .await
-            .map_err(|_e| Error::Send)
+            .map_err(|_e| Error::Send)?;
+        Ok(request_id)
     }
 
-    pub async fn import_key(&mut self, key_id: KeyId, data: &'data [u8]) -> Result<(), Error> {
+    pub async fn import_key(
+        &mut self,
+        key_id: KeyId,
+        data: &'data [u8],
+    ) -> Result<RequestId, Error> {
+        let request_id = self.next_request_id();
         self.requests
-            .send(Request::ImportKey { key_id, data })
+            .send(Request::ImportKey {
+                request_id,
+                key_id,
+                data,
+            })
             .await
-            .map_err(|_e| Error::Send)
+            .map_err(|_e| Error::Send)?;
+        Ok(request_id)
     }
 
     pub async fn encrypt(
@@ -60,11 +74,13 @@ impl<
         plaintext: &'data mut [u8],
         aad: &'data [u8],
         tag: &'data mut [u8],
-    ) -> Result<(), Error> {
+    ) -> Result<RequestId, Error> {
+        let request_id = self.next_request_id();
         match algorithm {
             SymmetricEncryptionAlgorithm::ChaCha20Poly1305 => self
                 .requests
                 .send(Request::EncryptChaChaPoly {
+                    request_id,
                     key_id,
                     nonce,
                     plaintext,
@@ -73,7 +89,8 @@ impl<
                 })
                 .await
                 .map_err(|_e| Error::Send),
-        }
+        }?;
+        Ok(request_id)
     }
 
     pub async fn encrypt_external_key(
@@ -84,11 +101,13 @@ impl<
         plaintext: &'data mut [u8],
         aad: &'data [u8],
         tag: &'data mut [u8],
-    ) -> Result<(), Error> {
+    ) -> Result<RequestId, Error> {
+        let request_id = self.next_request_id();
         match algorithm {
             SymmetricEncryptionAlgorithm::ChaCha20Poly1305 => self
                 .requests
                 .send(Request::EncryptChaChaPolyExternalKey {
+                    request_id,
                     key,
                     nonce,
                     plaintext,
@@ -97,7 +116,8 @@ impl<
                 })
                 .await
                 .map_err(|_e| Error::Send),
-        }
+        }?;
+        Ok(request_id)
     }
 
     pub async fn decrypt(
@@ -108,11 +128,13 @@ impl<
         ciphertext: &'data mut [u8],
         aad: &'data [u8],
         tag: &'data [u8],
-    ) -> Result<(), Error> {
+    ) -> Result<RequestId, Error> {
+        let request_id = self.next_request_id();
         match algorithm {
             SymmetricEncryptionAlgorithm::ChaCha20Poly1305 => self
                 .requests
                 .send(Request::DecryptChaChaPoly {
+                    request_id,
                     key_id,
                     nonce,
                     ciphertext,
@@ -121,7 +143,8 @@ impl<
                 })
                 .await
                 .map_err(|_e| Error::Send),
-        }
+        }?;
+        Ok(request_id)
     }
 
     pub async fn decrypt_external_key(
@@ -132,11 +155,13 @@ impl<
         ciphertext: &'data mut [u8],
         aad: &'data [u8],
         tag: &'data [u8],
-    ) -> Result<(), Error> {
+    ) -> Result<RequestId, Error> {
+        let request_id = self.next_request_id();
         match algorithm {
             SymmetricEncryptionAlgorithm::ChaCha20Poly1305 => self
                 .requests
                 .send(Request::DecryptChaChaPolyExternalKey {
+                    request_id,
                     key,
                     nonce,
                     ciphertext,
@@ -145,6 +170,13 @@ impl<
                 })
                 .await
                 .map_err(|_e| Error::Send),
-        }
+        }?;
+        Ok(request_id)
+    }
+
+    fn next_request_id(&mut self) -> RequestId {
+        let id = self.request_counter;
+        self.request_counter += 1;
+        id
     }
 }
diff --git a/heimlig/src/common/jobs.rs b/heimlig/src/common/jobs.rs
@@ -15,6 +15,8 @@ pub enum Error {
     Send,
 }
 
+pub type RequestId = u32;
+
 #[derive(Copy, Clone, Eq, PartialEq, Debug)]
 pub enum RequestType {
     ImportKey,
@@ -29,34 +31,40 @@ pub enum RequestType {
 #[derive(Eq, PartialEq, Debug)]
 pub enum Request<'a> {
     ImportKey {
+        request_id: RequestId,
         key_id: KeyId,
         data: &'a [u8],
     },
     GetRandom {
+        request_id: RequestId,
         output: &'a mut [u8],
     },
     EncryptChaChaPoly {
+        request_id: RequestId,
         key_id: KeyId,
         nonce: &'a [u8],
         plaintext: &'a mut [u8],
         aad: &'a [u8],
         tag: &'a mut [u8],
     },
     EncryptChaChaPolyExternalKey {
+        request_id: RequestId,
         key: &'a [u8],
         nonce: &'a [u8],
         plaintext: &'a mut [u8],
         aad: &'a [u8],
         tag: &'a mut [u8],
     },
     DecryptChaChaPoly {
+        request_id: RequestId,
         key_id: KeyId,
         nonce: &'a [u8],
         ciphertext: &'a mut [u8],
         aad: &'a [u8],
         tag: &'a [u8],
     },
     DecryptChaChaPolyExternalKey {
+        request_id: RequestId,
         key: &'a [u8],
         nonce: &'a [u8],
         ciphertext: &'a mut [u8],
@@ -85,16 +93,24 @@ impl<'data> Request<'data> {
 /// A response from the HSM containing the results of a cryptographic task.
 #[derive(Eq, PartialEq, Debug)]
 pub enum Response<'a> {
-    ImportKey,
-    Error(Error),
+    ImportKey {
+        request_id: RequestId,
+    },
+    Error {
+        request_id: RequestId,
+        error: Error,
+    },
     GetRandom {
+        request_id: RequestId,
         data: &'a mut [u8],
     },
     EncryptChaChaPoly {
+        request_id: RequestId,
         ciphertext: &'a mut [u8],
         tag: &'a mut [u8],
     },
     DecryptChaChaPoly {
+        request_id: RequestId,
         plaintext: &'a mut [u8],
     },
 }
diff --git a/heimlig/src/hsm/core.rs b/heimlig/src/hsm/core.rs
@@ -160,7 +160,11 @@ impl<
 
     async fn process(&mut self, request: Request<'data>) -> Result<(), Error> {
         match request {
-            Request::ImportKey { key_id, data } => {
+            Request::ImportKey {
+                request_id,
+                key_id,
+                data,
+            } => {
                 let response = {
                     if let Some(key_store) = self.key_store {
                         match key_store
@@ -169,11 +173,17 @@ impl<
                             .deref_mut()
                             .import(key_id, data)
                         {
-                            Ok(()) => Response::ImportKey,
-                            Err(e) => Response::Error(jobs::Error::KeyStore(e)),
+                            Ok(()) => Response::ImportKey { request_id },
+                            Err(e) => Response::Error {
+                                request_id,
+                                error: jobs::Error::KeyStore(e),
+                            },
                         }
                     } else {
-                        Response::Error(jobs::Error::NoKeyStore)
+                        Response::Error {
+                            request_id,
+                            error: jobs::Error::NoKeyStore,
+                        }
                     }
                 };
                 self.client