Browse Source

Use SmallRng in warp-rust (#6012)

Konrad Borowski 4 years ago
parent
commit
d8723f6293
2 changed files with 22 additions and 8 deletions
  1. 1 1
      frameworks/Rust/warp-rust/Cargo.toml
  2. 21 7
      frameworks/Rust/warp-rust/src/main.rs

+ 1 - 1
frameworks/Rust/warp-rust/Cargo.toml

@@ -6,7 +6,7 @@ edition = "2018"
 
 [dependencies]
 futures = "0.3.1"
-rand = "0.7.3"
+rand = { version = "0.7.3", features = ["small_rng"] }
 serde = { version = "1.0.103", features = ["derive"] }
 tokio = { version = "0.2.21", features = ["macros", "rt-threaded"] }
 tokio-postgres = "0.5.1"

+ 21 - 7
frameworks/Rust/warp-rust/src/main.rs

@@ -4,11 +4,24 @@ use crate::db::{Database, Fortune};
 use futures::stream::futures_unordered::FuturesUnordered;
 use futures::{FutureExt, StreamExt};
 use rand::distributions::{Distribution, Uniform};
+use rand::rngs::SmallRng;
+use rand::SeedableRng;
 use serde::Serialize;
+use std::cell::RefCell;
 use warp::http::header;
 use warp::{Filter, Rejection, Reply};
 use yarte::Template;
 
+// SmallRng is not a CSPRNG. It's used specifically to match performance of
+// benchmarks in other programming languages where the default RNG algorithm
+// is not a CSPRNG. Most Rust programs will likely want to use
+// rand::thread_rng instead which is more convenient to use and safer.
+thread_local!(static RNG: RefCell<SmallRng> = RefCell::new(SmallRng::from_entropy()));
+
+fn with_rng<T>(f: impl FnOnce(&mut SmallRng) -> T) -> T {
+    RNG.with(|rng| f(&mut rng.borrow_mut()))
+}
+
 #[derive(Serialize)]
 struct Message {
     message: &'static str,
@@ -29,7 +42,7 @@ fn plaintext() -> impl Filter<Extract = impl Reply, Error = Rejection> + Clone {
 fn db(database: &'static Database) -> impl Filter<Extract = impl Reply, Error = Rejection> + Clone {
     let between = Uniform::from(1..=10_000);
     warp::path!("db").and_then(move || async move {
-        let id = between.sample(&mut rand::thread_rng());
+        let id = with_rng(|rng| between.sample(rng));
         let world = database.get_world_by_id(id).await;
         Ok::<_, Rejection>(warp::reply::json(&world))
     })
@@ -43,12 +56,13 @@ fn queries(
     warp::path!("queries" / ..)
         .and(clamped.or(warp::any().map(|| 1)).unify())
         .and_then(move |queries| {
-            let mut rng = rand::thread_rng();
-            (0..queries)
-                .map(|_| database.get_world_by_id(between.sample(&mut rng)))
-                .collect::<FuturesUnordered<_>>()
-                .collect::<Vec<_>>()
-                .map(|worlds| Ok::<_, Rejection>(warp::reply::json(&worlds)))
+            with_rng(|rng| {
+                (0..queries)
+                    .map(|_| database.get_world_by_id(between.sample(rng)))
+                    .collect::<FuturesUnordered<_>>()
+                    .collect::<Vec<_>>()
+                    .map(|worlds| Ok::<_, Rejection>(warp::reply::json(&worlds)))
+            })
         })
 }