feat(stream): parse_features returns a good old tuple

1 files changed, 14 insertions, 17 deletions

diff --git a/lib/stream.ml b/lib/stream.ml
index 0944f8b..db4d5ee 100644
--- a/lib/stream.ml
+++ b/lib/stream.ml
@@ -6,12 +6,12 @@ exception InsufficientEncryption
 
 type feature =
   | Mechanisms of Sasl.auth_mechanism list
-  | STARTTLS of [`Required | `Optional]
+  | STARTTLS
   | Other of Xml.element
 
-(** [parse_features stanza] is a list of the features contained in the <features>
-    [stanza]. *)
-let parse_features (stanza : Xml.element) : feature list =
+(** [parse_features stanza] is a tuple of the list of all mandatory features and all
+    optional features described in the <features> [stanza]. *)
+let parse_features (stanza : Xml.element) : feature list * feature list =
   let open Xml in
   let open Either in
 
@@ -27,21 +27,18 @@ let parse_features (stanza : Xml.element) : feature list =
     | _ -> raise (InvalidStanza (element_to_string stanza))
   in
 
-  let parse_feature (stanza : Xml.element) : feature =
-    let parse_mechanisms ch = List.map parse_single_mechanism ch
-    and parse_starttls = function
-      | [Left {local_name="required"; _}] -> `Required
-      | [] -> `Optional
-      | _ -> raise (InvalidStanza (element_to_string stanza))
-    in match stanza.local_name with
-       | "mechanisms" -> Mechanisms (parse_mechanisms stanza.children)
-       | "starttls" -> STARTTLS (parse_starttls stanza.children)
-       | _ -> Other stanza
-  in let features = List.map parse_feature children
+  let parse_feature (stanza : Xml.element) : (feature, feature) Either.t =
+    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
      (* 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 `Required]
+        | [], [STARTTLS] -> [STARTTLS], []
         | _ -> features
 
 (** [negotiate domain portal auth] is a promise containing the features supported by the
@@ -62,7 +59,7 @@ let negotiate
       (portal : Portal.t)
       (auth : Sasl.auth_config) : feature list Lwt.t =
   (* Restart a stream: Send the usual business, ask for features. *)
-  let start_stream () : feature list Lwt.t =
+  let start_stream () : (feature list * feature list) Lwt.t =
     let* _id = Portal.header domain portal
     in Wire.get portal.stream >|= parse_features
   in