Add functions to add or get features out of taint representations

arthaud · facebook-github-bot · commit ba2d79b1dfdb · 2021-09-08T13:55:08.000-07:00
Summary:
A very recurrent pattern in our codebase is to gather or add features or breadcrumb on a taint abstraction. For now, this is usually done with a fold on the abstract domain.
Let's add helper functions to add or gather breadcrumbs and features.
This will make refactoring easier when changing the abstraction for the feature set, for instance splitting via-features and breadcrumbs, or implementing local features.

Reviewed By: dkgi

Differential Revision: D30761500

fbshipit-source-id: 5e7dc8e180b38b302b470cdb32ca40b9ffbfdb2a
diff --git a/source/interprocedural_analyses/taint/backwardAnalysis.ml b/source/interprocedural_analyses/taint/backwardAnalysis.ml
@@ -203,13 +203,7 @@ module AnalysisInstance (FunctionContext : FUNCTION_CONTEXT) = struct
                     (* No special path handling for side effect taint *)
                     [[]]
               in
-              let breadcrumbs =
-                BackwardTaint.fold
-                  Features.SimpleSet.Self
-                  element
-                  ~f:Features.gather_breadcrumbs
-                  ~init:Features.SimpleSet.bottom
-              in
+              let breadcrumbs = BackwardTaint.breadcrumbs element in
               let tito_depth =
                 BackwardTaint.fold
                   TraceLength.Self
@@ -238,8 +232,8 @@ module AnalysisInstance (FunctionContext : FUNCTION_CONTEXT) = struct
                 ~f:(fun taint extra_path ->
                   read_tree extra_path taint_to_propagate
                   |> BackwardState.Tree.collapse
-                       ~transform:(BackwardTaint.add_features Features.tito_broadening)
-                  |> BackwardTaint.transform Features.SimpleSet.Self Add ~f:breadcrumbs
+                       ~transform:(BackwardTaint.add_breadcrumbs Features.tito_broadening)
+                  |> BackwardTaint.add_breadcrumbs breadcrumbs
                   |> BackwardTaint.transform
                        TraceLength.Self
                        (Context (BackwardTaint.kind, Map))
@@ -254,8 +248,8 @@ module AnalysisInstance (FunctionContext : FUNCTION_CONTEXT) = struct
           in
           let add_tito_feature_and_position leaf_taint =
             leaf_taint
-            |> FlowDetails.transform Features.TitoPositionSet.Element Add ~f:location
-            |> FlowDetails.transform Features.SimpleSet.Element Add ~f:Features.tito
+            |> FlowDetails.add_tito_position location
+            |> FlowDetails.add_breadcrumb Features.tito
           in
           BackwardState.read
             ~transform_non_leaves
@@ -315,19 +309,16 @@ module AnalysisInstance (FunctionContext : FUNCTION_CONTEXT) = struct
           let state =
             match AccessPath.of_expression ~resolution argument with
             | Some { AccessPath.root; path } ->
-                let features_to_add =
+                let breadcrumbs_to_add =
                   BackwardState.Tree.filter_by_kind ~kind:Sinks.AddFeatureToArgument sink_taint
-                  |> BackwardTaint.fold
-                       Features.SimpleSet.Self
-                       ~f:Features.gather_breadcrumbs
-                       ~init:Features.SimpleSet.bottom
+                  |> BackwardTaint.breadcrumbs
                 in
-                if Features.SimpleSet.is_bottom features_to_add then
+                if Features.SimpleSet.is_bottom breadcrumbs_to_add then
                   state
                 else
                   let taint =
                     BackwardState.read state.taint ~root ~path
-                    |> BackwardState.Tree.transform Features.SimpleSet.Self Add ~f:features_to_add
+                    |> BackwardState.Tree.add_breadcrumbs breadcrumbs_to_add
                   in
                   { taint = BackwardState.assign ~root ~path taint state.taint }
             | None -> state
@@ -348,9 +339,9 @@ module AnalysisInstance (FunctionContext : FUNCTION_CONTEXT) = struct
               |> Features.SimpleSet.add Features.obscure
             in
             BackwardState.Tree.collapse
-              ~transform:(BackwardTaint.add_features Features.tito_broadening)
+              ~transform:(BackwardTaint.add_breadcrumbs Features.tito_broadening)
               call_taint
-            |> BackwardTaint.transform Features.SimpleSet.Self Add ~f:breadcrumbs
+            |> BackwardTaint.add_breadcrumbs breadcrumbs
             |> BackwardState.Tree.create_leaf
           else
             BackwardState.Tree.bottom
@@ -387,9 +378,9 @@ module AnalysisInstance (FunctionContext : FUNCTION_CONTEXT) = struct
           in
           let obscure_taint =
             BackwardState.Tree.collapse
-              ~transform:(BackwardTaint.add_features Features.tito_broadening)
+              ~transform:(BackwardTaint.add_breadcrumbs Features.tito_broadening)
               call_taint
-            |> BackwardTaint.transform Features.SimpleSet.Element Add ~f:Features.obscure
+            |> BackwardTaint.add_breadcrumb Features.obscure
             |> BackwardState.Tree.create_leaf
           in
           let state =
@@ -676,7 +667,7 @@ module AnalysisInstance (FunctionContext : FUNCTION_CONTEXT) = struct
             BackwardState.Tree.join
               taint
               (get_taint (Some (AccessPath.create (Root.Variable "$all") [])) state)
-            |> BackwardState.Tree.transform Features.SimpleSet.Element Add ~f:Features.lambda
+            |> BackwardState.Tree.add_breadcrumb Features.lambda
           in
           let all_assignee =
             Node.create
@@ -937,11 +928,7 @@ module AnalysisInstance (FunctionContext : FUNCTION_CONTEXT) = struct
             | Name
                 (Name.Attribute
                   { base = { Node.value = Expression.String _; _ }; attribute = "format"; _ }) ->
-                BackwardState.Tree.transform
-                  Features.SimpleSet.Element
-                  Add
-                  ~f:Features.format_string
-                  taint
+                BackwardState.Tree.add_breadcrumb Features.format_string taint
             | _ -> taint
           in
           match FunctionContext.get_callees ~location ~call:{ Call.callee; arguments } with
@@ -1003,13 +990,7 @@ module AnalysisInstance (FunctionContext : FUNCTION_CONTEXT) = struct
                     taint)
                 ~init:taint
           in
-          let taint =
-            BackwardState.Tree.transform
-              Features.SimpleSet.Element
-              Add
-              ~f:Features.format_string
-              taint
-          in
+          let taint = BackwardState.Tree.add_breadcrumb Features.format_string taint in
           List.fold
             expressions
             ~f:(fun state expression -> analyze_expression ~resolution ~taint ~state ~expression)
@@ -1034,9 +1015,7 @@ module AnalysisInstance (FunctionContext : FUNCTION_CONTEXT) = struct
           match ComparisonOperator.override ~location comparison with
           | Some override -> analyze_expression ~resolution ~taint ~state ~expression:override
           | None ->
-              let taint =
-                BackwardState.Tree.transform Features.SimpleSet.Self Add ~f:Features.type_bool taint
-              in
+              let taint = BackwardState.Tree.add_breadcrumbs Features.type_bool taint in
               analyze_expression ~resolution ~taint ~state ~expression:right
               |> fun state -> analyze_expression ~resolution ~taint ~state ~expression:left)
       | Call { callee; arguments } ->
@@ -1111,16 +1090,9 @@ module AnalysisInstance (FunctionContext : FUNCTION_CONTEXT) = struct
               let global_model = Model.get_global_model ~resolution ~expression ~location in
               let add_tito_features taint =
                 let attribute_features =
-                  global_model |> Model.GlobalModel.get_tito |> BackwardState.Tree.get_all_features
+                  global_model |> Model.GlobalModel.get_tito |> BackwardState.Tree.features
                 in
-                if not (Features.SimpleSet.is_bottom attribute_features) then
-                  BackwardState.Tree.transform
-                    Features.SimpleSet.Self
-                    Add
-                    ~f:attribute_features
-                    taint
-                else
-                  taint
+                BackwardState.Tree.add_breadcrumbs attribute_features taint
               in
 
               let apply_attribute_sanitizers taint =
@@ -1392,12 +1364,12 @@ let extract_tito_and_sink_models define ~is_constructor ~resolution ~existing_ba
         BackwardState.Tree.essential tree
     in
     BackwardState.Tree.shape
-      ~transform:(BackwardTaint.add_features Features.widen_broadening)
+      ~transform:(BackwardTaint.add_breadcrumbs Features.widen_broadening)
       tree
       ~mold:essential
-    |> BackwardState.Tree.transform Features.SimpleSet.Self Add ~f:type_breadcrumbs
+    |> BackwardState.Tree.add_breadcrumbs type_breadcrumbs
     |> BackwardState.Tree.limit_to
-         ~transform:(BackwardTaint.add_features Features.widen_broadening)
+         ~transform:(BackwardTaint.add_breadcrumbs Features.widen_broadening)
          ~width:maximum_model_width
     |> BackwardState.Tree.approximate_return_access_paths ~maximum_return_access_path_length
   in
@@ -1426,16 +1398,7 @@ let extract_tito_and_sink_models define ~is_constructor ~resolution ~existing_ba
             BackwardState.Tree.prune_maximum_length maximum_tito_depth candidate_tree
         | _ -> candidate_tree
       in
-      let candidate_tree =
-        if Features.SimpleSet.is_bottom features_to_attach then
-          candidate_tree
-        else
-          BackwardState.Tree.transform
-            Features.SimpleSet.Self
-            Add
-            ~f:features_to_attach
-            candidate_tree
-      in
+      let candidate_tree = BackwardState.Tree.add_breadcrumbs features_to_attach candidate_tree in
       let number_of_paths =
         BackwardState.Tree.fold
           BackwardState.Tree.Path
@@ -1445,7 +1408,7 @@ let extract_tito_and_sink_models define ~is_constructor ~resolution ~existing_ba
       in
       if number_of_paths > TaintConfiguration.maximum_tito_leaves then
         BackwardState.Tree.collapse_to
-          ~transform:(BackwardTaint.add_features Features.widen_broadening)
+          ~transform:(BackwardTaint.add_breadcrumbs Features.widen_broadening)
           ~depth:0
           candidate_tree
       else
@@ -1477,10 +1440,7 @@ let extract_tito_and_sink_models define ~is_constructor ~resolution ~existing_ba
           ~attach_to_kind:Sinks.Attach
           existing_backward.TaintResult.Backward.sink_taint
       in
-      if not (Features.SimpleSet.is_bottom features_to_attach) then
-        BackwardState.Tree.transform Features.SimpleSet.Self Add ~f:features_to_attach sink_taint
-      else
-        sink_taint
+      BackwardState.Tree.add_breadcrumbs features_to_attach sink_taint
     in
     TaintResult.Backward.
       {
diff --git a/source/interprocedural_analyses/taint/domains.ml b/source/interprocedural_analyses/taint/domains.ml
@@ -223,6 +223,14 @@ module FlowDetails = struct
     transform Features.TitoPositionSet.Self Map ~f:(fun _ -> Features.TitoPositionSet.bottom)
 
 
+  let add_tito_position position = transform Features.TitoPositionSet.Element Add ~f:position
+
+  let add_breadcrumb breadcrumb = transform Features.SimpleSet.Element Add ~f:breadcrumb
+
+  let add_breadcrumbs breadcrumbs = transform Features.SimpleSet.Self Add ~f:breadcrumbs
+
+  let features = get Slots.SimpleFeature
+
   let product_pp = pp (* shadow *)
 
   let pp formatter = Format.fprintf formatter "FlowDetails(%a)" product_pp
@@ -252,7 +260,13 @@ module type TAINT_DOMAIN = sig
 
   val trace_info : TraceInfo.t Abstract.Domain.part
 
-  val add_features : Features.SimpleSet.t -> t -> t
+  val add_breadcrumb : Features.Simple.t -> t -> t
+
+  val add_breadcrumbs : Features.SimpleSet.t -> t -> t
+
+  val breadcrumbs : t -> Features.SimpleSet.t
+
+  val features : t -> Features.SimpleSet.t
 
   val transform_on_widening_collapse : t -> t
 
@@ -384,12 +398,11 @@ end = struct
           |> Features.FirstField.to_json
         in
         ( List.concat [first_index_breadcrumbs; first_field_breadcrumbs; breadcrumbs],
-          FlowDetails.(
-            fold
-              Features.ReturnAccessPathSet.Element
-              ~f:gather_return_access_path
-              ~init:leaves
-              features) )
+          FlowDetails.fold
+            Features.ReturnAccessPathSet.Element
+            ~f:gather_return_access_path
+            ~init:leaves
+            features )
       in
       let tito_positions =
         FlowDetails.get FlowDetails.Slots.TitoPosition features
@@ -432,7 +445,27 @@ end = struct
     create_json ~trace_info_to_json:(TraceInfo.to_external_json ~filename_lookup)
 
 
-  let add_features features = transform Features.SimpleSet.Self Add ~f:features
+  let add_breadcrumb breadcrumb = transform Features.SimpleSet.Element Add ~f:breadcrumb
+
+  let add_breadcrumbs breadcrumbs taint =
+    if Features.SimpleSet.is_bottom breadcrumbs || Features.SimpleSet.is_empty breadcrumbs then
+      taint
+    else
+      transform Features.SimpleSet.Self Add ~f:breadcrumbs taint
+
+
+  let features taint =
+    let gather_features to_add features = Features.SimpleSet.add_set features ~to_add in
+    fold Features.SimpleSet.Self ~f:gather_features ~init:Features.SimpleSet.bottom taint
+
+
+  let breadcrumbs taint =
+    fold
+      Features.SimpleSet.Self
+      ~f:Features.gather_breadcrumbs
+      ~init:Features.SimpleSet.bottom
+      taint
+
 
   let transform_on_widening_collapse =
     (* using an always-feature here would break the widening invariant: a <= a widen b *)
@@ -444,7 +477,7 @@ end = struct
           { element = Simple.Breadcrumb Breadcrumb.IssueBroadening; in_under = false };
         ]
     in
-    add_features broadening
+    add_breadcrumbs broadening
 
 
   let prune_maximum_length maximum_length =
@@ -571,13 +604,31 @@ module MakeTaintTree (Taint : TAINT_DOMAIN) () = struct
     |> collapse ~transform:Fn.id
 
 
-  let get_all_features taint_tree =
-    let gather_features to_add features = Features.SimpleSet.add_set features ~to_add in
-    fold Features.SimpleSet.Self ~f:gather_features ~init:Features.SimpleSet.bottom taint_tree
+  let breadcrumbs taint_tree =
+    fold
+      Features.SimpleSet.Self
+      ~f:Features.gather_breadcrumbs
+      ~init:Features.SimpleSet.bottom
+      taint_tree
 end
 
 module MakeTaintEnvironment (Taint : TAINT_DOMAIN) () = struct
-  module Tree = MakeTaintTree (Taint) ()
+  module Tree = struct
+    include MakeTaintTree (Taint) ()
+
+    let add_breadcrumb breadcrumb = transform Features.SimpleSet.Element Add ~f:breadcrumb
+
+    let add_breadcrumbs breadcrumbs taint_tree =
+      if Features.SimpleSet.is_bottom breadcrumbs || Features.SimpleSet.is_empty breadcrumbs then
+        taint_tree
+      else
+        transform Features.SimpleSet.Self Add ~f:breadcrumbs taint_tree
+
+
+    let features taint_tree =
+      let gather_features to_add features = Features.SimpleSet.add_set features ~to_add in
+      fold Features.SimpleSet.Self ~f:gather_features ~init:Features.SimpleSet.bottom taint_tree
+  end
 
   include
     Abstract.MapDomain.Make
@@ -646,12 +697,20 @@ module MakeTaintEnvironment (Taint : TAINT_DOMAIN) () = struct
 
   let roots environment = fold Key ~f:List.cons ~init:[] environment
 
+  let add_breadcrumb breadcrumb = transform Features.SimpleSet.Element Add ~f:breadcrumb
+
+  let add_breadcrumbs breadcrumbs taint_tree =
+    if Features.SimpleSet.is_bottom breadcrumbs || Features.SimpleSet.is_empty breadcrumbs then
+      taint_tree
+    else
+      transform Features.SimpleSet.Self Add ~f:breadcrumbs taint_tree
+
+
   let extract_features_to_attach ~root ~attach_to_kind taint =
-    let gather_features to_add features = Features.SimpleSet.add_set features ~to_add in
     read ~root ~path:[] taint
     |> Tree.transform Taint.kind Filter ~f:(Taint.equal_kind attach_to_kind)
     |> Tree.collapse ~transform:Fn.id
-    |> Taint.fold Features.SimpleSet.Self ~f:gather_features ~init:Features.SimpleSet.bottom
+    |> Taint.features
 end
 
 module ForwardState = MakeTaintEnvironment (ForwardTaint) ()
diff --git a/source/interprocedural_analyses/taint/flow.ml b/source/interprocedural_analyses/taint/flow.ml
@@ -56,7 +56,8 @@ let generate_source_sink_matches ~location ~source_tree ~sink_tree =
   let make_source_sink_matches (path, sink_taint) matches =
     let source_taint =
       ForwardState.Tree.read path source_tree
-      |> ForwardState.Tree.collapse ~transform:(ForwardTaint.add_features Features.issue_broadening)
+      |> ForwardState.Tree.collapse
+           ~transform:(ForwardTaint.add_breadcrumbs Features.issue_broadening)
     in
     if ForwardTaint.is_bottom source_taint then
       matches
@@ -327,7 +328,7 @@ let code_metadata () =
 let compute_triggered_sinks ~triggered_sinks ~location ~source_tree ~sink_tree =
   let partial_sinks_to_taint =
     BackwardState.Tree.collapse
-      ~transform:(BackwardTaint.add_features Features.issue_broadening)
+      ~transform:(BackwardTaint.add_breadcrumbs Features.issue_broadening)
       sink_tree
     |> BackwardTaint.partition BackwardTaint.kind ByFilter ~f:(function
            | Sinks.PartialSink { Sinks.kind; label } -> Some { Sinks.kind; label }
diff --git a/source/interprocedural_analyses/taint/forwardAnalysis.ml b/source/interprocedural_analyses/taint/forwardAnalysis.ml
