From b6db33e54c617c6c57813076b67b5da096c29720 Mon Sep 17 00:00:00 2001
From: Ushitora Anqou <ushitora@anqou.net>
Date: Sun, 10 Mar 2024 18:45:07 +0900
Subject: [PATCH] cohttp-eio: server: return connection header

---
 cohttp-eio/src/server.ml | 16 +++++++++++++++-
 cohttp-eio/tests/test.ml |  8 ++++++--
 2 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/cohttp-eio/src/server.ml b/cohttp-eio/src/server.ml
index 364c83385..7906e19f6 100644
--- a/cohttp-eio/src/server.ml
+++ b/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
diff --git a/cohttp-eio/tests/test.ml b/cohttp-eio/tests/test.ml
index 2185b769e..5941cca30 100644
--- a/cohttp-eio/tests/test.ml
+++ b/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\