From 618cd8dc90f9485625e5beca1b2419a767c352ec Mon Sep 17 00:00:00 2001
From: Maxime Arthaud <arthaud@meta.com>
Date: Wed, 9 Aug 2023 02:24:16 -0700
Subject: [PATCH] Add tests for the parsing of ViaValueOf and ViaTypeOf

Summary: There are use cases of ViaValueOf and ViaTypeOf that we were not tested in our model parser tests. This diff fixes that.

Reviewed By: tianhan0

Differential Revision: D48160091

fbshipit-source-id: 8d1bdd969f64d41f24f5b8516920deb7d13a19d6
---
 .../taint/test/modelTest.ml                   | 88 ++++++++++++++++++-
 1 file changed, 86 insertions(+), 2 deletions(-)

diff --git a/source/interprocedural_analyses/taint/test/modelTest.ml b/source/interprocedural_analyses/taint/test/modelTest.ml
index 3b61446bc37..e7759046f5e 100644
--- a/source/interprocedural_analyses/taint/test/modelTest.ml
+++ b/source/interprocedural_analyses/taint/test/modelTest.ml
@@ -574,6 +574,26 @@ let test_source_models context =
     ~model_source:"def test.f() -> TaintSource[Test, ViaValueOf[x]]: ..."
     ~expect:[outcome ~kind:`Function ~returns:[Sources.NamedSource "Test"] "test.f"]
     ();
+  assert_model
+    ~source:"def f(x: int): ..."
+    ~model_source:{|def test.f(x) -> TaintSource[Test, ViaValueOf[x, WithTag["tag"]]]: ...|}
+    ~expect:[outcome ~kind:`Function ~returns:[Sources.NamedSource "Test"] "test.f"]
+    ();
+  assert_model
+    ~source:"def f(x: int): ..."
+    ~model_source:"def test.f(x) -> TaintSource[Test, ViaTypeOf[x]]: ..."
+    ~expect:[outcome ~kind:`Function ~returns:[Sources.NamedSource "Test"] "test.f"]
+    ();
+  assert_model
+    ~source:"def f(x: int): ..."
+    ~model_source:"def test.f() -> TaintSource[Test, ViaTypeOf[x]]: ..."
+    ~expect:[outcome ~kind:`Function ~returns:[Sources.NamedSource "Test"] "test.f"]
+    ();
+  assert_model
+    ~source:"def f(x: int): ..."
+    ~model_source:{|def test.f(x) -> TaintSource[Test, ViaTypeOf[x, WithTag["tag"]]]: ...|}
+    ~expect:[outcome ~kind:`Function ~returns:[Sources.NamedSource "Test"] "test.f"]
+    ();
   assert_model
     ~source:
       {|
@@ -1713,8 +1733,12 @@ let test_sink_models context =
     ();
   assert_model
     ~model_source:
-      "def test.sink(parameter0: TaintSink[Test], parameter1: TaintSink[Test, \
-       ViaValueOf[parameter0]]): ..."
+      {|
+      def test.sink(
+        parameter0: TaintSink[Test],
+        parameter1: TaintSink[Test, ViaValueOf[parameter0]]
+      ): ...
+    |}
     ~expect:
       [
         outcome
@@ -1727,6 +1751,54 @@ let test_sink_models context =
           "test.sink";
       ]
     ();
+  assert_model
+    ~model_source:
+      {|
+      def test.sink(
+        parameter0,
+        parameter1: TaintSink[Test, ViaValueOf[parameter0, WithTag["tag"]]]
+      ): ...
+    |}
+    ~expect:
+      [
+        outcome
+          ~kind:`Function
+          ~sink_parameters:[{ name = "parameter1"; sinks = [Sinks.NamedSink "Test"] }]
+          "test.sink";
+      ]
+    ();
+  assert_model
+    ~model_source:
+      {|
+      def test.sink(
+        parameter0,
+        parameter1: TaintSink[Test, ViaTypeOf[parameter0]]
+      ): ...
+    |}
+    ~expect:
+      [
+        outcome
+          ~kind:`Function
+          ~sink_parameters:[{ name = "parameter1"; sinks = [Sinks.NamedSink "Test"] }]
+          "test.sink";
+      ]
+    ();
+  assert_model
+    ~model_source:
+      {|
+      def test.sink(
+        parameter0,
+        parameter1: TaintSink[Test, ViaTypeOf[parameter0, WithTag["tag"]]]
+      ): ...
+    |}
+    ~expect:
+      [
+        outcome
+          ~kind:`Function
+          ~sink_parameters:[{ name = "parameter1"; sinks = [Sinks.NamedSink "Test"] }]
+          "test.sink";
+      ]
+    ();
   assert_model
     ~model_source:"def test.xss(parameter: TaintSink[XSS]): ..."
     ~expect:
@@ -4697,6 +4769,18 @@ Unexpected statement: `food(y)`
       test.C.x: TaintInTaintOut[ViaTypeOf] = ...
     |}
     ();
+  (* TODO(T159692443): support ViaTypeOf and WithTag on attributes *)
+  assert_invalid_model
+    ~source:{|
+      class C:
+        x: int = 0
+      |}
+    ~model_source:{|
+      test.C.x: ViaTypeOf[WithTag["tag"]] = ...
+    |}
+    ~expect:
+      {|`ViaTypeOf[WithTag["tag"]]` is an invalid taint annotation: Unsupported annotation for attributes|}
+    ();
   assert_valid_model
     ~source:{|
       class C: