1 files changed, 40 insertions, 37 deletions

diff --git a/lib/stream.ml b/lib/stream.ml
index 01b4930..ecfd31b 100644
--- a/lib/stream.ml
+++ b/lib/stream.ml
@@ -4,22 +4,40 @@ open Lwt.Infix
 exception ClosedStream
 exception InsufficientEncryption
 
-type feature =
-  | STARTTLS
-  | Mechanisms of Sasl.mechanism list
-  | Other of Xml.element
+module Feature = struct
+  open Either
+  open Xml
 
-type 'a requirement =
-  | Mandatory of 'a
-  | Optional of 'a
+  type t =
+    | STARTTLS
+    | Mechanisms of Sasl.mechanism list
+    | Other of Xml.element
 
-let unwrap = function
-  | Mandatory f -> f
-  | Optional f -> f
+  type requirement =
+    | Mandatory of t
+    | Optional of t
 
-let to_either = function
-  | Mandatory f -> Either.Left f
-  | Optional f -> Either.Right f
+  let unwrap = function
+    | Mandatory f -> f
+    | Optional f -> f
+
+  let to_either = function
+    | Mandatory f -> Left f
+    | Optional f -> Right f
+
+  let parse (stanza : element) : requirement =
+    let parse_single_mechanism = function
+      | Left {local_name = "mechanism"; children = [Right mechanism]; _} ->
+         Sasl.parse_mechanism mechanism
+      | _ -> raise (InvalidStanza (element_to_string stanza))
+    in
+    let parse_mechanisms mech_stanza = List.map parse_single_mechanism mech_stanza
+    in match stanza with
+       | {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)
+end
 
 type config = {
     starttls : Starttls.config;
@@ -27,34 +45,18 @@ type config = {
     other : (Markup.signal, Markup.sync) Markup.stream list;
   }
 
-type features = feature requirement list
+type feature = Feature.t
+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]. *)
 let parse_features (stanza : Xml.element) : features =
   let open Xml in
-  let open Either in
-
   let children =
-    if not (List.for_all is_left stanza.children)
+    if not (List.for_all Either.is_left stanza.children)
     then raise (InvalidStanza (element_to_string stanza))
-    else List.filter_map find_left stanza.children
-  in
-
-  let parse_single_mechanism = function
-    | Left {local_name = "mechanism"; children = [Right mechanism]; _} ->
-       Sasl.parse_mechanism mechanism
-    | _ -> raise (InvalidStanza (element_to_string stanza))
-  in
-
-  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"; _} -> 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
+    else List.filter_map Either.find_left stanza.children
+  in let features = List.map Feature.parse 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
@@ -89,11 +91,12 @@ let negotiate feature portal {starttls; sasl; _} : unit Lwt.t =
          | Error err -> Lwt.fail_with (parse_auth_error err)
          | Ok _ -> print_endline "Success!"; Lwt.return_unit
     else Lwt.fail InsufficientEncryption
-  in match feature with
+  in let open Feature in
+     match feature with
      | Mandatory STARTTLS -> Starttls.upgrade portal
      | Optional STARTTLS -> if starttls.prefer_starttls
-                           then Starttls.upgrade portal
-                           else Lwt.return_unit
+                            then Starttls.upgrade portal
+                            else Lwt.return_unit
      | f ->
         match unwrap f with
         | Mechanisms mechs -> authenticate mechs