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