Minimum managment of cookies

src/Tiny_httpd.ml

@@ -335,12 +335,12 @@ module Headers = struct
     try String.iter (fun c->if not (pred c) then raise Exit) s; true
     with Exit -> false
 
-  let parse_ ~buf (bs:byte_stream) : t =
+  let parse_ ~buf (bs:byte_stream) : t * t =
-    let rec loop acc =
+    let rec loop headers cookies =
       let line = Byte_stream.read_line ~buf bs in
       _debug (fun k->k  "parsed header line %S" line);
       if line = "\r" then (
-        acc
+        (headers,cookies)
       ) else (
         let k,v =
           try
@@ -352,10 +352,59 @@ module Headers = struct
             k,v
           with _ -> bad_reqf 400 "invalid header line: %S" line
         in
-        loop ((String.lowercase_ascii k,v)::acc)
+        let k = String.lowercase_ascii k in
-      )
+        let headers, cookies =
+          if k = "cookie" then
+            begin
+              let fn_eq s =
+                match String.split_on_char '=' s with
+                | k::v::_ when for_all is_tchar (String.trim k) ->
+                   Some String.(trim k,trim v)
+                | _       -> None
               in
-    loop []
+              let cookies =
+                List.filter_map fn_eq (String.split_on_char ';' v) @ cookies
+              in
+              (headers, cookies)
+            end
+          else
+            ((k,v)::headers, cookies)
+        in
+        loop headers cookies)
+    in
+    loop [] []
+
+end
+
+module SetCookie = struct
+  type sameSite = Strict | Lax | None
+  type t =
+    | MaxAge of int
+    | Expires of string (** FIXME: format date, but need computation
+                            of days of week *)
+    | Domain of string
+    | Path of string
+    | Secure
+    | HttpOnly
+    | SameSite of sameSite
+
+  let pp fmt = function
+    | MaxAge s -> Format.fprintf fmt "Max-Age=%d" s
+    | Expires d -> Format.fprintf fmt "Expires=%s" d
+    | Domain d -> Format.fprintf fmt "Domain=%s" d
+    | Path p -> Format.fprintf fmt "Path=%s" p
+    | Secure -> Format.fprintf fmt "Secure"
+    | HttpOnly -> Format.fprintf fmt "HttpOnly"
+    | SameSite Strict -> Format.fprintf fmt "Same=Strict"
+    | SameSite Lax -> Format.fprintf fmt "Same=Lax"
+    | SameSite None -> Format.fprintf fmt "Same=None"
+
+  let rec pps fmt = function
+    | [] -> ()
+    | prop::ls -> Format.fprintf fmt "; %a%a" pp prop pps ls
+
+  let set props key value =
+    Format.asprintf "%s=%s%a" key value pps props
 end
 
 module Request = struct
@@ -363,6 +412,7 @@ module Request = struct
     meth: Meth.t;
     host: string;
     headers: Headers.t;
+    cookies: Headers.t;
     http_version: int*int;
     path: string;
     path_components: string list;
@@ -372,6 +422,7 @@ module Request = struct
   }
 
   let headers self = self.headers
+  let cookies self = self.cookies
   let host self = self.host
   let meth self = self.meth
   let path self = self.path
@@ -384,6 +435,10 @@ module Request = struct
     | Some x -> (try Some (int_of_string x) with _ -> None)
     | None -> None
   let set_header k v self = {self with headers=Headers.set k v self.headers}
+  let get_cookie ?f self h = Headers.get ?f h self.cookies
+  let get_cookie_int self h = match get_cookie self h with
+    | Some x -> (try Some (int_of_string x) with _ -> None)
+    | None -> None
   let update_headers f self = {self with headers=f self.headers}
   let set_body b self = {self with body=b}
 
@@ -505,7 +560,7 @@ module Request = struct
       in
       let meth = Meth.of_string meth in
       _debug (fun k->k "got meth: %s, path %S" (Meth.to_string meth) path);
-      let headers = Headers.parse_ ~buf bs in
+      let headers, cookies = Headers.parse_ ~buf bs in
       let host =
         match Headers.get "Host" headers with
         | None -> bad_reqf 400 "No 'Host' header in request"
@@ -520,7 +575,7 @@ module Request = struct
       in
       let req = {
         meth; query; host; path; path_components;
-        headers; http_version=(1, version); body=(); start_time;
+        headers; cookies; http_version=(1, version); body=(); start_time;
       } in
       Ok (Some req)
     with
@@ -603,6 +658,10 @@ module Response = struct
   let update_headers f self = {self with headers=f self.headers}
   let set_header k v self = {self with headers = Headers.set k v self.headers}
   let set_code code self = {self with code}
+  let set_cookie ?(props=[]) k v self =
+      set_header "Set-Cookie" (SetCookie.set props k v) self
+  let unset_cookie k self =
+      set_cookie ~props:[MaxAge 0] k "UNSET" self
 
   let make_raw ?(headers=[]) ~code body : t =
     (* add content length to response *)

src/Tiny_httpd.mli

@@ -214,6 +214,22 @@ module Headers : sig
   (** Pretty print the headers. *)
 end
 
+(** {Set Cookie}
+
+    A module to set new cookies in the header *)
+module SetCookie : sig
+  type sameSite = Strict | Lax | None
+  type t =
+    | MaxAge of int
+    | Expires of string (** FIXME: format date, but need computation
+                            of days of week *)
+    | Domain of string
+    | Path of string
+    | Secure
+    | HttpOnly
+    | SameSite of sameSite
+end
+
 (** {2 Requests}
 
     Requests are sent by a client, e.g. a web browser or cURL. *)
@@ -223,6 +239,7 @@ module Request : sig
     meth: Meth.t;
     host: string;
     headers: Headers.t;
+    cookies: Headers.t;
     http_version: int*int;
     path: string;
     path_components: string list;
@@ -255,6 +272,18 @@ module Request : sig
 
   val get_header_int : _ t -> string -> int option
 
+  val cookies : _ t -> Headers.t
+  (** List of cookies of the request
+      @since 0.12 *)
+
+  val get_cookie : ?f:(string->string) -> _ t -> string -> string option
+  (** get a cookie
+      @since 0.12 *)
+
+  val get_cookie_int : _ t -> string -> int option
+  (** get a cookie as int
+      @since 0.12 *)
+
   val set_header : string -> string -> 'a t -> 'a t
   (** [set_header k v req] sets [k: v] in the request [req]'s headers. *)
 
@@ -350,6 +379,14 @@ module Response : sig
   (** Set a header.
       @since 0.11 *)
 
+  val set_cookie : ?props:SetCookie.t list -> string -> string -> t -> t
+  (** Set a cookie in the header
+      @since 0.12 *)
+
+  val unset_cookie : string -> t -> t
+  (** Unset a cookie in the header
+      @since 0.12 *)
+
   val update_headers : (Headers.t -> Headers.t) -> t -> t
   (** Modify headers
       @since 0.11 *)