Add Capsule.with_password

otherlibs/stdlib_alpha/capsule.ml

@@ -85,6 +85,8 @@ module Password : sig
      [void] can't be used for function argument and return types yet. *)
   type 'k t : value mod portable many unique uncontended
 
+  type packed = P : 'k t -> packed
+
   (* Can break the soundness of the API. *)
   val unsafe_mk : 'k Name.t -> 'k t @@ portable
   val name : 'k t @ local -> 'k Name.t @@ portable
@@ -103,6 +105,8 @@ module Password : sig
 end = struct
   type 'k t = 'k Name.t
 
+  type packed = P : 'k t -> packed
+
   let unsafe_mk name = name
   let name t = t
 
@@ -121,9 +125,6 @@ end
    it never returns is also [portable] *)
 external reraise : exn -> 'a @ portable @@ portable = "%reraise"
 
-external raise_with_backtrace :
-  exn -> Printexc.raw_backtrace -> 'a @ portable @@ portable = "%raise_with_backtrace"
-
 module Data = struct
   type ('a, 'k) t : value mod portable uncontended
 
@@ -377,14 +378,31 @@ let create_with_rwlock () =
 
 exception Protected : 'k Mutex.t * (exn, 'k) Data.t -> exn
 
-(* CR-soon mslater: replace with portable stdlib *)
-let get_raw_backtrace : unit -> Printexc.raw_backtrace @@ portable =
-  O.magic O.magic Printexc.get_raw_backtrace
-
-let protect f =
-  try f () with
-  | exn ->
-    let (P mut) = create_with_mutex () in
-    raise_with_backtrace (Protected (mut, Data.unsafe_mk exn)) (get_raw_backtrace ())
-  ;;
+let protect_local f = exclave_
+  let (P name) = Name.make () in
+  let password = Password.unsafe_mk name in
+  let reraise data = reraise (Protected ({ name; mutex = M.create (); poisoned = false }, data)) in
+  try f (Password.P password) with
+  | Encapsulated (inner, data) as exn ->
+    (match Name.equality_witness name inner with
+     | Some Equal -> reraise data
+     | None -> reraise (Data.unsafe_mk exn))
+  | exn -> reraise (Data.unsafe_mk exn)
+
+let with_password_local f = exclave_
+  let (P name) = Name.make () in
+  let password = Password.unsafe_mk name in
+  try f (Password.P password) with
+  | Encapsulated (inner, data) as exn ->
+    (match Name.equality_witness name inner with
+     | Some Equal -> reraise (Data.unsafe_get data)
+     | None -> reraise exn)
+  | exn -> reraise exn
+
+module Global = struct
+  type 'a t = { global : 'a @@ global } [@@unboxed]
+end
 
+open Global
+let protect f = (protect_local (fun password -> { global = f password })).global
+let with_password f = (with_password_local (fun password -> { global = f password })).global

otherlibs/stdlib_alpha/capsule.mli

@@ -104,6 +104,11 @@ module Password : sig
      mutex. This guarantees that uncontended access to the capsule is
      only granted to a single domain at once. *)
 
+  type packed = P : 'k t -> packed
+  (** [packed] is the type of a password for some unknown capsule.
+      Unpacking one provides a ['k t] together with a fresh existential
+      type brand for ['k]. *)
+
   val name : 'k t @ local -> 'k Name.t @@ portable
   (** [name t] identifies the capsule that [t] is associated with. *)
 
@@ -416,15 +421,27 @@ exception Encapsulated : 'k Name.t * (exn, 'k) Data.t -> exn
     the data. The [Name.t] can be used to associate the [Data.t] with a
     particular [Password.t] or [Mutex.t]. *)
 
+(* CR-soon mslater: ['k Key.t] instead of ['k Mutex.t]. *)
 exception Protected : 'k Mutex.t * (exn, 'k) Data.t -> exn
 (** If a function passed to [protect] raises an exception, it is wrapped
-    in [Protected] to provide access to the capsule in which the function ran. *)
-(* CR-soon mslater: this should return a key, not a mutex. *)
-
-val protect
-  :  (unit -> 'a @ portable contended) @ local portable
-  -> 'a @ portable contended
-  @@ portable
-(** [protect f] runs [f] in a fresh capsule. If [f] returns normally, [protect]
-    merges this capsule into the caller's capsule. If [f] raises, [protect]
-    raises [Protected], giving the caller access to the encapsulated exception. *)
+    in [Protected] to avoid leaking access to the data. The [Mutex.t] can
+    be used to access the [Data.t]. *)
+
+val protect : (Password.packed @ local -> 'a) @ local portable -> 'a @@ portable
+(** [protect f] runs [f password] in a fresh capsule represented by [password].
+    If [f] returns normally, [protect] merges the capsule into the caller's capsule.
+    If [f] raises an [Encapsulated] exception in the capsule represented by [password],
+    [protect] unwraps the exception and re-raises it as [Protected].
+    If [f] raises any other exception, [protect] re-raises it as [Protected]. *)
+
+val with_password : (Password.packed @ local -> 'a) @ local portable -> 'a @@ portable
+(** [with_password f] runs [f password] in a fresh capsule represented by [password].
+    If [f] returns normally, [with_password] merges the capsule into the caller's capsule.
+    If [f] raises an [Encapsulated] exception in the capsule represented by [password],
+    [with_password] unwraps the exception and re-raises it directly. *)
+
+val protect_local : (Password.packed @ local -> 'a @ local) @ local portable -> 'a @ local @@ portable
+(** See [protect]. *)
+
+val with_password_local : (Password.packed @ local -> 'a @ local) @ local portable -> 'a @ local @@ portable
+(** See [with_password]. *)

testsuite/tests/capsule-api/data.ml

@@ -176,24 +176,55 @@ let () =
   assert (Capsule.Data.project ptr' = 111)
 ;;
 
-
 (* [protect]. *)
 exception Exn of string
 
 let () =
-  match Capsule.protect (fun () -> "ok") with
+  match Capsule.protect (fun _password -> "ok") with
   | s -> assert (s = "ok")
   | exception _ -> assert false
 ;;
 
 let () =
-  match Capsule.protect (fun () -> Exn "ok") with
+  match Capsule.protect (fun _password -> Exn "ok") with
   | Exn s -> assert (s = "ok")
   | _ -> assert false
 ;;
 
 let () =
-  match Capsule.protect (fun () -> reraise (Exn "fail")) with
+  match Capsule.protect (fun _password -> reraise (Exn "fail")) with
+  | exception (Capsule.Protected (mut, exn)) ->
+    let s = Capsule.Mutex.with_lock mut (fun password ->
+      Capsule.Data.extract password (fun exn ->
+        match exn with
+        | Exn s -> s
+        | _ -> assert false) exn) in
+    assert (s = "fail")
+  | _ -> assert false
+;;
+
+let () =
+  match Capsule.protect (fun (Capsule.Password.P password) ->
+    let data = Capsule.Data.create (fun () -> "fail") in
+    let msg = Capsule.Data.extract password (fun s : string -> s) data in
+    reraise (Exn msg))
+  with
+  | exception (Capsule.Protected (mut, exn)) ->
+    let s = Capsule.Mutex.with_lock mut (fun password ->
+      Capsule.Data.extract password (fun exn ->
+        match exn with
+        | Exn s -> s
+        | _ -> assert false) exn) in
+    assert (s = "fail")
+  | _ -> assert false
+;;
+
+let () =
+  match Capsule.protect (fun (Capsule.Password.P password) ->
+    let data = Capsule.Data.create (fun () -> "fail") in
+    let () = Capsule.Data.extract password (fun s -> reraise (Exn s)) data in
+    ())
+  with
   | exception (Capsule.Protected (mut, exn)) ->
     let s = Capsule.Mutex.with_lock mut (fun password ->
       Capsule.Data.extract password (fun exn ->
@@ -203,3 +234,42 @@ let () =
     assert (s = "fail")
   | _ -> assert false
 ;;
+
+(* [with_password]. *)
+let () =
+  match Capsule.with_password (fun _password -> "ok") with
+  | s -> assert (s = "ok")
+  | exception _ -> assert false
+;;
+
+let () =
+  match Capsule.with_password (fun _password -> Exn "ok") with
+  | Exn s -> assert (s = "ok")
+  | _ -> assert false
+;;
+
+let () =
+  match Capsule.with_password (fun _password -> reraise (Exn "fail")) with
+  | exception (Exn s) -> assert (s = "fail")
+  | _ -> assert false
+;;
+
+let () =
+  match Capsule.with_password (fun (Capsule.Password.P password) ->
+    let data = Capsule.Data.create (fun () -> "fail") in
+    let msg = Capsule.Data.extract password (fun s : string -> s) data in
+    reraise (Exn msg))
+  with
+  | exception (Exn s) -> assert (s = "fail")
+  | _ -> assert false
+;;
+
+let () =
+  match Capsule.with_password (fun (Capsule.Password.P password) ->
+    let data = Capsule.Data.create (fun () -> "fail") in
+    let () = Capsule.Data.extract password (fun s -> reraise (Exn s)) data in
+    ())
+  with
+  | exception (Exn s) -> assert (s = "fail")
+  | _ -> assert false
+;;