From 62d583553edec221d1ff01f012fd5dffdfcb7fef Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Fri, 16 Jul 2021 10:56:04 -0400
Subject: [PATCH] allow extra headers in server-sent events

---
 examples/sse_clock.ml |  6 ++++++
 src/Tiny_httpd.ml     | 22 +++++++++++++++++-----
 src/Tiny_httpd.mli    |  6 ++++++
 3 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/examples/sse_clock.ml b/examples/sse_clock.ml
index 442ef088..92dfe961 100644
--- a/examples/sse_clock.ml
+++ b/examples/sse_clock.ml
@@ -12,9 +12,15 @@ let () =
     ]) (fun _ -> ()) "sse_clock [opt*]";
   let server = S.create ~port:!port () in
 
+  let extra_headers = [
+    "Access-Control-Allow-Origin", "*";
+    "Access-Control-Allow-Methods", "POST, GET, OPTIONS";
+  ] in
+
   S.add_route_server_sent_handler server S.Route.(exact "clock" @/ return)
     (fun _req (module EV : S.SERVER_SENT_GENERATOR) ->
        S._debug (fun k->k"new connection");
+       EV.set_headers extra_headers;
        let tick = ref true in
        while true do
          let now = Ptime_clock.now() in
diff --git a/src/Tiny_httpd.ml b/src/Tiny_httpd.ml
index 3c928757..ed4885aa 100644
--- a/src/Tiny_httpd.ml
+++ b/src/Tiny_httpd.ml
@@ -788,6 +788,7 @@ type cb_path_handler =
   unit
 
 module type SERVER_SENT_GENERATOR = sig
+  val set_headers : Headers.t -> unit
   val send_event :
     ?event:string ->
     ?id:string ->
@@ -901,15 +902,21 @@ let[@inline] _opt_iter ~f o = match o with
 let add_route_server_sent_handler ?accept self route f =
   let tr_req oc req ~resp f =
     let req = Request.read_body_full req in
+    let headers = ref Headers.(empty |> set "content-type" "text/event-stream") in
 
-    (* send 200 response now *)
-    let initial_resp =
-      let headers = Headers.(empty |> set "content-type" "text/event-stream") in
-      Response.make_raw ~headers ~code:200 ""
+    (* send response once *)
+    let resp_sent = ref false in
+    let send_response_idempotent_ () =
+      if not !resp_sent then (
+        resp_sent := true;
+        (* send 200 response now *)
+        let initial_resp = Response.make_raw ~headers:!headers ~code:200 "" in
+        resp initial_resp;
+      )
     in
-    resp initial_resp;
 
     let send_event ?event ?id ?retry ~data () : unit =
+      send_response_idempotent_();
       _opt_iter event ~f:(fun e -> Printf.fprintf oc "data: %s\n" e);
       _opt_iter id ~f:(fun e -> Printf.fprintf oc "id: %s\n" e);
       _opt_iter retry ~f:(fun e -> Printf.fprintf oc "retry: %s\n" e);
@@ -919,6 +926,11 @@ let add_route_server_sent_handler ?accept self route f =
       flush oc
     in
     let module SSG = struct
+      let set_headers h =
+        if not !resp_sent then (
+          headers := List.rev_append h !headers;
+          send_response_idempotent_()
+        )
       let send_event = send_event
     end in
     f req (module SSG : SERVER_SENT_GENERATOR);
diff --git a/src/Tiny_httpd.mli b/src/Tiny_httpd.mli
index 311209aa..90ca778b 100644
--- a/src/Tiny_httpd.mli
+++ b/src/Tiny_httpd.mli
@@ -576,6 +576,12 @@ val add_path_handler_stream :
     @since NEXT_RELEASE
   *)
 module type SERVER_SENT_GENERATOR = sig
+  val set_headers : Headers.t -> unit
+  (** Set headers of the response.
+      This is not mandatory but if used at all, it must be called before
+      any call to {!send_event} (once events are sent the response is
+      already sent too). *)
+
   val send_event :
     ?event:string ->
     ?id:string ->