refactor(stream): migrate to custom requirement wrapping type

1 files changed, 29 insertions, 15 deletions

diff --git a/lib/stream.ml b/lib/stream.ml
index bdf7987..01b4930 100644
--- a/lib/stream.ml
+++ b/lib/stream.ml
@@ -9,14 +9,25 @@ type feature =
   | Mechanisms of Sasl.mechanism list
   | Other of Xml.element
 
+type 'a requirement =
+  | Mandatory of 'a
+  | Optional of 'a
+
+let unwrap = function
+  | Mandatory f -> f
+  | Optional f -> f
+
+let to_either = function
+  | Mandatory f -> Either.Left f
+  | Optional f -> Either.Right f
+
 type config = {
     starttls : Starttls.config;
     sasl : Sasl.config;
     other : (Markup.signal, Markup.sync) Markup.stream list;
   }
 
-(** [features] is a tuple of features list, mandatory and optional. *)
-type features = (feature list * feature list)
+type features = feature requirement list
 
 (** [parse_features stanza] is a tuple of the list of all mandatory features and all
     optional features described in the <features> [stanza]. *)
@@ -36,18 +47,18 @@ let parse_features (stanza : Xml.element) : features =
     | _ -> raise (InvalidStanza (element_to_string stanza))
   in
 
-  let parse_feature (stanza : Xml.element) : (feature, feature) Either.t =
+  let parse_feature (stanza : Xml.element) : feature requirement =
     let parse_mechanisms mech_stanza = List.map parse_single_mechanism mech_stanza
     in match stanza with
-       | {local_name="mechanisms"; _} -> Left (Mechanisms (parse_mechanisms stanza.children))
-       | {local_name="starttls"; children=[Left {local_name="required"; _}]; _} -> Left STARTTLS
-       | {local_name="starttls"; children=[]; _} -> Right STARTTLS
-       | _ -> Right (Other stanza)
-  in let features = List.partition_map parse_feature children
+       | {local_name="mechanisms"; _} -> Mandatory (Mechanisms (parse_mechanisms stanza.children))
+       | {local_name="starttls"; children=[Left {local_name="required"; _}]; _} -> Mandatory STARTTLS
+       | {local_name="starttls"; children=[]; _} -> Optional STARTTLS
+       | _ -> Optional (Other stanza)
+  in let features = List.map parse_feature children
      (* The XMPP spec mandates that sending a features element that contains only a
         <starttls/> means the STARTTLS negotiation is required. *)
      in match features with
-        | [], [STARTTLS] -> [STARTTLS], []
+        | [Optional STARTTLS] -> [Mandatory STARTTLS]
         | _ -> features
 
 (** [start domain portal] is a promise to features that starts a stream negotiation with
@@ -60,7 +71,7 @@ let start (domain : string) (portal : Portal.t) : features Lwt.t =
     server at [portal].
 
     Some features can be sent as [mandatory] or not. *)
-let negotiate mandatory feature portal {starttls; sasl; _} : unit Lwt.t =
+let negotiate feature portal {starttls; sasl; _} : unit Lwt.t =
   (* authenticate using SASL with the XMPP server. *)
   let authenticate mechanisms =
     let open Sasl in
@@ -79,8 +90,11 @@ let negotiate mandatory feature portal {starttls; sasl; _} : unit Lwt.t =
          | Ok _ -> print_endline "Success!"; Lwt.return_unit
     else Lwt.fail InsufficientEncryption
   in match feature with
-     | STARTTLS -> if mandatory || starttls.prefer_starttls
-                   then Starttls.upgrade portal
-                   else Lwt.return_unit
-     | Mechanisms mechs -> authenticate mechs
-     | _ -> Lwt.return_unit
+     | Mandatory STARTTLS -> Starttls.upgrade portal
+     | Optional STARTTLS -> if starttls.prefer_starttls
+                           then Starttls.upgrade portal
+                           else Lwt.return_unit
+     | f ->
+        match unwrap f with
+        | Mechanisms mechs -> authenticate mechs
+        | _ ->  Lwt.return_unit