mte/unikernel/unikernel.ml

open Rresult
open Cmdliner
open Syntax

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 pgx_database =
  let doc = Arg.info ~doc:"database to use" [ "pg-database" ] in
  Arg.(value & opt string "postgres" doc)

let pgx_port =
  let doc = Arg.info ~doc:"port to use for postgresql" [ "pg-port" ] in
  Arg.(value & opt int 5432 doc)

let pgx_hostname =
  let doc = Arg.info ~doc:"host for postgres database" [ "pg-hostname" ] in
  Arg.(required & opt (some string) None doc)

let pgx_user =
  let doc = Arg.info ~doc:"postgres user" [ "pg-user" ] in
  Arg.(required & opt (some string) None doc)

let pgx_password =
  let doc = Arg.info ~doc:"postgres password" [ "pg-password" ] in
  Arg.(required & opt (some string) None doc)

type t = {
  pgx_database: string;
  pgx_port: int;
  pgx_hostname: string;
  pgx_user: string;
  pgx_password: string;
}

let pgx_setup =
  Term.(
    const (fun pgx_database pgx_port pgx_hostname pgx_user pgx_password ->
        { pgx_database; pgx_port; pgx_hostname; pgx_user; pgx_password })
    $ pgx_database
    $ pgx_port
    $ pgx_hostname
    $ pgx_user
    $ pgx_password)

module Make
    (Assets_ro : Mirage_kv.RO)
    (Certificates_ro : Mirage_kv.RO)
    (Keys_ro : Mirage_kv.RO)
    (Tcp : Tcpip.Tcp.S with type ipaddr = Ipaddr.t)
    (HTTP_server : Paf_mirage.S)
    (STACK : Tcpip.Stack.V4V6)
    (Happy_eyeballs :
      Happy_eyeballs_mirage.S
        with type stack = STACK.t
         and type flow = STACK.TCP.flow) =
struct
  module Pgx_mirage = Pgx_lwt_mirage.Make (STACK) (Happy_eyeballs)

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

  module Assets = struct
    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 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 =
        lwt_list_get_ok
        @@ Lwt_list.map_s
             (fun lang_dir ->
               let+? files = get_values ro lang_dir in
               (lang_dir, files))
             subdirs
      in
      Lwt.return
      @@
      let open Result.Syntax in
      let* l =
        list_get_ok
        @@ List.map
             (fun (lang, l) ->
               let lang = Mirage_kv.Key.basename lang in
               let* l =
                 List.map Mirage_kv.Key.basename l
                 |> List.map (fun s ->
                     match String.split_on_char '.' s with
                     | [] -> assert false
                     | [ etag; ext ] -> Ok (etag, ext)
                     | _ -> Error (`Msg (Fmt.str "invalid file name: `%s`" s)))
                 |> list_get_ok
               in
               let etag_l, ext_l = List.split l in
               Ok (lang, etag_l, ext_l))
             l
      in
      let lang_l =
        l |> List.map (fun (v, _, _) -> v) |> List.sort String.compare
      in
      let* etag =
        let etag_l = l |> List.map (fun (_, v, _) -> v) |> List.concat in
        match List.sort_uniq String.compare etag_l with
        | [ etag ] -> Ok etag
        | [] ->
            Fmt.error_msg "directory `%s` is empty"
              (Mirage_kv.Key.to_string dir)
        | l ->
            Fmt.error_msg "directory `%s` contains more than one etag: %a"
              (Mirage_kv.Key.to_string dir)
              (Fmt.list Fmt.string) l
      in
      let* ext_l =
        let* hd, tl =
          l
          |> List.map (fun (_, _, v) -> v)
          |> List.map (List.sort String.compare)
          |> function
          | [] ->
              Fmt.error_msg "directory `%s` is empty"
                (Mirage_kv.Key.to_string dir)
          | hd :: tl -> Ok (hd, tl)
        in
        match List.for_all (( = ) hd) tl with
        | false ->
            Fmt.error_msg
              "directory `%s` does not has the same set of file extensions for \
               each language"
              (Mirage_kv.Key.to_string dir)
        | true -> Ok hd
      in
      Ok (etag, lang_l, ext_l)

    let assets ro =
      let*? terms_etag, terms_lang_l, terms_ext_l = get ro "terms" in
      let*? privacy_etag, privacy_lang_l, privacy_ext_l = get ro "privacy" in
      Lwt.return
      @@
      let open Result.Syntax in
      let* lang_l =
        match terms_lang_l = privacy_lang_l with
        | false ->
            Fmt.error_msg
              "terms and privacy directories does not support the same set of \
               languages"
        | true -> Ok terms_lang_l
      in
      let ext_l =
        match terms_ext_l = privacy_ext_l with
        | false ->
            Fmt.error_msg
              "terms and privacy directories does not support the same set of \
               mimetype"
        | true -> Ok terms_ext_l
      in
      Ok (terms_etag, privacy_etag, lang_l, ext_l)
  end

  let tls certificate_ro key_ro =
    let*? keys =
      Keys_ro.list key_ro Mirage_kv.Key.empty
      |> map_err_to_string Keys_ro.pp_error
    in
    let keys = List.filter (fun (_, t) -> t = `Value) keys in
    let*? certificates =
      Certificates_ro.list certificate_ro Mirage_kv.Key.empty
      |> map_err_to_string Certificates_ro.pp_error
    in
    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
      | _ ->
          let*? data =
            Certificates_ro.get certificate_ro name
            |> map_err_to_string Certificates_ro.pp_error
          in
          let*? certificates =
            Lwt.return (X509.Certificate.decode_pem_multiple data)
          in
          let+? acc = Lwt.return acc in
          (name, certificates) :: acc
    in
    let*? certificates = Lwt_list.fold_left_s fold (Ok []) certificates in
    let fold acc (name, _) =
      match Mirage_kv.Key.basename name with
      | ".gitkeep" -> Lwt.return acc
      | _ ->
          let*? data =
            Keys_ro.get key_ro name |> map_err_to_string Keys_ro.pp_error
          in
          let*? key = Lwt.return (X509.Private_key.decode_pem data) in
          let+? acc = Lwt.return acc in
          (name, key) :: acc
    in
    let+? keys = Lwt_list.fold_left_s fold (Ok []) keys in
    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
    | [] -> `None
    | [ certchain ] -> `Single certchain
    | certchains -> `Multiple certchains

  let always x _ = x

  let http_1_1_request_handler _flow _edn reqd =
    Server.http_1_1_request_handler reqd

  let alpn_handler =
    {
      Alpn.error= Server.alpn_error_handler;
      Alpn.request=
        (fun _flow _edn reqd protocol ->
          Server.alpn_request_handler reqd protocol);
    }

  let run_with_tls ~tls http_server tls_port tcpv4v6 =
    let alpn_service = HTTP_server.alpn_service ~tls alpn_handler in
    let http_1_1_service =
      HTTP_server.http_service ~error_handler:Server.http_1_1_error_handler
        http_1_1_request_handler
    in
    let open Lwt.Syntax in
    let* server = HTTP_server.init ~port:tls_port tcpv4v6 in
    let (`Initialized th0) = Paf.serve alpn_service server in
    let (`Initialized th1) = Paf.serve http_1_1_service http_server in
    let+ (), () = Lwt.both th0 th1 in
    ()

  let run http_server =
    let http_1_1_service =
      HTTP_server.http_service ~error_handler:Server.http_1_1_error_handler
        http_1_1_request_handler
    in
    let (`Initialized th) = Paf.serve http_1_1_service http_server in
    th

  let start assets_ro certificate_ro key_ro tcpv4v6 http_server stack
      happy_eyeballs pgx_setup =
    Logs.set_reporter (Logs_fmt.reporter ());
    Logs.(set_level (Some Info));
    let open Lwt.Syntax in
    let pgx = Pgx_mirage.connect (stack, happy_eyeballs) in
    let module Pgx = (val pgx : Pgx_lwt.S) in
    let { pgx_database; pgx_port; pgx_hostname; pgx_user; pgx_password } =
      pgx_setup
    in
    Pgx.with_conn ~user:pgx_user ~host:pgx_hostname ~password:pgx_password
      ~port:pgx_port ~database:pgx_database
    @@ fun conn ->
    let* () =
      let+ alive = Pgx.alive conn in
      match alive with
      | false -> Fmt.failwith "Pgx connection failure: connection not alive"
      | true -> Logs.info (fun m -> m "Pgx connection success")
    in
    let* assets_res = Assets.assets assets_ro in
    match assets_res with
    | Error (`Msg m) -> Fmt.failwith "Assets configuration error: %s." m
    | Ok (_terms_etag, _privacy_etag, _lang_l, _ext_l) -> (
        let* tls_res = tls certificate_ro key_ro in
        match use_tls () with
        | false -> run http_server
        | true -> (
            match tls_res 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 ~tls http_server (tls_port ()) tcpv4v6)
            ))
end