unikernel/duniverse/ocaml-tls/async/io.ml

@@ -0,0 +1,200 @@
+open! Core
+open! Async
+include Io_intf
+
+module Tls_error = struct
+  module Alert = struct
+    type t = Tls.Packet.alert_type
+    let sexp_of_t a =
+      Sexplib.Sexp.Atom (Tls.Packet.alert_type_to_string a)
+  end
+  module Fail = struct
+    type t = Tls.Engine.failure
+    let sexp_of_t a =
+      Sexplib.Sexp.Atom (Tls.Engine.string_of_failure a)
+  end
+  type t =
+    | Tls_alert of Alert.t
+    (** [Tls_alert] exception received from the other endpoint *)
+    | Tls_failure of Fail.t
+    (** [Tls_failure] exception while processing incoming data *)
+    | Connection_closed
+    | Connection_not_ready
+    | Unexpected_eof
+    | Unable_to_renegotiate
+    | Unable_to_update_key
+  [@@deriving sexp_of]
+end
+
+module Make (Fd : Fd) : S with module Fd := Fd = struct
+  open Deferred.Or_error.Let_syntax
+
+  module State = struct
+    type t =
+      | Active of Tls.Engine.state
+      | Eof
+      | Error of Tls_error.t
+  end
+
+  type t =
+    { fd : Fd.t
+    ; mutable state : State.t
+    ; mutable linger : string option
+    ; recv_buf : bytes
+    }
+
+  let tls_error = Fn.compose Deferred.Or_error.error_s Tls_error.sexp_of_t
+
+  let rec read_react t =
+    let handle tls buf =
+      match Tls.Engine.handle_tls tls buf with
+      | Ok (state, eof, `Response resp, `Data data) ->
+        t.state
+        <- (match eof with
+          | None -> Active state
+          | Some `Eof -> Eof);
+        let%map () =
+          match resp with
+          | None -> return ()
+          | Some resp -> Fd.write_full t.fd resp
+        in
+        `Ok data
+      | Error (alert, `Response resp) ->
+        t.state <- Error (match alert with `Alert a -> Tls_alert a | f -> Tls_failure f);
+        let%bind () = Fd.write_full t.fd resp in
+        read_react t
+    in
+    match t.state with
+    | Error e -> tls_error e
+    | Eof -> return `Eof
+    | Active _ ->
+      let%bind n = Fd.read t.fd t.recv_buf in
+      (match t.state, n with
+       | Active _, `Eof ->
+         t.state <- Eof;
+         return `Eof
+       | Active tls, `Ok n -> handle tls (Stdlib.Bytes.sub_string t.recv_buf 0 n)
+       | Error e, _ -> tls_error e
+       | Eof, _ -> return `Eof)
+  ;;
+
+  let rec read t buf =
+    let writeout res =
+      let rlen = String.length res in
+      let n = min (Bytes.length buf) rlen in
+      Stdlib.Bytes.blit_string res 0 buf 0 n;
+      t.linger <- (if n < rlen then Some (Stdlib.String.sub res n (rlen - n)) else None);
+      return n
+    in
+    match t.linger with
+    | Some res -> writeout res
+    | None ->
+      (match%bind read_react t with
+       | `Eof -> return 0
+       | `Ok None -> read t buf
+       | `Ok (Some res) -> writeout res)
+  ;;
+
+  let writev t css =
+    match t.state with
+    | Error err -> tls_error err
+    | Eof -> tls_error Connection_closed
+    | Active tls ->
+      (match Tls.Engine.send_application_data tls css with
+       | Some (tls, tlsdata) ->
+         t.state <- Active tls;
+         Fd.write_full t.fd tlsdata
+       | None -> tls_error Connection_not_ready)
+  ;;
+
+  (*
+   * 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) -> return t
+    | _ ->
+      (match%bind read_react t with
+       | `Eof -> tls_error Unexpected_eof
+       | `Ok cs ->
+         push_linger t cs;
+         drain_handshake t)
+  ;;
+
+  let reneg ?authenticator ?acceptable_cas ?cert ?(drop = true) t =
+    match t.state with
+    | Error err -> tls_error err
+    | Eof -> tls_error Connection_closed
+    | Active tls ->
+      (match Tls.Engine.reneg ?authenticator ?acceptable_cas ?cert tls with
+       | None -> tls_error Unable_to_renegotiate
+       | Some (tls', buf) ->
+         if drop then t.linger <- None;
+         t.state <- Active tls';
+         let%bind () = Fd.write_full t.fd buf in
+         let%bind _ = drain_handshake t in
+         return ())
+  ;;
+
+  let key_update ?request t =
+    match t.state with
+    | Error err -> tls_error err
+    | Eof -> tls_error Connection_closed
+    | Active tls ->
+      (match Tls.Engine.key_update ?request tls with
+       | Error _ -> tls_error Unable_to_update_key
+       | Ok (tls', buf) ->
+         t.state <- Active tls';
+         Fd.write_full t.fd buf)
+  ;;
+
+  let close_tls t =
+    match t.state with
+    | Active tls ->
+      let _, buf = Tls.Engine.send_close_notify tls in
+      t.state <- Eof;
+      Fd.write_full t.fd buf
+    | _ -> return ()
+  ;;
+
+  let server_of_fd config fd =
+    drain_handshake
+      { state = Active (Tls.Engine.server config)
+      ; fd
+      ; linger = None
+      ; recv_buf = Bytes.create 4096
+      }
+  ;;
+
+  let client_of_fd config ?host fd =
+    let config' =
+      match host with
+      | None -> config
+      | Some host -> Tls.Config.peer config host
+    in
+    let t = { state = Eof; fd; linger = None; recv_buf = Bytes.create 4096 } in
+    let tls, init = Tls.Engine.client config' in
+    let t = { t with state = Active tls } in
+    let%bind () = Fd.write_full t.fd init in
+    drain_handshake t
+  ;;
+
+  let epoch t =
+    match t.state with
+    | Active tls -> (match Tls.Engine.epoch tls with
+       | Ok _ as o -> o
+       | Error () -> Or_error.error_string "no TLS state available yet")
+    | Eof -> Or_error.error_string "TLS state is end of file"
+    | Error _ -> Or_error.error_string "TLS state is error"
+  ;;
+end