[sledge] Protect Trace.parse from invalid user input

Summary:
Trace.parse is exposed across a module boundary, so handle its own
exceptions and return an `error`.

Reviewed By: mbouaziz

Differential Revision: D14081602

fbshipit-source-id: 4087c5d5b

sledge/src/config.ml

@@ -9,8 +9,9 @@
 
 let trace_conv =
   let parse s =
-    try Ok (Trace.parse s).trace_mods_funs with _ ->
-      Error (`Msg ("Invalid trace spec: " ^ s))
+    match Trace.parse s with
+    | Ok c -> Ok c.trace_mods_funs
+    | Error err -> Error err
   in
   let print fs c =
     let pf fmt = Format.fprintf fs fmt in

sledge/src/trace/trace.ml

@@ -24,53 +24,55 @@ let all = {none with trace_all= true}
 let config : config ref = ref none
 
 let parse s =
-  let default = Map.empty (module String) in
-  let index_from s i =
-    Option.merge ~f:min
-      (String.index_from s i '+')
-      (String.index_from s i '-')
-  in
-  let rec split s rev_parts i =
-    match index_from s (i + 1) with
-    | Some j when j = i -> split s rev_parts j
-    | Some j -> split s (String.sub s ~pos:i ~len:(j - i) :: rev_parts) j
-    | _ -> List.rev (String.subo s ~pos:i :: rev_parts)
-  in
-  let parts = split s [] 0 in
-  let trace_mods_funs =
-    List.fold parts ~init:default ~f:(fun m part ->
-        let parse_part part =
-          let sign, rest =
-            match part.[0] with
-            | '-' -> (false, String.subo part ~pos:1)
-            | '+' -> (true, String.subo part ~pos:1)
-            | _ -> (true, part)
-          in
-          assert (not (String.is_empty rest)) ;
-          assert (Char.is_uppercase rest.[0]) ;
-          match String.lsplit2 rest ~on:'.' with
-          | Some (mod_name, fun_name) ->
-              assert (Char.is_lowercase fun_name.[0]) ;
-              (mod_name, Some fun_name, sign)
-          | None -> (rest, None, sign)
-        in
-        match parse_part part with
-        | mod_name, Some fun_name, enabled ->
-            let {trace_mod; trace_funs} =
-              match Map.find m mod_name with
-              | Some c -> c
-              | None -> {trace_mod= None; trace_funs= default}
+  try
+    let default = Map.empty (module String) in
+    let index_from s i =
+      Option.merge ~f:min
+        (String.index_from s i '+')
+        (String.index_from s i '-')
+    in
+    let rec split s rev_parts i =
+      match index_from s (i + 1) with
+      | Some j when j = i -> split s rev_parts j
+      | Some j -> split s (String.sub s ~pos:i ~len:(j - i) :: rev_parts) j
+      | _ -> List.rev (String.subo s ~pos:i :: rev_parts)
+    in
+    let parts = split s [] 0 in
+    let trace_mods_funs =
+      List.fold parts ~init:default ~f:(fun m part ->
+          let parse_part part =
+            let sign, rest =
+              match part.[0] with
+              | '-' -> (false, String.subo part ~pos:1)
+              | '+' -> (true, String.subo part ~pos:1)
+              | _ -> (true, part)
             in
-            Map.set m ~key:mod_name
-              ~data:
-                { trace_mod
-                ; trace_funs= Map.set trace_funs ~key:fun_name ~data:enabled
-                }
-        | mod_name, None, enabled ->
-            Map.set m ~key:mod_name
-              ~data:{trace_mod= Some enabled; trace_funs= default} )
-  in
-  {none with trace_mods_funs}
+            assert (not (String.is_empty rest)) ;
+            assert (Char.is_uppercase rest.[0]) ;
+            match String.lsplit2 rest ~on:'.' with
+            | Some (mod_name, fun_name) ->
+                assert (Char.is_lowercase fun_name.[0]) ;
+                (mod_name, Some fun_name, sign)
+            | None -> (rest, None, sign)
+          in
+          match parse_part part with
+          | mod_name, Some fun_name, enabled ->
+              let {trace_mod; trace_funs} =
+                match Map.find m mod_name with
+                | Some c -> c
+                | None -> {trace_mod= None; trace_funs= default}
+              in
+              Map.set m ~key:mod_name
+                ~data:
+                  { trace_mod
+                  ; trace_funs=
+                      Map.set trace_funs ~key:fun_name ~data:enabled }
+          | mod_name, None, enabled ->
+              Map.set m ~key:mod_name
+                ~data:{trace_mod= Some enabled; trace_funs= default} )
+    in
+    Ok {none with trace_mods_funs}
+  with Assert_failure _ -> Error (`Msg ("Invalid trace spec: " ^ s))
 
 let init ?(margin = 300) ~config:c () =
   Format.set_margin margin ;

sledge/src/trace/trace.mli

@@ -23,7 +23,7 @@ type config =
 
 val none : config
 val all : config
-val parse : string -> config
+val parse : string -> (config, [> `Msg of string]) result
 
 val init : ?margin:int -> config:config -> unit -> unit
 (** Initialize the configuration of debug tracing. *)