@salesforce/afv-skills 1.14.0 → 1.16.0

package/skills/investigating-agentforce-architecture/scripts/tests/test_soql_loader.py

@@ -0,0 +1,501 @@
+"""Tests for load_soql revalidates every substituted string."""
+from __future__ import annotations
+
+import tempfile
+import unittest
+from pathlib import Path
+from unittest import mock
+
+from . import _bootstrap  # noqa: F401 — sys.path setup
+
+
+class LoadSoqlValidationTests(unittest.TestCase):
+    """every param must be revalidated at the substitution boundary."""
+
+    def setUp(self) -> None:
+        # Use a throwaway SOQL_DIR populated per test so we don't depend on
+        # the shipped assets having a specific template name.
+        self._tmpdir = tempfile.TemporaryDirectory()
+        self.soql_dir = Path(self._tmpdir.name)
+        # Patch config.SOQL_DIR BEFORE importing soql_loader so the module
+        # constant reflects the tmpdir. Because soql_loader binds SOQL_DIR
+        # at import, we re-import it fresh under the patch.
+        self._patch = mock.patch("config.SOQL_DIR", self.soql_dir)
+        self._patch.start()
+        # Force a fresh import so load_soql reads the patched SOQL_DIR.
+        import importlib
+        import soql_loader  # type: ignore
+        importlib.reload(soql_loader)
+        self.soql_loader = soql_loader
+
+    def tearDown(self) -> None:
+        self._patch.stop()
+        self._tmpdir.cleanup()
+
+    def _write_template(self, name: str, body: str) -> None:
+        (self.soql_dir / f"{name}.soql").write_text(body)
+
+    # ---- happy path -------------------------------------------------------
+
+    def test_valid_param_substitutes(self):
+        self._write_template(
+            "bot_lookup",
+            "SELECT Id FROM BotDefinition WHERE DeveloperName = '{{NAME}}'",
+        )
+        out = self.soql_loader.load_soql("bot_lookup", NAME="MyAgent")
+        self.assertEqual(
+            out,
+            "SELECT Id FROM BotDefinition WHERE DeveloperName = 'MyAgent'",
+        )
+
+    def test_multiple_valid_params(self):
+        self._write_template(
+            "lookup",
+            "SELECT Id FROM Obj WHERE A='{{A}}' AND B='{{B}}'",
+        )
+        out = self.soql_loader.load_soql("lookup", A="Foo", B="Bar_v2")
+        self.assertIn("A='Foo'", out)
+        self.assertIn("B='Bar_v2'", out)
+
+    # ---- injection attempts must raise ------------------------------------
+
+    def test_injection_quote_or_clause_raises(self):
+        self._write_template("q", "SELECT Id FROM X WHERE Name='{{NAME}}'")
+        with self.assertRaises(self.soql_loader.SoqlParamError) as ctx:
+            self.soql_loader.load_soql("q", NAME="x' OR Id!=null--")
+        self.assertEqual(ctx.exception.key, "NAME")
+
+    def test_injection_drop_table_raises(self):
+        self._write_template("q", "SELECT Id FROM X WHERE Name='{{NAME}}'")
+        with self.assertRaises(self.soql_loader.SoqlParamError):
+            self.soql_loader.load_soql("q", NAME="'; DROP TABLE x;--")
+
+    def test_injection_or_1eq1_raises(self):
+        self._write_template("q", "SELECT Id FROM X WHERE Name='{{NAME}}'")
+        with self.assertRaises(self.soql_loader.SoqlParamError):
+            self.soql_loader.load_soql("q", NAME="x OR 1=1")
+
+    def test_whitespace_rejected(self):
+        self._write_template("q", "SELECT Id FROM X WHERE Name='{{NAME}}'")
+        with self.assertRaises(self.soql_loader.SoqlParamError):
+            self.soql_loader.load_soql("q", NAME="x OR y")  # space
+
+    def test_dash_rejected(self):
+        # Common SOQL-injection payload prefix.
+        self._write_template("q", "SELECT Id FROM X WHERE Name='{{NAME}}'")
+        with self.assertRaises(self.soql_loader.SoqlParamError):
+            self.soql_loader.load_soql("q", NAME="x-y")
+
+    # ---- type errors ------------------------------------------------------
+
+    def test_non_string_value_raises(self):
+        self._write_template("q", "SELECT {{N}}")
+        with self.assertRaises(self.soql_loader.SoqlParamError):
+            self.soql_loader.load_soql("q", N=42)
+
+    def test_none_value_raises(self):
+        self._write_template("q", "SELECT {{N}}")
+        with self.assertRaises(self.soql_loader.SoqlParamError):
+            self.soql_loader.load_soql("q", N=None)
+
+    def test_empty_string_raises(self):
+        self._write_template("q", "SELECT {{N}}")
+        with self.assertRaises(self.soql_loader.SoqlParamError):
+            self.soql_loader.load_soql("q", N="")
+
+    # ---- single-pass substitution guarantee -------------------------------
+
+    def test_value_containing_other_placeholder_does_not_retrigger(self):
+        """A valid value that contains `{{OTHER}}` must NOT trigger a
+        second substitution pass on OTHER. str.replace is single-pass
+        by contract; we assert that the validator's regex prevents any
+        value from containing `{`, `}`, or whitespace in the first place
+        — so the injection path is closed at validation, not at
+        substitution.
+        """
+        self._write_template("q", "SELECT {{A}} AND {{B}}")
+        # Reject values that carry placeholder syntax — the regex forbids
+        # `{` and `}` (they're not in [A-Za-z0-9_]).
+        with self.assertRaises(self.soql_loader.SoqlParamError):
+            self.soql_loader.load_soql("q", A="{{B}}", B="safe")
+
+    def test_raw_replace_is_single_pass(self):
+        """Belt-and-braces: even if a value slipped through (it cannot,
+        per the regex), Python's str.replace does not recursively scan
+        the output. Simulate by bypassing the validator and confirming
+        str.replace behavior directly.
+        """
+        template = "SELECT {{A}} AND {{B}}"
+        # Substitute A first with a value that "looks like" the B
+        # placeholder. str.replace for B should then replace the
+        # template's own `{{B}}` — but A's embedded `{{B}}` should stay
+        # (Python replaces left-to-right in a single pass; the output
+        # of the A replacement is the NEW string and B's pass runs on
+        # that new string). The assertion below therefore confirms
+        # that B's replacement hits BOTH occurrences (the original and
+        # the one inside A's substituted value) — demonstrating that
+        # str.replace is NOT recursive on the ORIGINAL template alone,
+        # but IS a single linear scan of the full post-A string.
+        step1 = template.replace("{{A}}", "{{B}}")
+        step2 = step1.replace("{{B}}", "replaced")
+        # Both occurrences of `{{B}}` in step1 are replaced in a single
+        # left-to-right pass. This is the property tested for.
+        self.assertEqual(step2, "SELECT replaced AND replaced")
+        # The safety net: if a caller EVER skipped revalidation and
+        # allowed a `{{X}}` to reach substitution, the above behavior
+        # could matter — which is PRECISELY why revalidates at
+        # the boundary. The regex denies `{`, `}`, whitespace, making
+        # this scenario unreachable in production.
+
+
+class LoadSoqlNameValidationTests(unittest.TestCase):
+    """the `name` argument is validated before any filesystem access.
+
+    Without this, a caller that sources `name` from data (config file, user
+    argument, discovered string) could read arbitrary files via traversal
+    (`../../../etc/passwd`). The regex gate closes that before `read_text()`.
+    """
+
+    def setUp(self) -> None:
+        self._tmpdir = tempfile.TemporaryDirectory()
+        self.soql_dir = Path(self._tmpdir.name)
+        self._patch = mock.patch("config.SOQL_DIR", self.soql_dir)
+        self._patch.start()
+        import importlib
+        import soql_loader  # type: ignore
+        importlib.reload(soql_loader)
+        self.soql_loader = soql_loader
+
+    def tearDown(self) -> None:
+        self._patch.stop()
+        self._tmpdir.cleanup()
+
+    # ---- traversal attempts must raise before any file read ---------------
+
+    def test_parent_traversal_raises(self):
+        """`../../../etc/passwd` must be caught by validation, not by the
+        filesystem. The error must be SoqlParamError (clearly labeled) —
+        NOT a bare FileNotFoundError or a ValidationError with raw path.
+        """
+        with self.assertRaises(self.soql_loader.SoqlParamError) as ctx:
+            self.soql_loader.load_soql("../../../etc/passwd")
+        self.assertEqual(ctx.exception.key, "soql_template_name")
+        # The absolute SOQL_DIR must NOT appear in the surfaced message.
+        self.assertNotIn(str(self.soql_dir), str(ctx.exception))
+
+    def test_dotdot_alone_raises(self):
+        with self.assertRaises(self.soql_loader.SoqlParamError):
+            self.soql_loader.load_soql("..")
+
+    def test_slash_in_name_raises(self):
+        """`/` is not in [A-Za-z0-9_], so `plugins/by_planner` must be
+        caught at validation — never reach the filesystem."""
+        with self.assertRaises(self.soql_loader.SoqlParamError):
+            self.soql_loader.load_soql("plugins/by_planner")
+
+    def test_backslash_in_name_raises(self):
+        """Windows-style separators — belt and braces."""
+        with self.assertRaises(self.soql_loader.SoqlParamError):
+            self.soql_loader.load_soql("plugins\\by_planner")
+
+    def test_absolute_path_raises(self):
+        """A caller passing an absolute path (a different traversal variant)
+        must still hit validation, not a false-negative file read."""
+        with self.assertRaises(self.soql_loader.SoqlParamError):
+            self.soql_loader.load_soql("/etc/passwd")
+
+    def test_empty_name_raises(self):
+        with self.assertRaises(self.soql_loader.SoqlParamError):
+            self.soql_loader.load_soql("")
+
+    def test_none_name_raises(self):
+        with self.assertRaises(self.soql_loader.SoqlParamError):
+            self.soql_loader.load_soql(None)  # type: ignore[arg-type]
+
+    def test_whitespace_in_name_raises(self):
+        with self.assertRaises(self.soql_loader.SoqlParamError):
+            self.soql_loader.load_soql("bot lookup")
+
+    def test_traversal_never_reads_filesystem(self):
+        """Validation runs before any I/O — verify read_text is never
+        called when the name is invalid. If the order ever flips, a bad
+        name could still leak a FileNotFoundError with raw path info.
+        """
+        with mock.patch.object(Path, "read_text") as mock_read:
+            with self.assertRaises(self.soql_loader.SoqlParamError):
+                self.soql_loader.load_soql("../../../evil")
+            mock_read.assert_not_called()
+
+
+class LoadSoqlTemplateNotFoundTests(unittest.TestCase):
+    """FileNotFoundError is translated into SoqlTemplateNotFound
+    whose message is free of filesystem-path leakage.
+    """
+
+    def setUp(self) -> None:
+        self._tmpdir = tempfile.TemporaryDirectory()
+        self.soql_dir = Path(self._tmpdir.name)
+        self._patch = mock.patch("config.SOQL_DIR", self.soql_dir)
+        self._patch.start()
+        import importlib
+        import soql_loader  # type: ignore
+        importlib.reload(soql_loader)
+        self.soql_loader = soql_loader
+
+    def tearDown(self) -> None:
+        self._patch.stop()
+        self._tmpdir.cleanup()
+
+    def test_missing_template_raises_custom_exception(self):
+        with self.assertRaises(self.soql_loader.SoqlTemplateNotFound) as ctx:
+            self.soql_loader.load_soql("nonexistent_template")
+        self.assertEqual(ctx.exception.name, "nonexistent_template")
+
+    def test_missing_template_message_excludes_soql_dir(self):
+        """The SOQL_DIR absolute path must NOT appear in the surfaced error
+        string — information-disclosure hygiene. Attackers don't need to know
+        where the skill install lives on disk.
+        """
+        with self.assertRaises(self.soql_loader.SoqlTemplateNotFound) as ctx:
+            self.soql_loader.load_soql("nonexistent_template")
+        msg = str(ctx.exception)
+        self.assertNotIn(str(self.soql_dir), msg)
+        self.assertNotIn(".soql", msg)
+        # Template name IS allowed in the message — that's the triage signal.
+        self.assertIn("nonexistent_template", msg)
+
+    def test_missing_template_does_not_leak_via_cause_chain(self):
+        """`raise ... from None` is load-bearing: without it, the
+        FileNotFoundError (with its raw `filename` attribute) would be
+        reachable via `exception.__cause__`. Verify the chain is severed.
+        """
+        try:
+            self.soql_loader.load_soql("nonexistent_template")
+        except self.soql_loader.SoqlTemplateNotFound as e:
+            # `from None` sets __cause__ = None AND __suppress_context__ = True.
+            # Either alone suppresses traceback rendering of the underlying
+            # FileNotFoundError.
+            self.assertIsNone(e.__cause__)
+            self.assertTrue(e.__suppress_context__)
+        else:
+            self.fail("expected SoqlTemplateNotFound")
+
+    def test_not_found_exception_is_distinct_from_file_not_found(self):
+        """Callers should be able to tell 'template missing' apart from
+        'permission denied / I/O error' at the except-clause layer.
+        """
+        self.assertFalse(
+            issubclass(
+                self.soql_loader.SoqlTemplateNotFound,
+                FileNotFoundError,
+            ),
+            "SoqlTemplateNotFound must not subclass FileNotFoundError",
+        )
+
+    def test_valid_name_with_params_unchanged(self):
+        """Regression: must not break the happy path."""
+        (self.soql_dir / "lookup.soql").write_text(
+            "SELECT Id FROM X WHERE Name = '{{NAME}}'"
+        )
+        out = self.soql_loader.load_soql("lookup", NAME="MyAgent")
+        self.assertEqual(
+            out,
+            "SELECT Id FROM X WHERE Name = 'MyAgent'",
+        )
+
+
+class LoadSoqlInListParamTests(unittest.TestCase):
+    """`load_soql_in` renders `WHERE X IN (...)` list placeholders.
+
+    Same validation surface as `load_soql` — every list element passes
+    through `fs_guard.validate_api_name`. Empty lists fail fast (SOQL
+    `WHERE X IN ()` is invalid). Dedup + sort are load-bearing for
+    stable cache keys.
+    """
+
+    def setUp(self) -> None:
+        self._tmpdir = tempfile.TemporaryDirectory()
+        self.soql_dir = Path(self._tmpdir.name)
+        self._patch = mock.patch("config.SOQL_DIR", self.soql_dir)
+        self._patch.start()
+        import importlib
+        import soql_loader  # type: ignore
+        importlib.reload(soql_loader)
+        self.soql_loader = soql_loader
+
+    def tearDown(self) -> None:
+        self._patch.stop()
+        self._tmpdir.cleanup()
+
+    def _write_template(self, name: str, body: str) -> None:
+        (self.soql_dir / f"{name}.soql").write_text(body)
+
+    # ---- happy path ------------------------------------------------------
+
+    def test_list_params_render_single_quoted_comma_joined(self):
+        self._write_template(
+            "apex_by_names",
+            "SELECT Id FROM ApexClass WHERE Name IN ({{NAMES_LIST}})",
+        )
+        out = self.soql_loader.load_soql_in(
+            "apex_by_names",
+            list_params={"NAMES_LIST": ["ClassA", "ClassB"]},
+        )
+        self.assertIn("WHERE Name IN ('ClassA','ClassB')", out)
+
+    def test_mixed_string_and_list_params(self):
+        self._write_template(
+            "functions_q",
+            "SELECT Id FROM GenAiFunctionDefinition "
+            "WHERE PlannerId = '{{PLANNER_ID}}' OR PluginId IN ({{PLUGIN_IDS}})",
+        )
+        out = self.soql_loader.load_soql_in(
+            "functions_q",
+            string_params={"PLANNER_ID": "X"},
+            list_params={"PLUGIN_IDS": ["P1", "P2"]},
+        )
+        self.assertIn("PlannerId = 'X'", out)
+        self.assertIn("PluginId IN ('P1','P2')", out)
+
+    def test_string_params_optional(self):
+        """string_params defaults to None — list_params alone should work."""
+        self._write_template(
+            "flow_by_names",
+            "SELECT Id FROM FlowDefinition WHERE DeveloperName IN ({{NAMES_LIST}})",
+        )
+        out = self.soql_loader.load_soql_in(
+            "flow_by_names",
+            list_params={"NAMES_LIST": ["Flow_A"]},
+        )
+        self.assertIn("IN ('Flow_A')", out)
+
+    # ---- validation: list elements must match api_name regex ---------------
+
+    def test_injection_in_list_element_raises_with_list_key(self):
+        """A SOQL injection attempt inside a list element must raise
+        SoqlParamError whose `key` is the LIST key (not the element
+        index or a synthetic name) — so the caller can log / mark
+        `_unresolved[NAMES_LIST]` at the upstream boundary.
+        """
+        self._write_template(
+            "q", "SELECT Id FROM ApexClass WHERE Name IN ({{NAMES_LIST}})",
+        )
+        with self.assertRaises(self.soql_loader.SoqlParamError) as ctx:
+            self.soql_loader.load_soql_in(
+                "q",
+                list_params={"NAMES_LIST": ["ClassA", "ClassB'; DROP TABLE x;--"]},
+            )
+        self.assertEqual(ctx.exception.key, "NAMES_LIST")
+
+    def test_non_string_element_raises(self):
+        self._write_template("q", "SELECT Id FROM X WHERE Id IN ({{IDS}})")
+        with self.assertRaises(self.soql_loader.SoqlParamError) as ctx:
+            self.soql_loader.load_soql_in(
+                "q", list_params={"IDS": ["ClassA", 42]},
+            )
+        self.assertEqual(ctx.exception.key, "IDS")
+
+    def test_whitespace_in_list_element_raises(self):
+        self._write_template("q", "SELECT Id FROM X WHERE Id IN ({{IDS}})")
+        with self.assertRaises(self.soql_loader.SoqlParamError):
+            self.soql_loader.load_soql_in(
+                "q", list_params={"IDS": ["ClassA", "x OR 1=1"]},
+            )
+
+    # ---- empty list fails fast --------------------------------------------
+
+    def test_empty_list_raises(self):
+        """SOQL `WHERE X IN ()` is a syntax error; fail at the loader,
+        not at the CLI. The reason string must mention empty so the
+        `_unresolved` bucket can be tagged distinctly from injection.
+        """
+        self._write_template("q", "SELECT Id FROM X WHERE Id IN ({{IDS}})")
+        with self.assertRaises(self.soql_loader.SoqlParamError) as ctx:
+            self.soql_loader.load_soql_in("q", list_params={"IDS": []})
+        self.assertEqual(ctx.exception.key, "IDS")
+        self.assertIn("empty", ctx.exception.reason.lower())
+
+    def test_list_params_not_a_list_raises(self):
+        """Defensive: a dict or string passed in `list_params[KEY]`
+        must be rejected — silently iterating a string would produce
+        one-char-per-element SOQL, which is worse than an explicit error.
+        """
+        self._write_template("q", "SELECT Id FROM X WHERE Id IN ({{IDS}})")
+        with self.assertRaises(self.soql_loader.SoqlParamError):
+            # type: ignore[arg-type]
+            self.soql_loader.load_soql_in("q", list_params={"IDS": "notalist"})
+
+    # ---- dedupe + deterministic order -------------------------------------
+
+    def test_dedupe_eliminates_duplicates(self):
+        self._write_template("q", "SELECT Id FROM X WHERE Name IN ({{NS}})")
+        out = self.soql_loader.load_soql_in(
+            "q", list_params={"NS": ["ClassA", "ClassA", "ClassB"]},
+        )
+        self.assertEqual(out.count("'ClassA'"), 1)
+        self.assertEqual(out.count("'ClassB'"), 1)
+
+    def test_output_is_sorted_for_deterministic_order(self):
+        """Stable cache-key requirement: input order MUST NOT affect
+        output. `sorted(set(...))` lands ClassA before ClassB regardless
+        of input order.
+        """
+        self._write_template("q", "SELECT Id FROM X WHERE Name IN ({{NS}})")
+        out1 = self.soql_loader.load_soql_in(
+            "q", list_params={"NS": ["ClassB", "ClassA"]},
+        )
+        out2 = self.soql_loader.load_soql_in(
+            "q", list_params={"NS": ["ClassA", "ClassB"]},
+        )
+        self.assertEqual(out1, out2)
+        # And the order is alphabetical, not input-dependent.
+        self.assertLess(out1.index("'ClassA'"), out1.index("'ClassB'"))
+
+    # ---- scalar validation path shared with load_soql ---------------------
+
+    def test_scalar_injection_still_raises(self):
+        """string_params go through the same validator as `load_soql` —
+        no shortcut. A SOQL-injection attempt in a scalar must still
+        surface SoqlParamError.
+        """
+        self._write_template(
+            "q",
+            "SELECT Id FROM X WHERE P = '{{PID}}' OR Q IN ({{LIST}})",
+        )
+        with self.assertRaises(self.soql_loader.SoqlParamError) as ctx:
+            self.soql_loader.load_soql_in(
+                "q",
+                string_params={"PID": "x' OR Id!=null--"},
+                list_params={"LIST": ["A"]},
+            )
+        self.assertEqual(ctx.exception.key, "PID")
+
+    # ---- template-name validation reused ----------------------------------
+
+    def test_template_traversal_raises(self):
+        with self.assertRaises(self.soql_loader.SoqlParamError) as ctx:
+            self.soql_loader.load_soql_in(
+                "../../../evil", list_params={"IDS": ["A"]},
+            )
+        self.assertEqual(ctx.exception.key, "soql_template_name")
+
+    def test_missing_template_raises_template_not_found(self):
+        with self.assertRaises(self.soql_loader.SoqlTemplateNotFound):
+            self.soql_loader.load_soql_in(
+                "nonexistent_for_in", list_params={"IDS": ["A"]},
+            )
+
+    # ---- existing load_soql unchanged -------------------------------------
+
+    def test_load_soql_signature_untouched(self):
+        """contract: `load_soql(name, **params)` keeps its original
+        signature — no kwargs-only, no extra params. A caller that still
+        uses the old form must keep working.
+        """
+        self._write_template("q", "SELECT Id FROM X WHERE Name = '{{NAME}}'")
+        out = self.soql_loader.load_soql("q", NAME="Foo")
+        self.assertIn("'Foo'", out)
+
+
+if __name__ == "__main__":
+    unittest.main()