Merge pull request #1025 from ushitora-anqou/return-connection-header

cohttp-{lwt,eio}: server: return connection header

CHANGES.md

@@ -1,5 +1,6 @@
 ## Unreleased
 
+- cohttp-{lwt,eio}: server: add connection header to response if not present (ushitora-anqou #1025)
 - cohttp-curl: Curl no longer prepends the first HTTP request header to the output. (jonahbeckford #1030, #987)
 - cohttp-eio: client: use permissive argument type for make_generic
 - cohttp-eio: Improve error handling in example server (talex5 #1023)

cohttp-eio/src/server.ml

@@ -116,7 +116,21 @@ let callback { conn_closed; handler } ((_, peer_address) as conn) input output =
         let () =
           try
             match handler (conn, id) request body with
-            | `Response (response, body) -> write output response body
+            | `Response (response, body) ->
+                let keep_alive =
+                  Http.Request.is_keep_alive request
+                  && Http.Response.is_keep_alive response
+                in
+                let response =
+                  let headers =
+                    Http.Header.add_unless_exists
+                      (Http.Response.headers response)
+                      "connection"
+                      (if keep_alive then "keep-alive" else "close")
+                  in
+                  { response with Http.Response.headers }
+                in
+                write output response body
             | `Expert (response, handler) ->
                 let () = Io.Response.write_header response output in
                 handler input output

cohttp-eio/tests/test.ml

@@ -44,7 +44,8 @@ let () =
         [ Cstruct.of_string "GET / HTTP/1.1\r\nconnection: close\r\n\r\n" ]
     in
     Alcotest.(check ~here:[%here] string)
-      "response" "HTTP/1.1 200 OK\r\ncontent-length: 4\r\n\r\nroot"
+      "response"
+      "HTTP/1.1 200 OK\r\nconnection: close\r\ncontent-length: 4\r\n\r\nroot"
       Eio.Buf_read.(of_flow ~max_size:max_int socket |> take_all)
   and missing socket =
     let () =
@@ -54,7 +55,8 @@ let () =
         ]
     in
     Alcotest.(check ~here:[%here] string)
-      "response" "HTTP/1.1 404 Not Found\r\ncontent-length: 0\r\n\r\n"
+      "response"
+      "HTTP/1.1 404 Not Found\r\nconnection: close\r\ncontent-length: 0\r\n\r\n"
       Eio.Buf_read.(of_flow ~max_size:max_int socket |> take_all)
   and streaming_response socket =
     let () =
@@ -66,6 +68,7 @@ let () =
     Alcotest.(check ~here:[%here] string)
       "response"
       "HTTP/1.1 200 OK\r\n\
+       connection: close\r\n\
        transfer-encoding: chunked\r\n\
        \r\n\
        5\r\n\
@@ -90,6 +93,7 @@ let () =
     Alcotest.(check ~here:[%here] string)
       "response"
       "HTTP/1.1 200 OK\r\n\
+       connection: close\r\n\
        transfer-encoding: chunked\r\n\
        \r\n\
        c\r\n\

cohttp-lwt/src/server.ml

@@ -146,6 +146,15 @@ module Make (IO : S.IO) = struct
             let keep_alive =
               Http.Request.is_keep_alive req && Http.Response.is_keep_alive res
             in
+            let res =
+              let headers =
+                Http.Header.add_unless_exists
+                  (Http.Response.headers res)
+                  "connection"
+                  (if keep_alive then "keep-alive" else "close")
+              in
+              { res with Http.Response.headers }
+            in
             handle_response ~keep_alive oc res body
               (fun () -> spec.conn_closed conn)
               (fun oc -> handle_client ic oc conn spec)