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
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
let log_src = Logs.Src.create "sihl.middleware.flash"
module Logs = (val Logs.src_log log_src : Logs.LOG)
module Flash = struct
open Sexplib.Conv
type t =
{ alert : string option
; notice : string option
; custom : (string * string) list
}
[@@deriving yojson, sexp]
let equals f1 f2 =
Option.equal String.equal f1.alert f2.alert
&& Option.equal String.equal f1.notice f2.notice
&& CCList.equal (CCPair.equal String.equal String.equal) f1.custom f2.custom
;;
let of_json (json : string) : t option =
try Some (of_yojson (Yojson.Safe.from_string json) |> Result.get_ok) with
| _ -> None
;;
let to_json (flash : t) : string = flash |> to_yojson |> Yojson.Safe.to_string
end
module Env = struct
let key : Flash.t Opium.Context.key =
Opium.Context.Key.create ("flash", Flash.sexp_of_t)
;;
end
let find' req = Opium.Context.find Env.key req.Opium.Request.env
let find_alert req = Option.bind (find' req) (fun flash -> flash.alert)
let find_notice req = Option.bind (find' req) (fun flash -> flash.notice)
let find key req =
Option.bind (find' req) (fun flash ->
flash.custom
|> List.find_opt (fun (k, _) -> String.equal key k)
|> Option.map snd)
;;
let set_alert alert resp =
let flash = Opium.Context.find Env.key resp.Opium.Response.env in
let flash =
match flash with
| None -> Flash.{ alert = Some alert; notice = None; custom = [] }
| Some flash -> Flash.{ flash with alert = Some alert }
in
let env = resp.Opium.Response.env in
let env = Opium.Context.add Env.key flash env in
{ resp with env }
;;
let set_notice notice resp =
let flash = Opium.Context.find Env.key resp.Opium.Response.env in
let flash =
match flash with
| None -> Flash.{ alert = None; notice = Some notice; custom = [] }
| Some flash -> Flash.{ flash with notice = Some notice }
in
let env = resp.Opium.Response.env in
let env = Opium.Context.add Env.key flash env in
{ resp with env }
;;
let set values resp =
let flash = Opium.Context.find Env.key resp.Opium.Response.env in
let flash =
match flash with
| None -> Flash.{ alert = None; notice = None; custom = values }
| Some flash -> Flash.{ flash with custom = values }
in
let env = resp.Opium.Response.env in
let env = Opium.Context.add Env.key flash env in
{ resp with env }
;;
type decode_status =
| No_cookie_found
| Parse_error
| Found of Flash.t
let decode_flash cookie_key req =
match Opium.Request.cookie cookie_key req with
| None -> No_cookie_found
| Some cookie_value ->
(match Flash.of_json cookie_value with
| None ->
Logs.err (fun m ->
m
"Failed to parse value found in flash cookie '%s': '%s'"
cookie_key
cookie_value);
Logs.info (fun m ->
m
"Maybe the cookie key '%s' collides with a cookie issued by \
someone else. Try to change the cookie key."
cookie_key);
Parse_error
| Some flash -> Found flash)
;;
let persist_flash ?old_flash ?(delete_if_not_set = false) cookie_key resp =
let flash = Opium.Context.find Env.key resp.Opium.Response.env in
match flash with
| None ->
if delete_if_not_set
then
Opium.Response.add_cookie_or_replace
~expires:(`Max_age Int64.zero)
~scope:(Uri.of_string "/")
(cookie_key, "")
resp
else resp
| Some flash ->
(match old_flash with
| Some old_flash ->
if Flash.equals old_flash flash
then
resp
else (
let cookie_value = Flash.to_json flash in
let cookie = cookie_key, cookie_value in
let resp =
Opium.Response.add_cookie_or_replace
~scope:(Uri.of_string "/")
cookie
resp
in
resp)
| None ->
let cookie_value = Flash.to_json flash in
let cookie = cookie_key, cookie_value in
let resp =
Opium.Response.add_cookie_or_replace
~scope:(Uri.of_string "/")
cookie
resp
in
resp)
;;
let middleware ?(cookie_key = "_flash") () =
let filter handler req =
match decode_flash cookie_key req with
| No_cookie_found ->
let%lwt resp = handler req in
Lwt.return @@ persist_flash cookie_key resp
| Parse_error ->
let%lwt resp = handler req in
Lwt.return @@ persist_flash ~delete_if_not_set:true cookie_key resp
| Found flash ->
let env = req.Opium.Request.env in
let env = Opium.Context.add Env.key flash env in
let req = { req with env } in
let%lwt resp = handler req in
Lwt.return
@@ persist_flash ~delete_if_not_set:true ~old_flash:flash cookie_key resp
in
Rock.Middleware.create ~name:"flash" ~filter
;;