From 72d4c4877add5f9d112210f93b859c2f67af15ef Mon Sep 17 00:00:00 2001
From: Hannes Mehnert <hannes@mehnert.org>
Date: Mon, 11 Mar 2024 12:44:29 +0100
Subject: [PATCH] rng: provide a generate_into : ?g -> bytes -> ?off:int -> int
 -> unit

and reimplement the generate in terms of generate_into

this keeps the allocation at the API boundary if desired
---
 bench/speed.ml                         |  5 ++++-
 rng/fortuna.ml                         | 26 ++++++++++++++------------
 rng/hmac_drbg.ml                       | 23 ++++++++++++++++-------
 rng/mirage_crypto_rng.mli              | 13 ++++++++++---
 rng/rng.ml                             | 23 +++++++++++++++--------
 tests/test_eio_entropy_collection.ml   |  2 +-
 tests/test_entropy_collection.ml       |  2 +-
 tests/test_entropy_collection_async.ml |  2 +-
 tests/test_rsa.ml                      |  7 +++----
 9 files changed, 65 insertions(+), 38 deletions(-)

diff --git a/bench/speed.ml b/bench/speed.ml
index cb6b24ef..c344f230 100644
--- a/bench/speed.ml
+++ b/bench/speed.ml
@@ -27,6 +27,8 @@ let burn_period = 2.0
 let sizes = [16; 64; 256; 1024; 8192]
 (* let sizes = [16] *)
 
+let big_b = Bytes.create List.(hd (rev sizes))
+
 let burn f n =
   let cs = Cstruct.of_string (Mirage_crypto_rng.generate n) in
   let (t1, i1) =
@@ -410,7 +412,8 @@ let benchmarks = [
     let open Mirage_crypto_rng.Fortuna in
     let g = create () in
     reseed ~g "abcd" ;
-    throughput name (fun cs -> generate ~g (Cstruct.length cs))) ;
+    throughput name (fun cs ->
+        generate_into ~g big_b ~off:0 (Cstruct.length cs))) ;
 
   bm "md5"    (fun name -> throughput name MD5.digest) ;
   bm "sha1"   (fun name -> throughput name SHA1.digest) ;
diff --git a/rng/fortuna.ml b/rng/fortuna.ml
index 66daf849..6cd696df 100644
--- a/rng/fortuna.ml
+++ b/rng/fortuna.ml
@@ -64,15 +64,14 @@ let iter1 a     f = f a
 
 let reseed ~g cs = reseedi ~g (iter1 cs)
 
-let generate_rekey ~g bytes =
-  let b  = bytes // block + 2 in
+let generate_rekey ~g buf ~off len =
+  let b  = len // block + 2 in
   let n  = b * block in
   let r  = Cstruct.to_string (AES_CTR.stream ~key:g.key ~ctr:g.ctr n) in
-  let r1 = String.sub r 0 bytes
-  and r2 = String.sub r (n - 32) 32 in
+  Bytes.blit_string r 0 buf off len;
+  let r2 = String.sub r (n - 32) 32 in
   set_key ~g r2 ;
-  g.ctr <- AES_CTR.add_ctr g.ctr (Int64.of_int b);
-  r1
+  g.ctr <- AES_CTR.add_ctr g.ctr (Int64.of_int b)
 
 let add_pool_entropy g =
   if g.pool0_size > min_pool_size then
@@ -94,14 +93,17 @@ let add_pool_entropy g =
       done
     end
 
-let generate ~g bytes =
+let generate_into ~g buf ~off len =
   add_pool_entropy g;
   if not (seeded ~g) then raise Rng.Unseeded_generator ;
-  let rec chunk acc = function
-    | i when i <= 0 -> acc
-    | n -> let n' = imin n 0x10000 in
-           chunk (generate_rekey ~g n' :: acc) (n - n') in
-  String.concat "" @@ chunk [] bytes
+  let rec chunk off = function
+    | i when i <= 0 -> ()
+    | n ->
+      let n' = imin n 0x10000 in
+      generate_rekey ~g buf ~off n';
+      chunk (off + n') (n - n')
+  in
+  chunk off len
 
 let _buf = Bytes.create 2
 
diff --git a/rng/hmac_drbg.ml b/rng/hmac_drbg.ml
index 8092eec4..bb81b0d1 100644
--- a/rng/hmac_drbg.ml
+++ b/rng/hmac_drbg.ml
@@ -24,17 +24,26 @@ module Make (H : Digestif.S) = struct
     let v = H.hmac_string ~key:k v |> H.to_raw_string in
     g.k <- k ; g.v <- v ; g.seeded <- true
 
-  let generate ~g bytes =
+  let generate_into ~g buf ~off len =
     if not g.seeded then raise Rng.Unseeded_generator ;
-    let rec go acc k v = function
-      | 0 -> (v, String.concat "" @@ List.rev acc)
+    let rec go off k v = function
+      | 0 -> v (* unlikely this happens *)
+      | 1 ->
+        let v = H.hmac_string ~key:k v |> H.to_raw_string in
+        let len =
+          let rem = len mod H.digest_size in
+          if rem = 0 then H.digest_size else rem
+        in
+        Bytes.blit_string v 0 buf off len;
+        v
       | i ->
         let v = H.hmac_string ~key:k v |> H.to_raw_string in
-        go (v::acc) k v (pred i) in
-    let (v, buf) = go [] g.k g.v Mirage_crypto.Uncommon.(bytes // H.digest_size) in
+        Bytes.blit_string v 0 buf off H.digest_size;
+        go (off + H.digest_size) k v (pred i)
+    in
+    let v = go off g.k g.v Mirage_crypto.Uncommon.(len // H.digest_size) in
     g.k <- H.hmac_string ~key:g.k (v ^ bx00) |> H.to_raw_string;
-    g.v <- H.hmac_string ~key:g.k v |> H.to_raw_string;
-    String.sub buf 0 (Mirage_crypto.Uncommon.imax 0 bytes)
+    g.v <- H.hmac_string ~key:g.k v |> H.to_raw_string
 
   (* XXX *)
   let accumulate ~g:_ = invalid_arg "Implement Hmac_drbg.accumulate..."
diff --git a/rng/mirage_crypto_rng.mli b/rng/mirage_crypto_rng.mli
index 6a14fa2a..453ee193 100644
--- a/rng/mirage_crypto_rng.mli
+++ b/rng/mirage_crypto_rng.mli
@@ -156,9 +156,12 @@ module type Generator = sig
   val create : ?time:(unit -> int64) -> unit -> g
   (** Create a new, unseeded {{!g}g}. *)
 
-  val generate : g:g -> int -> string
-  (** [generate ~g n] produces [n] uniformly distributed random bytes,
-      updating the state of [g]. *)
+  val generate_into : g:g -> bytes -> off:int -> int -> unit
+  (** [generate_into ~g buf ~off n] produces [n] uniformly distributed random
+      bytes into [buf] at offset [off], updating the state of [g].
+
+      @raise Invalid_argument if buffer is too small (it must be: Bytes.length buf - off >= n)
+  *)
 
   val reseed : g:g -> string -> unit
   (** [reseed ~g bytes] directly updates [g]. Its new state depends both on
@@ -231,6 +234,10 @@ val generate : ?g:g -> int -> string
 (** Invoke {{!Generator.generate}generate} on [g] or
     {{!generator}default generator}. *)
 
+val generate_into : ?g:g -> bytes -> ?off:int -> int -> unit
+(** Invoke {{!Generator.generate}generate} on [g] or
+    {{!generator}default generator}. The offset [off] defaults to 0. *)
+
 val block : g option -> int
 (** {{!Generator.block}Block} size of [g] or
     {{!generator}default generator}. *)
diff --git a/rng/rng.ml b/rng/rng.ml
index f7e295c2..2ddaa4cd 100644
--- a/rng/rng.ml
+++ b/rng/rng.ml
@@ -36,13 +36,13 @@ let () = Printexc.register_printer (function
 
 module type Generator = sig
   type g
-  val block      : int
-  val create     : ?time:(unit -> int64) -> unit -> g
-  val generate   : g:g -> int -> string
-  val reseed     : g:g -> string -> unit
+  val block : int
+  val create : ?time:(unit -> int64) -> unit -> g
+  val generate_into : g:g -> bytes -> off:int -> int -> unit
+  val reseed : g:g -> string -> unit
   val accumulate : g:g -> source -> [`Acc of string -> unit]
-  val seeded     : g:g -> bool
-  val pools      : int
+  val seeded : g:g -> bool
+  val pools : int
 end
 
 type 'a generator = (module Generator with type g = 'a)
@@ -67,8 +67,15 @@ let default_generator () =
 
 let get = function Some g -> g | None -> default_generator ()
 
-let generate ?(g = default_generator ()) n =
-  let Generator (g, _, m) = g in let module M = (val m) in M.generate ~g n
+let generate_into ?(g = default_generator ()) b ?(off = 0) n =
+  let Generator (g, _, m) = g in
+  let module M = (val m) in
+  M.generate_into ~g b ~off n
+
+let generate ?g n =
+  let data = Bytes.create n in
+  generate_into ?g data ~off:0 n;
+  Bytes.unsafe_to_string data
 
 let reseed ?(g = default_generator ()) cs =
   let Generator (g, _, m) = g in let module M = (val m) in M.reseed ~g cs
diff --git a/tests/test_eio_entropy_collection.ml b/tests/test_eio_entropy_collection.ml
index 0f18284d..15e69af3 100644
--- a/tests/test_eio_entropy_collection.ml
+++ b/tests/test_eio_entropy_collection.ml
@@ -3,7 +3,7 @@ module Printing_rng = struct
 
   let block = 16
   let create ?time:_ () = ()
-  let generate ~g:_ _n = assert false
+  let generate_into ~g:_ _buf ~off:_ _len = assert false
   let seeded ~g:_ = true
   let pools = 1
 
diff --git a/tests/test_entropy_collection.ml b/tests/test_entropy_collection.ml
index 99e152be..4905af8a 100644
--- a/tests/test_entropy_collection.ml
+++ b/tests/test_entropy_collection.ml
@@ -7,7 +7,7 @@ module Printing_rng = struct
 
   let create ?time:_ () = ()
 
-  let generate ~g:_ _n = assert false
+  let generate_into ~g:_ _buf ~off:_ _len = assert false
 
   let reseed ~g:_ data =
     Format.printf "reseeding: %a@.%!" Cstruct.hexdump_pp (Cstruct.of_string data)
diff --git a/tests/test_entropy_collection_async.ml b/tests/test_entropy_collection_async.ml
index d7908317..f2beeccb 100644
--- a/tests/test_entropy_collection_async.ml
+++ b/tests/test_entropy_collection_async.ml
@@ -8,7 +8,7 @@ module Printing_rng = struct
 
   let create ?time:_ () = ()
 
-  let generate ~g:_ _n = assert false
+  let generate_into ~g:_ _buf ~off:_ _len = assert false
 
   let reseed ~g:_ data =
     Format.printf "reseeding: %a@.%!" Cstruct.hexdump_pp (Cstruct.of_string data)
diff --git a/tests/test_rsa.ml b/tests/test_rsa.ml
index 5ee8ab63..f512ba4d 100644
--- a/tests/test_rsa.ml
+++ b/tests/test_rsa.ml
@@ -16,11 +16,10 @@ module Null = struct
 
   let create ?time:_ () = ref ""
 
-  let generate ~g n =
+  let generate_into ~g buf ~off n =
     try
-      let (a, b) = String.sub !g 0 n, String.sub !g n (String.length !g - n) in
-      g := b;
-      a
+      Bytes.blit_string !g 0 buf off n;
+      g := String.sub !g n (String.length !g - n)
     with Invalid_argument _ -> raise Mirage_crypto_rng.Unseeded_generator
 
   let reseed ~g buf = g := !g ^ buf