mte/unikernel/unikernel.ml

open Rresult
open Lwt.Infix
open Cmdliner

let port =
  let doc = Arg.info ~doc:"Port of HTTP service." [ "p"; "port" ] in
  Arg.(value & opt int 3434 doc)

let use_tls =
  let doc =
    Arg.info ~doc:"Start an HTTP server with a TLS certificate." [ "tls" ]
  in
  Mirage_runtime.register_arg Arg.(value & flag doc)

let tls_port =
  let doc = Arg.info ~doc:"Port of HTTPS service." [ "tls-port" ] in
  Mirage_runtime.register_arg Arg.(value & opt int 4343 doc)

let alpn =
  let alpns = [ "h2"; "http/1.1" ] in
  let doc =
    Printf.sprintf "Protocols handled by the HTTP server. Must be %s."
      (Arg.doc_alts alpns)
  in
  let doc = Arg.info ~doc [ "alpn" ] in
  Mirage_runtime.register_arg
    Arg.(value & opt_all (enum (List.map (fun v -> (v, v)) alpns)) alpns doc)

let ( <.> ) f g x = f (g x)
let always x _ = x

let map_err_to_string pp_err res =
  Lwt.map (R.reword_error (R.msgf "%a" pp_err)) res

module Make
    (Certificates_ro : Mirage_kv.RO)
    (Keys_ro : Mirage_kv.RO)
    (Tcp : Tcpip.Tcp.S with type ipaddr = Ipaddr.t)
    (Connect : Connect.S)
    (HTTP_server : Paf_mirage.S) =
struct
  (* WIP
     read directories *)
  module Assets = struct
    module Assets_ro = Keys_ro

    let ( let*? ) = Lwt_result.bind
    let ( let+? ) x f = Lwt_result.map f x
    let map_err_to_string = map_err_to_string Assets_ro.pp_error

    let get_subdirs ro k =
      let+? keys = Assets_ro.list ro k |> map_err_to_string in
      List.filter (fun (_, t) -> t = `Dictionary) keys |> List.map fst

    let get_values ro k =
      let+? keys = Assets_ro.list ro k |> map_err_to_string in
      List.filter (fun (_, t) -> t = `Value) keys |> List.map fst

    let find keys name =
      keys
      |> List.find_opt Mirage_kv.Key.(equal (v name))
      |> Option.to_result ~none:(`Msg (Fmt.str "missing `%s` directory" name))
      |> Lwt_result.lift

    let lwt_list_get_ok l =
      let open Lwt.Syntax in
      let+ l = l in
      let err = ref None in
      try
        l
        |> List.map (function
          | Error _e as e ->
              err := Some e;
              raise Exit
          | Ok v -> v)
        |> Result.ok
      with Exit -> ( match !err with None -> assert false | Some v -> v)

    let get ro s =
      let*? base_dir = get_subdirs ro Mirage_kv.Key.empty in
      let*? dir = find base_dir s in
      let*? subdirs = get_subdirs ro dir in
      let+? l =
        subdirs
        |> Lwt_list.map_s (fun lang_dir ->
            let+? files = get_values ro lang_dir in
            (lang_dir, files))
        |> lwt_list_get_ok
      in
      l

    let assets ro =
      let*? terms_assoc = get ro "terms" in
      let+? privacy_assoc = get ro "privacy" in
      (terms_assoc, privacy_assoc)
  end

  let tls certificate_ro key_ro =
    let ( >>= ) = Lwt_result.bind in
    Keys_ro.list key_ro Mirage_kv.Key.empty
    |> map_err_to_string Keys_ro.pp_error
    >>= fun keys ->
    let keys = List.filter (fun (_, t) -> t = `Value) keys in
    Certificates_ro.list certificate_ro Mirage_kv.Key.empty
    |> map_err_to_string Certificates_ro.pp_error
    >>= fun certificates ->
    let certificates = List.filter (fun (_, t) -> t = `Value) certificates in
    let fold acc (name, _) =
      match Mirage_kv.Key.basename name with
      | ".gitkeep" -> Lwt.return acc
      | _ ->
          Certificates_ro.get certificate_ro name
          |> map_err_to_string Certificates_ro.pp_error
          >>= (Lwt.return <.> X509.Certificate.decode_pem_multiple)
          >>= fun certificates ->
          Lwt.return acc >>= fun acc ->
          Lwt.return_ok ((name, certificates) :: acc)
    in
    Lwt_list.fold_left_s fold (Ok []) certificates >>= fun certificates ->
    let fold acc (name, _) =
      match Mirage_kv.Key.basename name with
      | ".gitkeep" -> Lwt.return acc
      | _ ->
          Keys_ro.get key_ro name
          |> map_err_to_string Keys_ro.pp_error
          >>= (Lwt.return <.> X509.Private_key.decode_pem)
          >>= fun key ->
          Lwt.return acc >>= fun acc -> Lwt.return_ok ((name, key) :: acc)
    in
    Lwt_list.fold_left_s fold (Ok []) keys >>= fun keys ->
    let tbl = Hashtbl.create 0x10 in
    List.iter
      (fun (name, certificates) ->
        match List.assoc_opt name keys with
        | Some key -> Hashtbl.add tbl name (certificates, key)
        | None -> ())
      certificates;
    match Hashtbl.fold (fun _ certchain acc -> certchain :: acc) tbl [] with
    | [] -> Lwt.return_ok `None
    | [ certchain ] -> Lwt.return_ok (`Single certchain)
    | certchains -> Lwt.return_ok (`Multiple certchains)

  let http_1_1_request_handler ~ctx ~authenticator flow _edn =
    let module R = (val Mimic.repr HTTP_server.tcp_protocol) in
    fun reqd ->
      match (H1.Reqd.request reqd).H1.Request.meth with
      | `CONNECT ->
          HTTP_server.TCP.no_close flow;
          let to_close = function
            | R.T flow -> HTTP_server.TCP.to_close flow
            | _ -> ()
          in
          Server.http_1_1_request_handler ~ctx ~authenticator ~to_close
            (R.T flow) reqd
      | _ ->
          Server.http_1_1_request_handler ~ctx ~authenticator
            ~to_close:(always ()) (R.T flow) reqd

  let alpn_handler ~ctx ~authenticator =
    let module R = (val Mimic.repr HTTP_server.tls_protocol) in
    let to_close = function
      | R.T flow -> HTTP_server.TLS.to_close flow
      | _ -> ()
    in
    {
      Alpn.error= Server.alpn_error_handler;
      Alpn.request=
        (fun flow edn reqd protocol ->
          Server.alpn_request_handler ~ctx ~authenticator ~to_close (R.T flow)
            edn reqd protocol);
    }

  let run_with_tls ~ctx ~authenticator ~tls http_server tls_port tcpv4v6 =
    let alpn_service =
      HTTP_server.alpn_service ~tls (alpn_handler ~ctx ~authenticator)
    in
    let http_1_1_service =
      HTTP_server.http_service ~error_handler:Server.http_1_1_error_handler
        (http_1_1_request_handler ~ctx ~authenticator)
    in
    HTTP_server.init ~port:tls_port tcpv4v6 >|= Paf.serve alpn_service
    >>= fun (`Initialized th0) ->
    Paf.serve http_1_1_service http_server |> fun (`Initialized th1) ->
    Lwt.both th0 th1 >>= fun ((), ()) -> Lwt.return_unit

  let run ~ctx ~authenticator http_server =
    let http_1_1_service =
      HTTP_server.http_service ~error_handler:Server.http_1_1_error_handler
        (http_1_1_request_handler ~ctx ~authenticator)
    in
    Paf.serve http_1_1_service http_server |> fun (`Initialized th) -> th

  let start certificate_ro key_ro tcpv4v6 ctx http_server =
    let open Lwt.Infix in
    let authenticator = Connect.authenticator in
    tls certificate_ro key_ro >>= fun tls ->
    match use_tls () with
    | false -> run ~ctx ~authenticator http_server
    | true -> (
        match tls with
        | Error (`Msg m) ->
            Fmt.failwith
              "A TLS server requires, at least, one certificate and one \
               private key. Received error %s."
              m
        | Ok certificates -> (
            let alpn_protocols = alpn () in
            match Tls.Config.server ~certificates ~alpn_protocols () with
            | Error (`Msg m) -> Fmt.failwith "TLS configuration error: %s." m
            | Ok tls ->
                run_with_tls ~ctx ~authenticator ~tls http_server (tls_port ())
                  tcpv4v6))
end