rng.ml

type source = int * string

exception Unseeded_generator

exception No_default_generator

let setup_rng =
  "\nPlease setup your default random number generator. On Unix, the best \
   path is to call [Mirage_crypto_rng_unix.use_default ()].\
   \nBut you can use Fortuna (or any other RNG) and setup the seeding \
   (done by default in MirageOS): \
   \n\
   \nTo initialize the RNG with a default generator, and set up entropy \
   collection and periodic reseeding as a background task, do the \
   following:\
   \n  If you are using MirageOS, use the random device in config.ml: \
   `let main = Mirage.main \"Unikernel.Main\" (random @-> job)`, \
   and `let () = register \"my_unikernel\" [main $ default_random]`. \
   \n  If you are using Lwt, execute \
   `Mirage_crypto_rng_lwt.initialize (module Mirage_crypto_rng.Fortuna)` \
   at startup. \
   \n  If you are using Async, execute \
   `Mirage_crypto_rng_async.initialize (module Mirage_crypto_rng.Fortuna)` \
   at startup. \
   \n  If you are using Eio, execute in one of the fibers \
   `Mirage_crypto_rng_eio.run (module Fortuna) env` (`env` from `Eio_main.run`).
   \n  Otherwise, there is no periodic reseeding. For an initial seed from \
   getrandom(), execute \
   `Mirage_crypto_rng_unix.initialize (module Mirage_crypto_rng.Fortuna)`. \
   You can use `Mirage_crypto_rng.accumulate` and `Mirage_crypto_rng.reseed` \
   to reseed the RNG manually."

let () = Printexc.register_printer (function
    | Unseeded_generator ->
      Some ("The RNG has not been seeded." ^ setup_rng)
    | No_default_generator ->
      Some ("The default generator is not yet initialized. " ^ setup_rng)
    | _ -> None)

module type Generator = sig
  type g
  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
end

type 'a generator = (module Generator with type g = 'a)
type g = Generator : ('a * bool * 'a generator) -> g

let create (type a) ?g ?seed ?(strict=false) ?time (m : a generator) =
  let module M = (val m) in
  let g = Option.value g ~default:(M.create ?time ()) in
  Option.iter (M.reseed ~g) seed;
  Generator (g, strict, m)

let _default_generator = Atomic.make None

let set_default_generator g = Atomic.set _default_generator (Some g)

let unset_default_generator () = Atomic.set _default_generator None

let default_generator () =
  match Atomic.get _default_generator with
  | None -> raise No_default_generator
  | Some g -> g

let get = function Some g -> g | None -> default_generator ()

let generate_into ?(g = default_generator ()) b ?(off = 0) n =
  let Generator (g, _, m) = g in
  let module M = (val m) in
  if Bytes.length b - off < n then
    invalid_arg "buffer too short";
  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

let accumulate g source =
  let Generator (g, _, m) = get g in
  let module M = (val m) in
  M.accumulate ~g source

let seeded g =
  let Generator (g, _, m) = get g in let module M = (val m) in M.seeded ~g

let block g =
  let Generator (_, _, m) = get g in let module M = (val m) in M.block

let pools g =
  let Generator (_, _, m) = get g in let module M = (val m) in M.pools

let strict g =
  let Generator (_, s, _) = get g in s