feat(stream): add function initiate

2 files changed, 18 insertions, 2 deletions

diff --git a/lib/stream.ml b/lib/stream.ml
index 01d6a55..ee33575 100644
--- a/lib/stream.ml
+++ b/lib/stream.ml
@@ -47,3 +47,20 @@ let negotiate (domain : string) (portal : Portal.t) : features Lwt.t =
   let* _id = Portal.header domain portal
   in let+ features = Wire.get portal.stream >|= parse_features
      in features
+
+(** [initiate domain] initiates a stream with the XMPP server [domain].
+
+    Once [None] is pushed into the stream, the receiving stream is drained and the
+    socket is closed. *)
+let initiate (domain : string) : (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
+        in (portal, features)
diff --git a/test/hello.ml b/test/hello.ml
index f9d429b..e55b44a 100644
--- a/test/hello.ml
+++ b/test/hello.ml
@@ -18,8 +18,7 @@ let main =
       preferred_mechanisms = [Sasl.PLAIN]
     }
   in let domain = (List.nth (String.split_on_char '@' config.jid) 1) in
-     let* portal = Portal.connect domain in
-     let* features = Stream.negotiate domain portal in
+     let* portal, features = Stream.initiate domain in
      Lwt.catch
        (fun () -> program portal config features >|= (fun () -> portal.push None))
        (fun exn ->