unikernel/duniverse/ocaml-tls/lwt/tls_lwt.ml

@@ -0,0 +1,364 @@
+open Lwt.Infix
+
+exception Tls_alert   of Tls.Packet.alert_type
+exception Tls_failure of Tls.Engine.failure
+
+(* This really belongs just about anywhere else: generic unix name resolution. *)
+let resolve host service =
+  let open Lwt_unix in
+  getprotobyname "tcp" >>= fun tcp ->
+  getaddrinfo host service [AI_PROTOCOL tcp.p_proto] >>= function
+  | []    ->
+      let msg = Printf.sprintf "no address for %s:%s" host service in
+      Lwt.reraise (Invalid_argument msg)
+  | ai::_ -> Lwt.return ai.ai_addr
+
+module Lwt_cs = struct
+
+  let naked ~name f fd cs off len =
+    f fd cs off len >>= fun res ->
+    match Lwt_unix.getsockopt_error fd with
+    | None     -> Lwt.return res
+    | Some err -> Lwt.reraise @@ Unix.Unix_error (err, name, "")
+
+  let write = naked ~name:"Tls_lwt.write" Lwt_unix.write
+  and read  = naked ~name:"Tls_lwt.read"  Lwt_unix.read
+
+  let rec write_full ?(off = 0) ?len fd buf =
+    let len = Option.value ~default:(String.length buf - off) len in
+    if len = 0 then
+      Lwt.return_unit
+    else
+      write fd (Bytes.unsafe_of_string buf) off len >>= fun written ->
+      write_full ~off:(off + written) ~len:(len - written) fd buf
+
+  let read fd buf = read fd buf 0 (Bytes.length buf)
+end
+
+module Lwt_fd = struct
+  type t = {
+    read  : bytes -> int Lwt.t ;
+    write : string -> unit Lwt.t ;
+    close : unit -> unit Lwt.t ;
+  }
+
+  let read t cs = t.read cs
+  let write t cs = t.write cs
+  let close t = t.close ()
+
+  let of_fd fd =
+    let close () =
+      (* (partially) avoid double-closes by checking if the fd has already been closed *)
+      match Lwt_unix.state fd with
+      | Lwt_unix.Closed -> Lwt.return_unit
+      | Lwt_unix.Opened | Lwt_unix.Aborted _ -> Lwt_unix.close fd
+    in
+    {
+      read  = Lwt_cs.read fd ;
+      write = Lwt_cs.write_full fd ;
+      close = close ;
+    }
+
+  let of_channels ic oc =
+    {
+      read  = (fun bs -> Lwt_io.read_into ic bs 0 (Bytes.length bs)) ;
+      write = (Lwt_io.write oc) ;
+      close = (fun () -> Lwt_io.close oc <&> Lwt_io.close ic) ;
+    }
+end
+
+module Unix = struct
+
+  type t = {
+    fd             : Lwt_fd.t ;
+    mutable state  : [ `Active of Tls.Engine.state
+                     | `Read_closed of Tls.Engine.state
+                     | `Write_closed of Tls.Engine.state
+                     | `Closed
+                     | `Error of exn ] ;
+    mutable linger : string option ;
+    recv_buf       : bytes ;
+  }
+
+  let half_close state mode =
+    match state, mode with
+    | `Active tls, `read -> `Read_closed tls
+    | `Active tls, `write -> `Write_closed tls
+    | `Active _, `read_write -> `Closed
+    | `Read_closed tls, `read -> `Read_closed tls
+    | `Read_closed _, (`write | `read_write) -> `Closed
+    | `Write_closed tls, `write -> `Write_closed tls
+    | `Write_closed _, (`read | `read_write) -> `Closed
+    | (`Closed | `Error _) as e, (`read | `write | `read_write) -> e
+
+  let inject_state tls = function
+    | `Active _ -> `Active tls
+    | `Read_closed _ -> `Read_closed tls
+    | `Write_closed _ -> `Write_closed tls
+    | (`Closed | `Error _) as e -> e
+
+  let safely th =
+    Lwt.catch
+      (fun () -> th >>= fun _ -> Lwt.return_unit)
+      (function
+        | Out_of_memory -> raise Out_of_memory
+        | _ -> Lwt.return_unit)
+
+  let (read_t, write_t) =
+    let recording_errors op t cs =
+      Lwt.catch
+        (fun () -> op t.fd cs)
+        (function
+          | Out_of_memory -> raise Out_of_memory
+          | exn -> (match t.state with
+              | `Error _ -> ()
+              | _ -> t.state <- `Error exn) ;
+            Lwt.reraise exn)
+    in
+    (recording_errors Lwt_fd.read, recording_errors Lwt_fd.write)
+
+  let when_some f = function None -> Lwt.return_unit | Some x -> f x
+
+  let rec read_react t =
+
+    let handle tls buf =
+      match Tls.Engine.handle_tls tls buf with
+      | Ok (state', eof, `Response resp, `Data data) ->
+          let state' = inject_state state' t.state in
+          let state' = Option.(value ~default:state' (map (fun `Eof -> half_close state' `read) eof)) in
+          t.state <- state' ;
+          safely (resp |> when_some (write_t t)) >|= fun () ->
+          `Ok data
+
+      | Error (fail, `Response resp) ->
+          t.state <- `Error (match fail with
+            | `Alert a -> Tls_alert a
+            | f -> Tls_failure f);
+          write_t t resp >>= fun () -> read_react t
+    in
+
+    match t.state with
+    | `Error e -> Lwt.reraise e
+    | `Closed
+    | `Read_closed _ -> Lwt.return `Eof
+    | _ ->
+        read_t t t.recv_buf >>= function
+        | 0 ->
+          t.state <- half_close t.state `read;
+          Lwt.return `Eof
+        | n ->
+          match t.state with
+          | `Error e -> Lwt.reraise e
+          | `Active tls | `Read_closed tls | `Write_closed tls ->
+            handle tls (String.sub (Bytes.unsafe_to_string t.recv_buf) 0 n)
+          | `Closed -> Lwt.return `Eof
+
+  let rec read t ?(off = 0) buf =
+    if off < 0 || off >= Bytes.length buf then
+      invalid_arg "offset must be >= 0 and < Bytes.length buf";
+
+    let writeout res =
+      let rlen = String.length res in
+      let n    = min (Bytes.length buf - off) rlen in
+      Bytes.blit_string res 0 buf off n ;
+      t.linger <-
+        (if n < rlen then Some (String.sub res n (rlen - n)) else None) ;
+      Lwt.return n in
+
+    match t.linger with
+    | Some res -> writeout res
+    | None     ->
+        read_react t >>= function
+          | `Eof           -> Lwt.return 0
+          | `Ok None       -> read t ~off buf
+          | `Ok (Some res) -> writeout res
+
+  let writev t css =
+    match t.state with
+    | `Error err  -> Lwt.reraise err
+    | `Write_closed _ | `Closed -> Lwt.reraise @@ Invalid_argument "tls: closed socket"
+    | `Active tls | `Read_closed tls ->
+        match Tls.Engine.send_application_data tls css with
+        | Some (tls, tlsdata) ->
+            ( t.state <- inject_state tls t.state ; write_t t tlsdata )
+        | None -> Lwt.reraise @@ Invalid_argument "tls: write: socket not ready"
+
+  let write t cs = writev t [cs]
+
+  (*
+   * XXX bad XXX
+   * This is a point that should particularly be protected from concurrent r/w.
+   * Doing this before a `t` is returned is safe; redoing it during rekeying is
+   * not, as the API client already sees the `t` and can mistakenly interleave
+   * writes while this is in progress.
+   * *)
+  let rec drain_handshake t =
+    let push_linger t mcs =
+      match (mcs, t.linger) with
+      | (None, _)         -> ()
+      | (scs, None)       -> t.linger <- scs
+      | (Some cs, Some l) -> t.linger <- Some (l ^ cs)
+    in
+    match t.state with
+    | `Active tls when not (Tls.Engine.handshake_in_progress tls) ->
+        Lwt.return t
+    | _ ->
+        read_react t >>= function
+          | `Eof     -> Lwt.reraise End_of_file
+          | `Ok cs   -> push_linger t cs ; drain_handshake t
+
+  let reneg ?authenticator ?acceptable_cas ?cert ?(drop = true) t =
+    match t.state with
+    | `Error err  -> Lwt.reraise err
+    | `Closed | `Read_closed _ | `Write_closed _ ->
+        Lwt.reraise @@ Invalid_argument "tls: closed socket"
+    | `Active tls ->
+        match Tls.Engine.reneg ?authenticator ?acceptable_cas ?cert tls with
+        | None -> Lwt.reraise @@ Invalid_argument "tls: can't renegotiate"
+        | Some (tls', buf) ->
+           if drop then t.linger <- None ;
+           t.state <- inject_state tls' t.state ;
+           write_t t buf >>= fun () ->
+           drain_handshake t >>= fun _ ->
+           Lwt.return_unit
+
+  let key_update ?request t =
+    match t.state with
+    | `Error err -> Lwt.reraise err
+    | `Write_closed _ | `Closed -> Lwt.reraise @@ Invalid_argument "tls: closed socket"
+    | `Active tls | `Read_closed tls ->
+      match Tls.Engine.key_update ?request tls with
+      | Error f -> Lwt.reraise @@ Invalid_argument (Format.asprintf "tls: can't update key: %a" Tls.Engine.pp_failure f)
+      | Ok (tls', buf) ->
+        t.state <- inject_state tls' t.state ;
+        write_t t buf
+
+  let shutdown t mode =
+    (match mode with
+     | `read -> Lwt.return_unit
+     | `write | `read_write ->
+       match t.state with
+       | `Active tls | `Read_closed tls ->
+         let tls', buf = Tls.Engine.send_close_notify tls in
+         t.state <- inject_state tls' (half_close t.state `write) ;
+         write_t t buf
+       | _ -> Lwt.return_unit) >>= fun () ->
+    t.state <- half_close t.state mode;
+    match t.state with
+    | `Closed | `Error _ -> safely (Lwt_fd.close t.fd)
+    | _ -> Lwt.return_unit
+
+  let close t = shutdown t `read_write
+
+  let server_of_fd config fd =
+    drain_handshake {
+      state    = `Active (Tls.Engine.server config) ;
+      fd       = fd ;
+      linger   = None ;
+      recv_buf = Bytes.create 4096
+    }
+
+  let server_of_channels config (ic, oc) =
+    server_of_fd config (Lwt_fd.of_channels ic oc)
+
+  let server_of_fd config fd =
+    server_of_fd config (Lwt_fd.of_fd fd)
+
+  let client_of_fd config ?host fd =
+    let config' = match host with
+      | None -> config
+      | Some host -> Tls.Config.peer config host
+    in
+    let (tls, init) = Tls.Engine.client config' in
+    let t = {
+      state    = `Active tls ;
+      fd       = fd ;
+      linger   = None ;
+      recv_buf = Bytes.create 4096
+    }
+    in
+    write_t t init >>= fun () ->
+    drain_handshake t
+
+  let client_of_channels config ?host (ic, oc) =
+    client_of_fd config ?host (Lwt_fd.of_channels ic oc)
+
+  let client_of_fd config ?host fd =
+    client_of_fd config ?host (Lwt_fd.of_fd fd)
+
+  let accept conf fd =
+    Lwt_unix.accept fd >>= fun (fd', addr) ->
+    Lwt.catch (fun () -> server_of_fd conf fd' >|= fun t -> (t, addr))
+      (function
+        | Out_of_memory -> raise Out_of_memory
+        | exn -> safely (Lwt_unix.close fd') >>= fun () -> Lwt.reraise exn)
+
+  let connect conf (host, port) =
+    resolve host (string_of_int port) >>= fun addr ->
+    let fd = Lwt_unix.(socket (Unix.domain_of_sockaddr addr) SOCK_STREAM 0) in
+    Lwt.catch (fun () ->
+      let host =
+        Result.to_option
+          (Result.bind (Domain_name.of_string host) Domain_name.host)
+      in
+      Lwt_unix.connect fd addr >>= fun () -> client_of_fd conf ?host fd)
+      (function
+        | Out_of_memory -> raise Out_of_memory
+        | exn -> safely (Lwt_unix.close fd) >>= fun () -> Lwt.reraise exn)
+
+  let read_bytes t bs off len =
+    let buf = Bytes.create len in
+    read t buf >|= fun n ->
+    let to_copy = min n len in
+    Lwt_bytes.blit_from_bytes buf 0 bs off to_copy;
+    to_copy
+
+  let write_bytes t bs off len =
+    let buf = Bytes.create len in
+    Lwt_bytes.blit_to_bytes bs off buf 0 len;
+    write t (Bytes.unsafe_to_string buf)
+
+  let epoch t =
+    match t.state with
+    | `Active tls | `Read_closed tls | `Write_closed tls -> Tls.Engine.epoch tls
+    | `Closed | `Error _ -> Error ()
+end
+
+type ic = Lwt_io.input_channel
+type oc = Lwt_io.output_channel
+
+let of_t ?close t =
+  let close = match close with
+    | Some f -> (fun () -> Unix.safely (f ()))
+    | None   -> (fun () -> Unix.(safely (close t)))
+  in
+  (Lwt_io.make ~close ~mode:Lwt_io.Input (Unix.read_bytes t)),
+  (Lwt_io.make ~close ~mode:Lwt_io.Output @@
+    fun a b c -> Unix.write_bytes t a b c >>= fun () -> Lwt.return c)
+
+let accept_ext conf fd =
+  Unix.accept conf fd >|= fun (t, peer) -> (of_t t, peer)
+
+and connect_ext conf addr =
+  Unix.connect conf addr >|= of_t
+
+let accept certificate fd =
+  match Tls.Config.server ~certificates:certificate () with
+  | Ok config -> accept_ext config fd >|= fun w -> Ok w
+  | Error _ as e -> Lwt.return e
+
+and connect authenticator addr =
+  match Tls.Config.client ~authenticator () with
+  | Ok config -> connect_ext config addr >|= fun w -> Ok w
+  | Error _ as e -> Lwt.return e
+
+(* Boot the entropy loop at module init time. *)
+let () = Mirage_crypto_rng_unix.use_default ()
+
+let () =
+  Printexc.register_printer (function
+      | Tls_alert typ ->
+        Some ("TLS alert from peer: " ^ Tls.Packet.alert_type_to_string typ)
+      | Tls_failure f ->
+        Some ("TLS failure: " ^ Tls.Engine.string_of_failure f)
+      | _ -> None)