From a953f0799b1a4b08b76df8057a941d33e894511d Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Sun, 12 Dec 2021 20:51:45 -0500
Subject: [PATCH] store http version in request, and close connections when
 needed

cases include:
- http 1.0 and no "connection: keep-alive"
- http 1.1 and "connection: close"
---
 src/Tiny_httpd.ml  | 25 +++++++++++++++++++------
 src/Tiny_httpd.mli |  3 ++-
 2 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/src/Tiny_httpd.ml b/src/Tiny_httpd.ml
index 58ed49bd..f058512e 100644
--- a/src/Tiny_httpd.ml
+++ b/src/Tiny_httpd.ml
@@ -363,6 +363,7 @@ module Request = struct
     meth: Meth.t;
     host: string;
     headers: Headers.t;
+    http_version: int*int;
     path: string;
     path_components: string list;
     query: (string*string) list;
@@ -384,6 +385,13 @@ module Request = struct
     | None -> None
   let set_header self k v = {self with headers=Headers.set k v self.headers}
 
+  (** Should we close the connection after this request? *)
+  let close_after_req (self:_ t) : bool =
+    match self.http_version with
+    | 1, 1 -> get_header self "connection" =Some"close"
+    | 1, 0 -> not (get_header self "connection"=Some"keep-alive")
+    | _ -> false
+
   let pp_comp_ out comp =
     Format.fprintf out "[%s]"
       (String.concat ";" @@ List.map (Printf.sprintf "%S") comp)
@@ -483,11 +491,11 @@ module Request = struct
   let parse_req_start ~buf (bs:byte_stream) : unit t option resp_result =
     try
       let line = Byte_stream.read_line ~buf bs in
-      let meth, path =
+      let meth, path, version =
         try
-          let m, p, v = Scanf.sscanf line "%s %s HTTP/1.%d\r" (fun x y z->x,y,z) in
-          if v != 0 && v != 1 then raise Exit;
-          m, p
+          let meth, path, version = Scanf.sscanf line "%s %s HTTP/1.%d\r" (fun x y z->x,y,z) in
+          if version != 0 && version != 1 then raise Exit;
+          meth, path, version
         with _ ->
           _debug (fun k->k "invalid request line: `%s`" line);
           raise (Bad_req (400, "Invalid request line"))
@@ -507,8 +515,11 @@ module Request = struct
         | Ok l -> l
         | Error e -> bad_reqf 400 "invalid query: %s" e
       in
-      Ok (Some {meth; query; host; path; path_components;
-                headers; body=()})
+      let req = {
+        meth; query; host; path; path_components;
+        headers; http_version=(1, version); body=()
+      } in
+      Ok (Some req)
     with
     | End_of_file | Sys_error _ -> Ok None
     | Bad_req (c,s) -> Error (c,s)
@@ -978,6 +989,8 @@ let handle_client_ (self:t) (client_sock:Unix.file_descr) : unit =
     | Ok (Some req) ->
       _debug (fun k->k "req: %s" (Format.asprintf "@[%a@]" Request.pp_ req));
 
+      if Request.close_after_req req then continue := false;
+
       try
         (* is there a handler for this path? *)
         let handler =
diff --git a/src/Tiny_httpd.mli b/src/Tiny_httpd.mli
index 21752f7c..c27951e0 100644
--- a/src/Tiny_httpd.mli
+++ b/src/Tiny_httpd.mli
@@ -219,10 +219,11 @@ end
     Requests are sent by a client, e.g. a web browser or cURL. *)
 
 module Request : sig
-  type 'body t = {
+  type 'body t = private {
     meth: Meth.t;
     host: string;
     headers: Headers.t;
+    http_version: int*int;
     path: string;
     path_components: string list;
     query: (string*string) list;