open Syntax

let src = Logs.Src.create "server"

module Log = (val Logs.src_log src : Logs.LOG)

let is_digit = function '0' .. '9' -> true | _ -> false

let root =
  {plain|Hello fellows!

This unikernel is a simple example of a website that supports HTTP/1.1, H2 and
the TLS security layer. This is an example of how to make a website with
MirageOS.

The website has several endpoints:
- `http{,s}://$hostname/` the page you are about to read.
- `http{,s}://$hostname/transmit` a page that copies what is sent (no matter
  how big it is)
- `http{,s}://$hostname/hash` calculates a hash from a random content
  generated by a seed `x-seed`
- `http{,s}://$hostname/random` a page that waits in the request for a size
  (`x-length`) and generates a random (optionally seeded by `x-seed`) content
  encoded in base64

These different endpoints allow performance tests of what MirageOS can offer.
Have fun, and hack it!
|plain}

let random_cstruct ~g buf len =
  for i = 0 to Cstruct.length buf - 1 do
    let v = Random.State.bits g land 0xff in
    Cstruct.set_uint8 buf i v
  done;
  Cstruct.sub buf 0 len

let random_state_of_seed str =
  match Base64.decode str with
  | Error _ -> None
  | Ok seed ->
      let res = Array.make (String.length seed / 2) 0 in
      for i = 0 to (String.length seed / 2) - 1 do
        res.(i) <-
          (Char.code seed.[i * 2] lsl 8) lor Char.code seed.[(i * 2) + 1]
      done;
      Some (Random.State.make res)

let transmit_random ~write_string ~flush ~close_writer
    ?(g = Random.State.make_self_init ()) length body =
  let tmp = Bytes.create (0x1000 * 4) in
  let rnd = Cstruct.create 0x1000 in
  let ctx = Digestif.SHA256.empty in

  let encoder = Base64_rfc2045.encoder `Manual in

  let rec go ctx length cs = function
    | `Ok when Cstruct.length cs = 0 -> encode ctx length
    | `Ok ->
        go ctx length (Cstruct.shift cs 1)
          (Base64_rfc2045.encode encoder (`Char (Cstruct.get_char cs 0)))
    | `Partial ->
        let len = Bytes.length tmp - Base64_rfc2045.dst_rem encoder in
        write_string body (Bytes.sub_string tmp 0 len);
        Base64_rfc2045.dst encoder tmp 0 (Bytes.length tmp);
        let next () = go ctx length cs (Base64_rfc2045.encode encoder `Await) in
        flush body next
  and encode ctx = function
    | 0L -> finalize ctx (Base64_rfc2045.encode encoder `End)
    | length ->
        let len = min (Int64.of_int (Cstruct.length rnd)) length in
        let ({ Cstruct.buffer; off; len= buffer_len } as rnd) =
          random_cstruct ~g rnd (Int64.to_int len)
        in
        Log.debug (fun m ->
            m "@[<hov>%a@]"
              (Hxd_string.pp Hxd.default)
              (Bigstringaf.substring buffer ~off ~len:buffer_len));
        let ctx =
          Digestif.SHA256.feed_bigstring ctx ~off ~len:buffer_len buffer
        in
        go ctx (Int64.sub length len) rnd `Ok
  and finalize ctx = function
    | `Partial ->
        let len = Bytes.length tmp - Base64_rfc2045.dst_rem encoder in
        write_string body (Bytes.sub_string tmp 0 len);
        Base64_rfc2045.dst encoder tmp 0 (Bytes.length tmp);
        let next () = finalize ctx (Base64_rfc2045.encode encoder `Await) in
        flush body next
    | `Ok ->
        Logs.debug (fun m ->
            m "%a" Digestif.SHA256.pp (Digestif.SHA256.get ctx));
        close_writer body
  in

  Base64_rfc2045.dst encoder tmp 0 (Bytes.length tmp);
  encode ctx length

let hash ~version ~create ?seed length =
  let rec go g rnd ctx = function
    | 0L ->
        let headers =
          [
            ("content-length", string_of_int (Digestif.SHA256.digest_size * 2));
            ("content-type", "text/plain");
          ]
        in
        (* XXX(dinosaure): Connection: Close header is only true for http/1.1.
           For h2, curl complains and abruptely closes the connection. The most
           important seems Content-Length which closes properly the connection
           on both sides. *)
        let headers =
          match version with
          | `HTTP_1_1 -> ("connection", "close") :: headers
          | `HTTP_2_0 -> headers
        in
        let response = create headers `OK in
        (response, Digestif.SHA256.(to_hex (get ctx)))
    | length ->
        let len = min (Int64.of_int (Cstruct.length rnd)) length in
        let { Cstruct.buffer; off; len= buffer_len } =
          random_cstruct ~g rnd (Int64.to_int len)
        in
        Log.debug (fun m ->
            m "@[<hov>%a@]"
              (Hxd_string.pp Hxd.default)
              (Bigstringaf.substring buffer ~off ~len:buffer_len));
        let ctx =
          Digestif.SHA256.feed_bigstring ctx ~off ~len:buffer_len buffer
        in
        go g rnd ctx (Int64.sub length len)
  in
  match
    ( Option.bind seed random_state_of_seed,
      Option.bind length Int64.of_string_opt )
  with
  | Some g, Some length ->
      go g (Cstruct.create 0x1000) Digestif.SHA256.empty length
  | _ ->
      let contents = "Invalid seed." in
      let headers =
        [
          ("content-length", string_of_int (String.length contents));
          ("content-type", "text/plain");
        ]
      in
      let headers =
        match version with
        | `HTTP_1_1 -> ("connection", "close") :: headers
        | `HTTP_2_0 -> headers
      in
      let response = create headers `Bad_request in
      (response, contents)

module Cache = Ephemeron.K1.Make (struct
  type t = string option * string option

  let equal = ( = )
  let hash = Hashtbl.hash
end)

module type S = sig
  type response
  type request

  val version : [ `HTTP_1_1 | `HTTP_2_0 ]
  val create : (string * string) list -> H2.Status.t -> response
  val with_etag : string -> response -> response
  val with_status : H2.Status.t -> response -> response
  val get : request -> string -> string option
end

let hash_of_seed : type response request.
    (module S with type response = response and type request = request) ->
    respond:(response -> string -> unit) ->
    request ->
    unit =
 fun (module S) ->
  ();
  let tbl = Cache.create 0x100 in
  fun ~respond request ->
    let seed = S.get request "x-seed" and length = S.get request "x-length" in

    match
      (Cache.find_opt tbl (seed, length), S.get request "if-none-match")
    with
    | None, _ | Some _, None ->
        let response, contents =
          hash ~version:S.version ~create:S.create ?seed length
        in
        Cache.add tbl (seed, length) (response, contents);
        respond response contents
    | Some (response, contents), Some etags ->
        let hash = Digestif.SHA256.digest_string contents in
        let hash = Digestif.SHA256.to_hex hash in

        if List.mem hash (String.split_on_char ',' etags) then
          let response =
            S.with_etag hash (S.with_status `Not_modified response)
          in
          respond response String.empty
        else respond response contents

let transmit_over_http : to_close:_ -> Mimic.flow -> Mimic.flow -> unit Lwt.t =
 fun ~to_close src dst ->
  let closed = Lwt_mvar.create_empty () in
  let rec loop ~src ~dst () =
    let* res =
      Lwt.pick
        [
          (let+? v = Mimic.read src in
           (v :> [ `Closed | _ Mirage_flow.or_eof ]));
          (let+ v = Lwt_mvar.take closed in
           Ok v);
        ]
    in
    match res with
    | Error err ->
        Log.err (fun m ->
            m "Got an error while we reading the source (CONNECT): %a."
              Mimic.pp_error err);
        if Lwt_mvar.is_empty closed then Lwt_mvar.put closed `Closed
        else Lwt.return_unit
    | Ok `Closed -> Lwt.return_unit
    | Ok `Eof ->
        if Lwt_mvar.is_empty closed then Lwt_mvar.put closed `Closed
        else Lwt.return_unit
    | Ok (`Data cs) -> (
        Log.debug (fun m -> m "Transfer over HTTP:");
        Log.debug (fun m ->
            m "@[<hov>%a@]." (Hxd_string.pp Hxd.default) (Cstruct.to_string cs));
        let* res = Mimic.write dst cs in
        match res with
        | Ok () ->
            let* () = Lwt.pause () in
            loop ~src ~dst ()
        | Error err ->
            Log.err (fun m ->
                m
                  "Got an error while we writing into the destination \
                   (CONNECT): %a."
                  Mimic.pp_write_error err);
            if Lwt_mvar.is_empty closed then Lwt_mvar.put closed `Closed
            else Lwt.return_unit)
  in
  let* () = Lwt.join [ loop ~src ~dst (); loop ~src:dst ~dst:src () ] in
  to_close src;
  let* () = Lwt.join [ Mimic.close src; Mimic.close dst ] in
  Log.debug (fun m -> m "Connection closed properly on both side.");
  Lwt.return_unit

(***** HTTP/1.1 *****)

module S_HTTP_1_1 = struct
  type request = H1.Request.t
  type response = H1.Response.t

  let version = `HTTP_1_1
  let get request name = H1.Headers.get request.H1.Request.headers name

  let create headers = function
    | #H1.Status.t as status ->
        H1.Response.create ~headers:(H1.Headers.of_list headers) status
    | _ -> assert false

  let with_etag etag response =
    let headers = response.H1.Response.headers in
    { response with H1.Response.headers= H1.Headers.add headers "etag" etag }

  let with_status (status : H2.Status.t) response =
    match status with
    | #H1.Status.t as status -> { response with H1.Response.status }
    | _ -> assert false
end

let transmit src dst =
  let rec on_eof () = H1.Body.Reader.close src; H1.Body.Writer.close dst
  and on_read buf ~off ~len =
    H1.Body.Writer.write_bigstring dst ~off ~len buf;
    H1.Body.Reader.schedule_read src ~on_eof ~on_read
  in
  H1.Body.Reader.schedule_read src ~on_eof ~on_read

let connect_http_1_1 ~ctx ~authenticator ~to_close flow reqd =
  let request = H1.Reqd.request reqd in
  match H1.Headers.get request.H1.Request.headers "host" with
  | Some uri ->
      let uri = "http://" ^ uri in
      Lwt.async (fun () ->
          let* res = Connect.create_connection ~ctx ~authenticator uri in
          match res with
          | Ok dst ->
              let headers = H1.Headers.of_list [ ("connection", "close") ] in
              let response =
                H1.Response.create ~reason:"CONNECT" ~headers `OK
              in
              H1.Reqd.respond_with_string reqd response "";
              H1.Body.Reader.close (H1.Reqd.request_body reqd);
              transmit_over_http ~to_close flow dst
          | Error err ->
              Log.err (fun m ->
                  m "Got an error while connection to %S: %a." uri
                    Mimic.pp_error err);
              let contents = Fmt.str "Invalid URI: %S" uri in
              let headers =
                H1.Headers.of_list
                  [
                    ("content-length", string_of_int (String.length contents));
                    ("connection", "close"); ("content-type", "text/plain");
                  ]
              in
              let response =
                H1.Response.create ~reason:"CONNECT" ~headers `Bad_request
              in
              H1.Reqd.respond_with_string reqd response contents;
              Lwt.return_unit)
  | None ->
      let contents = "Missing Host field." in
      let headers =
        H1.Headers.of_list
          [
            ("content-length", string_of_int (String.length contents));
            ("connection", "close"); ("content-type", "text/plain");
          ]
      in
      let response =
        H1.Response.create ~reason:"CONNECT" ~headers `Bad_request
      in
      H1.Reqd.respond_with_string reqd response contents

let http_1_1_request_handler ~ctx ~authenticator ~to_close =
  let hash_of_seed = hash_of_seed (module S_HTTP_1_1) in
  fun flow reqd ->
    let request = H1.Reqd.request reqd in
    Log.debug (fun m ->
        m "(HTTP/1.1) request-handler: %S" request.H1.Request.target);
    match request.H1.Request.meth with
    | `CONNECT ->
        Log.debug (fun m -> m "Start to transmit data over HTTP/1.1.");
        connect_http_1_1 ~ctx ~authenticator ~to_close flow reqd
    | _meth -> (
        match String.split_on_char '/' request.H1.Request.target with
        | [ ""; "" ] ->
            let headers =
              H1.Headers.of_list
                [
                  ("content-length", string_of_int (String.length root));
                  ("connection", "close"); ("content-type", "text/plain");
                ]
            in
            let response = H1.Response.create ~reason:"root" ~headers `OK in
            H1.Reqd.respond_with_string reqd response root
        | [ ""; "transmit" ] ->
            let content_type =
              H1.Headers.get request.H1.Request.headers "content-type"
            in
            let content_type =
              Option.value ~default:"application/octet-stream" content_type
            in
            let headers =
              H1.Headers.of_list
                [
                  ("transfer-encoding", "chunked");
                  ("content-type", content_type);
                ]
            in
            let response = H1.Response.create ~reason:"transmit" ~headers `OK in
            let src = H1.Reqd.request_body reqd in
            let dst = H1.Reqd.respond_with_streaming reqd response in
            transmit src dst
        | [ ""; "hash" ] ->
            let respond response contents =
              H1.Reqd.respond_with_string reqd response contents
            in
            hash_of_seed ~respond request
        | [ ""; "random" ] -> (
            match H1.Headers.get request.H1.Request.headers "x-length" with
            | Some v when String.for_all is_digit v ->
                let length = Int64.of_string v in
                let g =
                  Option.bind
                    (H1.Headers.get request.H1.Request.headers "x-seed")
                    random_state_of_seed
                in
                let headers =
                  H1.Headers.of_list
                    [
                      ("content-type", "text/plain");
                      ("transfer-encoding", "chunked");
                    ]
                in
                let response =
                  H1.Response.create ~reason:"random" ~headers `OK
                in
                let body = H1.Reqd.respond_with_streaming reqd response in
                transmit_random
                  ~write_string:(fun body str ->
                    H1.Body.Writer.write_string body str)
                  ~flush:H1.Body.Writer.flush ~close_writer:H1.Body.Writer.close
                  ?g length body
            | _ ->
                let contents = "Invalid length." in
                let headers =
                  H1.Headers.of_list
                    [
                      ("content-length", string_of_int (String.length contents));
                      ("connection", "close"); ("content-type", "text/plain");
                    ]
                in
                let response =
                  H1.Response.create ~reason:"random" ~headers `Bad_request
                in
                H1.Reqd.respond_with_string reqd response contents)
        | _ ->
            let contents = "Not found." in
            let headers =
              H1.Headers.of_list
                [
                  ("content-type", "text/plain"); ("connection", "close");
                  ("content-length", string_of_int (String.length contents));
                ]
            in
            let response =
              H1.Response.create ~reason:"not-found" ~headers `Not_found
            in
            H1.Reqd.respond_with_string reqd response contents)

(***** H2 *****)

module S_HTTP_2_0 = struct
  type request = H2.Request.t
  type response = H2.Response.t

  let version = `HTTP_2_0
  let get request name = H2.Headers.get request.H2.Request.headers name

  let create headers status =
    H2.Response.create ~headers:(H2.Headers.of_list headers) status

  let with_etag etag response =
    let headers = response.H2.Response.headers in
    { response with H2.Response.headers= H2.Headers.add headers "etag" etag }

  let with_status status response = { response with H2.Response.status }
end

let transmit src dst =
  let rec on_eof () = H2.Body.Reader.close src; H2.Body.Writer.close dst
  and on_read buf ~off ~len =
    H2.Body.Writer.write_bigstring dst ~off ~len buf;
    H2.Body.Reader.schedule_read src ~on_eof ~on_read
  in
  H2.Body.Reader.schedule_read src ~on_eof ~on_read

let connect_http_2_0 ~ctx ~authenticator ~to_close flow reqd =
  let request = H2.Reqd.request reqd in
  match H2.Headers.get request.H2.Request.headers "host" with
  | Some uri ->
      let uri = "http://" ^ uri in
      Lwt.async (fun () ->
          let* res = Connect.create_connection ~ctx ~authenticator uri in
          match res with
          | Ok dst ->
              let response = H2.Response.create `OK in
              H2.Reqd.respond_with_string reqd response "";
              H2.Body.Reader.close (H2.Reqd.request_body reqd);
              transmit_over_http ~to_close flow dst
          | Error err ->
              Log.err (fun m ->
                  m "Got an error while connection to %S: %a." uri
                    Mimic.pp_error err);
              let contents = Fmt.str "Invalid URI: %S" uri in
              let headers =
                H2.Headers.of_list
                  [
                    ("content-length", string_of_int (String.length contents));
                    ("content-type", "text/plain");
                  ]
              in
              let response = H2.Response.create ~headers `Bad_request in
              H2.Reqd.respond_with_string reqd response contents;
              Lwt.return_unit)
  | None ->
      let contents = "Missing Host field." in
      let headers =
        H2.Headers.of_list
          [
            ("content-length", string_of_int (String.length contents));
            ("content-type", "text/plain");
          ]
      in
      let response = H2.Response.create ~headers `Bad_request in
      H2.Reqd.respond_with_string reqd response contents

let http_2_0_request_handler ~ctx ~authenticator ~to_close =
  let hash_of_seed = hash_of_seed (module S_HTTP_2_0) in
  fun flow reqd ->
    let request = H2.Reqd.request reqd in
    Log.debug (fun m -> m "(H2) request-handler: %S" request.H2.Request.target);
    match request.H2.Request.meth with
    | `CONNECT ->
        Log.debug (fun m -> m "Start to transmit data over H2.");
        connect_http_2_0 ~ctx ~authenticator ~to_close flow reqd
    | _meth -> (
        match String.split_on_char '/' request.H2.Request.target with
        | [ ""; "" ] ->
            let headers =
              H2.Headers.of_list
                [
                  ("content-length", string_of_int (String.length root));
                  ("content-type", "text/plain");
                ]
            in
            let response = H2.Response.create ~headers `OK in
            H2.Reqd.respond_with_string reqd response root
        | [ ""; "transmit" ] ->
            let content_type =
              H2.Headers.get request.H2.Request.headers "content-type"
            in
            let content_type =
              Option.value ~default:"application/octet-stream" content_type
            in
            let headers =
              H2.Headers.of_list [ ("content-type", content_type) ]
            in
            let response = H2.Response.create ~headers `OK in
            let src = H2.Reqd.request_body reqd in
            let dst = H2.Reqd.respond_with_streaming reqd response in
            transmit src dst
        | [ ""; "hash" ] ->
            let respond response contents =
              H2.Reqd.respond_with_string reqd response contents
            in
            hash_of_seed ~respond request
        | [ ""; "random" ] -> (
            match H2.Headers.get request.H2.Request.headers "x-length" with
            | Some v when String.for_all is_digit v ->
                let length = Int64.of_string v in
                let g =
                  Option.bind
                    (H2.Headers.get request.H2.Request.headers "x-seed")
                    random_state_of_seed
                in
                let headers =
                  H2.Headers.of_list [ ("content-type", "text/plain") ]
                in
                let response = H2.Response.create ~headers `OK in
                let body = H2.Reqd.respond_with_streaming reqd response in
                let flush body next =
                  H2.Body.Writer.flush body (fun _reason -> next ())
                in
                transmit_random
                  ~write_string:(fun body str ->
                    H2.Body.Writer.write_string body str)
                  ~flush ~close_writer:H2.Body.Writer.close ?g length body
            | _ ->
                let contents = "Invalid length." in
                let headers =
                  H2.Headers.of_list
                    [
                      ("content-length", string_of_int (String.length contents));
                      ("content-type", "text/plain");
                    ]
                in
                let response = H2.Response.create ~headers `Bad_request in
                H2.Reqd.respond_with_string reqd response contents)
        | _ ->
            let contents = "Not found." in
            let headers =
              H2.Headers.of_list
                [
                  ("content-type", "text/plain");
                  ("content-length", string_of_int (String.length contents));
                ]
            in
            let response = H2.Response.create ~headers `Not_found in
            H2.Reqd.respond_with_string reqd response contents)

let alpn_request_handler : type reqd headers request response ro wo.
    ctx:_ ->
    authenticator:_ ->
    to_close:_ ->
    ?shutdown:_ ->
    _ ->
    _ ->
    reqd ->
    (reqd, headers, request, response, ro, wo) Alpn.protocol ->
    unit =
 fun ~ctx ~authenticator ~to_close ?shutdown:_ flow _edn reqd -> function
  | Alpn.HTTP_1_1 _ ->
      http_1_1_request_handler ~ctx ~authenticator ~to_close flow reqd
  | Alpn.H2 _ ->
      http_2_0_request_handler ~ctx ~authenticator ~to_close flow reqd

let headers_of_list : type reqd headers request response ro wo.
    (reqd, headers, request, response, ro, wo) Alpn.protocol ->
    (string * string) list ->
    headers =
 fun protocol lst ->
  match protocol with
  | Alpn.HTTP_1_1 _ -> H1.Headers.of_list lst
  | Alpn.H2 _ -> H2.Headers.of_list lst

let respond_with_string : type reqd headers request response ro wo.
    (reqd, headers, request, response, ro, wo) Alpn.protocol ->
    headers:headers ->
    respond:(headers -> wo) ->
    string ->
    unit =
 fun protocol ~headers ~respond str ->
  let body = respond headers in
  match protocol with
  | Alpn.HTTP_1_1 _ ->
      H1.Body.Writer.write_string body str;
      H1.Body.Writer.close body
  | Alpn.H2 _ ->
      H2.Body.Writer.write_string body str;
      H2.Body.Writer.close body

let alpn_error_handler : type reqd headers request response ro wo.
    _ ->
    (reqd, headers, request, response, ro, wo) Alpn.protocol ->
    ?request:_ ->
    _ ->
    (headers -> wo) ->
    unit =
 fun _edn protocol ?request:_ error respond ->
  let contents =
    match error with
    | `Bad_gateway -> {plain|Bad gateway.|plain}
    | `Bad_request -> {plain|Bad request.|plain}
    | `Exn (Paf.Flow err) | `Exn (Paf.Flow_write err) ->
        Fmt.str {plain|I/O error: %s.|plain} err
    | `Exn exn ->
        Fmt.str {plain|Unknown error: %S.|plain} (Printexc.to_string exn)
    | `Internal_server_error -> {plain|Internal server error.|plain}
  in
  let headers =
    headers_of_list protocol
      [
        ("content-type", "text/plain");
        ("content-length", string_of_int (String.length contents));
      ]
  in
  respond_with_string protocol ~respond ~headers contents

let http_1_1_error_handler edn ?request error respond =
  alpn_error_handler edn ?request Alpn.http_1_1
    (error :> Alpn.server_error)
    respond

let http_2_0_error_handler edn ?request error respond =
  alpn_error_handler edn ?request Alpn.h2 (error :> Alpn.server_error) respond