Add kind-specific breadcrumbs for via:transform-tito-depth:N

Summary: # Problem The `via:transform-tito-depth:N` is kind-specific: it should only be attached to flows that went through the transform. However, right now, breadcrumbs are stored in `LocalTaint`, which means that breadcrumbs for a given call info are all mixed together, regardless of taint kind. This means we attach `via:transform-tito-depth:N` to taint that didn't go through the transform, making it confusing. # Solution We could move local breadcrumbs down in the taint representation (in `Frame`), so that their are tied to a specific kind. We have tried this and it lead to a performance regression of 6%. Instead, what this diff does is have a special set for kind-specific breadcrumbs that are stored in `Frame`. This is then used for `via:transform-tito-depth:N`. This way, we get the right semantic, and preserve performance. Reviewed By: tianhan0 Differential Revision: D58468250 fbshipit-source-id: 4889f78c5d57b211ec466ac7794fd0661751bb0f
facebook · Jun 21, 2024 · 9be972d · 9be972d
1 parent b0a10c0
commit 9be972d
Show file tree

Hide file tree

Showing 9 changed files with 253 additions and 193 deletions.
diff --git a/source/interprocedural_analyses/taint/backwardAnalysis.ml b/source/interprocedural_analyses/taint/backwardAnalysis.ml
@@ -517,9 +517,10 @@ module State (FunctionContext : FUNCTION_CONTEXT) = struct
             else
               let breadcrumb = CallModel.transform_tito_depth_breadcrumb tito_taint in
               let taint_to_propagate =
-                BackwardState.Tree.add_local_breadcrumb
-                  ~add_on_tito:false
-                  breadcrumb
+                BackwardState.Tree.transform_non_tito
+                  Features.LocalKindSpecificBreadcrumbSet.Self
+                  Map
+                  ~f:(Features.BreadcrumbSet.add breadcrumb)
                   taint_to_propagate
               in
               add_extra_traces_for_transforms

diff --git a/source/interprocedural_analyses/taint/callModel.ml b/source/interprocedural_analyses/taint/callModel.ml
@@ -293,7 +293,7 @@ let taint_in_taint_out_mapping_for_argument
       else if SanitizeTransformSet.is_empty obscure_sanitize then
         let tito =
           Domains.local_return_frame ~output_path:[] ~collapse_depth:0
-          |> Frame.update Frame.Slots.Breadcrumb obscure_breadcrumbs
+          |> Frame.update Frame.Slots.PropagatedBreadcrumb obscure_breadcrumbs
           |> BackwardTaint.singleton CallInfo.declaration output_kind
           |> BackwardState.Tree.create_leaf
         in
@@ -309,7 +309,7 @@ let taint_in_taint_out_mapping_for_argument
         in
         let tito =
           Domains.local_return_frame ~output_path:[] ~collapse_depth:0
-          |> Frame.update Frame.Slots.Breadcrumb obscure_breadcrumbs
+          |> Frame.update Frame.Slots.PropagatedBreadcrumb obscure_breadcrumbs
           |> BackwardTaint.singleton CallInfo.Tito output_kind
           |> BackwardState.Tree.create_leaf
         in

diff --git a/source/interprocedural_analyses/taint/domains.ml b/source/interprocedural_analyses/taint/domains.ml
@@ -403,7 +403,8 @@ module Frame = struct
     let name = "frame"
 
     type 'a slot =
-      | Breadcrumb : Features.PropagatedBreadcrumbSet.t slot
+      | PropagatedBreadcrumb : Features.PropagatedBreadcrumbSet.t slot
+      | LocalKindSpecificBreadcrumb : Features.LocalKindSpecificBreadcrumbSet.t slot
       | ViaFeature : Features.ViaFeatureSet.t slot
       | ReturnAccessPath : Features.ReturnAccessPathTree.t slot
       | TraceLength : TraceLength.t slot
@@ -413,11 +414,12 @@ module Frame = struct
       | ExtraTraceFirstHopSet : ExtraTraceFirstHop.Set.t slot
 
     (* Must be consistent with above variants *)
-    let slots = 8
+    let slots = 9
 
     let slot_name (type a) (slot : a slot) =
       match slot with
-      | Breadcrumb -> "Breadcrumb"
+      | PropagatedBreadcrumb -> "PropagatedBreadcrumb"
+      | LocalKindSpecificBreadcrumb -> "LocalKindSpecificBreadcrumb"
       | ViaFeature -> "ViaFeature"
       | ReturnAccessPath -> "ReturnAccessPath"
       | TraceLength -> "TraceLength"
@@ -429,7 +431,10 @@ module Frame = struct
 
     let slot_domain (type a) (slot : a slot) =
       match slot with
-      | Breadcrumb -> (module Features.PropagatedBreadcrumbSet : Abstract.Domain.S with type t = a)
+      | PropagatedBreadcrumb ->
+          (module Features.PropagatedBreadcrumbSet : Abstract.Domain.S with type t = a)
+      | LocalKindSpecificBreadcrumb ->
+          (module Features.LocalKindSpecificBreadcrumbSet : Abstract.Domain.S with type t = a)
       | ViaFeature -> (module Features.ViaFeatureSet : Abstract.Domain.S with type t = a)
       | ReturnAccessPath ->
           (module Features.ReturnAccessPathTree : Abstract.Domain.S with type t = a)
@@ -449,6 +454,7 @@ module Frame = struct
     create
       [
         Part (Features.PropagatedBreadcrumbSet.Self, Features.BreadcrumbSet.empty);
+        Part (Features.LocalKindSpecificBreadcrumbSet.Self, Features.BreadcrumbSet.empty);
         Part (TraceLength.Self, 0);
       ]
 
@@ -798,12 +804,12 @@ end = struct
         (key, `List list) :: assoc
     in
     let open Features in
+    let breadcrumb_to_json { Abstract.OverUnderSetDomain.element; in_under } breadcrumbs =
+      let element = BreadcrumbInterned.unintern element in
+      let json = Breadcrumb.to_json element ~on_all_paths:in_under in
+      json :: breadcrumbs
+    in
     let breadcrumbs_to_json ~breadcrumbs ~first_indices ~first_fields =
-      let breadcrumb_to_json { Abstract.OverUnderSetDomain.element; in_under } breadcrumbs =
-        let element = BreadcrumbInterned.unintern element in
-        let json = Breadcrumb.to_json element ~on_all_paths:in_under in
-        json :: breadcrumbs
-      in
       let breadcrumbs =
         BreadcrumbSet.fold BreadcrumbSet.ElementAndUnder ~f:breadcrumb_to_json ~init:[] breadcrumbs
       in
@@ -883,12 +889,20 @@ end = struct
 
         let breadcrumbs =
           breadcrumbs_to_json
-            ~breadcrumbs:(Frame.get Frame.Slots.Breadcrumb frame)
+            ~breadcrumbs:(Frame.get Frame.Slots.PropagatedBreadcrumb frame)
             ~first_indices:(Frame.get Frame.Slots.FirstIndex frame)
             ~first_fields:(Frame.get Frame.Slots.FirstField frame)
         in
+        (* Those are "propagated breadcrumbs", but stored as "features" for backward
+           compatibility. *)
         let json = cons_if_non_empty "features" breadcrumbs json in
 
+        let local_breadcrumbs =
+          Frame.get Frame.Slots.LocalKindSpecificBreadcrumb frame
+          |> BreadcrumbSet.fold BreadcrumbSet.ElementAndUnder ~f:breadcrumb_to_json ~init:[]
+        in
+        let json = cons_if_non_empty "local_features" local_breadcrumbs json in
+
         let extra_traces =
           Frame.get Frame.Slots.ExtraTraceFirstHopSet frame
           |> ExtraTraceFirstHop.Set.to_json ~expand_overrides ~is_valid_callee
@@ -982,9 +996,21 @@ end = struct
       taint
 
 
-  let get_features ~frame_slot ~local_slot ~bottom ~join ~sequence_join taint =
+  let get_features
+      ~propagated_slot
+      ~local_kind_specific_slot
+      ~local_slot
+      ~bottom
+      ~join
+      ~sequence_join
+      taint
+    =
     let local_taint_features local_taint sofar =
-      let frame_features frame sofar = Frame.get frame_slot frame |> join sofar in
+      let frame_features frame sofar =
+        let propagated_features = Frame.get propagated_slot frame in
+        let local_kind_specific_features = Frame.get local_kind_specific_slot frame in
+        sequence_join propagated_features local_kind_specific_features |> join sofar
+      in
       let features = LocalTaintDomain.fold Frame.Self local_taint ~init:bottom ~f:frame_features in
       let features = LocalTaintDomain.get local_slot local_taint |> sequence_join features in
       join sofar features
@@ -994,7 +1020,8 @@ end = struct
 
   let accumulated_breadcrumbs taint =
     get_features
-      ~frame_slot:Frame.Slots.Breadcrumb
+      ~propagated_slot:Frame.Slots.PropagatedBreadcrumb
+      ~local_kind_specific_slot:Frame.Slots.LocalKindSpecificBreadcrumb
       ~local_slot:LocalTaintDomain.Slots.Breadcrumb
       ~bottom:Features.BreadcrumbSet.bottom
       ~join:Features.BreadcrumbSet.sequence_join
@@ -1004,19 +1031,20 @@ end = struct
 
   let joined_breadcrumbs taint =
     get_features
-      ~frame_slot:Frame.Slots.Breadcrumb
+      ~propagated_slot:Frame.Slots.PropagatedBreadcrumb
+      ~local_kind_specific_slot:Frame.Slots.LocalKindSpecificBreadcrumb
       ~local_slot:LocalTaintDomain.Slots.Breadcrumb
       ~bottom:Features.BreadcrumbSet.bottom
       ~join:Features.BreadcrumbSet.join
       ~sequence_join:Features.BreadcrumbSet.sequence_join
       taint
 
 
-  let get_first ~frame_slot ~local_slot ~bottom ~join ~sequence_join taint =
+  let get_first ~propagated_slot ~local_slot ~bottom ~join ~sequence_join taint =
     let local_taint_first local_taint sofar =
       let local_first = LocalTaintDomain.get local_slot local_taint in
       let frame_first frame sofar =
-        Frame.get frame_slot frame |> sequence_join local_first |> join sofar
+        Frame.get propagated_slot frame |> sequence_join local_first |> join sofar
       in
       LocalTaintDomain.fold Frame.Self local_taint ~init:sofar ~f:frame_first
     in
@@ -1025,7 +1053,7 @@ end = struct
 
   let first_indices taint =
     get_first
-      ~frame_slot:Frame.Slots.FirstIndex
+      ~propagated_slot:Frame.Slots.FirstIndex
       ~local_slot:LocalTaintDomain.Slots.FirstIndex
       ~bottom:Features.FirstIndexSet.bottom
       ~join:Features.FirstIndexSet.join
@@ -1035,7 +1063,7 @@ end = struct
 
   let first_fields taint =
     get_first
-      ~frame_slot:Frame.Slots.FirstField
+      ~propagated_slot:Frame.Slots.FirstField
       ~local_slot:LocalTaintDomain.Slots.FirstField
       ~bottom:Features.FirstFieldSet.bottom
       ~join:Features.FirstFieldSet.join
@@ -1224,10 +1252,15 @@ end = struct
               frame
           | _ -> Frame.update Frame.Slots.ExtraTraceFirstHopSet ExtraTraceFirstHop.Set.bottom frame
         in
+        let local_breadcrumbs =
+          Frame.get Frame.Slots.LocalKindSpecificBreadcrumb frame
+          |> Features.BreadcrumbSet.sequence_join local_breadcrumbs
+        in
         frame
         |> Frame.update Frame.Slots.ViaFeature Features.ViaFeatureSet.bottom
+        |> Frame.update Frame.Slots.LocalKindSpecificBreadcrumb Features.BreadcrumbSet.empty
         |> Frame.transform
-             Features.BreadcrumbSet.Self
+             Features.PropagatedBreadcrumbSet.Self
              Map
              ~f:(Features.BreadcrumbSet.sequence_join local_breadcrumbs)
         |> Frame.transform
@@ -1851,7 +1884,8 @@ let local_return_frame ~output_path ~collapse_depth =
     [
       Part (TraceLength.Self, 0);
       Part (Features.ReturnAccessPathTree.Path, (output_path, collapse_depth));
-      Part (Features.BreadcrumbSet.Self, Features.BreadcrumbSet.empty);
+      Part (Features.PropagatedBreadcrumbSet.Self, Features.BreadcrumbSet.empty);
+      Part (Features.LocalKindSpecificBreadcrumbSet.Self, Features.BreadcrumbSet.empty);
     ]
 
 

diff --git a/source/interprocedural_analyses/taint/features.ml b/source/interprocedural_analyses/taint/features.ml
@@ -410,12 +410,23 @@ module BreadcrumbSet = Abstract.OverUnderSetDomain.MakeWithSet (struct
   let element_name = BreadcrumbInterned.name
 end)
 
+(* Stores local breadcrumbs (also called trace breadcrumbs) inferred during the analysis of the
+   current function. *)
 module LocalBreadcrumbSet = Abstract.WrapperDomain.Make (struct
   include BreadcrumbSet
 
   let name = "local breadcrumbs"
 end)
 
+(* Stores local breadcrumbs that are kind-specific, unlike `LocalBreadcrumbSet`. For instance,
+   `via:transform-tito-depth:N` should only be emitted on transform kinds. *)
+module LocalKindSpecificBreadcrumbSet = Abstract.WrapperDomain.Make (struct
+  include BreadcrumbSet
+
+  let name = "local kind-specific breadcrumbs"
+end)
+
+(* Stores breadcrumbs propagated from the callee. *)
 module PropagatedBreadcrumbSet = Abstract.WrapperDomain.Make (struct
   include BreadcrumbSet
 

diff --git a/source/interprocedural_analyses/taint/forwardAnalysis.ml b/source/interprocedural_analyses/taint/forwardAnalysis.ml
@@ -552,9 +552,10 @@ module State (FunctionContext : FUNCTION_CONTEXT) = struct
             else
               let breadcrumb = CallModel.transform_tito_depth_breadcrumb tito_taint in
               let taint_to_propagate =
-                ForwardState.Tree.add_local_breadcrumb
-                  ~add_on_tito:false
-                  breadcrumb
+                ForwardState.Tree.transform_non_tito
+                  Features.LocalKindSpecificBreadcrumbSet.Self
+                  Map
+                  ~f:(Features.BreadcrumbSet.add breadcrumb)
                   taint_to_propagate
               in
               add_extra_traces_for_transforms

diff --git a/source/interprocedural_analyses/taint/modelParser.ml b/source/interprocedural_analyses/taint/modelParser.ml
@@ -1234,7 +1234,7 @@ let introduce_sink_taint
     let leaf_taint =
       Frame.initial
       |> Frame.transform Features.LeafNameSet.Self Add ~f:leaf_names
-      |> Frame.transform Features.BreadcrumbSet.Self Add ~f:breadcrumbs
+      |> Frame.transform Features.PropagatedBreadcrumbSet.Self Map ~f:(fun _ -> breadcrumbs)
       |> Frame.transform Features.ViaFeatureSet.Self Add ~f:via_features
       |> transform_trace_length
       |> BackwardTaint.singleton CallInfo.declaration taint_sink_kind
@@ -1299,7 +1299,7 @@ let introduce_taint_in_taint_out
   >>= fun output_path ->
   let tito_result_taint =
     Domains.local_return_frame ~output_path ~collapse_depth
-    |> Frame.transform Features.BreadcrumbSet.Self Add ~f:breadcrumbs
+    |> Frame.transform Features.PropagatedBreadcrumbSet.Self Map ~f:(fun _ -> breadcrumbs)
     |> Frame.transform Features.ViaFeatureSet.Self Add ~f:via_features
     |> BackwardTaint.singleton CallInfo.Tito taint_sink_kind
     |> BackwardState.Tree.create_leaf
@@ -1343,7 +1343,7 @@ let introduce_taint_in_taint_out
     | Sinks.Attach ->
         let attach_taint =
           Frame.initial
-          |> Frame.transform Features.BreadcrumbSet.Self Add ~f:breadcrumbs
+          |> Frame.transform Features.PropagatedBreadcrumbSet.Self Map ~f:(fun _ -> breadcrumbs)
           |> Frame.transform Features.ViaFeatureSet.Self Add ~f:via_features
           |> BackwardTaint.singleton CallInfo.declaration taint_sink_kind
           |> BackwardState.Tree.create_leaf
@@ -1425,7 +1425,7 @@ let introduce_source_taint
       in
       Frame.initial
       |> Frame.transform Features.LeafNameSet.Self Add ~f:leaf_names
-      |> Frame.transform Features.BreadcrumbSet.Self Add ~f:breadcrumbs
+      |> Frame.transform Features.PropagatedBreadcrumbSet.Self Map ~f:(fun _ -> breadcrumbs)
       |> Frame.transform Features.ViaFeatureSet.Self Add ~f:via_features
       |> transform_trace_length
       |> ForwardTaint.singleton CallInfo.declaration taint_source_kind

diff --git a/source/interprocedural_analyses/taint/test/integration/closure_models.py.models b/source/interprocedural_analyses/taint/test/integration/closure_models.py.models
@@ -788,6 +788,9 @@
           {
             "kinds": [
               {
+                "local_features": [
+                  { "always-via": "transform-tito-depth:1" }
+                ],
                 "features": [ { "always-via": "special_source" } ],
                 "leaves": [
                   { "name": "_test_source", "port": "leaf:return" }
@@ -796,7 +799,6 @@
               }
             ],
             "local_features": [
-              { "always-via": "transform-tito-depth:1" },
               { "always-via-feature1-value": "breadcrumb1" },
               { "always-via": "tito" },
               { "always-via": "lambda" }
@@ -876,6 +878,9 @@
           {
             "kinds": [
               {
+                "local_features": [
+                  { "always-via": "transform-tito-depth:1" }
+                ],
                 "features": [ { "always-via": "special_source" } ],
                 "leaves": [
                   { "name": "_test_source", "port": "leaf:return" }
@@ -884,7 +889,6 @@
               }
             ],
             "local_features": [
-              { "always-via": "transform-tito-depth:1" },
               { "always-via-feature1-value": "breadcrumb1" },
               { "always-via": "tito" },
               { "always-via": "lambda" }