unikernel/duniverse/ocaml-h1/lib/reqd.ml

@@ -0,0 +1,294 @@
+(*----------------------------------------------------------------------------
+    Copyright (c) 2017 Inhabited Type LLC.
+    Copyright (c) 2025 Robur Cooperative
+
+    All rights reserved.
+
+    Redistribution and use in source and binary forms, with or without
+    modification, are permitted provided that the following conditions
+    are met:
+
+    1. Redistributions of source code must retain the above copyright
+       notice, this list of conditions and the following disclaimer.
+
+    2. Redistributions in binary form must reproduce the above copyright
+       notice, this list of conditions and the following disclaimer in the
+       documentation and/or other materials provided with the distribution.
+
+    3. Neither the name of the author nor the names of his contributors
+       may be used to endorse or promote products derived from this software
+       without specific prior written permission.
+
+    THIS SOFTWARE IS PROVIDED BY THE CONTRIBUTORS ``AS IS'' AND ANY EXPRESS
+    OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+    WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+    DISCLAIMED.  IN NO EVENT SHALL THE AUTHORS OR CONTRIBUTORS BE LIABLE FOR
+    ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+    DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
+    OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
+    HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
+    STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
+    ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+    POSSIBILITY OF SUCH DAMAGE.
+  ----------------------------------------------------------------------------*)
+
+open Httpun_types
+
+type error =
+  [ `Bad_request | `Bad_gateway | `Internal_server_error | `Exn of exn ]
+
+module Response_state = struct
+  type t =
+    | Waiting
+    | Upgrade   of Response.t
+    | Fixed     of Response.t
+    | Streaming of Response.t * Body.Writer.t
+end
+
+module Input_state = struct
+  type t =
+    | Waiting
+    | Ready
+    | Complete
+    | Upgraded
+end
+
+module Output_state = struct
+  type t =
+    | Waiting
+    | Ready
+    | Complete
+    | Upgraded
+end
+
+type error_handler =
+  ?request:Request.t -> error -> (Headers.t -> Body.Writer.t) -> unit
+
+module Writer = Serialize.Writer
+
+(* XXX(seliopou): The current design assumes that a new [Reqd.t] will be
+ * allocated for each new request/response on a connection. This is wasteful,
+ * as it creates garbage on persistent connections. A better approach would be
+ * to allocate a single [Reqd.t] per connection and reuse it across
+ * request/responses. This would allow a single [Faraday.t] to be allocated for
+ * the body and reused. The [response_state] type could then be inlined into
+ * the [Reqd.t] record, with dummy values occuping the fields for [response].
+ * Something like this:
+ *
+ * {[
+ *   type 'handle t =
+ *     { mutable request        : Request.t
+ *     ; mutable request_body   : Response.Body.Reader.t
+ *     ; mutable response       : Response.t (* Starts off as a dummy value,
+ *                                            * using [(==)] to identify it when
+ *                                            * necessary *)
+ *     ; mutable response_body  : Response.Body.Writer.t
+ *     ; mutable persistent     : bool
+ *     ; mutable response_state : [ `Waiting | `Started | `Streaming ]
+ *     }
+ *  ]}
+ *
+ * *)
+type t = {
+  request : Request.t;
+  request_body : Body.Reader.t;
+  writer : Writer.t;
+  response_body_buffer : Bstr.t;
+  error_handler : error_handler;
+  mutable persistent : bool;
+  mutable response_state : Response_state.t;
+  mutable error_code : [ `Ok | error ];
+}
+
+let create error_handler request request_body writer response_body_buffer =
+  {
+    request;
+    request_body;
+    writer;
+    response_body_buffer;
+    error_handler;
+    persistent = Request.persistent_connection request;
+    response_state = Waiting;
+    error_code = `Ok;
+  }
+
+let request { request; _ } = request
+let request_body { request_body; _ } = request_body
+
+let response { response_state; _ } =
+  match response_state with
+  | Waiting -> None
+  | Streaming (response, _)
+  | Upgrade response
+  | Fixed response -> Some response
+
+let response_exn { response_state; _ } =
+  match response_state with
+  | Waiting -> failwith "H1.Reqd.response_exn: response has not started"
+  | Streaming (response, _)
+  | Upgrade response
+  | Fixed response -> response
+
+let respond_with_string t response str =
+  if t.error_code <> `Ok then
+    failwith
+      "H1.Reqd.respond_with_string: invalid state, currently handling error";
+  match t.response_state with
+  | Waiting ->
+      (* XXX(seliopou): check response body length *)
+      Writer.write_response t.writer response;
+      Writer.write_string t.writer str;
+      if t.persistent then
+        t.persistent <- Response.persistent_connection response;
+      t.response_state <- Fixed response;
+      Writer.wakeup t.writer
+  | Streaming _ ->
+    failwith "H1.Reqd.respond_with_string: response already started"
+  | Upgrade _
+  | Fixed _ ->
+    failwith "H1.Reqd.respond_with_string: response already complete"
+
+let respond_with_bigstring t response (bstr : Bstr.t) =
+  if t.error_code <> `Ok then
+    failwith
+      "H1.Reqd.respond_with_bigstring: invalid state, currently handling error";
+  match t.response_state with
+  | Waiting ->
+      (* XXX(seliopou): check response body length *)
+      Writer.write_response t.writer response;
+      Writer.schedule_bigstring t.writer bstr;
+      if t.persistent then
+        t.persistent <- Response.persistent_connection response;
+      t.response_state <- Fixed response;
+      Writer.wakeup t.writer
+  | Streaming _ ->
+    failwith "H1.Reqd.respond_with_bigstring: response already started"
+  | Upgrade _
+  | Fixed _ ->
+      failwith "H1.Reqd.respond_with_bigstring: response already complete"
+
+let unsafe_respond_with_streaming ~flush_headers_immediately t response =
+  match t.response_state with
+  | Waiting ->
+      let encoding =
+        match Response.body_length ~request_method:t.request.meth response with
+        | (`Fixed _ | `Close_delimited | `Chunked) as encoding -> encoding
+        | `Error (`Bad_gateway | `Internal_server_error) ->
+            failwith
+              "H1.Reqd.respond_with_streaming: invalid response body length"
+      in
+      let response_body =
+        Body.Writer.create t.response_body_buffer t.writer ~encoding
+      in
+      Writer.write_response t.writer response;
+      if t.persistent then
+        t.persistent <- Response.persistent_connection response;
+      t.response_state <- Streaming (response, response_body);
+      if flush_headers_immediately then Writer.wakeup t.writer;
+      response_body
+  | Streaming _ ->
+    failwith "H1.Reqd.respond_with_streaming: response already started"
+  | Upgrade _
+  | Fixed _ ->
+      failwith "H1.Reqd.respond_with_streaming: response already complete"
+
+let respond_with_streaming ?(flush_headers_immediately = false) t response =
+  if t.error_code <> `Ok then
+    failwith
+      "H1.Reqd.respond_with_streaming: invalid state, currently handling error";
+  unsafe_respond_with_streaming ~flush_headers_immediately t response
+
+let respond_with_upgrade ?reason t headers =
+  match t.response_state with
+  | Waiting ->
+    if not (Request.is_upgrade t.request) then
+      failwith "H1.Reqd.respond_with_upgrade: request was not an upgrade request"
+    else (
+      let response = Response.create ?reason ~headers `Switching_protocols in
+      t.response_state <- Upgrade response;
+      (* The parser ensures it only passes empty bodies in the case of an
+         upgrade request *)
+      assert (Body.Reader.is_closed t.request_body);
+      Writer.write_response t.writer response;
+      Writer.wakeup t.writer);
+  | Streaming _ ->
+    failwith "H1.Reqd.respond_with_upgrade: response already started"
+  | Upgrade _
+  | Fixed _ ->
+    failwith "H1.Reqd.respond_with_upgrade: response already complete"
+
+let report_error t error =
+  t.persistent <- false;
+  Body.Reader.close t.request_body;
+  match (t.response_state, t.error_code) with
+  | Waiting, `Ok ->
+      t.error_code <- (error :> [ `Ok | error ]);
+      let status =
+        match (error :> [ error | Status.standard ]) with
+        | `Exn _ -> `Internal_server_error
+        | #Status.standard as status -> status
+      in
+      t.error_handler ~request:t.request error (fun headers ->
+          unsafe_respond_with_streaming ~flush_headers_immediately:true t
+            (Response.create ~headers status))
+  | Waiting, `Exn _ ->
+      (* XXX(seliopou): Decide what to do in this unlikely case. There is an
+       * outstanding call to the [error_handler], but an intervening exception
+       * has been reported as well. *)
+      failwith "H1.Reqd.report_exn: NYI"
+  | Streaming (_response, response_body), `Ok -> Body.Writer.close response_body
+  | Streaming (_response, response_body), `Exn _ ->
+    Body.Writer.close response_body;
+    Writer.close_and_drain t.writer
+  | (Fixed _ | Streaming _ | Waiting | Upgrade _) , _ ->
+    (* XXX(seliopou): Once additional logging support is added, log the error
+     * in case it is not spurious. *)
+    ()
+
+let report_exn t exn = report_error t (`Exn exn)
+
+let try_with t f : (unit, exn) result =
+  try
+    f ();
+    Ok ()
+  with exn ->
+    report_exn t exn;
+    Error exn
+
+(* Private API, not exposed to the user through h1.mli *)
+
+let close_request_body { request_body; _ } = Body.Reader.close request_body
+
+let error_code t =
+  match t.error_code with #error as error -> Some error | `Ok -> None
+
+let persistent_connection t = t.persistent
+
+let input_state t : Input_state.t =
+  match t.response_state with
+  | Upgrade _ -> Upgraded
+  | Waiting when Request.is_upgrade t.request -> Waiting
+  | Waiting | Fixed _ | Streaming _ ->
+    if Body.Reader.is_closed t.request_body
+    then Complete
+    else Ready
+;;
+
+let output_state t : Output_state.t =
+  match t.response_state with
+  | Upgrade _ -> Upgraded
+  | Fixed _ -> Complete
+  | Streaming (_, response_body) ->
+      if Body.Writer.has_pending_output response_body then Ready
+      else if Body.Writer.is_closed response_body then Complete
+      else Waiting
+  | Waiting -> Waiting
+
+let flush_request_body t =
+  if Body.Reader.has_pending_output t.request_body then
+    try Body.Reader.execute_read t.request_body with exn -> report_exn t exn
+
+let flush_response_body t =
+  match t.response_state with
+  | Streaming (_, response_body) -> Body.Writer.transfer_to_writer response_body
+  | _ -> ()