From 01c4c413d524a862987f8234bf65c01e7486ed58 Mon Sep 17 00:00:00 2001
From: Norbert Fabritius <norbert.fabritius@accenture.com>
Date: Mon, 8 Jul 2024 16:57:43 +0200
Subject: [PATCH] Use try_fill_bytes() to handle RNG errors in RngWorker

---
 heimlig/src/crypto/mod.rs             |  2 ++
 heimlig/src/hsm/workers/rng_worker.rs | 25 +++++++++++++++++++------
 heimlig/src/integration/raw_errors.rs |  3 +++
 3 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/heimlig/src/crypto/mod.rs b/heimlig/src/crypto/mod.rs
index 00fc3fdc..056cfe2f 100644
--- a/heimlig/src/crypto/mod.rs
+++ b/heimlig/src/crypto/mod.rs
@@ -12,6 +12,8 @@ pub mod x25519;
 /// Common errors.
 #[derive(Copy, Clone, Debug, Eq, PartialEq)]
 pub enum Error {
+    /// Error during random number generation.
+    Random,
     /// Error during encryption.
     Encrypt,
     /// Error during decryption.
diff --git a/heimlig/src/hsm/workers/rng_worker.rs b/heimlig/src/hsm/workers/rng_worker.rs
index 1d59451f..af5738b6 100644
--- a/heimlig/src/hsm/workers/rng_worker.rs
+++ b/heimlig/src/hsm/workers/rng_worker.rs
@@ -1,5 +1,6 @@
 use crate::common::jobs::{ClientId, Error, Request, RequestId, Response};
 use crate::common::limits::MAX_RANDOM_SIZE;
+use crate::crypto;
 use crate::hsm::keystore::{self, KeyId};
 use embassy_sync::blocking_mutex::raw::RawMutex;
 use embassy_sync::mutex::Mutex;
@@ -82,11 +83,17 @@ impl<
                 error: Error::RequestTooLarge,
             };
         }
-        self.rng.lock().await.fill_bytes(output);
-        Response::GetRandom {
-            client_id,
-            request_id,
-            data: output,
+        match self.rng.lock().await.try_fill_bytes(output) {
+            Ok(()) => Response::GetRandom {
+                client_id,
+                request_id,
+                data: output,
+            },
+            Err(_) => Response::Error {
+                client_id,
+                request_id,
+                error: Error::Crypto(crypto::Error::Random),
+            },
         }
     }
 
@@ -109,7 +116,13 @@ impl<
             Ok(key_info) => {
                 let mut key = [0u8; keystore::KeyType::MAX_SYMMETRIC_KEY_SIZE];
                 let key = &mut key[0..key_info.ty.key_size()];
-                self.rng.lock().await.fill_bytes(key);
+                if self.rng.lock().await.try_fill_bytes(key).is_err() {
+                    return Response::Error {
+                        client_id,
+                        request_id,
+                        error: Error::Crypto(crypto::Error::Random),
+                    };
+                };
                 let mut locked_key_store = key_store.lock().await;
 
                 // Check overwrite permission
diff --git a/heimlig/src/integration/raw_errors.rs b/heimlig/src/integration/raw_errors.rs
index 79ef4725..d6624606 100644
--- a/heimlig/src/integration/raw_errors.rs
+++ b/heimlig/src/integration/raw_errors.rs
@@ -28,6 +28,8 @@ pub enum JobErrorRaw {
 #[repr(u8)]
 #[derive(Copy, Clone, Debug, Eq, PartialEq)]
 pub enum CryptoErrorRaw {
+    /// Error during random number generation
+    Random,
     /// Error during encryption.
     Encrypt,
     /// Error during decryption.
@@ -98,6 +100,7 @@ impl From<jobs::Error> for JobErrorRaw {
 impl From<crypto::Error> for CryptoErrorRaw {
     fn from(value: crypto::Error) -> Self {
         match value {
+            crypto::Error::Random => CryptoErrorRaw::Random,
             crypto::Error::Encrypt => CryptoErrorRaw::Encrypt,
             crypto::Error::Decrypt => CryptoErrorRaw::Decrypt,
             crypto::Error::Sign => CryptoErrorRaw::Sign,