Source file restorer.ml

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
open Lwt.Infix
open Capnp_rpc_lwt

module Core_types = Private.Capnp_core.Core_types
module Log = Capnp_rpc.Debug.Log

module Id = struct
  type t = string

  let generate () =
    Mirage_crypto_rng.generate 20

  let public x = x

  let derived ~secret name =
    Digestif.SHA256.hmac_string ~key:secret name
    |> Digestif.SHA256.to_raw_string

  let digest alg t =
    let alg = (alg :> Digestif.hash') in
    let module H = (val Digestif.module_of_hash' alg : Digestif.S) in
    H.digest_string t |> H.to_raw_string

  let to_string x = x

  let equal = ( = )
  let pp = Fmt.string
end

type resolution = (Core_types.cap, Capnp_rpc.Exception.t) result

module type LOADER = sig
  type t
  val hash : t -> Auth.hash
  val make_sturdy : t -> Id.t -> Uri.t
  val load : t -> 'a Sturdy_ref.t -> string -> resolution Lwt.t
end

type t = Id.t -> resolution Lwt.t

let grant x : resolution = Ok (Cast.cap_to_raw x)
let reject ex = Error ex

let unknown_service_id = reject (Capnp_rpc.Exception.v "Unknown persistent service ID")

let fn (r:t) =
  fun k object_id ->
    Lwt.async (fun () ->
        Lwt.try_bind
          (fun () -> r object_id)
          (fun r -> k r; Lwt.return_unit)
          (fun ex ->
             Log.err (fun f -> f "Uncaught exception restoring object: %a" Fmt.exn ex);
             k (reject (Capnp_rpc.Exception.v "Internal error restoring object"));
             Lwt.return_unit
          )
      )

let restore (f:t) x = f x |> Lwt_result.map Cast.cap_of_raw

let none : t = fun _ ->
  Lwt.return @@ Error (Capnp_rpc.Exception.v "This vat has no restorer")

let single id cap =
  let cap = Cast.cap_to_raw cap in
  (* Hash the ID to prevent timing attacks. *)
  let id = Digestif.SHA256.digest_string id |> Digestif.SHA256.to_raw_string in
  fun requested_id ->
    let requested_id = Digestif.SHA256.digest_string requested_id |> Digestif.SHA256.to_raw_string in
    if String.equal id requested_id then (
      Core_types.inc_ref cap;
      Lwt.return (Ok cap)
    ) else Lwt.return unknown_service_id

module Table = struct
  type digest = string

  type entry =
    | Cached of resolution Lwt.t
    | Manual of Core_types.cap          (* We hold a ref on the cap *)

  type t = {
    hash : Digestif.hash';
    cache : (digest, entry) Hashtbl.t;
    load : Id.t -> digest -> resolution Lwt.t;
    make_sturdy : Id.t -> Uri.t;
  }

  (* [cache] contains promises or capabilities with positive ref-counts. *)

  let create make_sturdy =
    let hash = `SHA256 in
    let cache = Hashtbl.create 53 in
    let load _ _ = Lwt.return unknown_service_id in
    { hash; cache; load; make_sturdy }

  let hash t id =
    Id.digest t.hash id

  let resolve t id =
    let digest = hash t id in
    match Hashtbl.find t.cache digest with
    | Manual cap ->
      Core_types.inc_ref cap;
      Lwt.return @@ Ok cap
    | Cached res ->
      begin res >>= function
        | Error _ as e -> Lwt.return e
        | Ok cap ->
          Core_types.inc_ref cap;
          Lwt.pause () >|= fun () ->
          Ok cap
      end
    | exception Not_found ->
      let cap = t.load id digest in
      Hashtbl.add t.cache digest (Cached cap);
      Lwt.try_bind
        (fun () -> cap)
        (fun result ->
           begin match result with
             | Error _ -> Hashtbl.remove t.cache digest
             | Ok cap -> cap#when_released (fun () -> Hashtbl.remove t.cache digest)
           end;
           (* Ensure all [inc_ref]s are done before handing over to the user. *)
           Lwt.pause () >|= fun () ->
           result
        )
        (fun ex ->
           Hashtbl.remove t.cache digest;
           Lwt.fail ex
        )

  let of_loader (type l) (module L : LOADER with type t = l) loader =
    let hash = (L.hash loader :> Digestif.hash') in
    let cache = Hashtbl.create 53 in
    let rec load id digest =
      let sr : Private.Capnp_core.sturdy_ref = object
        method connect = resolve t id
        method to_uri_with_secrets = L.make_sturdy loader id
      end in
      L.load loader (Cast.sturdy_of_raw sr) digest
    and t = { hash; cache; load; make_sturdy = L.make_sturdy loader } in
    t

  let add t id cap =
    let cap = Cast.cap_to_raw cap in
    let id = hash t id in
    assert (not (Hashtbl.mem t.cache id));
    Hashtbl.add t.cache id (Manual cap)

  let sturdy_ref t id =
    Cast.sturdy_of_raw @@ object
      method connect = resolve t id
      method to_uri_with_secrets = t.make_sturdy id
    end

  let release = function
    | Manual cap -> Core_types.dec_ref cap;
    | Cached _ -> ()

  let remove t id =
    let id = hash t id in
    match Hashtbl.find t.cache id with
    | exception Not_found -> failwith "Service ID not in restorer table"
    | value ->
      release value;
      Hashtbl.remove t.cache id

  let clear t =
    Hashtbl.iter (fun _ v -> release v) t.cache;
    Hashtbl.clear t.cache
end

let of_table = Table.resolve