Handle Content-Length at the HTTP/1 adapter

src/dream.ml

@@ -8,7 +8,6 @@
 module Catch = Dream__server.Catch
 module Cipher = Dream__cipher.Cipher
 module Cookie = Dream__server.Cookie
-module Content_length = Dream__server.Content_length
 module Csrf = Dream__server.Csrf
 module Echo = Dream__server.Echo
 module Error_handler = Dream__http.Error_handler
@@ -354,7 +353,6 @@ let catch = Catch.catch
 let run = Http.run
 let serve = Http.serve
 let lowercase_headers = Lowercase_headers.lowercase_headers
-let content_length = Content_length.content_length
 let with_site_prefix = Site_prefix.with_site_prefix
 
 
@@ -391,8 +389,7 @@ let request = Helpers.request_with_body
    corresponding tests. *)
 let test ?(prefix = "") handler request =
   let app =
-    Content_length.content_length
-    @@ Site_prefix.with_site_prefix prefix
+    Site_prefix.with_site_prefix prefix
     @@ handler
   in
 

src/dream.mli

@@ -2297,7 +2297,6 @@ val serve :
     {[
       Dream.run ~builtins:false
       @@ Dream.lowercase_headers
-      @@ Dream.content_length
       @@ Dream.catch ~error_handler
       @@ my_app
     ]}
@@ -2308,19 +2307,6 @@ val serve :
 val lowercase_headers : middleware
 (** Lowercases response headers for HTTP/2 requests. *)
 
-val content_length : middleware
-(** If the request has {!Dream.version} [(1, _)], then...
-
-    - if the response does not have [Content-Length:] and the body is a string,
-      sets [Content-Length:] to the string's length, or
-    - if the response does not have [Transfer-Encoding:] and the body is a
-      stream, sets [Transfer-Encoding: chunked].
-
-    This is built in because an application cannot be expected to decide
-    including these headers in the face of transparent HTTP/2 upgrades. The
-    headers are necessary in HTTP/1, and forbidden or redundant and difficult to
-    use in HTTP/2. *)
-
 val with_site_prefix : string -> middleware
 (** Removes the given prefix from the path in each request, and adds it to the
     request prefix. Responds with [502 Bad Gateway] if the path does not have

src/http/http.ml

@@ -6,7 +6,6 @@
 
 
 module Catch = Dream__server.Catch
-module Content_length = Dream__server.Content_length
 module Helpers = Dream__server.Helpers
 module Log = Dream__server.Log
 module Lowercase_headers = Dream__server.Lowercase_headers
@@ -123,18 +122,13 @@ let wrap_handler
            2. Upon failure to establish a WebSocket, the function is called to
               transmit the resulting error response. *)
         let forward_response response =
+          Message.set_content_length_headers response;
+
           let headers =
             Httpaf.Headers.of_list (Message.all_headers response) in
 
-          (* let version =
-            match Dream.version_override response with
-            | None -> None
-            | Some (major, minor) -> Some Httpaf.Version.{major; minor}
-          in *)
           let status =
             to_httpaf_status (Message.status response) in
-          (* let reason =
-            Dream.reason_override response in *)
 
           let httpaf_response =
             Httpaf.Response.create ~headers status in
@@ -246,6 +240,7 @@ let wrap_handler_h2
         (* Extract the Dream response's headers. *)
 
         let forward_response response =
+          Message.drop_content_length_headers response;
           let headers =
             H2.Headers.of_list (Message.all_headers response) in
           let status =
@@ -386,7 +381,6 @@ let ocaml_tls = {
 let built_in_middleware error_handler =
   Message.pipeline [
     Lowercase_headers.lowercase_headers;
-    Content_length.content_length;
     Catch.catch (Error_handler.app error_handler);
   ]
 

src/pure/message.ml

@@ -213,6 +213,28 @@ let set_body message body =
   | Request -> message.server_stream <- Stream.string body
   | Response -> message.client_stream <- Stream.string body
 
+let set_content_length_headers message =
+  if has_header message "Content-Length" then
+    ()
+  else
+    if has_header message "Transfer-Encoding" then
+      ()
+    else
+      match message.body with
+      | None ->
+        add_header message "Transfer-Encoding" "chunked"
+      | Some body_promise ->
+        match Lwt.poll body_promise with
+        | None ->
+          add_header message "Transfer-Encoding" "chunked"
+        | Some body ->
+          let length = string_of_int (String.length body) in
+          add_header message "Content-Length" length
+
+let drop_content_length_headers message =
+  drop_header message "Content-Length";
+  drop_header message "Transfer-Encoding"
+
 
 
 (* Streams *)

src/pure/message.mli

@@ -66,6 +66,8 @@ val sort_headers : (string * string) list -> (string * string) list
 
 val body : 'a message -> string promise
 val set_body : 'a message -> string -> unit
+val set_content_length_headers : 'a message -> unit
+val drop_content_length_headers : 'a message -> unit
 
 
 

src/server/helpers.ml

@@ -60,25 +60,25 @@ let request_with_body ?method_ ?target ?version ?headers body =
 
 
 
-let html ?status ?code ?headers body =
+let response_with_body ?status ?code ?headers body =
   let response =
-    Message.response ?status ?code ?headers (Stream.string body) Stream.null in
+    Message.response ?status ?code ?headers Stream.null Stream.null in
+  Message.set_body response body;
+  response
+
+let respond ?status ?code ?headers body =
+  Lwt.return (response_with_body ?status ?code ?headers body)
+
+let html ?status ?code ?headers body =
+  let response = response_with_body ?status ?code ?headers body in
   Message.set_header response "Content-Type" Formats.text_html;
   Lwt.return response
 
 let json ?status ?code ?headers body =
-  let response =
-    Message.response ?status ?code ?headers (Stream.string body) Stream.null in
+  let response = response_with_body ?status ?code ?headers body in
   Message.set_header response "Content-Type" Formats.application_json;
   Lwt.return response
 
-let response_with_body ?status ?code ?headers body =
-  Message.response ?status ?code ?headers (Stream.string body) Stream.null
-
-let respond ?status ?code ?headers body =
-  Message.response ?status ?code ?headers (Stream.string body) Stream.null
-  |> Lwt.return
-
 (* TODO Actually use the request and extract the site prefix. *)
 let redirect ?status ?code ?headers _request location =
   let status = (status :> Status.redirection option) in
@@ -87,8 +87,7 @@ let redirect ?status ?code ?headers _request location =
     | None, None -> Some (`See_Other)
     | _ -> status
   in
-  let response =
-    Message.response ?status ?code ?headers Stream.empty Stream.null in
+  let response = response_with_body ?status ?code ?headers "" in
   Message.set_header response "Location" location;
   Lwt.return response