refactor(stream): move handle_feature outside of negotiate

1 files changed, 13 insertions, 22 deletions

diff --git a/lib/stream.ml b/lib/stream.ml
index 92aebd4..ac8c08b 100644
--- a/lib/stream.ml
+++ b/lib/stream.ml
@@ -50,20 +50,9 @@ let parse_features (stanza : Xml.element) : features =
         | [], [STARTTLS] -> [STARTTLS], []
         | _ -> features
 
-(** [starttls mandatory portal config] negotiates STARTTLS and establishes a TLS
-    handshake with the XMPP server [portal], following the stream config [config].
-
-    If STARTLS is required during negotiation, [mandatory] is true. *)
-let starttls mandatory portal {starttls=config; _} =
-  if (mandatory || config.prefer_starttls)
-  then Starttls.upgrade portal
-  else Lwt.return_unit
-
-(** [sasl mechanisms _mandatory portal config] authenticates using SASL with the XMPP
-    server [portal], following the stream config [config].
-
-    [_mandatory] has no effect (SASL negotiation is always mandatory, if present). *)
-let sasl mechanisms _ (portal : Portal.t) {sasl=config; _} =
+(** [authenticate mechanisms _mandatory portal config] authenticates using SASL with the
+    XMPP server [portal], following the stream config [config]. *)
+let authenticate mechanisms (portal : Portal.t) config =
   let open Sasl in
   let allow_auth () =
     Portal._encrypted portal._socket ||
@@ -80,6 +69,15 @@ let sasl mechanisms _ (portal : Portal.t) {sasl=config; _} =
        | Ok _ -> print_endline "Success!"; Lwt.return_unit
   else Lwt.fail InsufficientEncryption
 
+let handle_feature (mandatory : bool) (feat : feature) (portal : Portal.t)
+      ({starttls; sasl; _} : config) : unit Lwt.t =
+  match feat with
+  | STARTTLS -> if mandatory || starttls.prefer_starttls
+                then Starttls.upgrade portal
+                else Lwt.return_unit
+  | Mechanisms mechs -> authenticate mechs portal sasl
+  | _ -> Lwt.return_unit
+
 (** [negotiate domain portal auth] is a promise containing the features supported by the
     XMPP server [portal], after eventual STARTTLS negotiation and authentication using
     the auth config [auth].
@@ -103,16 +101,9 @@ let negotiate
     let* _id = Portal.header domain portal
     in Wire.get portal.stream >|= parse_features
   in
-  (* Handle a single feature. Mandatory is whether the feature is mandatory. *)
-  let handle_feature (mandatory : bool) (f : feature) : unit Lwt.t =
-    match f with
-    | STARTTLS -> starttls mandatory portal config
-    | Mechanisms m -> sasl m mandatory portal config
-    | _ -> Lwt.return_unit
-  in
   let rec handle_features (f : features) : features Lwt.t =
     match f with
-    | m :: mandatory, optional -> let* () = handle_feature true m
+    | m :: mandatory, optional -> let* () = handle_feature true m portal config
                                   in if needs_restart m
                                      then start_stream () >>= handle_features
                                      else handle_features (mandatory, optional)