mte/unikernel/duniverse/ocaml-h1/lib/server_connection.ml
2025-11-11 02:07:51 +01:00

307 lines
11 KiB
OCaml

(*----------------------------------------------------------------------------
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
module Queue = struct
include Queue
let peek_exn = peek
let peek t = if is_empty t then None else Some (peek_exn t)
end
module Reader = Parse.Reader
module Writer = Serialize.Writer
type request_handler = Reqd.t -> unit
type error =
[ `Bad_gateway | `Bad_request | `Internal_server_error | `Exn of exn ]
type error_handler =
?request:Request.t -> error -> (Headers.t -> Body.Writer.t) -> unit
type t = {
reader : Reader.request;
writer : Writer.t;
response_body_buffer : Bstr.t;
request_handler : request_handler;
error_handler : error_handler;
request_queue : Reqd.t Queue.t;
(* invariant: If [request_queue] is not empty, then the head of the queue
has already had [request_handler] called on it. *)
mutable is_errored : bool;
(* if there is a parse or connection error, we invoke the [error_handler]
and set [is_errored] to indicate we should not close the writer yet. *)
mutable wakeup_reader : Optional_thunk.t;
}
let is_closed t = Reader.is_closed t.reader && Writer.is_closed t.writer
let is_active t = not (Queue.is_empty t.request_queue)
let current_reqd_exn t = Queue.peek_exn t.request_queue
let yield_reader t k =
if is_closed t then failwith "yield_reader on closed conn"
else if Optional_thunk.is_some t.wakeup_reader then
failwith "yield_reader: only one callback can be registered at a time"
else t.wakeup_reader <- Optional_thunk.some k
let wakeup_reader t =
let f = t.wakeup_reader in
t.wakeup_reader <- Optional_thunk.none;
Optional_thunk.call_if_some f
let yield_writer t k =
if Writer.is_closed t.writer then k () else Writer.on_wakeup t.writer k
let wakeup_writer t = Writer.wakeup t.writer
let default_error_handler ?request:_ error handle =
let message =
match error with
| `Exn exn -> Printexc.to_string exn
| (#Status.client_error | #Status.server_error) as error ->
Status.to_string error
in
let body = handle Headers.empty in
Body.Writer.write_string body message;
Body.Writer.close body
let create ?(config = Config.default) ?(error_handler = default_error_handler)
request_handler =
let { Config.response_buffer_size; response_body_buffer_size; _ } = config in
let writer = Writer.create ~buffer_size:response_buffer_size () in
let request_queue = Queue.create () in
let response_body_buffer = Bstr.create response_body_buffer_size in
let handler request request_body =
let reqd =
Reqd.create error_handler request request_body writer response_body_buffer
in
Queue.push reqd request_queue
in
{
reader = Reader.request handler;
writer;
response_body_buffer;
request_handler;
error_handler;
request_queue;
is_errored = false;
wakeup_reader = Optional_thunk.none;
}
let shutdown_reader t =
if is_active t then Reqd.close_request_body (current_reqd_exn t);
Reader.force_close t.reader;
wakeup_reader t
let shutdown_writer t =
if is_active t then (
let reqd = current_reqd_exn t in
(* XXX(dpatti): I'm not sure I understand why we close the *request* body
here. Maybe we can write a test such that removing this line causes it to
fail? *)
Reqd.close_request_body reqd;
Reqd.flush_response_body reqd);
Writer.close t.writer;
wakeup_writer t
let error_code t =
if is_active t then Reqd.error_code (current_reqd_exn t) else None
let shutdown t =
shutdown_reader t;
shutdown_writer t
let set_error_and_handle ?request t error =
if is_active t then (
assert (request = None);
let reqd = current_reqd_exn t in
Reqd.report_error reqd error)
else (
t.is_errored <- true;
let status =
match (error :> [ error | Status.standard ]) with
| `Exn _ -> `Internal_server_error
| #Status.standard as status -> status
in
shutdown_reader t;
let writer = t.writer in
t.error_handler ?request error (fun headers ->
let response = Response.create ~headers status in
Writer.write_response writer response;
let encoding =
(* If we haven't parsed the request method, just use GET as a standard
placeholder. The method is only used for edge cases, like HEAD or
CONNECT. *)
let request_method =
match request with None -> `GET | Some request -> request.meth
in
match Response.body_length ~request_method response with
| (`Fixed _ | `Close_delimited) as encoding -> encoding
| `Chunked ->
(* XXX(dpatti): Because we pass the writer's faraday directly to the
new body, we don't write the chunked encoding. A client won't be
able to interpret this. *)
`Close_delimited
| `Error (`Bad_gateway | `Internal_server_error) ->
failwith
"H1.Server_connection.error_handler: invalid response body \
length"
in
Body.Writer.of_faraday (Writer.faraday writer) writer ~encoding
)
)
let report_exn t exn = set_error_and_handle t (`Exn exn)
let advance_request_queue t =
ignore (Queue.take t.request_queue);
if not (Queue.is_empty t.request_queue) then
t.request_handler (Queue.peek_exn t.request_queue)
let rec _next_read_operation t =
if not (is_active t)
then (
(* If the request queue is empty, there is no connection error, and the
reader is closed, then we can assume that no more user code will be able
to write. *)
if Reader.is_closed t.reader && not t.is_errored then shutdown_writer t;
Reader.next t.reader)
else
let reqd = current_reqd_exn t in
match Reqd.input_state reqd with
| Waiting -> _yield_reader t
| Ready -> Reader.next t.reader
| Complete -> _final_read_operation_for t reqd
| Upgraded -> `Upgrade
and _final_read_operation_for t reqd =
if not (Reqd.persistent_connection reqd) then (
shutdown_reader t;
Reader.next t.reader;
) else (
match Reqd.output_state reqd with
| Waiting | Ready -> _yield_reader t
| Upgraded ->
(* If the input state is not [Upgraded], the output state cannot be
either. *)
assert false
| Complete ->
advance_request_queue t;
_next_read_operation t;
)
and _yield_reader t =
(* XXX(dpatti): This is a way in which the reader and writer are not
parallel -- we tell the writer when it needs to yield but the reader is
always asking for more data. This is the only branch in either
operation function that does not return `(Reader|Writer).next`, which
means there are surprising states you can get into. For example, we ask
the runtime to yield but then raise when it tries to because the reader
is closed. I think this can be avoided if we allow this module to tell the
reader when it should yield/resume, then we'd just do an inlined
`Reader.next` call instead. I put this function here to describe why this
is subtle. *)
if Reader.is_closed t.reader
then Reader.next t.reader
else `Yield
;;
let next_read_operation t =
match _next_read_operation t with
| `Error (`Parse _) -> set_error_and_handle t `Bad_request; `Close
| `Error (`Bad_request request) -> set_error_and_handle ~request t `Bad_request; `Close
| (`Read | `Yield | `Close | `Upgrade) as operation -> operation
let rec read_with_more t bs ~off ~len more =
let call_handler = Queue.is_empty t.request_queue in
let consumed = Reader.read_with_more t.reader bs ~off ~len more in
if is_active t then (
let reqd = current_reqd_exn t in
if call_handler then t.request_handler reqd;
Reqd.flush_request_body reqd);
(* Keep consuming input as long as progress is made and data is
available, in case multiple requests were received at once. *)
if consumed > 0 && consumed < len then
let off = off + consumed and len = len - consumed in
consumed + read_with_more t bs ~off ~len more
else consumed
let read t bs ~off ~len = read_with_more t bs ~off ~len Incomplete
let read_eof t bs ~off ~len = read_with_more t bs ~off ~len Complete
let rec _next_write_operation t =
if not (is_active t)
then Writer.next t.writer
else (
let reqd = current_reqd_exn t in
match Reqd.output_state reqd with
| Waiting ->
(* XXX(dpatti): I don't think we should need to call this, but it is
necessary in the case of a streaming, non-chunked body so that you can
set the appropriate flag. *)
Reqd.flush_response_body reqd;
Writer.next t.writer
| Ready ->
Reqd.flush_response_body reqd;
Writer.next t.writer
| Complete -> _final_write_operation_for t reqd
| Upgraded ->
wakeup_reader t;
(* Even in the Upgrade case, we're still responsible for writing the
response header, so we might have work to do. *)
if Writer.has_pending_output t.writer
then Writer.next t.writer
else `Upgrade)
and _final_write_operation_for t reqd =
let next =
if not (Reqd.persistent_connection reqd) then (
shutdown_writer t;
Writer.next t.writer)
else
match Reqd.input_state reqd with
| Waiting -> `Yield
| Ready -> Writer.next t.writer;
| Upgraded -> `Upgrade
| Complete ->
advance_request_queue t;
_next_write_operation t
in
wakeup_reader t;
next
let next_write_operation t = _next_write_operation t
let report_write_result t result = Writer.report_result t.writer result