feat(portal_tcp): send closing element when pushing None

2 files changed, 13 insertions, 10 deletions

diff --git a/lib/stream.ml b/lib/stream.ml
index e6a46b8..b997a25 100644
--- a/lib/stream.ml
+++ b/lib/stream.ml
@@ -79,12 +79,6 @@ let negotiate
 let initiate (domain : string) (auth : Sasl.auth_config) : (Portal.t * features) Lwt.t =
   let open Portal in
   let* p = connect domain
-  in let push = function
-       | Some n -> p.push (Some n)
-       | None ->
-          p.push (Some close);
-          (* Empty the stream completely, then close the socket. *)
-          Lwt.async (fun () -> let+ () = Markup_lwt.drain p.stream in p.push None)
-     in let portal = {p with push} in
-        let+ features = negotiate domain portal auth
-        in (portal, features)
+  in let+ features = negotiate domain p auth
+     in (p, features)
+
diff --git a/portal/tcp/portal.ml b/portal/tcp/portal.ml
index 31c45c0..c767f3e 100644
--- a/portal/tcp/portal.ml
+++ b/portal/tcp/portal.ml
@@ -123,8 +123,17 @@ let socket_to_stream (sock : socket) =
   in
   let outbound_stream, outbound_push = Lwt_stream.create () in
   let push = function
-    | None -> outbound_push None
     | Some signals -> Markup.iter (fun f -> outbound_push (Some f)) signals
+    | None -> begin
+        (* XMPP streams are one long XML document, so naturally ending the document
+           closes the stream. *)
+        outbound_push (Some `End_element);
+        Lwt.async
+          (fun () ->
+            (* We drain completely the stream when closing, so the socket can close. *)
+            let+ () = Markup_lwt.drain stream
+            in outbound_push None)
+      end
   in Lwt.async begin
          fun () ->
          let* _ =