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 open Lwt.Infix in
    let closed = Lwt_mvar.create_empty () in
    let rec loop ~src ~dst () =
      Lwt.pick [ Lwt_result.Infix.(Mimic.read src >|= fun v -> (v :> [ `Closed | _ Mirage_flow.or_eof ]))
               ; (Lwt_mvar.take closed >|= fun v -> Ok v) ] >>= function
      | 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)) ;
        Mimic.write dst cs >>= function
        | Ok () -> Lwt.pause () >>= 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
    Lwt.join
      [ loop ~src:src ~dst:dst ()
      ; loop ~src:dst ~dst:src () ] >>= fun () ->
    to_close src ;
    Lwt.join [ Mimic.close src; Mimic.close dst ] >>= fun () ->
    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 = Httpaf.Request.t
  type response = Httpaf.Response.t

  let version = `HTTP_1_1
  let get request name = Httpaf.Headers.get request.Httpaf.Request.headers name
  let create headers = function
    | #Httpaf.Status.t as status -> Httpaf.Response.create ~headers:(Httpaf.Headers.of_list headers) status
    | _ -> assert false
  let with_etag etag response =
    let headers = response.Httpaf.Response.headers in
    { response with Httpaf.Response.headers= Httpaf.Headers.add headers "etag" etag }
  let with_status (status : H2.Status.t) response = match status with
    | #Httpaf.Status.t as status -> { response with Httpaf.Response.status }
    | _ -> assert false
end

let transmit src dst =
  let rec on_eof () =
    Httpaf.Body.close_reader src ;
    Httpaf.Body.close_writer dst
  and on_read buf ~off ~len =
    Httpaf.Body.write_bigstring dst ~off ~len buf ;
    Httpaf.Body.schedule_read src ~on_eof ~on_read in
  Httpaf.Body.schedule_read src ~on_eof ~on_read

let connect_http_1_1 ~ctx ~authenticator ~to_close flow reqd =
  let request = Httpaf.Reqd.request reqd in
  ( match Httpaf.Headers.get request.Httpaf.Request.headers "host" with
  | Some uri ->
    let uri = "http://" ^ uri in
    let open Lwt.Infix in
    Lwt.async (fun () -> Connect.create_connection ~ctx ~authenticator uri >>= function
    | Ok dst ->
      let headers = Httpaf.Headers.of_list
        [ "connection", "close" ] in
      let response = Httpaf.Response.create ~reason:"CONNECT" ~headers `OK in
      Httpaf.Reqd.respond_with_string reqd response "" ;
      Httpaf.Body.close_reader (Httpaf.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 = Httpaf.Headers.of_list
        [ "content-length", string_of_int (String.length contents)
        ; "connection", "close"
        ; "content-type", "text/plain" ] in
      let response = Httpaf.Response.create ~reason:"CONNECT" ~headers `Bad_request in
      Httpaf.Reqd.respond_with_string reqd response contents ;
      Lwt.return_unit )
  | None ->
    let contents = "Missing Host field." in
    let headers = Httpaf.Headers.of_list
      [ "content-length", string_of_int (String.length contents)
      ; "connection", "close"
      ; "content-type", "text/plain" ] in
    let response = Httpaf.Response.create ~reason:"CONNECT" ~headers `Bad_request in
    Httpaf.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 = Httpaf.Reqd.request reqd in
  Log.debug (fun m -> m "(HTTP/1.1) request-handler: %S" request.Httpaf.Request.target) ;
  match request.Httpaf.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.Httpaf.Request.target with
    | [ ""; "" ] ->
      let headers = Httpaf.Headers.of_list
        [ "content-length", string_of_int (String.length root)
        ; "connection", "close"
        ; "content-type", "text/plain" ] in
      let response = Httpaf.Response.create ~reason:"root" ~headers `OK in
      Httpaf.Reqd.respond_with_string reqd response root
    | [ ""; "transmit" ] ->
      let content_type = Httpaf.Headers.get request.Httpaf.Request.headers "content-type" in
      let content_type = Option.value ~default:"application/octet-stream" content_type in
      let headers = Httpaf.Headers.of_list
        [ "transfer-encoding", "chunked"
        ; "content-type", content_type ] in     
      let response = Httpaf.Response.create ~reason:"transmit" ~headers `OK in
      let src = Httpaf.Reqd.request_body reqd in
      let dst = Httpaf.Reqd.respond_with_streaming reqd response in
      transmit src dst
    | [ ""; "hash" ] ->
      let respond response contents = Httpaf.Reqd.respond_with_string reqd response contents in
      hash_of_seed ~respond request
    | [ ""; "random" ] ->
      ( match Httpaf.Headers.get request.Httpaf.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 (Httpaf.Headers.get request.Httpaf.Request.headers "x-seed") random_state_of_seed in
        let headers = Httpaf.Headers.of_list
          [ "content-type", "text/plain"
          ; "transfer-encoding", "chunked" ] in
        let response = Httpaf.Response.create ~reason:"random" ~headers `OK in
        let body = Httpaf.Reqd.respond_with_streaming reqd response in
        transmit_random
          ~write_string:(fun body str -> Httpaf.Body.write_string body str)
          ~flush:Httpaf.Body.flush
          ~close_writer:Httpaf.Body.close_writer
          ?g length body
      | _ ->
        let contents = "Invalid length." in
        let headers = Httpaf.Headers.of_list
          [ "content-length", string_of_int (String.length contents)
          ; "connection", "close"
          ; "content-type", "text/plain" ] in
        let response = Httpaf.Response.create ~reason:"random" ~headers `Bad_request in
        Httpaf.Reqd.respond_with_string reqd response contents )
    | _ ->
      let contents = "Not found." in
      let headers = Httpaf.Headers.of_list
        [ "content-type", "text/plain"
        ; "connection", "close"
        ; "content-length", string_of_int (String.length contents) ] in
      let response = Httpaf.Response.create ~reason:"not-found" ~headers `Not_found in
      Httpaf.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
    let open Lwt.Infix in
    Lwt.async (fun () -> Connect.create_connection ~ctx ~authenticator uri >>= function
    | 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
        transmit_random
          ~write_string:(fun body str -> H2.Body.Writer.write_string body str)
          ~flush:H2.Body.Writer.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 _ -> Httpaf.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 _ ->
    Httpaf.Body.write_string body str ;
    Httpaf.Body.close_writer 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
