ORAM.ml

(** Oblivious random-access memory *)

(* Copyright Xavier Leroy.
   License: LGPL 2.1 or later with OCaml LGPL Linking Exception *)

(* Goldreich-Ostrovsky "square root" ORAM *)

let rng = Cryptokit.Random.(pseudo_rng (string secure_rng 32))

(* Randomized, symmetric encryption.  Cleartext is (int, string) pairs.  *)

type key = string

let keygen () : key = Cryptokit.Random.string rng 16

let encrypt (k: key) ((i, v) : int * string) : string =
  let b = Bytes.create (8 + String.length v) in
  Bytes.set_int64_be b 0 (Int64.of_int i);
  Bytes.blit_string v 0 b 8 (String.length v);
  let iv = Cryptokit.Random.string rng 8 in
  iv ^
  Cryptokit.(transform_string
    (Cipher.chacha20 ~iv k Cipher.Encrypt)
    (Bytes.to_string b))

let decrypt (k: key) (c: string) : int * string =
  let iv = String.sub c 0 8 in
  let m =
    Cryptokit.(transform_string
      (Cipher.chacha20 ~iv k Cipher.Decrypt)
      (String.sub c 8 (String.length c - 8))) in
  (Int64.to_int (String.get_int64_be m 0), String.sub m 8 (String.length m - 8))

(* Pseudo-random functions from integers to bitstrings *)

type prf = Cryptokit.Block.block_cipher

let new_prf () : prf =
  let k = Cryptokit.Random.string rng 16 in
  new Cryptokit.Block.aes_encrypt k

let prf (h: prf) (i: int) : string =
  let b = Bytes.make 16 '\000' in
  Bytes.set_int64_be b 8 (Int64.of_int i);
  h#transform b 0 b 0;
  Bytes.to_string b

(* The physical RAM is modeled as a pair of read, write functions.
   RAM contents are strings (ciphertexts). *)

type phys_ram = {
  size: int;
  read: int -> string;
  write: int -> string -> unit
  }

let make_ram (size: int) : phys_ram =
  let a = Array.make size "" in
  { size;
    read = (fun i -> a.(i));
    write = (fun i s -> a.(i) <- s)
  }

(* Initial encryption of the RAM *)

let setup (ram: phys_ram) (k: key) =
  for i = 0 to ram.size - 1 do
    ram.write i (encrypt k (i, ram.read i))
  done

(* Sorting network, using Batcher's odd-even mergesort *)

let batcher_network f n =
  let p = ref 1 in
  while !p < n do
    let k = ref !p in
    while !k >= 1 do
      let j = ref (!k mod !p) in
      while !j <= n - 1 - !k do
        for i = 0 to Int.min (!k - 1) (n - !j - !k - 1) do
          let x = i + !j
          and y = i + !j + !k in
          if x / (2 * !p) = y / (2 * !p) then f x y
        done;
        j := !j + 2 * !k
      done;
      k := !k / 2
    done;
    p := !p * 2
  done

(* Sorting the physical RAM by increasing PRF values *)

let compare_swap (ram: phys_ram) (k: key) (h: prf) (u: int) (v: int) =
  let (i, vi) = decrypt k (ram.read u) in
  let (j, vj) = decrypt k (ram.read v) in
  if prf h i <= prf h j then begin
    ram.write u (encrypt k (i, vi));
    ram.write v (encrypt k (j, vj))
  end else begin
    ram.write u (encrypt k (j, vj));
    ram.write v (encrypt k (i, vi))
  end

let sort (ram: phys_ram) (k: key) (h: prf) =
  batcher_network
    (compare_swap ram k h)
    ram.size

(* Oblivious lookup of cell number i in the physical RAM *)

let lookup (ram: phys_ram) (k: key) (h: prf) (i: int) : int * string =
  let rec bsearch lo hi =
    if lo >= hi then raise Not_found;
    let p = lo + (hi - lo) / 2 in
    let (j, v) = decrypt k (ram.read p) in
    if j = i then (p, v) else
    if prf h i <= prf h j
    then bsearch lo p
    else bsearch (p + 1) hi
  in bsearch 0 ram.size

(* The state of the oblivious RAM *)

type oblivious_ram = {
    key: key;
    mutable hash: prf;
    data_size: int;
    cache_size: int;
    cache: (int, string) Hashtbl.t;
    mutable trail: (int * int) list;
    mutable pos: int
  }

let make_oram (ram: phys_ram) data_size cache_size =
  assert (ram.size = data_size + cache_size);
  let key = keygen() in
  let hash = new_prf() in
  setup ram key;
  sort ram key hash;
  { key; hash; data_size; cache_size;
    cache = Hashtbl.create (2 * cache_size);
    trail = [];
    pos = 0 }

let perform_write (ram: phys_ram) (o: oblivious_ram) (p, i) =
  let v =
    match Hashtbl.find_opt o.cache i with
    | Some v -> v
    | None -> "" in
  ram.write p (encrypt o.key (i, v))

let rmw (ram: phys_ram) (o: oblivious_ram) (i: int) (f: string -> string) : string =
  if o.pos >= o.cache_size then begin
    List.iter (perform_write ram o) (List.rev o.trail);
    o.trail <- [];
    o.hash <- new_prf();
    sort ram o.key o.hash;
    o.pos <- 0;
    Hashtbl.clear o.cache
  end;
  match Hashtbl.find_opt o.cache i with
  | None ->
      let (p, v0) = lookup ram o.key o.hash i in
      let v1 = f v0 in
      Hashtbl.replace o.cache i v1;
      o.trail <- (p, i) :: o.trail;
      o.pos <- o.pos + 1;
      v0
  | Some v0 ->
      let j = o.data_size + o.pos in
      let (p, _) = lookup ram o.key o.hash j in
      let v1 = f v0 in
      Hashtbl.replace o.cache i v1;
      o.trail <- (p, i) :: o.trail;
      o.pos <- o.pos + 1;
      v0

let read (ram: phys_ram) (o: oblivious_ram) (i: int) : string =
  rmw ram o i (fun v -> v)

let write (ram: phys_ram) (o: oblivious_ram) (i: int) (v: string) : unit =
  ignore (rmw ram o i (fun _ -> v))