|
| 1 | +open Mirage_crypto_rng |
| 2 | + |
| 3 | +module Pfortuna = Pfortuna |
| 4 | + |
| 5 | +type _ Effect.t += Spawn : (unit -> unit) -> unit Effect.t |
| 6 | +external reraise : exn -> 'a = "%reraise" |
| 7 | + |
| 8 | +let periodic fn delta = |
| 9 | + let rec one () = |
| 10 | + fn (); |
| 11 | + Miou_unix.sleep (Duration.to_f delta); |
| 12 | + one () in |
| 13 | + Effect.perform (Spawn one) |
| 14 | + |
| 15 | +let getrandom delta source = |
| 16 | + let fn () = |
| 17 | + let per_pool = 8 in |
| 18 | + let size = per_pool * pools None in |
| 19 | + let random = Mirage_crypto_rng_unix.getrandom size in |
| 20 | + let idx = ref 0 in |
| 21 | + let fn () = incr idx; String.sub random (per_pool * (pred !idx)) per_pool in |
| 22 | + Entropy.feed_pools None source fn in |
| 23 | + periodic fn delta |
| 24 | + |
| 25 | +let getrandom_init i = |
| 26 | + let data = Mirage_crypto_rng_unix.getrandom 128 in |
| 27 | + Entropy.header i data |
| 28 | + |
| 29 | +let rdrand delta = |
| 30 | + match Entropy.cpu_rng with |
| 31 | + | Error `Not_supported -> () |
| 32 | + | Ok cpu_rng -> periodic (cpu_rng None) delta |
| 33 | + |
| 34 | +let running = Atomic.make false |
| 35 | + |
| 36 | +let switch fn = |
| 37 | + let orphans = Miou.orphans () in |
| 38 | + let open Effect.Deep in |
| 39 | + let retc = Fun.id in |
| 40 | + let exnc = reraise in |
| 41 | + let effc : type c. c Effect.t -> ((c, 'r) continuation -> 'r) option |
| 42 | + = function |
| 43 | + | Spawn fn -> |
| 44 | + ignore (Miou.async ~orphans fn); |
| 45 | + Some (fun k -> continue k ()) |
| 46 | + | _ -> None in |
| 47 | + match_with fn orphans { retc; exnc; effc } |
| 48 | + |
| 49 | +let default_generator_already_set = |
| 50 | + "Mirage_crypto_rng.default_generator has already \ |
| 51 | + been set (but not via Mirage_crypto_rng_miou). Please check \ |
| 52 | + that this is intentional" |
| 53 | + |
| 54 | +let miou_generator_already_launched = |
| 55 | + "Mirage_crypto_rng_miou.initialize has already been launched \ |
| 56 | + and a task is already seeding the RNG." |
| 57 | + |
| 58 | +type rng = unit Miou.t |
| 59 | + |
| 60 | +let rec compare_and_set ?(backoff= Miou_backoff.default) t a b = |
| 61 | + if Atomic.compare_and_set t a b = false |
| 62 | + then compare_and_set ~backoff:(Miou_backoff.once backoff) t a b |
| 63 | + |
| 64 | +let rec clean_up sleep orphans = match Miou.care orphans with |
| 65 | + | Some None | None -> Miou_unix.sleep (Duration.to_f sleep); clean_up sleep orphans |
| 66 | + | Some (Some prm) -> Miou.await_exn prm; clean_up sleep orphans |
| 67 | + |
| 68 | +let call_if_domain_available fn = |
| 69 | + let available = Miou.Domain.available () in |
| 70 | + let current = (Stdlib.Domain.self () :> int) in |
| 71 | + if current = 0 && available > 0 |
| 72 | + || current <> 0 && available > 1 |
| 73 | + then Miou.call fn |
| 74 | + else Miou.async fn |
| 75 | + |
| 76 | +let initialize (type a) ?g ?(sleep= Duration.of_sec 1) (rng : a generator) = |
| 77 | + if Atomic.compare_and_set running false true |
| 78 | + then begin |
| 79 | + let seed = |
| 80 | + let init = Entropy.[ bootstrap; whirlwind_bootstrap; bootstrap; getrandom_init ] in |
| 81 | + List.mapi (fun i fn -> fn i) init |> String.concat "" in |
| 82 | + let () = |
| 83 | + try let _ = default_generator () in |
| 84 | + Logs.warn (fun m -> m "%s" default_generator_already_set) |
| 85 | + with No_default_generator -> () in |
| 86 | + let rng = create ?g ~seed ~time:Mtime_clock.elapsed_ns rng in |
| 87 | + set_default_generator rng; |
| 88 | + call_if_domain_available @@ fun () -> switch @@ fun orphans -> |
| 89 | + rdrand sleep; |
| 90 | + let source = Entropy.register_source "getrandom" in |
| 91 | + getrandom (Int64.mul sleep 10L) source; |
| 92 | + clean_up sleep orphans |
| 93 | + end else invalid_arg miou_generator_already_launched |
| 94 | + |
| 95 | +let kill prm = |
| 96 | + Miou.cancel prm; |
| 97 | + compare_and_set running true false; |
| 98 | + unset_default_generator () |
0 commit comments