Source file `cookie.ml`

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
open Misc
open Sexplib.Std
module Co = Cohttp

let encode x = Uri.pct_encode ~component:`Query_key x

let decode = Uri.pct_decode

module Env = struct
  type cookie = (string * string) list

  let key : cookie Hmap0.key =
    Hmap0.Key.create ("cookie", [%sexp_of: (string * string) list])
end

module Env_resp = struct
  type cookie = Co.Cookie.Set_cookie_hdr.t list

  let key : cookie Hmap0.key =
    Hmap0.Key.create ("cookie_res", [%sexp_of: Co.Cookie.Set_cookie_hdr.t list])
end

let current_cookies env record =
  Option.value ~default:[] (Hmap0.find Env.key (env record))

let current_cookies_resp env record =
  Option.value ~default:[] (Hmap0.find Env_resp.key (env record))

let cookies_raw req =
  req |> Rock.Request.request |> Co.Request.headers
  |> Co.Cookie.Cookie_hdr.extract

let cookies req =
  req |> cookies_raw
  |> List.filter_map ~f:(fun (k, v) ->
         (* ignore bad cookies *)
         Option.try_with (fun () -> (k, decode v)))

let get req ~key =
  let cookie1 =
    let env = current_cookies (fun r -> r.Rock.Request.env) req in
    List.find_map env ~f:(fun (k, v) -> if k = key then Some v else None)
  in
  match cookie1 with
  | Some cookie -> Some cookie
  | None ->
      let cookies = cookies_raw req in
      cookies
      |> List.find_map ~f:(fun (k, v) ->
             if k = key then Some (decode v) else None)

(* Path defaulted to "/" as otherwise the default is the path of the request's
   URI *)
let set_cookies ?expiration ?(path = "/") ?domain ?secure ?http_only resp
    cookies =
  let env = Rock.Response.env resp in
  let current_cookies =
    current_cookies_resp (fun r -> r.Rock.Response.env) resp
  in
  let cookies' =
    List.map cookies ~f:(fun (key, data) ->
        Co.Cookie.Set_cookie_hdr.make ~path ?domain ?expiration ?secure
          ?http_only
          (key, encode data))
  in
  (* WRONG cookies cannot just be concatenated *)
  let all_cookies = current_cookies @ cookies' in
  {resp with Rock.Response.env= Hmap0.add Env_resp.key all_cookies env}

let set ?expiration ?path ?domain ?secure ?http_only resp ~key ~data =
  set_cookies ?expiration ?path ?domain ?secure ?http_only resp [(key, data)]

let m =
  (* TODO: "optimize" *)
  let filter handler req =
    handler req
    >>| fun response ->
    let cookie_headers =
      let module Cookie = Co.Cookie.Set_cookie_hdr in
      response
      |> current_cookies_resp (fun r -> r.Rock.Response.env)
      |> List.map ~f:Cookie.serialize
    in
    let old_headers = Rock.Response.headers response in
    { response with
      Rock.Response.headers=
        List.fold_left cookie_headers ~init:old_headers
          ~f:(fun headers (k, v) -> Co.Header.add headers k v) }
  in
  Rock.Middleware.create ~filter ~name:"Cookie"