From 9b7a1a339fce75334b5ec8015b16ce102d24d064 Mon Sep 17 00:00:00 2001
From: Egor Lazarchuk <yegorlz@amazon.co.uk>
Date: Sun, 28 May 2023 11:53:49 +0100
Subject: [PATCH] feat: added selection of entropy crate

Added ability to select between `rand` and
`aws-lc-rs` crates for entropy device.

Signed-off-by: Egor Lazarchuk <yegorlz@amazon.co.uk>
---
 Cargo.lock                               |  1 +
 src/cpu-template-helper/Cargo.toml       |  2 +-
 src/firecracker/Cargo.toml               |  7 ++++++-
 src/vmm/Cargo.toml                       |  8 +++++++-
 src/vmm/build.rs                         |  5 +++++
 src/vmm/src/devices/virtio/rng/device.rs | 14 ++++++++++++--
 6 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 4f77cecfe62..8f24598ef64 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1391,6 +1391,7 @@ dependencies = [
  "mmds",
  "net_gen",
  "proptest",
+ "rand",
  "rate_limiter",
  "seccompiler",
  "serde",
diff --git a/src/cpu-template-helper/Cargo.toml b/src/cpu-template-helper/Cargo.toml
index 1aee2cb1ffc..3fa955a08d5 100644
--- a/src/cpu-template-helper/Cargo.toml
+++ b/src/cpu-template-helper/Cargo.toml
@@ -18,7 +18,7 @@ serde = { version = "1.0.136", features = ["derive"] }
 serde_json = "1.0.78"
 thiserror = "1.0.32"
 
-vmm = { path = "../vmm" }
+vmm = { path = "../vmm", features = ["rng-rand"] }
 
 [dev-dependencies]
 utils = { path = "../utils" }
diff --git a/src/firecracker/Cargo.toml b/src/firecracker/Cargo.toml
index 930557adc1c..703d5cf1bee 100644
--- a/src/firecracker/Cargo.toml
+++ b/src/firecracker/Cargo.toml
@@ -12,6 +12,11 @@ license = "Apache-2.0"
 name = "firecracker"
 bench = false
 
+[features]
+default = ["rng-aws-lc-rs"]
+rng-aws-lc-rs = ["vmm/rng-aws-lc-rs"]
+rng-rand = ["vmm/rng-rand"]
+
 [dependencies]
 event-manager = "0.3.0"
 libc = "0.2.117"
@@ -24,7 +29,7 @@ mmds = { path = "../mmds" }
 seccompiler = { path = "../seccompiler" }
 snapshot = { path = "../snapshot" }
 utils = { path = "../utils" }
-vmm = { path = "../vmm" }
+vmm = { path = "../vmm", default-features = false }
 
 [dev-dependencies]
 cargo_toml = "0.15.2"
diff --git a/src/vmm/Cargo.toml b/src/vmm/Cargo.toml
index d33e37f3e56..b94775dc54e 100644
--- a/src/vmm/Cargo.toml
+++ b/src/vmm/Cargo.toml
@@ -8,8 +8,14 @@ license = "Apache-2.0"
 [lib]
 bench = false
 
+[features]
+default = []
+rng-aws-lc-rs = ["dep:aws-lc-rs"]
+rng-rand = ["dep:rand"]
+
 [dependencies]
-aws-lc-rs = "1.0.2"
+aws-lc-rs = { version = "1.0.2", optional = true }
+rand = { version = "0.8.5", optional = true }
 bitflags = "2.0.2"
 derive_more = { version = "0.99.17", default-features = false, features = ["from", "display"] }
 event-manager = "0.3.0"
diff --git a/src/vmm/build.rs b/src/vmm/build.rs
index 839440a81aa..f8c83c612c4 100644
--- a/src/vmm/build.rs
+++ b/src/vmm/build.rs
@@ -127,4 +127,9 @@ fn cpuid() {
         all(target_arch = "x86_64", not(target_env = "sgx"))
     ))]
     println!("cargo:rustc-cfg=cpuid");
+
+    #[cfg(
+        not(any(feature="rng-aws-lc-rs", feature="rng-rand")) // If neither are enabled
+    )]
+    compile_error!("Please enable the feature \"rng-aws-lc-rs\" OR the feature \"rng-rand\".");
 }
diff --git a/src/vmm/src/devices/virtio/rng/device.rs b/src/vmm/src/devices/virtio/rng/device.rs
index 7d4ba53f6ce..ae596a8bc91 100644
--- a/src/vmm/src/devices/virtio/rng/device.rs
+++ b/src/vmm/src/devices/virtio/rng/device.rs
@@ -5,8 +5,11 @@ use std::io;
 use std::sync::atomic::AtomicUsize;
 use std::sync::Arc;
 
-use aws_lc_rs::rand;
+#[cfg(feature = "rng-aws-lc-rs")]
+use aws_lc_rs::{error::Unspecified as RandomError, rand};
 use logger::{debug, error, IncMetric, METRICS};
+#[cfg(all(feature = "rng-rand", not(feature = "rng-aws-lc-rs")))]
+use rand::{rngs::OsRng, Error as RandomError, RngCore};
 use rate_limiter::{RateLimiter, TokenType};
 use utils::eventfd::EventFd;
 use utils::vm_memory::{GuestMemoryError, GuestMemoryMmap};
@@ -27,7 +30,7 @@ pub enum Error {
     #[error("Bad guest memory buffer: {0}")]
     GuestMemory(#[from] GuestMemoryError),
     #[error("Could not get random bytes: {0}")]
-    Random(#[from] aws_lc_rs::error::Unspecified),
+    Random(#[from] RandomError),
 }
 
 type Result<T> = std::result::Result<T, Error>;
@@ -109,10 +112,17 @@ impl Entropy {
         }
 
         let mut rand_bytes = vec![0; iovec.len()];
+
+        #[cfg(feature = "rng-aws-lc-rs")]
         rand::fill(&mut rand_bytes).map_err(|err| {
             METRICS.entropy.host_rng_fails.inc();
             err
         })?;
+        #[cfg(all(feature = "rng-rand", not(feature = "rng-aws-lc-rs")))]
+        OsRng.try_fill_bytes(&mut rand_bytes).map_err(|err| {
+            METRICS.entropy.host_rng_fails.inc();
+            err
+        })?;
 
         // It is ok to unwrap here. We are writing `iovec.len()` bytes at offset 0.
         Ok(iovec.write_at(&rand_bytes, 0).unwrap().try_into().unwrap())