support first-index for tito

Reviewed By: dkgi Differential Revision: D23876747 fbshipit-source-id: babf337593f4de0f233a4e826c877242b552011e
facebook · Sep 24, 2020 · 9f7ad43 · 9f7ad43
1 parent d2debe6
commit 9f7ad43
Show file tree

Hide file tree

Showing 14 changed files with 334 additions and 74 deletions.
diff --git a/interprocedural_analyses/taint/backwardAnalysis.ml b/interprocedural_analyses/taint/backwardAnalysis.ml
@@ -50,6 +50,15 @@ module AnalysisInstance (FunctionContext : FUNCTION_CONTEXT) = struct
       Log.log ~section:`Taint format
 
 
+  let add_first_index index indices =
+    if Features.FirstIndexSet.is_bottom indices then
+      Features.to_first_name index
+      >>| Features.FirstIndexSet.singleton
+      |> Option.value ~default:Features.FirstIndexSet.bottom
+    else
+      indices
+
+
   (* This is where we can observe access paths reaching into LocalReturn and record the extraneous
      paths for more precise tito. *)
   let initial_taint =
@@ -544,11 +553,26 @@ module AnalysisInstance (FunctionContext : FUNCTION_CONTEXT) = struct
               else
                 indirect_targets, taint
             in
+            let arguments = Option.to_list receiver_option @ arguments in
+            (* Add index breadcrumb if appropriate. *)
+            let taint =
+              if not (String.equal attribute "get") then
+                taint
+              else
+                match arguments with
+                | _ :: index :: _ ->
+                    let label = get_index index.value in
+                    BackwardState.Tree.transform
+                      BackwardTaint.first_indices
+                      Abstract.Domain.(Map (add_first_index label))
+                      taint
+                | _ -> taint
+            in
             apply_call_targets
               ~resolution
               ~call_expression:(Expression.Call call_expression)
               location
-              (Option.to_list receiver_option @ arguments)
+              arguments
               state
               taint
               indirect_targets
@@ -692,7 +716,13 @@ module AnalysisInstance (FunctionContext : FUNCTION_CONTEXT) = struct
        arguments = [{ Call.Argument.value = argument_value; _ }];
       } ->
           let index = AccessPath.get_index argument_value in
-          let taint = BackwardState.Tree.prepend [index] taint in
+          let taint =
+            BackwardState.Tree.prepend [index] taint
+            |> BackwardState.Tree.transform
+                 BackwardTaint.first_indices
+                 Abstract.Domain.(Map (add_first_index index))
+          in
+
           analyze_expression ~resolution ~taint ~state ~expression:base
       (* Special case x.__next__() as being a random index access (this pattern is the desugaring of
          `for element in x`). *)

diff --git a/interprocedural_analyses/taint/features.ml b/interprocedural_analyses/taint/features.ml
@@ -220,3 +220,15 @@ let gather_breadcrumbs feature breadcrumbs =
 let is_breadcrumb = function
   | { Abstract.OverUnderSetDomain.element = Simple.Breadcrumb _; _ } -> true
   | _ -> false
+
+
+let number_regexp = Str.regexp "[0-9]+"
+
+let is_numeric name = Str.string_match number_regexp name 0
+
+let to_first_name label =
+  match label with
+  | Abstract.TreeDomain.Label.Field name when is_numeric name -> Some "<numeric>"
+  | Field name -> Some name
+  | DictionaryKeys -> None
+  | Any -> Some "<unknown>"
diff --git a/interprocedural_analyses/taint/flow.ml b/interprocedural_analyses/taint/flow.ml
@@ -97,19 +97,41 @@ let get_issue_features { source_taint; sink_taint } =
     Features.SimpleSet.sequence_join source_features sink_features
   in
   let first_indices =
-    ForwardTaint.fold
-      Features.FirstIndexSet.Self
-      ~f:Features.FirstIndexSet.join
-      ~init:Features.FirstIndexSet.bottom
-      source_taint
+    let source_indices =
+      ForwardTaint.fold
+        Features.FirstIndexSet.Self
+        ~f:Features.FirstIndexSet.join
+        ~init:Features.FirstIndexSet.bottom
+        source_taint
+    in
+    let sink_indices =
+      BackwardTaint.fold
+        Features.FirstIndexSet.Self
+        ~f:Features.FirstIndexSet.join
+        ~init:Features.FirstIndexSet.bottom
+        sink_taint
+    in
+
+    Features.FirstIndexSet.join source_indices sink_indices
   in
   let first_fields =
-    ForwardTaint.fold
-      Features.FirstFieldSet.Self
-      ~f:Features.FirstFieldSet.join
-      ~init:Features.FirstFieldSet.bottom
-      source_taint
+    let source_fields =
+      ForwardTaint.fold
+        Features.FirstFieldSet.Self
+        ~f:Features.FirstFieldSet.join
+        ~init:Features.FirstFieldSet.bottom
+        source_taint
+    in
+    let sink_fields =
+      BackwardTaint.fold
+        Features.FirstFieldSet.Self
+        ~f:Features.FirstFieldSet.join
+        ~init:Features.FirstFieldSet.bottom
+        sink_taint
+    in
+    Features.FirstFieldSet.join source_fields sink_fields
   in
+
   { simple; first_indices; first_fields }
 
 

diff --git a/interprocedural_analyses/taint/forwardAnalysis.ml b/interprocedural_analyses/taint/forwardAnalysis.ml
@@ -56,10 +56,6 @@ module type FUNCTION_CONTEXT = sig
   val add_triggered_sinks : location:Location.t -> triggered_sinks:(Root.t * Sinks.t) list -> unit
 end
 
-let number_regexp = Str.regexp "[0-9]+"
-
-let is_numeric name = Str.string_match number_regexp name 0
-
 let ( |>> ) (taint, state) f = f taint, state
 
 module AnalysisInstance (FunctionContext : FUNCTION_CONTEXT) = struct
@@ -99,12 +95,9 @@ module AnalysisInstance (FunctionContext : FUNCTION_CONTEXT) = struct
 
     let add_first_index index indices =
       if Features.FirstIndexSet.is_bottom indices then
-        match index with
-        | Abstract.TreeDomain.Label.Field name when is_numeric name ->
-            Features.FirstIndexSet.singleton "<numeric>"
-        | Field name -> Features.FirstIndexSet.singleton name
-        | DictionaryKeys -> Features.FirstIndexSet.bottom
-        | Any -> Features.FirstIndexSet.singleton "<unknown>"
+        Features.to_first_name index
+        >>| Features.FirstIndexSet.singleton
+        |> Option.value ~default:Features.FirstIndexSet.bottom
       else
         indices
 
@@ -1491,7 +1484,6 @@ let run ~environment ~qualifier ~define ~existing_model =
       let triggered, candidates =
         Flow.compute_triggered_sinks ~triggered_sinks ~location ~source_tree ~sink_tree
       in
-
       List.iter triggered ~f:(fun sink ->
           Hash_set.add triggered_sinks (Sinks.show_partial_sink sink));
       List.iter candidates ~f:add_flow_candidate

diff --git a/interprocedural_analyses/taint/test/integration/constructors.py.models b/interprocedural_analyses/taint/test/integration/constructors.py.models
@@ -629,7 +629,11 @@
             "decl": null,
             "tito": [ { "line": 32, "start": 41, "end": 50 } ],
             "leaves": [ { "kind": "LocalReturn", "name": "[oauth_token]" } ],
-            "features": [ { "always-via": "tito" } ]
+            "features": [
+              { "has": "first-index" },
+              { "first-index": "<numeric>" },
+              { "always-via": "tito" }
+            ]
           }
         ]
       },
@@ -642,7 +646,11 @@
             "leaves": [
               { "kind": "LocalReturn", "name": "[oauth_token_secret]" }
             ],
-            "features": [ { "always-via": "tito" } ]
+            "features": [
+              { "has": "first-index" },
+              { "first-index": "<numeric>" },
+              { "always-via": "tito" }
+            ]
           }
         ]
       }

diff --git a/interprocedural_analyses/taint/test/integration/dictionary.py b/interprocedural_analyses/taint/test/integration/dictionary.py
@@ -218,3 +218,13 @@ def test_service_with_dict():
 def test_service_with_mapping():
     service = Service()
     __test_sink(service.async_get_mapping(__test_source()))
+
+
+def tito_with_index(d: Dict[str, str]) -> str:
+    result = d["a"]
+    return result
+
+
+def test_index_from_tito():
+    d = {"a": __test_source(), "b": __test_source()}
+    __test_sink(tito_with_index(d))
diff --git a/interprocedural_analyses/taint/test/integration/dictionary.py.cg b/interprocedural_analyses/taint/test/integration/dictionary.py.cg
@@ -31,10 +31,12 @@ dictionary.return_comprehension_with_untainted_keys (fun) -> [__test_source (fun
 dictionary.return_tito_literally (fun) -> []
 dictionary.sink_dictionary_through_keys (fun) -> [__test_sink (fun)]
 dictionary.tainted_setitem (fun) -> [__test_source (fun) dictionary.SpecialSetitemDict::__setitem__ (method)]
+dictionary.test_index_from_tito (fun) -> [__test_sink (fun) __test_source (fun) dictionary.tito_with_index (fun)]
 dictionary.test_keys_and_values (fun) -> [__test_sink (fun) __test_source (fun) tuple::__getitem__ (method) typing.Mapping::keys (method) typing.Mapping::keys (method) typing.Mapping::values (method) typing.Mapping::values (method)]
 dictionary.test_service_with_dict (fun) -> [__test_sink (fun) __test_source (fun) dictionary.Service::async_get_dict (method) object::__init__ (method) object::__new__ (method)]
 dictionary.test_service_with_mapping (fun) -> [__test_sink (fun) __test_source (fun) dictionary.Service::async_get_mapping (method) object::__init__ (method) object::__new__ (method)]
 dictionary.test_with_issue_in_dict_comprehension (fun) -> [__test_sink (fun) __test_source (fun)]
+dictionary.tito_with_index (fun) -> [dict::__getitem__ (method)]
 dictionary.to_map (fun) -> []
 dictionary.update_dictionary_indirectly (fun) -> [dict::update (method)]
 dictionary.update_parameter (fun) -> []

diff --git a/interprocedural_analyses/taint/test/integration/dictionary.py.models b/interprocedural_analyses/taint/test/integration/dictionary.py.models
@@ -381,6 +381,68 @@
     ]
   }
 }
+{
+  "kind": "issue",
+  "data": {
+    "callable": "dictionary.test_index_from_tito",
+    "callable_line": 228,
+    "code": 5002,
+    "line": 230,
+    "start": 16,
+    "end": 34,
+    "filename": "dictionary.py",
+    "message": "Data from [Test] source(s) may reach [Test] sink(s)",
+    "traces": [
+      {
+        "name": "forward",
+        "roots": [
+          {
+            "root": {
+              "filename": "dictionary.py",
+              "line": 229,
+              "start": 14,
+              "end": 29
+            },
+            "tito": [ { "line": 230, "start": 32, "end": 33 } ],
+            "leaves": [
+              {
+                "kind": "Test",
+                "name": "__test_source",
+                "on_all_flows": true
+              }
+            ],
+            "features": [
+              { "always-via": "tito" },
+              { "always-via": "special_source" }
+            ]
+          }
+        ]
+      },
+      {
+        "name": "backward",
+        "roots": [
+          {
+            "root": {
+              "filename": "dictionary.py",
+              "line": 230,
+              "start": 16,
+              "end": 34
+            },
+            "leaves": [
+              { "kind": "Test", "name": "__test_sink", "on_all_flows": true }
+            ],
+            "features": [ { "always-via": "special_sink" } ]
+          }
+        ]
+      }
+    ],
+    "features": [
+      { "always-via": "tito" },
+      { "always-via": "special_source" },
+      { "always-via": "special_sink" }
+    ]
+  }
+}
 {
   "kind": "issue",
   "data": {
@@ -880,7 +942,12 @@
               { "line": 207, "start": 16, "end": 58 }
             ],
             "leaves": [ { "kind": "LocalReturn", "name": "" } ],
-            "features": [ { "always-via": "tito" }, { "via": "obscure" } ]
+            "features": [
+              { "has": "first-index" },
+              { "first-index": "<unknown>" },
+              { "always-via": "tito" },
+              { "via": "obscure" }
+            ]
           }
         ]
       }
@@ -925,7 +992,12 @@
               { "line": 210, "start": 15, "end": 65 }
             ],
             "leaves": [ { "kind": "LocalReturn", "name": "" } ],
-            "features": [ { "always-via": "tito" }, { "via": "obscure" } ]
+            "features": [
+              { "has": "first-index" },
+              { "first-index": "<unknown>" },
+              { "always-via": "tito" },
+              { "via": "obscure" }
+            ]
           }
         ]
       }
@@ -1516,6 +1588,26 @@
     ]
   }
 }
+{
+  "kind": "model",
+  "data": {
+    "callable": "dictionary.tito_with_index",
+    "sources": [],
+    "sinks": [],
+    "tito": [
+      {
+        "port": "formal(d)[a]",
+        "taint": [
+          {
+            "decl": null,
+            "leaves": [ { "kind": "LocalReturn", "name": "" } ],
+            "features": [ { "has": "first-index" }, { "first-index": "a" } ]
+          }
+        ]
+      }
+    ]
+  }
+}
 {
   "kind": "model",
   "data": {

diff --git a/interprocedural_analyses/taint/test/integration/long_access_path_taint.py.models b/interprocedural_analyses/taint/test/integration/long_access_path_taint.py.models
@@ -33,7 +33,12 @@
               { "kind": "LocalReturn", "name": "[timestamp]" },
               { "kind": "LocalReturn", "name": "[app_id]" }
             ],
-            "features": [ { "always-via": "tito" } ]
+            "features": [
+              { "has": "first-index" },
+              { "first-index": "app_id" },
+              { "first-index": "timestamp" },
+              { "always-via": "tito" }
+            ]
           }
         ]
       },
@@ -93,7 +98,12 @@
               { "kind": "LocalReturn", "name": "[timestamp]" },
               { "kind": "LocalReturn", "name": "[app_id]" }
             ],
-            "features": [ { "always-via": "tito" } ]
+            "features": [
+              { "has": "first-index" },
+              { "first-index": "app_id" },
+              { "first-index": "timestamp" },
+              { "always-via": "tito" }
+            ]
           }
         ]
       },
@@ -157,7 +167,12 @@
               { "kind": "LocalReturn", "name": "[timestamp]" },
               { "kind": "LocalReturn", "name": "[app_id]" }
             ],
-            "features": [ { "always-via": "tito" } ]
+            "features": [
+              { "has": "first-index" },
+              { "first-index": "app_id" },
+              { "first-index": "timestamp" },
+              { "always-via": "tito" }
+            ]
           }
         ]
       },