feat(stream): return an optional socket on stream negotiation

this allows encoding whether we need a stream restart in negotiate, bypassing some
ugliness later :)

1 files changed, 12 insertions, 9 deletions

diff --git a/lib/stream.ml b/lib/stream.ml
index d84dd57..91ac9da 100644
--- a/lib/stream.ml
+++ b/lib/stream.ml
@@ -71,8 +71,11 @@ let start (portal : Portal.t) : features Lwt.t =
   Segment.get portal.stream >|= parse_features
 
 (** [negotiate mandatory feature portal] negotiates the feature [feature] with the XMPP
-    server at [portal]. *)
-let negotiate feature portal {starttls; sasl; _} : unit Lwt.t =
+    server at [portal].
+
+    The returned promise contains an optional socket. If the negotiation leads to a
+    stream restart, this option is a socket, that can be used to restart the stream. *)
+let negotiate feature portal {starttls; sasl; _} : Portal.socket option Lwt.t =
   (* authenticate using SASL with the XMPP server. *)
   let authenticate mechanisms =
     let open Sasl in
@@ -88,17 +91,17 @@ let negotiate feature portal {starttls; sasl; _} : unit Lwt.t =
       let* auth_result = authenticate portal sasl mechanisms
       in match auth_result with
          | Error err -> Lwt.fail_with (parse_auth_error err)
-         | Ok _ -> Lwt.return_unit
+         | Ok _ -> Lwt.return_some portal._socket
     else Lwt.fail InsufficientEncryption
   in
   let open Feature in
   match feature with
-  | Mandatory STARTTLS -> Starttls.upgrade portal
+  | Mandatory STARTTLS -> Starttls.upgrade portal >|= Option.some
   | Optional STARTTLS -> if starttls.prefer
-                         then Starttls.upgrade portal
-                         else Lwt.return_unit
+                         then Starttls.upgrade portal >|= Option.some
+                         else Lwt.return_none
   | Mandatory (Mechanisms mechs) -> authenticate mechs
-  | Mandatory Bind -> Lwt.return_unit  (* TODO: binding *)
-  | Mandatory (Other _) -> Lwt.return_unit
+  | Mandatory Bind -> Lwt.return_none (* TODO: binding *)
+  | Mandatory (Other _) -> Lwt.return_none
   (* So far, for convenience, when something is optional, we just don't do it. *)
-  | Optional _ -> Lwt.return_unit
+  | Optional _ -> Lwt.return_none