Add alias tracking for self attributes to prevent trust bypasses

Author: pjwerneck

packages/syft-restrict/src/syft_restrict/verifier.py

@@ -197,6 +197,11 @@ def __init__(self, policy: Policy, scan: FileScan, ranges):
         # _mark_annotation_subtrees; exempt from the name/container checks in _check_name and
         # _check_container_literal because annotations are never invoked (see visit).
         self._annotation_nodes: set[int] = set()
+        # One {local_name: self_attr_name} dict per enclosing class/function/lambda scope (mirrors
+        # _scope_stack's push/pop), tracking which local variables currently alias a self.<attr> --
+        # so `tmp = self.fn; tmp(x)` is checked the same way `self.fn(x)` directly is; see
+        # _track_self_attr_alias and its use in _check_call.
+        self._alias_stack: list[dict[str, str]] = []
 
     def report(self, node: ast.AST, code: ViolationCode, message: str) -> None:
         self.violations.append(
@@ -212,13 +217,15 @@ def visit(self, node: ast.AST) -> None:
         is_scope = isinstance(node, (ast.ClassDef, ast.FunctionDef, ast.Lambda))
         if is_scope:
             self._scope_stack.append(node)
+            self._alias_stack.append({})
         self._mark_annotation_subtrees(node)
 
         for child in ast.iter_child_nodes(node):
             self.visit(child)
 
         if is_scope:
             self._scope_stack.pop()
+            self._alias_stack.pop()
 
     def _enclosing_class(self) -> ast.ClassDef | None:
         """The nearest enclosing ClassDef, skipping any FunctionDef/Lambda frames above it."""
@@ -408,6 +415,18 @@ def _check_call(self, node: ast.Call) -> None:
                         f"call to {self._resolve(func.id)!r} is not allow-listed",
                     )
                 return
+            aliased_attr = self._current_aliases().get(func.id)
+            if aliased_attr is not None and not self._self_attr.is_safe(
+                aliased_attr, self._enclosing_class()
+            ):
+                self.report(
+                    node,
+                    "attr-on-value",
+                    f"{func.id!r} was assigned from self.{aliased_attr!r}, which isn't an "
+                    f"allow-listed constructor or a locally-defined class; calling it is not "
+                    f"allowed",
+                )
+                return
             # Otherwise a bare-name call (local var / hidden or visible def / safe builtin) is allowed.
             return
         if isinstance(func, ast.Attribute):
@@ -608,6 +627,49 @@ def _check_assign_targets(self, node) -> None:
                     "banned-construct",
                     "storing a reference to a banned builtin into a container slot is not allowed",
                 )
+        self._track_self_attr_alias(node)
+
+    def _track_self_attr_alias(self, node) -> None:
+        # Detect `tmp = self.<attr>` / `tmp = self.<attr>[i]` / `tmp = <already-tracked alias>` so a
+        # later `tmp(...)` call in _check_call is checked the same way calling self.<attr> directly
+        # would be -- otherwise reading an unsafe self.<attr> into a local variable trivially
+        # defeats _SelfAttrTrust (self.<attr> reads are always allowed; calling a local variable is
+        # always allowed). Any other assignment to a previously-tracked name clears its alias -- it
+        # no longer refers to that self.<attr>. AugAssign always clears (no value to trace here).
+        if not self._alias_stack:
+            return
+        aliases = self._alias_stack[-1]
+        if isinstance(node, ast.Assign):
+            targets, value = node.targets, node.value
+        elif isinstance(node, ast.AnnAssign):
+            targets, value = ([node.target] if node.value is not None else []), node.value
+        else:  # AugAssign
+            targets, value = [node.target], None
+        for t in targets:
+            if not isinstance(t, ast.Name):
+                continue
+            attr = self._self_attr_source(value) if value is not None else None
+            if attr is not None:
+                aliases[t.id] = attr
+            else:
+                aliases.pop(t.id, None)
+
+    def _self_attr_source(self, value: ast.AST) -> str | None:
+        """The self.<attr> a plain expression traces back to: a direct self.<attr> read, a
+        single-level self.<attr>[i] subscript, or a copy of an already-tracked local alias."""
+        attr = self_attr_name(value)
+        if attr is not None:
+            return attr
+        if isinstance(value, ast.Subscript) and isinstance(value.value, ast.Attribute):
+            attr = self_attr_name(value.value)
+            if attr is not None:
+                return attr
+        if isinstance(value, ast.Name):
+            return self._current_aliases().get(value.id)
+        return None
+
+    def _current_aliases(self) -> dict[str, str]:
+        return self._alias_stack[-1] if self._alias_stack else {}
 
     def _check_reserved_target(self, target: ast.AST) -> None:
         # Guards against: rebinding self/cls (which would forge the identifier-based trust — see

packages/syft-restrict/tests/verify/test_bypasses.py

@@ -568,3 +568,81 @@ def test_self_attr_for_loop_target_must_not_grant_trust(verify_all):
         "        return x",
     ]
     assert "attr-on-value" in error_codes(verify_all("\n".join(src)))
+
+
+def test_self_attr_local_alias_must_not_grant_trust(verify_all):
+    # Copying self.<attr> into a local variable before calling it must not defeat _SelfAttrTrust --
+    # reading self.<attr> is unconditionally allowed, and calling a local variable is
+    # unconditionally allowed, so without alias tracking this trivially bypasses the same check
+    # that correctly rejects calling self.fn directly.
+    src = [
+        "class M(object):",
+        "    def setup(self):",
+        "        self.fn = lambda x: x + 1",
+        "    def __call__(self, x):",
+        "        tmp = self.fn",
+        "        return tmp(x)",
+    ]
+    assert "attr-on-value" in error_codes(verify_all("\n".join(src)))
+
+
+def test_self_attr_subscript_local_alias_must_not_grant_trust(verify_all):
+    # Same bypass via the Flax self.layer[i](x) idiom: copy the element to a local first. Uses an
+    # opaque parameter (not a BANNED_NAMES literal like `open`) so the only thing standing between
+    # this and a violation is the self-attr trust check the alias is meant to defeat.
+    src = [
+        "class M(object):",
+        "    def setup(self, cb):",
+        "        self.layer = [cb]",
+        "    def __call__(self, x):",
+        "        tmp = self.layer[0]",
+        "        return tmp(x)",
+    ]
+    assert "attr-on-value" in error_codes(verify_all("\n".join(src)))
+
+
+def test_self_attr_two_hop_local_alias_must_not_grant_trust(verify_all):
+    # A second-hop copy (tmp2 = tmp) must not escape alias tracking either, mirroring the
+    # multi-hop rigor already applied to BANNED_NAMES aliasing (test_homoglyph_two_hop_alias).
+    src = [
+        "class M(object):",
+        "    def setup(self):",
+        "        self.fn = lambda x: x + 1",
+        "    def __call__(self, x):",
+        "        tmp = self.fn",
+        "        tmp2 = tmp",
+        "        return tmp2(x)",
+    ]
+    assert "attr-on-value" in error_codes(verify_all("\n".join(src)))
+
+
+def test_self_attr_reassigned_local_alias_is_not_stuck_tainted(verify_all):
+    # Reassigning a variable that once aliased an unsafe self.<attr> to something else entirely
+    # must clear the taint -- tracking must reflect the CURRENT value, not the first one ever seen.
+    src = [
+        "class M(object):",
+        "    def setup(self):",
+        "        self.fn = lambda x: x + 1",
+        "    def __call__(self, x):",
+        "        tmp = self.fn",
+        "        tmp = x",
+        "        return tmp",
+    ]
+    result = verify_all("\n".join(src))
+    assert result.ok, [(v.code, v.message) for v in result.violations]
+
+
+def test_self_attr_local_alias_of_safe_source_is_still_allowed(verify_all):
+    # A local alias of a self.<attr> that WAS assigned a vetted-safe source (or is presumed
+    # inherited, since the class never assigns it) must still be callable -- alias tracking must
+    # track safety, not blanket-ban aliasing self.<attr> outright (see
+    # test_calling_a_value_is_allowed in test_whitelist.py, which already covers this shape and
+    # must keep passing).
+    src = [
+        "class M(object):",
+        "    def __call__(self, x):",
+        "        layer = self.layers[0]",
+        "        return layer(x)",
+    ]
+    result = verify_all("\n".join(src))
+    assert result.ok, [(v.code, v.message) for v in result.violations]