diff --git a/packages/syft-restrict/README.md b/packages/syft-restrict/README.md index 9d7be27e75b..9b20381f86c 100644 --- a/packages/syft-restrict/README.md +++ b/packages/syft-restrict/README.md @@ -1,263 +1,61 @@ # syft-restrict -## 1. Overview - -`syft-restrict` is inspired by [**RestrictedPython**](https://github.com/zopefoundation/RestrictedPython): it **default-denies** every dynamic part of Python. It differs in a few ways: - -- **Syft-restrict analyzes, it doesn't run.** syft-restrict statically analyzes a source file and **raises if the analyzed part uses any dynamic Python**. Its output is a **report** (the violations, or — on success — an attestation) plus an **obfuscated copy** of the code. -- **syft-restrict is currently tailored for analysis of machine learning inference code** -- **Public / private split.** The user marks some lines of the file _private_ and leaves the rest _public_. restrict analyzes only the **AST of the private part**, and only those private lines are **hidden (obfuscated)** in the emitted artifact; the public lines are copied through and read directly. -- **Imports are allowed (and therefore trusted).** Unlike RestrictedPython, the analyzed code may **import allowlisted libraries** and call into them — which means the reader of the output **needs to trust those imported classes/functions** (only their allow-listed paths; see §3). -- **Dynamic Python lives in the public part.** Anything dynamic the author needs (file/tokenizer loading, the generation loop, wrappers around library methods) must be written in the **public** region — where it is read directly — and may be **called by** the private part. - -## 2. Usage +A **static analyzer** for Python source files that **default-denies** dynamic +Python constructs. + +## Overview + +`syft-restrict` is inspired by +[**RestrictedPython**](https://github.com/zopefoundation/RestrictedPython), but +differs in a few ways: + +- **Syft-restrict analyzes, it doesn't run.** syft-restrict statically analyzes + a source file and **fails if the analyzed part uses any dynamic Python**. Its + output is a **report** (the violations, or — on success — an attestation) plus + an **obfuscated copy** of the code. +- **Public / private split.** The user marks some lines of the file _private_ + and leaves the rest _public_. restrict analyzes only the **private regions**, + and only those are **obfuscated** in the emitted artifact; the public lines + are copied through and read directly. +- The public code may **import allowlisted libraries** and call into them. + Imports are not allowed in private code, and any library call must be + **explicitly allow-listed**. +- **Dynamic Python lives in the public part.** Anything dynamic the author needs + (file/tokenizer loading, the generation loop, wrappers around library methods) + must be written in the **public** region -- where it is reviewed directly --and + may be **called by** the private part. + +## Usage ```python import syft_restrict as restrict result = restrict.run( "gemma_inference.py", - obfuscate=[[22, 93], [99, 280]], # 1-based ranges: identifiers renamed, constants blanked - hide=[], # 1-based ranges: whole line replaced with ■■■■■■■■ - allow_functions=["jax.*", "flax.linen.*"], # things callable BY NAME (path-resolved) - allow_methods=["arithmetic", "indexing", "comparison"], # operators allowed ON A VALUE + obfuscate=[[22, 93], [99, 280]], # 1-based ranges: identifiers renamed, constants blanked + hide=[], # 1-based ranges: whole line replaced with ■■■■■■■■ + allow_functions=["jax.*", "flax.linen.*"], # functions callable BY NAME (path-resolved) + allow_operators=["arithmetic", "indexing", "comparison"], # operators allowed ON A VALUE ) -# Verified region = obfuscate ∪ hide (both are private code; both are checked). # On success: writes gemma_inference.obfuscated.py and returns result.certificate. -# On a policy violation: raises PolicyViolation naming each offending line (strict=True, the default). +# On a policy violation: raises PolicyViolation naming each offending line. ``` -This command reads **[examples/gemma_inference.py](examples/gemma_inference.py)** and generates **[examples/gemma_inference.obfuscated.py](examples/gemma_inference.obfuscated.py)**. If we can assume that syft-restrict was executed on the true inference file and the library was not modified, for instance because it was executed in a [TEE](https://en.wikipedia.org/wiki/Trusted_execution_environment), this proves: - -1. this code is an inference pipelines for a jax model, where the model architecture is hidden in the report -2. this code does not steal the inputs - -Use `restrict.verify(...)` for the check alone (it returns violations instead of raising), or pass `strict=False` to `run` to get a `RunResult` with `.ok` / `.violations` and no exception. - -## 3. The whitelist (for the private region) - -What the private code legitimately needs to do: **tie together JAX operations** — define Flax modules (classes with `setup`/`__call__`), call JAX functions, do arithmetic, build the shape/dtype plumbing (lists, dicts, tuples, slices, f-strings for einsum equations), and run data-independent control flow. Nothing that can reach the host machine, the filesystem, the network, the Python interpreter's internals, or build-and-run code from a string. - -Default rule, exactly like RestrictedPython: **any node type not in the ALLOW column is REJECTED.** Future new Python syntax is denied until a human reviews it. - -For **examples of disallowed constructs** (and why each is rejected), see **[disallowed-ast-examples.md](disallowed-ast-examples.md)**. Those examples are illustrative only — safety comes from the default-deny whitelist below, not from any deny list. - -### 3.1 ALLOW — always-on whitelist - -These structural nodes are always permitted in the private region. Beyond them, the author enables a **configurable set per-file** — allowlisted call/attribute _paths_ (`allow_functions`) and operator _bundles_ on values (`allow_methods`) — listed in **§3.2**. A node is permitted only if it appears here or in an _enabled_ §3.2 entry; everything else is rejected (default-deny). - -| Category | Allowed `ast` nodes | Example (from the private region of `gemma_inference.py`) | Notes / constraints | -| --------------------------------- | ------------------------------------------------------------------ | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| Module structure | `Module`, `Expr` | the file itself (`Module`); the module docstring `"""Gemma 3 IT — Flax Inference Module ..."""` is a bare-expression statement (`Expr`) | top level | -| Definitions | `FunctionDef`, `ClassDef`, `arguments`, `arg`, `Return`, `Lambda` | `class RMSNorm(nn.Module):` (`ClassDef`); `def __call__(self, x):` (`FunctionDef` + `arguments`/`arg`); `return x * jax.lax.rsqrt(...)` (`Return`); `lambda: None` in `_get(...)` (`Lambda`) | class bases restricted to allowlisted module attrs (e.g. `nn.Module`); decorators restricted to allowlist (`@nn.compact`, `@jax.jit` — see §3.2.1) | -| Names — `Name` | `Name` | any bare identifier: `x`, `half`, `cfg`, `Attention` (a variable, parameter, function, or class reference) | **underscore ban applies to the identifier text** — see the note below this row | -| Names — read context `Load` | `Load` | the `x` and `half` in `x * half` (every value _read_) | not a node — it's the `ctx` flag on a `Name`/`Attribute`/`Subscript` marking a _read_; always allowed | -| Names — write context `Store` | `Store` | the `half` in `half = dim // 2` (the name being _bound_) | same `ctx` flag marking a _write_ (vs `Load`'s read); allowed only for a local `Name` or `self.` — no global or foreign-attribute writes | -| Constants | `Constant` | `2` and `1e-6` (numbers); `"bsd,ndh->bsnh"` (str literal); `None`, `True`, `False` | any literal scalar — pure, no constraint | -| Assignment | `Assign`, `AugAssign`, `AnnAssign` | `x = x + h` (`Assign`); `attn_type: str = "local"` (`AnnAssign`). _(`AugAssign` e.g. `h += ...` — not used in the private region, illustrative.)_ | targets: local `Name` or `self.`; **no global writes**, no attribute writes to foreign objects | -| f-strings | `JoinedStr`, `FormattedValue` | none in the private region — the einsum equations here are plain `Constant` strings (`"bsd,ndh->bsnh"`). _(A real `JoinedStr` is the public `format_chat`'s `f"user\n{prompt}..."`.)_ | needed if einsum equations are built as f-strings; the format-string `__format__` trick must be neutralised — only allow formatting of `Constant`/`Name`/allowlisted exprs, and ban `_`-prefixed format specs (a historical RestrictedPython bug class) | -| Type hints | `Subscript`/`Name` in annotations | `cfg: dict` (`Transformer`); `attn_type: str = "local"` (`Block`) | `cfg: dict` etc. | -| Containers | `List`, `Tuple`, `Dict`, `Set` | `dict(num_layers=18, embed_dim=640, ...)` (`Dict`); `{"local": ..., "global": ...}` in `make_masks` (`Dict`); `("local",) * 5 + ("global",)` in `_attn_types` (`Tuple`) | literal data plumbing; elements/keys must themselves be allowed nodes — a denied call or attr inside (e.g. `[open(f)]`) still fails | -| Comprehensions | `ListComp`, `DictComp`, `SetComp`, `GeneratorExp`, `comprehension` | `self.layer = [Block(cfg=self.cfg, attn_type=attn_types[i]) for i in range(num_layers)]` (`ListComp` over `range(...)`) | **only if** the iterable is a pure expression (`range(...)`, a literal, a local). Banned if the iterable touches I/O. No `async` comprehensions. | -| Calls | `Call`, `keyword`, `Starred` | `Attention(cfg=self.cfg)` (`Call` + `keyword` `cfg=`); `jnp.einsum("bsd,ndh->bsnh", x)` (`Call`). *(`Starred` e.g. `jnp.stack([*xs])` — not used here, illustrative.)\* | the function being called **must resolve to an allowlisted symbol** (§3.2.1); `Starred` allowed only into allowlisted calls | -| Attribute — `self.` (read) | `Attribute` (`Load`) | `self.scale`, `self.cfg`, `self.layer` | reading the module's own attributes — receiver is the class being defined, not an opaque value, so always safe. (Reads _on a value_, even `x.shape`, are **denied** — route them through a visible wrapper, §3.2.) | -| Attribute — `self.` (write) | `Attribute` (`Store`) | `self.w = ...` inside `setup` | the **only** attribute write allowed; Flax needs it. Writes to any other object (`obj.x = ...`) are denied | -| Control flow | `If`, `For`, `While`, `Break`, `Continue`, `Pass`, `IfExp` | `for i in range(num_layers):` (`For`); `if cache is None:` (`If`); `cache[i] if cache is not None else None` (`IfExp`) | data-independent loops fine; `For` iterable restricted like comprehensions | - -> **The underscore rule (carried over from RestrictedPython) — the most delicate part of the policy.** Any identifier starting with `_` — the `id` of a `Name`, the `attr` of an `Attribute`, an `arg` name, or an import `alias` — is in principle rejected; this single rule kills the introspection escape ladder (`obj.__class__.__bases__[0].__subclasses__()` …). But the private model uses leading-underscore names heavily (`_get`, `_attn_types`, `_query_norm`, `__call__`), so a **curated relaxation** is required: a leading-underscore identifier is **rejected** when it is (a) **read off a foreign object** (`obj._secret`, `obj.__class__`) or (b) used as a **call target** (`_f(...)`), but **allowed** when it is (c) a _local definition_ (`def _get(...)`, `_tmp = ...`) or (d) a `self.`-attribute (`self._query_norm`, `self._key_norm`). (Defining `_`/dunder **hook methods** on a model class is separately restricted to `setup` / `__call__` / `__post_init__` — §3.2.1 #6.) This carve-out is where review effort concentrates. - -> **Attribute — what is NOT allowed (the two `self.` rows above, plus the allowlisted-path row in §3.2, are the whole allowed set).** If a read doesn't match one of those, it's denied. Denied **reads** (`Load`), with examples: -> -> - **A dotted path that resolves to a non-allowlisted symbol** — `jnp.save`, `jax.experimental.io_callback`, `jax.numpy.load` (denylist beats allow, §3.2.1). -> - **Any attribute read on an opaque value** — `x.shape`, `x.dtype`, `embed_table.T`, `x.real`, `x.device` (we can't pin the receiver's type, so _every_ attribute name on a value is denied — including `.shape`/`.ndim`/`.dtype`; route each through a visible wrapper function, §3.2). -> - **Any underscore/dunder attribute read off a foreign object** — `obj._secret`, `x.__class__`, `module.__dict__` (the reflection ladder; §3.1). -> -> And the one denied **write** (`Store`): an attribute write to any non-`self` object — `obj.send = data`. - -> **The one allowed-syntax exception — class keyword arguments (`metaclass=`).** A `keyword` node is on the allow-list (Calls row) for ordinary call kwargs like `Attention(cfg=...)`. But a `keyword` attached to a **`ClassDef`** — `class C(metaclass=evil)`, or _any_ class keyword — is **denied**, because building a class with a `metaclass=` (or other class kwarg) runs attacker-chosen code at class-creation time. So the rule is _position-dependent_: the same `keyword` node is allowed inside a `Call` but rejected inside a `ClassDef`. (Base classes are separately restricted to allow-listed symbols — see the Definitions row and §3.2.1 #5.) - -### 3.2 Configurable to allow - -Two channels the author sets per-file. Anything not enabled in one of them is rejected like any other non-allowlisted node. - -**Paths called/read by name — `allow_functions` (§3.2.1).** Dotted paths resolved exactly against the import bindings; the denylist beats the allow. - -| Form | AST node | Example (from the private region) | Constraint | -| ------------------------------ | -------------------- | -------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------- | -| Allowlisted module path (read) | `Attribute` (`Load`) | `jnp.einsum`, `jax.lax.rsqrt` — the dotted path in front of a `Call` | allowed only when the **whole path resolves to an allowlisted symbol** (§3.2.1); the underscore ban applies to each `attr` | - -**Operator bundles on a value — `allow_methods` (§3.2.2).** Type-agnostic-safe generic operators, never library-specific named methods — so each can be toggled on without re-opening an escape. - -| Bundle (`allow_methods`) | AST nodes / operators | Example (from the private region) | -| ------------------------ | ------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------- | -| `arithmetic` | `BinOp`, `UnaryOp`; operators `Add Sub Mult Div FloorDiv Mod Pow MatMult LShift RShift BitOr BitAnd BitXor`, `USub UAdd Invert` | `var + 1e-6` (`Add`); `base_freq ** freq_exp` (`Pow`); `x @ transpose(embed_table)` (`MatMult`); `causal & window` (`BitAnd`) | -| `comparison` | `Compare`, `BoolOp`; operators `Eq NotEq Lt LtE Gt GtE Is IsNot In NotIn`, `And Or`, `Not` | `attn_type == "local"` (`Eq`); `cache is None` (`Is`); `causal and window` (`And`) | -| `indexing` | `Subscript`, `Slice` | `x[..., :half]` (`Subscript` + `Slice`); `kv[0]`, `cache[i]`, `shape_of(x)[-1]` | - -#### 3.2.1 Constraining how allow-listed functions and classes are used - -The private region still imports and calls JAX, and **an import is the only way the code can name anything outside its own AST** — while _most JAX is safe pure math, a handful of functions can reach the host_. So once §3.2 has said _which_ functions and methods may be named, these rules fix **how those allow-listed names may then be used** — so nothing defeats the path resolver or smuggles execution into a `def`/`class`. Rule 1 is the resolution mechanism itself; rules 2–6 keep it sound. (**Rebinding**, below, means pointing a name at a different object — `jnp = something_else` — innocuous normally, but a weapon here because the resolver assumes `jnp` still means the allow-listed module.) - -**1. We resolve every call/attribute to a dotted path and require it to be allow-listed (and not deny-listed).** - -```python -import jax.numpy as jnp # binding recorded: jnp → jax.numpy -jnp.einsum("...", x) # resolves to jax.numpy.einsum → allow-listed ✓ -jnp.save(x) # resolves to jax.numpy.save → deny-listed ✗ -``` - -We don't allow-list a module wholesale; we allow-list **(module, attribute-path)** pairs, enforced by _reading_ the code — which works precisely because dynamic attribute access is banned. Mechanism: (a) `Import`/`ImportFrom` may only name allow-listed modules, and each binding is recorded (`jnp → jax.numpy`); (b) every `Attribute`/`Call` is traced to a fully-qualified dotted path via those bindings; (c) the resolved path must match an allow pattern (`jax.numpy.*`, `jax.lax.*`, …) **and not** match a deny pattern — _the denylist beats the allow_, even inside an otherwise-allowed module (so `jax.numpy.einsum` passes while `jax.numpy.save` is rejected). Because there's no `getattr`/`exec`/underscore-reflection, this static resolution is **sound**: there is no other runtime path from `jnp` to a denied symbol. The dangerous-API denylist (host callbacks, `numpy.save/load`, FFI/dlpack, profiler/distributed, serialization, …) is catalogued in **[../../research/verifuscate/approach-B-jax-denylist-reference.md](../../research/verifuscate/approach-B-jax-denylist-reference.md)**. - -**2. We deny reassigning a trusted module name (`jnp`, `nn`, `lax`, …).** - -```python -jnp = SomethingEvil() -jnp.save(x) # checker sees allow-listed path "jnp.save" → approves - # at runtime jnp is the evil object, .save() runs attacker code -``` - -The static path resolver (rule 1) trusts the binding table built from the imports. If the code later _rebinds_ `jnp`, the table would be a lie and every "allow-listed path" through that name attacker-controlled — so this vector would defeat the path allow-list if left unaddressed. -_Mitigation:_ treat the imported-module alias names as **RESERVED**. The reserved set is exactly the aliases the import resolver recorded — `jax`, `jnp` (`jax.numpy`), `lax` (`jax.lax`), `nn` (`flax.linen`), `jrandom`/`random` as imported, `functools`, etc. A reserved name may **not** be reassigned, used as a function/lambda parameter, used as a `for`/comprehension/`with` target, or otherwise rebound _anywhere_ in the private region. Any node that would store to a reserved name ⇒ REJECT. (This is cheap to enforce: it's just "is this `Store`/`arg`/loop-target one of the reserved names?") - -**3. We deny referencing an allow-listed callable without calling it inline.** - -```python -f = jnp.einsum # path resolver sees a *reference*, not a Call -... -f(equation, x) # later call goes through plain Name `f` — unchecked -d = {"s": jnp.save}; d["s"](x) # same trick via a container -``` - -The rule-1 resolver only follows _literal dotted paths that are the function of a Call_. Pull an allow-listed callable out into a variable, list, or dict and the eventual `f(...)` is just a call on an ordinary local name — outside the path allow-list entirely. -_Mitigation:_ **require allow-listed callables to be called inline.** An allow-listed attribute path (`jnp.einsum`, `jax.lax.rsqrt`, …) may appear _only_ as the function position of a `Call` — never referenced and stored, returned, or passed as an argument. `jnp.einsum(...)` is fine; `g = jnp.einsum`, `return jnp.einsum`, `vmap(jnp.einsum)` (passing it as an arg) are all REJECT. This keeps every allow-listed call statically resolvable at its call site. - -**4. We deny any decorator that isn't on the allow-list.** - -```python -@evil # `evil` is called when the def/class is created, before any "math" runs -def block(...): ... -@a.b(...) # a.b(...) is called, then its result is called on the function -class Block(...): ... -``` - -A decorator is an ordinary call that executes the instant the `def`/`class` is reached. -_Mitigation:_ **allow-list the decorators** — `@nn.compact`, the `@jax.jit` family (incl. `@functools.partial(jax.jit, …)` when `functools.partial` is allow-listed and its first arg resolves to an allow-listed transform), `@jax.named_scope` (review). Any decorator expression not on that list ⇒ REJECT. The decorated body is still fully walked, so an allow-listed decorator can't smuggle denied ops. - -**5. We deny non-allow-listed base classes.** - -```python -class Block(EvilBase): ... # EvilBase.__init_subclass__ / metaclass runs at class creation -``` - -Creating a class consults its bases — `EvilBase.__init_subclass__` (and any metaclass it carries) runs attacker code. -_Mitigation:_ **base classes must be allow-listed** (`nn.Module`, `object`); a `ClassDef` with a non-allow-listed base ⇒ REJECT. (The related `metaclass=` / class-keyword route is handled as a syntax exception in §3.1 — a `keyword` on a `ClassDef` is denied.) - -**6. We deny defining magic/hook methods on a model class (only `setup`/`__call__` allowed).** - -```python -class Block(nn.Module): - def __getattr__(self, name): ... # runs on every attribute miss - def __init_subclass__(cls, **kw): ... # runs at class creation - def __reduce__(self): ... # runs at pickle time - def __getitem__(self, k): ... # runs on subscripting -``` - -These are hooks: Python calls them automatically at attribute-access, class-creation, pickling, or subscript time. The body is attacker code that fires _without an explicit call in the math_ — so even a clean-looking `block[0]` could detonate. -_Mitigation:_ inside a class body, **only allow defining the handful of method names Flax actually needs** — `setup`, `__call__`, and `__post_init__` if Flax requires it. Defining _any other_ underscore/dunder method ⇒ REJECT. (Note this is stricter than the §3.1 underscore relaxation, and it should be: §3.1 lets you _read_ `self._foo` and _define_ `_foo`-style helpers; this rule additionally forbids _defining magic hook methods_ on a model class.) - -#### 3.2.2 The method-call gap: constraining `value.method(...)` calls on returned objects - -For a call of the form `value.method(...)` we do **not** know the receiver's type — so, unlike a resolved `module.func(...)` path (§3.2.1), we cannot tell what invoking a method (or a `getattr`) on it actually does: - -```python -x.reshape(8, -1) # x's type is unknown — is this jax.Array.reshape, or a .format-style escape? -``` - -**The fix: only generic operator _methods_ on values; library ops become _functions_.** - -Two rules: - -1. **On an opaque value, allow only generic, language-level operator methods — never a library-specific named method.** These are the dunder operators (`__add__`, `__getitem__`, comparisons, …): they're safe on _any_ type, and — crucially — they aren't named-method calls, so the `.format`-style escape has nowhere to hide. They are enabled as toggleable **bundles**: - - | Bundle (pass in `allow_methods`) | Covers | Example from the private region | - | -------------------------------- | -------------------------------------- | ------------------------------------ | ------------------------------------------ | - | `arithmetic` | `+ - * / // % ** @`, unary, bitwise `& | ^ ~` | `var + 1e-6`, `x @ transpose(embed_table)` | - | `indexing` | subscript + slice | `x[..., :half]`, `kv[0]`, `cache[i]` | - | `comparison` | `== != < <= > >=`, `and`/`or`/`not` | `attn_type == "local"` | - - There is deliberately **no metadata bundle**: a `.shape`/`.ndim`/`.dtype` read is a _named attribute access_ on a value whose type we can't pin (it could be the `.format`-style escape just as easily as a real array), so it is treated like any other attribute-on-value and rejected — it must be wrapped (rule 2). - -2. **Every library-specific method or attribute read is extracted into a visible wrapper function** — `.T`, `reshape`, `astype`, `sum`, `x.at[i].set(v)`, a Flax `Variable.value`, a `.shape`/`.ndim`/`.dtype` read. Instead of `x.transpose()` in the private code, the author writes the wrapper in the **public** region and the private code calls it _by name_; the data owner reads it, where an `isinstance` guard pins the receiver type the checker couldn't: - - ```python - # in the public (non-private) region — the data owner reads this - def transpose(x): - if not isinstance(x, jax.Array): # explicit raise, NOT assert (asserts vanish under python -O) - raise TypeError - return x.T - ``` - - So the private code writes `transpose(embed_table)`, not `embed_table.T`. The wrapper names are **reserved** — the private region can't rebind them (§3.2.1 #2), so `transpose = evil` is rejected. (Optionally, restrict template-checks each wrapper — a type guard plus a single delegated call — so they're machine-checked, not only read.) - -Net effect: **no named method is ever called on an opaque value.** The only thing done _to a value_ is a generic operator; everything library-specific is a named, type-guarded function. - -#### 3.2.3 The full call-target rule - -Combining the §3.2.1 rules with the method rules above, what may sit in a `Call`'s function position (or be done to a value) is a **four-way OR**: - -A call target is allowed iff it is one of: - -- **(a)** an external allow-listed dotted path called inline (§3.2.1 — a _function_); -- **(b)** a name defined-and-checked in the private region (transitively safe — e.g. `Attention(...)`, `apply_rope(...)`); -- **(c)** a **wrapper function defined in the public region** (read by the DO, its name reserved against rebinding); -- **(d)** — for an operation _on a value_ — a generic operator **method** from an enabled bundle (§3.2.2). - -**No library-specific named method on a value is ever allowed.** - -Anything else ⇒ REJECT. Cases (b) keep the rule from being too strict: the model's own classes/functions are calls on names defined _inside_ the private lines, so calling them can't introduce anything the checker hasn't already walked. - -## 4. Order of operations — verify the private lines BEFORE obfuscating them - -``` -restrict.run("file.py", obfuscate=[[84, 280]], hide=[], allow_functions=["jax.*"]) - │ - ├─ 1. PARSE the full source ──► ast.parse(source) - │ (restrict sees everything; the DO will only get the public part) - │ - ├─ 2. RESOLVE imports → binding table (import jax.numpy as jnp ⇒ jnp→jax.numpy) - │ - ├─ 3. WALK the PRIVATE lines (84–280) with a default-deny visitor: - │ • node type ∈ ALLOW set (§3.1)? else REJECT - │ • every Name/Attribute/arg/alias: underscore rule (§3.1) - │ • Call target ∈ {external allowlisted path called INLINE, - │ locally-defined-and-checked name, allowlisted method - │ name on a value} (§3.2, §3.2.3)? else REJECT - │ • resolved external path ∉ denylist (§3.2)? else REJECT - │ • reserved module-alias names never rebound (§3.2.1#2)? else REJECT - │ • decorators/bases/class-keywords/own-dunder-methods - │ ∈ allowlist (§3.2.1#4-6)? else REJECT - │ • no Global/Nonlocal/foreign-attr-Store else REJECT - │ • no exec/eval/getattr/open/try/with/... else REJECT - │ (call-checking is POSITION-INDEPENDENT: defaults, - │ annotations, class-body stmts are walked too) - │ ── if ANY private node fails ⇒ restrict ABORTS, emits nothing. - │ - ├─ 4. Only now: MANGLE allowlisted jax calls per allow_functions, - │ and OBFUSCATE lines 84–280 → ■■■■■■. - │ - └─ 5. EMIT (artifact = public glue + obfuscated math, - attestation{source_hash, policy_hash, PASS}). -``` +This command reads +**[examples/gemma_inference.py](examples/gemma_inference.py)** and generates +**[examples/gemma_inference.obfuscated.py](examples/gemma_inference.obfuscated.py)**. -"Attestation" (step 5) is a cryptographically signed statement the secure enclave produces: a machine-verifiable receipt saying _"I ran exactly this check over exactly these bytes and it passed."_ The DO checks the signature instead of trusting anyone's word. +If syft-restrict was succesfully executed on the true inference file and the +library was not modified, this proves the code does not exfiltrate the inputs, and +the obfuscated file can be safely shared with a third party. -The enclave runs steps 1–5 and attests: _"I, the enclave, parsed bytes with sha256 = H, ran whitelist policy version P over the PRIVATE lines, it PASSED, and the obfuscated artifact you're reading was derived from exactly those bytes."_ +Use `restrict.verify(...)` for the check alone (it returns violations instead of +raising), or pass `strict=False` to `run` to get a `RunResult` with `.ok` / +`.violations` and no exception. -## 5. What this whitelist does NOT stop -Everything in §3 is a _closed_ vector: no file/socket/host-callback/dynamic-code node can pass, including all the subtle dynamic-Python escapes of §3.2.1. The problems below are the ones the static whitelist structurally cannot solve. Each needs an orthogonal control (output discipline, a runtime sandbox, resource bounds, attestation) — none is fixed by tightening the AST checker. +## Documentation -1. **The output is itself a leak channel.** Inference must return logits/tokens, and a malicious model can encode private data in them (low-order bits, token choices); "only pure math ran" doesn't imply "the output carries no private bits." → output rate-limiting / quantization / DP-noise / reviewed schema. -2. **Timing & cache side channels.** Execution time, memory-access and CPU-cache patterns leak information regardless of which ops ran; the whitelist can't enforce constant-time execution. → enclave-level side-channel mitigations. -3. **A future JAX host-callback API under an already-allowed prefix.** Default-deny closes unknown _new_ paths automatically, but a newly added dangerous symbol inside an allowed prefix (e.g. a hypothetical `jax.numpy.`) is a gap until the denylist is updated. → keep the policy versioned and reviewed against each pinned JAX release. -4. **Bugs or supply-chain compromise in the trusted libraries.** The model only constrains the _caller's_ code; a malicious jax/flax/orbax build defeats it. → attest the exact library versions/hashes; the DO must trust those releases. +- [docs/verify.md](docs/verify.md) — how verification works and what private code may do (allow side). +- [docs/blacklist.md](docs/blacklist.md) — default-deny catalog and violation codes (deny side). +- [docs/code-layout.md](docs/code-layout.md) — source modules and test layout. diff --git a/packages/syft-restrict/disallowed-ast-examples.md b/packages/syft-restrict/disallowed-ast-examples.md deleted file mode 100644 index 15d94f07f9e..00000000000 --- a/packages/syft-restrict/disallowed-ast-examples.md +++ /dev/null @@ -1,27 +0,0 @@ -# Disallowed AST — illustrative examples - -> Companion to **[restrict.md](restrict.md)**. **Safety does not come from this list.** The enforcement mechanism is the **default-deny whitelist** (restrict.md §3): any AST node type not explicitly in the ALLOW set is rejected — including constructs nobody has listed and syntax that doesn't exist yet. This table is therefore **not exhaustive and not load-bearing**: it names the data-theft and escape constructs that matter most, purely to make the threat surface legible to a reviewer. Deleting a row here changes nothing — the construct stays rejected unless it _also_ appears in the allow set. Node types absent from the whitelist (the walrus operator `NamedExpr`, the `match`/`case` pattern family, `Assert`, `except*`/`TryStar`, PEP 695 type-parameter nodes, …) are denied by default-deny — the "unlisted modern syntax" row below lists a representative set so a reviewer can confirm they were considered, not because the checker special-cases them. - -| Banned construct (example, not exhaustive) | `ast` node / form | Why | -| ----------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| **Build-and-run code from a string** | `exec`, `eval`, `compile`, `__import__` as call targets | arbitrary code / import escape | -| `Exec` / `eval` | (py2 `Exec`; `Call` to `eval`) | same | -| **Non-allowlisted imports** | `Import`, `ImportFrom` to any module not in the allowlist | an import is the _only_ way the code can name anything outside itself; this is the crux (restrict.md §3.2) | -| **Reflection / dunder** | any `Name`/`Attribute`/`arg`/alias starting with `_` | the `__class__` → `__globals__` → `__subclasses__` → `__builtins__` escape ladder | -| `getattr`/`setattr`/`delattr`/`hasattr`/`vars`/`globals`/`locals`/`dir` | `Call` to these names | dynamic attribute access defeats the static attribute-path allowlist | -| `getattr`/`setattr` with a computed name | — | even a guarded `getattr` would break static provability; deny outright | -| **`.format` / any named method on a value** | `Call` whose func is an `Attribute` on a non-allowlisted value, e.g. `"{0.__class__.__init__.__globals__}".format(x)` | a format string walks attributes at runtime, _invisibly_ in the AST; subsumed by the no-named-method-on-a-value rule (restrict.md §3.2.2) — only generic operator bundles are allowed on a value | -| **Host I/O** | `open`, real `print`, `input`, `file` | filesystem / stdout exfiltration channels | -| **OS / process** | any `os`, `sys`, `subprocess`, `shutil`, `pathlib`, `socket`, `ssl`, `http`, `urllib`, `requests`, `ctypes`, `cffi`, `mmap`, `multiprocessing`, `threading`, `asyncio`, `signal` import or attr | direct exfiltration / native code / escape | -| **Pickle / marshal** | `pickle`, `marshal`, `dill`, `shelve`, `joblib` | code-execution on load + serialization exfiltration | -| **Global mutation** | `Global`, `Nonlocal`; assignment to module-level names from inside functions | prevents stashing data in module state for later read-out | -| **Attribute write to a foreign object** | `Store` on `obj.` where `obj` isn't `self` (e.g. `some_obj.send = data`) | only `self.` writes are allowed (restrict.md §3.1); writing onto another object is an exfiltration channel | -| **Async** | `AsyncFunctionDef`, `Await`, `AsyncFor`, `AsyncWith`, `Yield`, `YieldFrom` | concurrency/escape, not needed for inference | -| **Context managers** | `With` | `open(...) as f`, `socket(...)` — deny by default (none needed in pure inference) | -| **Exceptions reaching the host** | `Try`, `Raise` | allow _optionally_ but only `raise ValueError(...)` shapes; danger is `except` swallowing a guard or walking traceback frames — default DENY, review per-file | -| **Arbitrary decorators** | `decorator_list` entries not in allowlist | a decorator is an arbitrary call wrapping the function — must be on the allowlist (restrict.md §3.2.1) | -| **`@property` / descriptors** | `@property` (and any descriptor) — a decorator not on the allowlist | runs code on a _bare attribute access_ (`obj.w`), a hook that fires with no explicit call; rejected by the decorator allowlist (restrict.md §3.2.1). Pure inference needs only `setup`/`__call__` | -| **Loop/comprehension over I/O** | `For`/`*Comp` whose iterable is a denied call | `[x for x in open(f)]` reintroduces I/O | -| **Denied call in a "passive" position** | a `Call` to a denied target in a default arg (`def f(x=evil())`), an annotation (`def g(x: evil())`), or a bare class-body statement (`class C: evil()`) | these run early, at def/class creation; call-checking is position-independent — the checker walks every node, so a denied target is caught wherever it sits | -| `del` of names that are guards | `Delete` | could un-define a guard; deny `Delete` entirely | -| **Unlisted modern syntax** | walrus `NamedExpr` `(x := …)`; `Match`/`case`; `Assert`; `except*`/`TryStar`; PEP 695 type params (`def f[T]`, `type X = …`) | none appear in the allowlist, so default-deny rejects them — listed so a reviewer can confirm they were considered (asserts also vanish under `python -O`, so they must never be load-bearing for safety) | diff --git a/packages/syft-restrict/docs/blacklist.md b/packages/syft-restrict/docs/blacklist.md new file mode 100644 index 00000000000..c0554057a2e --- /dev/null +++ b/packages/syft-restrict/docs/blacklist.md @@ -0,0 +1,261 @@ +# What is not allowed + +This document lists everything the checker rejects in the **private** region. + +- The model is **default-deny**. Tables below name common rejections with clear +codes. +- Anything not listed as allowed in [verify.md](verify.md) is also +rejected. +- Each entry includes the **violation code** returned by `verify()`. + + +--- + +## Index of violation codes + +| Code | Meaning | +| ------------------ | ------------------------------------------------------------------------ | +| `node-type` | Syntax not on the allow list (walrus, `match`, …) | +| `banned-name` | Reference or call to a banned builtin (`open`, `eval`, …) | +| `banned-construct` | Forbidden statement/expression form (import, `with`, f-string, …) | +| `call-not-allowed` | Library path not allow-listed (or hit by `disallow_functions`) | +| `call-unresolved` | Call target not provably safe | +| `method-on-value` | Named method on an unknown value (`x.reshape`) | +| `attr-on-value` | Attribute read/write on an unknown value, or bad `self` chain | +| `attr-not-allowed` | Library attribute path not allow-listed | +| `dunder-attr` | Dunder attribute (`.__class__`, …) | +| `dunder-name` | Bare dunder name (`__class__`) | +| `dunder-def` | Defining a forbidden magic method | +| `decorator` | Decorator not on the allow list | +| `class-keyword` | e.g. `metaclass=` | +| `class-base` | Base class not allow-listed | +| `reserved-name` | Rebinding a trusted name (`jnp`, public wrapper, private def, `self`, …) | +| `operator-disabled` | Operator used without enabling its group in `allow_operators` | + +--- + +## Forbidden constructs + +These constructs are banned outright in private code. + +Code: **`banned-construct`**. + +| Construct | Example | Why | +| --------------------- | ---------------------------------- | ------------------------------------------------------ | +| Import | `import os`, `from os import path` | Imports belong in public code | +| `with` | `with open(f) as g: ...` | Runs enter/exit hooks; often I/O | +| `try` / `raise` | `try: ... finally: ...` | Exception tricks / host surfaces | +| `global` / `nonlocal` | `global x` | Escape local naming rules | +| `del` | `del x` | Can remove names the policy relies on | +| `assert` | `assert cond` | Disappears under `python -O` | +| Async | `async def`, `await`, … | Out of scope for pure inference | +| Generators | `yield`, `yield from` | Suspended execution | +| F-strings | `f"hi"`, `f"{x}"` | Interpolation runs formatting with no normal call site | + +```python +# rejected +import os +y = f"value={x}" +with ctx() as g: + pass +``` + +### Unknown / future syntax + +Any syntax not on the allow list is rejected. + +Code: **`node-type`**. + +| Example | Notes | +| ---------------------- | ---------------- | +| `y = (z := 1)` | Walrus | +| `match x: case 1: ...` | Pattern matching | + +New Python syntax stays blocked until reviewed and explicitly allowed. + +--- + +## Banned builtins + +These names may never be **used** (i.e, loaded) in private code, even if they are +**not called**. + +Code: **`banned-name`**. + +### Dynamic code / reflection / I/O + +- `eval` +- `exec` +- `compile` +- `__import__` +- `getattr` +- `setattr` +- `delattr` +- `hasattr` +- `vars` +- `globals` +- `locals` +- `dir` +- `open` +- `input` +- `breakpoint` +- `memoryview` +- `type` +- `__build_class__` +- `print` + +### Formatting / buffer builtins + +Same idea as calling dunders on a value (`x.__repr__()`), spelled as a bare call: + +- `repr` +- `str` +- `ascii` +- `format` +- `bytes` + +> [!WARNING] +> `bytes(x)` can dump raw buffer contents. `print` is a stdout channel. F-strings are banned +> separately as constructs (above), including with no `{...}` at all. + +```python +# all rejected (banned-name) +f = open +d = {"o": open} +def run(op=open): ... +open("/etc/passwd") # call reports banned-name once +y = [v for v in open(path)] # passive positions still checked +``` + +--- + +## Calls and attributes + +| What | Example | Code | +| -------------------------------------------- | ------------------------------------------ | ------------------ | +| Library path not allowed | `np.dot(a, b)` | `call-not-allowed` | +| Disallow list hits an otherwise-allowed path | `jnp.save(...)` under `disallow_functions` | `call-not-allowed` | +| Call target not proven safe | `fn(x)` where `fn` is a parameter | `call-unresolved` | +| Named method on a value | `x.reshape(8, -1)`, `items.append(1)` | `method-on-value` | +| Attribute on a value | `x.shape`, `x.T`, `obj.send = data` | `attr-on-value` | +| Deep `self` chain | `self.a.b`, `self.sub.evil(...)` | `attr-on-value` | +| Unsafe `self.(...)` | stashed `open` on `self` in `setup` | `attr-on-value` | +| Dunder attribute | `obj.__class__`, `self.__dict__` | `dunder-attr` | +| Bare dunder name | `c = __class__` | `dunder-name` | +| Library attr not allowed | `np.pi` when numpy is not allowed | `attr-not-allowed` | + +```python +# rejected +a = x.reshape(8, -1) # method-on-value +b = x.shape # attr-on-value +c = obj.__class__ # dunder-attr + +def apply(fn, x): + return fn(x) # call-unresolved + +``` + +> [!IMPORTANT] +> **`call-unresolved` is default-deny for callees.** It catches dangerous callables that never +> mention a banned builtin by name (parameters, opaque locals, `d["k"](x)`). + +Only single-level `self.` / `cls.` is special-cased. Rules for when `self.x(...)` is +safe are in [verify.md](verify.md#self-and-flax-style-modules). + +--- + +## Classes, decorators, definitions + +| What | Example | Code | +| -------------------------- | ------------------------------------- | --------------- | +| Non-allow-listed decorator | `@evil`, `@property`, `@staticmethod` | `decorator` | +| Class keywords | `class M(nn.Module, metaclass=Meta)` | `class-keyword` | +| Bad base class | `class M(SomeLib)`, `class M(object)` | `class-base` | +| Forbidden magic method | `def __getattr__`, `def __reduce__` | `dunder-def` | + +Allowed hooks only: `setup`, `__call__`, `__post_init__`. +Allowed decorators only: `nn.compact`, `jax.jit`, `jax.named_scope`, `flax.linen.compact`. + +```python +# rejected +class M(object): # class-base + @property # decorator + def w(self): + return 1 + + def __getattr__(self, name): # dunder-def + return None +``` + +--- + +## Reserved names + +Rebinding a name the checker trusts reports **`reserved-name`**. + +| Reserved | Example rebind | Why | +| -------------------- | --------------------------------------------- | -------------------------------------- | +| Import aliases | `jnp = make_evil()` | Would lie about resolved library paths | +| Public wrappers | nested `def transpose` shadowing a public one | Defeats the reviewed wrapper | +| Private defs/classes | `helper = evil` after `def helper` | Bare calls trust the def by name | +| Safe builtins | `list = None` | Bare calls trust `list(...)` by name | +| `self` / `cls` | `self = x`, nested `def f(self):` | `self.*` is trusted by spelling | + +Applies to assignment, `for` / comprehension targets, parameters, and nested defs. +Private def names are also protected against rebind in **public** glue between private ranges. + +```python +# rejected +import jax.numpy as jnp +jnp = 1 # reserved-name + +def helper(x): + return x +def f(evil): + helper = evil # reserved-name + return helper(1) +``` + +--- + +## Operator bundles + +If a group of operators is not in `allow_operators`, using any of its operators +reports **`operator-disabled`**. + +| Bundle | If missing, these fail | +| ------------ | ---------------------- | +| `arithmetic` | `a + b`, `-a`, … | +| `comparison` | `a < b`, `a and b`, … | +| `indexing` | `x[0]`, `x[1:3]`, … | + +--- + +## Optional `disallow_functions` + +There is no built-in library denylist. Everything not explicitly allow-listed is +rejected. Safety comes from **`allow_functions`**. + +When you use a broad allow (`jax.*`), pass `disallow_functions=[...]` for a hard floor. Hits report +**`call-not-allowed`**. + +Useful patterns under broad JAX/Flax allows: + +- Host / debug / experimental: `jax.experimental.*`, `jax.debug.*`, `jax.pure_callback`, `*.io_callback`, `*.host_callback*` +- FFI / interop: `jax.dlpack.*`, `jax.ffi*` +- Array ↔ disk: `jax.numpy.save`, `load`, `tofile`, `fromfile`, `memmap`, … +- Checkpointing: `flax.serialization.*`, `flax.training.checkpoints.*`, `orbax.*` + +```python +# public import, private use — still resolved and checked +from jax.numpy import save as persist +persist(x, "out.npz") # call-not-allowed if disallow includes jax.numpy.save +``` + +--- + +## See also + +- [verify.md](verify.md) — how checking works and what is allowed +- `tests/verify/test_disallowed.py` — disallowed tests +- `tests/verify/test_bypasses.py` — multi-step attack regressions diff --git a/packages/syft-restrict/docs/code-layout.md b/packages/syft-restrict/docs/code-layout.md new file mode 100644 index 00000000000..b10cf159a18 --- /dev/null +++ b/packages/syft-restrict/docs/code-layout.md @@ -0,0 +1,23 @@ +# Code layout + +Map of `src/syft_restrict/`: + +| Module | Role | +| --- | --- | +| `astutil.py` | Line ranges, import scan, small syntax helpers. No policy. | +| `policy.py` | Allow lists, safe/banned builtins, decorator/hook lists, `Policy`. | +| `verifier.py` | Static checker: walks private lines default-deny, reports violations. | +| `obfuscator.py` | After a clean verify: rename private identifiers, blank constants. | +| `runner.py` | `run()` — verify → obfuscate → certificate. | +| `errors.py` | `RestrictError`, `PolicyViolation`. | +| `__init__.py` | Public exports. | + +Tests under `tests/verify/`: + +| File | Role | +| --- | --- | +| `test_whitelist.py` | Green path (aligns with [verify.md](verify.md)) | +| `test_whitelisted_lib.py` | Library paths + operator bundles | +| `test_disallowed.py` | Default-deny catalog (aligns with [blacklist.md](blacklist.md)) | +| `test_bypasses.py` | Multi-step escape regressions | +| `test_ranges.py` | Invalid private ranges | diff --git a/packages/syft-restrict/docs/verify.md b/packages/syft-restrict/docs/verify.md new file mode 100644 index 00000000000..7b272aee8c1 --- /dev/null +++ b/packages/syft-restrict/docs/verify.md @@ -0,0 +1,305 @@ +# How verification works + +`syft-restrict` checks the **private** lines of a Python file. Those lines must only do trusted +inference math: no host I/O, no dynamic code, no reaching into interpreter internals. + +The tool **never runs** your code. It reads the source, checks the private region, and returns a +list of violations (empty on success). + +> [!NOTE] +> This page is the **allow side**: how checking works and what private code may do. +> +> For everything that is default-denied, see [blacklist.md](blacklist.md). + +Tests mirror this split: + +| Doc | Test module | Role | +| ---------------------------- | ----------------------------------------------------------- | ---------------------- | +| This page | `tests/verify/test_whitelist.py`, `test_whitelisted_lib.py` | Green path | +| [blacklist.md](blacklist.md) | `tests/verify/test_disallowed.py` | Default-deny catalog | +| Edge / attack shapes | `tests/verify/test_bypasses.py` | Multi-step regressions | + +--- + +## Public vs private + +Every file has two regions: + +| Region | Typical contents | Checked? | +| ----------- | ------------------------------------------------------- | ----------------- | +| **Public** | Imports, data loading, wrappers the data owner can read | No (human review) | +| **Private** | Hidden model math (`setup` / `__call__`, layers) | Yes | + +You mark private code with 1-based line ranges (the union of `obfuscate` and `hide` in `run()`). + +```python +# public — data owner reads this +import jax.numpy as jnp +from flax import linen as nn + +def transpose(x): + if not isinstance(x, jax.Array): + raise TypeError + return x.T + +# private — checked + obfuscated +class Net(nn.Module): + def setup(self): + self.dense = nn.Dense(8) + + def __call__(self, x): + return self.dense(transpose(x)) +``` + +> [!WARNING] +> Private code may **call** public wrappers by name. Public code is trusted, not verified. +> +> A clean `verify()` means the private region cannot escape on its own, not that +> the whole file is safe to execute. + +--- + +## What the checker does + +The checker walks the AST of the private region and verifies each construct against the allow list. + +1. **Parse** the whole file. +2. **Record imports** from the public region (`import jax.numpy as jnp` → `jnp` means `jax.numpy`). +3. **Walk every private construct.** For each piece of syntax, either an + explicitly rule allows it or it fails. +4. **Collect violations** with a line number and a short code (e.g. `banned-name`). + +The default-deny semantics means that if nothing explicitly allows a construct, +it is rejected. New Python syntax stays blocked until someone reviews it. + +--- + +## Always allowed syntax + +These shapes are fine in private code **when** nested pieces also obey the rules: + +| Category | Examples | +| ----------------- | ---------------------------------------------------------- | +| Definitions | `def`, `class`, `lambda`, `return` | +| Names & constants | variables, numbers, strings (not f-strings) | +| Assignment | `=`, `+=`, annotated assigns | +| Containers | `list`, `tuple`, `dict`, `set`, comprehensions | +| Calls | `f(...)` when the callee is allowed (below) | +| Control flow | `if`, `for`, `while`, `break`, `continue`, `pass`, ternary | +| Operators | only if the matching **bundle** is enabled (below) | + +> [!NOTE] +> **Not** always allowed: imports, `with`, `try`/`raise`, `async`, generators, `assert`, `del`, +> f-strings, walrus `:=`, `match`/`case`. Those are listed under [blacklist.md](blacklist.md). + +--- + +## Per-file knobs + +### `allow_functions` — library paths callable by name + +Dotted paths are resolved through the import table. An optional +`disallow_functions` list has priority over the allow list (hard floor over a +broad glob). + +```python +import jax.numpy as jnp + +jnp.einsum("ij->i", x) # → jax.numpy.einsum — must match allow_functions +jnp.save(path, x) # → jax.numpy.save — fails if in disallow_functions +``` + +| Rule | Meaning | +| ---------------------------- | ------------------------------------------------------- | +| Use paths, not whole modules | `jax.numpy.einsum` is named; `x.method()` is not a path | +| Call libraries **inline** | `jnp.sin(x)` yes; `f = jnp.sin; f(x)` no | +| Do not rebind import aliases | `jnp = evil` is rejected | +| Prefer specific allows | `jax.numpy.einsum` over bare `jax.*` | + +### `allow_operators` — operator bundles on a value + +Named methods on an unknown value (`x.reshape(...)`, `x.T`) are never allowed: the checker cannot +prove what they do. Instead you enable **language operators** as bundles: + +| Bundle | Operators | Example | +| ------------ | ----------------------------------- | --------------- | +| `arithmetic` | `+ - * / // % ** @`, unary | `x + 1e-6` | +| `comparison` | `== != < <= > >=`, `and`/`or`/`not` | `t == "local"` | +| `indexing` | `[]` and slices | `x[..., :half]` | + +Anything else should go through a **public wrapper**: + +```python +# public wrapper — data owner can review this +def shape_of(x): + if not isinstance(x, jax.Array): + raise TypeError + return x.shape + +# private +h = shape_of(x) # allowed if shape_of is a public def +``` + +Disabled groups report `operator-disabled`. + +--- + +## What you may call + +A call is allowed only if the checker can **prove** the callee. Otherwise it reports +`call-unresolved` (even when nothing looks “banned”). + +| Allowed callee | Example | +| ---------------------------------------------------------- | ------------------------------------ | +| Allow-listed import path, called inline | `jnp.sin(x)`, `nn.Dense(8)` | +| Function or class defined in this file (private or public) | `Attention(cfg)`, `transpose(w)` | +| Safe builtin (fixed list below) | `len(xs)`, `list(range(n))` | +| Local traced to a safe source | `block = Attn(); block(x)` | +| Vetted `self.` / `self.[i]` | `self.dense(x)`, `self.layers[i](x)` | + +| Not allowed | Example | Why | +| ------------------------ | ------------------- | ---------------------- | +| Named method on a value | `x.reshape(8, -1)` | Type unknown | +| Stashed library function | `f = jnp.sin; f(x)` | Must call paths inline | +| Opaque subscript call | `d["k"](x)` | Callee not identified | + +### Safe builtins + +These may be called by bare name. They may **not** be reassigned. + +- `int` +- `float` +- `bool` +- `len` +- `range` +- `enumerate` +- `zip` +- `min` +- `max` +- `sum` +- `abs` +- `round` +- `all` +- `any` +- `tuple` +- `list` +- `dict` +- `set` +- `sorted` +- `reversed` +- `isinstance` +- `super` + +### Names you may not reassign + +Because call sites trust some names by spelling alone: + +| Name kind | Example | +| -------------------- | -------------------------------------------- | +| Import aliases | `jnp`, `nn` | +| Public wrappers | `transpose` | +| Private defs/classes | `Attention`, `helper` | +| Safe builtins | `list`, `range` | +| `self` / `cls` | only as the real first parameter of a method | + +Rebind is rejected even in **public** glue if the name is a private def, otherwise a public +`helper = evil` between private chunks could be used to evade the checker. + +--- + +## `self` and Flax-style modules + +Private models usually look like: + +```python +class Net(nn.Module): # base must be allow-listed (e.g. flax.linen.Module) + def setup(self): # often public (data owner can read wiring) + self.dense = nn.Dense(8) # allow-listed constructor → safe to call later + self.layers = [Block() for _ in range(3)] + + def __call__(self, x): # often private + x = self.dense(x) + block = self.layers[0] + return block(x) +``` + +| Rule | Meaning | +| ------------------------------------------------- | --------------------------------------------------------------------- | +| Only first method param may be named `self`/`cls` | Nested `def helper(self):` is rejected | +| Do not rebind `self`/`cls` | `self = x` is rejected | +| Single-level only | `self.dense` yes; `self.sub.evil` no | +| Call only if every assignment was a vetted source | allow-listed constructor, file-local class/def, or list/comp of those | +| Inherited attrs (never assigned) | Treated as safe (e.g. `self.param` on `nn.Module`) | +| `self.x += ...` | Always unsafe for later calls | + +> [!TIP] +> Put dangerous or opaque wiring in **public** `setup` if the data owner should read it; keep +> private `__call__` as pure math. +> +> The checker still uses public `setup` assignments when deciding +> whether `self.(...)` is safe. + +### Classes and hooks + +| Allowed | Not allowed | +| ---------------------------------------------------------------------------- | ----------------------------------------------------- | +| Bases that resolve to an allow-listed path (e.g. `nn.Module`) | `object`, random private bases, non-allow-listed libs | +| Decorators: `nn.compact`, `jax.jit`, `jax.named_scope`, `flax.linen.compact` | `@property`, `@staticmethod`, arbitrary functions | +| Defining `setup`, `__call__`, `__post_init__` | `__getattr__`, `__reduce__`, other magic methods | +| — | `metaclass=` / other class keywords | + +--- + +## Whitelist examples + +```python +# control flow + safe builtins +def f(xs): + acc = [abs(v) for v in xs if v] + for v in xs: + if v: + break + return acc +``` + +```python +# public wrapper + private call +def transpose(x): + return x # data owner reviews this + +def private_math(w): + return transpose(w) +``` + +```python +# library path (imports public, use private) +import jax.numpy as jnp +# private: +r = jnp.einsum("ij,jk->ik", a, b) +``` + +```python +# annotations may mention banned type names; they are not executed +def f(x: list[str]) -> dict[str, bytes]: + return {} +``` + +--- + +## Limits of static checking + +Even a clean verify does **not** stop: + +1. **Encoding secrets in model outputs** (logits/tokens) +2. **Timing / cache side channels** — enclave-level concern. +3. **New dangerous APIs under a broad allow** (`jax.*`) — prefer tight allows + `disallow_functions`. +4. **Compromised JAX/Flax builds** — attest library versions separately. +5. **Malicious public wrappers** — still depend on human review of the public region. + +--- + +## See also + +- [blacklist.md](blacklist.md) — default-deny catalog and violation codes +- [code-layout.md](code-layout.md) — source modules +- `tests/verify/` — executable examples of allow vs deny diff --git a/packages/syft-restrict/examples/gemma_inference.certificate.json b/packages/syft-restrict/examples/gemma_inference.certificate.json index 0c61ae7c906..327620e36fe 100644 --- a/packages/syft-restrict/examples/gemma_inference.certificate.json +++ b/packages/syft-restrict/examples/gemma_inference.certificate.json @@ -1,6 +1,6 @@ { - "source_sha256": "11452760f1bd895b55214caa7a07e314174e75c4664af02bd177d18f4b008101", - "policy_id": "a73438a4b78b4438", + "source_sha256": "3b462b5a6490b9dba8a93cf1a06da37601b002ab3a6f7497c01ee6698a671a18", + "policy_id": "b4a557058704eed6", "restrict_version": "0.1.0", "private_ranges": [ [24, 88], diff --git a/packages/syft-restrict/examples/gemma_inference.obfuscated.py b/packages/syft-restrict/examples/gemma_inference.obfuscated.py index 57798b95bd8..ab009defb83 100644 --- a/packages/syft-restrict/examples/gemma_inference.obfuscated.py +++ b/packages/syft-restrict/examples/gemma_inference.obfuscated.py @@ -138,12 +138,12 @@ def _get(module, name): def shape_of(x): - """Visible wrapper: read an array's shape — an attribute read on a value, not allowed in the hidden region.""" + """Public wrapper: read an array's shape — an attribute read on a value, not allowed in the private region.""" return x.shape def append_to(lst, item): - """Visible wrapper: append to a Python list (a named method on a value).""" + """Public wrapper: append to a Python list (a named method on a value).""" lst.append(item) return lst @@ -317,7 +317,7 @@ def load_params(weights_dir, cfg): # ── Setup (convenience entry point) ─────────────────────────────────────── -def setup(weights_dir): +def setup_model(weights_dir): """Configure model, load weights and tokenizer. Returns (model, tokenizer, params). diff --git a/packages/syft-restrict/examples/gemma_inference.py b/packages/syft-restrict/examples/gemma_inference.py index ce9160aed02..db075c9bc63 100644 --- a/packages/syft-restrict/examples/gemma_inference.py +++ b/packages/syft-restrict/examples/gemma_inference.py @@ -138,12 +138,12 @@ def _get(module, name): def shape_of(x): - """Visible wrapper: read an array's shape — an attribute read on a value, not allowed in the hidden region.""" + """Public wrapper: read an array's shape — an attribute read on a value, not allowed in the private region.""" return x.shape def append_to(lst, item): - """Visible wrapper: append to a Python list (a named method on a value).""" + """Public wrapper: append to a Python list (a named method on a value).""" lst.append(item) return lst @@ -317,7 +317,7 @@ def load_params(weights_dir, cfg): # ── Setup (convenience entry point) ─────────────────────────────────────── -def setup(weights_dir): +def setup_model(weights_dir): """Configure model, load weights and tokenizer. Returns (model, tokenizer, params). diff --git a/packages/syft-restrict/examples/generate.py b/packages/syft-restrict/examples/generate.py index 6b476e8c004..0f2d65d0c82 100644 --- a/packages/syft-restrict/examples/generate.py +++ b/packages/syft-restrict/examples/generate.py @@ -2,7 +2,7 @@ Run from the package root: uv run python examples/generate.py -The policy here lists the *exact* lower-level JAX/Flax symbols the hidden region uses, +The policy here lists the *exact* lower-level JAX/Flax symbols the private region uses, rather than broad globs (`jax.*`). Each call leaf (`jax.numpy.einsum`, …) is enforced individually; `jax.lax` / `jax.nn` are included only because the calls are written as deep attribute paths (`jax.lax.rsqrt`), so the checker also evaluates those module @@ -88,7 +88,7 @@ "jax.lax", "jax.nn", ], - allow_methods=["arithmetic", "indexing", "comparison"], + allow_operators=["arithmetic", "indexing", "comparison"], ) cert_path = EX / "gemma_inference.certificate.json" diff --git a/packages/syft-restrict/src/syft_restrict/__init__.py b/packages/syft-restrict/src/syft_restrict/__init__.py index c4501d0e1f9..ac90467b52d 100644 --- a/packages/syft-restrict/src/syft_restrict/__init__.py +++ b/packages/syft-restrict/src/syft_restrict/__init__.py @@ -1,8 +1,8 @@ """syft-restrict — verify + obfuscate JAX/Flax inference code. `run` is the entry point: it statically proves the private model-definition lines only do trusted math -(no data theft), then obfuscates them so the model architecture stays secret. See README and the design -under `research/verifuscate/`. +(no data theft), then obfuscates them so the model architecture stays secret. See the README for the +pitch, docs/verify.md for how verification works, and docs/blacklist.md for what is rejected. """ from __future__ import annotations diff --git a/packages/syft-restrict/src/syft_restrict/astutil.py b/packages/syft-restrict/src/syft_restrict/astutil.py new file mode 100644 index 00000000000..c8e65d0a79c --- /dev/null +++ b/packages/syft-restrict/src/syft_restrict/astutil.py @@ -0,0 +1,153 @@ +"""Pure AST + line-range helpers shared by the verifier, obfuscator, and runner.""" + +from __future__ import annotations + +import ast + +from pydantic import BaseModel + + +# ── line ranges ──────────────────────────────────────────────────────────────────────────── +# The caller marks the private region as a list of [start, end] 1-based inclusive line ranges. +def normalize_ranges(private) -> list[tuple[int, int]]: + """Coerce the caller's ``[[lo, hi], ...]`` into a list of ``(int, int)`` tuples. + + Raises ``ValueError`` on a malformed range (``hi < lo``) rather than silently matching no + lines -- an inverted range must never be mistaken for "nothing to check here". + """ + ranges = [(int(lo), int(hi)) for lo, hi in private] + for lo, hi in ranges: + if hi < lo: + raise ValueError(f"invalid range [{lo}, {hi}]: end must be >= start") + return ranges + + +def row_in_ranges(row: int, ranges) -> bool: + """True if a 1-based line number falls inside any range.""" + return any(lo <= row <= hi for lo, hi in ranges) + + +def node_in_ranges(node: ast.AST, ranges) -> bool: + """True if a node has a line number and it falls inside any range. + + Some nodes (e.g. ``ast.comprehension``) carry no position; those are never "in range". + """ + line = getattr(node, "lineno", None) + return line is not None and row_in_ranges(line, ranges) + + +# ── reading names and dotted paths off the tree ────────────────────────────────────────────── +def dotted_name(node: ast.AST) -> str | None: + """The dotted path for a pure Name/Attribute chain (``jnp.numpy.einsum``), else None. + + Returns None the moment the chain hits anything dynamic (a call, a subscript), because such a + chain can no longer be resolved to a static import path. + """ + parts: list[str] = [] + cur = node + while isinstance(cur, ast.Attribute): + parts.append(cur.attr) + cur = cur.value + if isinstance(cur, ast.Name): + parts.append(cur.id) + return ".".join(reversed(parts)) + return None + + +def is_dunder(name: str) -> bool: + """True for any identifier that starts with ``__`` (the reflection/hook surface).""" + return name.startswith("__") + + +def self_attr_name(node: ast.AST) -> str | None: + """The attr for a single-level ``self.`` / ``cls.`` access, else None.""" + if ( + isinstance(node, ast.Attribute) + and isinstance(node.value, ast.Name) + and node.value.id in ("self", "cls") + ): + return node.attr + return None + + +def rooted_in_self(node: ast.AST) -> bool: + """True if an Attribute/Subscript chain's ultimate base is ``self`` / ``cls``.""" + cur = node + while isinstance(cur, (ast.Attribute, ast.Subscript)): + cur = cur.value + return isinstance(cur, ast.Name) and cur.id in ("self", "cls") + + +def describe(node: ast.AST) -> str: + """A human-readable label for a node in a violation message.""" + return dotted_name(node) or type(node).__name__ + + +# ── whole-file scan ────────────────────────────────────────────────────────────────────────── +class FileScan(BaseModel): + """Names harvested from the whole file, used to classify calls in the private region.""" + + import_bindings: dict[ + str, str + ] # import alias -> fully-qualified module path (jnp -> jax.numpy) + private_defs: set[str] # class/func names defined inside the private region + public_defs: set[str] # class/func names defined in the public region (non-methods) + private_def_ids: set[ + int + ] # id() of the first (canonical) node behind each private_defs name + method_ids: set[int] # id() of direct FunctionDef children of any ClassDef body + + +def scan_file(tree: ast.Module, private_ranges) -> FileScan: + """Collect import bindings and the class/func names defined inside vs. outside the private region. + + Imports live in the public region (they're banned inside the private one), but their bindings are + what the checker resolves private-region calls against — so we scan the whole file, not just the + private lines. + """ + import_bindings: dict[str, str] = {} # import alias -> fully-qualified module path + private_defs: set[str] = set() # class/func names defined inside the private region + public_defs: set[str] = set() # class/func names defined in the public region + private_def_ids: set[int] = ( + set() + ) # id() of the first node claiming each private_defs name + # Methods (direct FunctionDef children of a class body) are called via `self.`, never + # by bare name, so they're excluded here -- otherwise unrelated classes sharing a hook name + # (`setup`, `__call__`) would look like one shadowing the other. + method_ids = { + id(child) + for node in ast.walk(tree) + if isinstance(node, ast.ClassDef) + for child in node.body + if isinstance(child, ast.FunctionDef) + } + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + # `import jax.numpy as jnp` -> bindings["jnp"] = "jax.numpy"; `import os.path` -> bindings["os"] = "os.path" + import_bindings[alias.asname or alias.name.split(".")[0]] = alias.name + elif isinstance(node, ast.ImportFrom) and node.module: + for alias in node.names: + import_bindings[alias.asname or alias.name] = ( + f"{node.module}.{alias.name}" + ) + elif isinstance(node, ast.FunctionDef) and id(node) in method_ids: + continue + elif isinstance(node, (ast.FunctionDef, ast.ClassDef)): + if node_in_ranges(node, private_ranges): + # `class Net:` / `def _secret():` on a private line -> private_defs.add("Net" / "_secret") + # ast.walk is breadth-first, so the first node to claim a name is always the + # shallowest (outermost) one -- any later def/class reusing the name is a shadow. + if node.name not in private_defs: + private_def_ids.add(id(node)) + private_defs.add(node.name) + else: + # public FunctionDef *or* ClassDef — both may be called by bare name from private + public_defs.add(node.name) + return FileScan( + import_bindings=import_bindings, + private_defs=private_defs, + public_defs=public_defs, + private_def_ids=private_def_ids, + method_ids=method_ids, + ) diff --git a/packages/syft-restrict/src/syft_restrict/obfuscator.py b/packages/syft-restrict/src/syft_restrict/obfuscator.py index 80c1d2b4d14..f9bfe838bea 100644 --- a/packages/syft-restrict/src/syft_restrict/obfuscator.py +++ b/packages/syft-restrict/src/syft_restrict/obfuscator.py @@ -1,4 +1,4 @@ -"""The display transform (research approach A). +"""The display transform — turn the verified private lines into a readable-but-secret artifact. ``obfuscate(source, private, scan)`` returns a copy of ``source`` in which only the lines inside the private ranges are transformed — identifiers renamed to neutral placeholders, constant values and @@ -16,8 +16,15 @@ import keyword import tokenize +from .astutil import ( + FileScan, + dotted_name, + is_dunder, + node_in_ranges, + normalize_ranges, + row_in_ranges, +) from .policy import DEFAULT_KEEP -from .verifier import FileScan, _dotted, _is_dunder, _normalize_ranges _BLANK = "■" # ■ — replaces a single constant/string token (obfuscate mode) _HIDE = "■■■■■■■■" # replaces a whole line's code, indentation kept (hide mode) @@ -25,6 +32,13 @@ "# hidden/obfuscated lines can only execute restricted python, " "see restrict docs for more details" ) +# Layout-only tokens: never renamed/blanked, and don't count as "last token was an operator dot" +# for the attribute-vs-value rename decision in obfuscate() below. +_LAYOUT_TOKENS = (tokenize.NL, tokenize.NEWLINE, tokenize.INDENT, tokenize.DEDENT) +# PEP 701 (Python 3.12+) tokenizes an f-string's literal text as its own FSTRING_MIDDLE token(s) +# instead of folding the whole f-string into one STRING token -- detect it by feature rather than +# a hardcoded version number, since that's what actually determines which token carries the text. +_FSTRING_MIDDLE = getattr(tokenize, "FSTRING_MIDDLE", None) # Builtins kept readable (they reveal nothing about the architecture). _KEEP_BUILTINS = frozenset( @@ -60,33 +74,32 @@ ) -def obfuscate(source: str, obfuscate_ranges, hide_ranges, scan: FileScan) -> str: - ranges = _normalize_ranges(obfuscate_ranges) - tree = ast.parse(source) - value_map, attr_map = _build_maps(tree, ranges, scan) - - keep_values = ( - DEFAULT_KEEP - | set(scan.bindings) - | set(scan.visible_defs) - | _KEEP_BUILTINS +def _keep_values(scan: FileScan) -> set[str]: + """Names left readable in the output: import aliases, public wrapper names, safe builtins, keywords.""" + return ( + set(DEFAULT_KEEP) + | set(scan.import_bindings) + | set(scan.public_defs) + | set(_KEEP_BUILTINS) | set(keyword.kwlist) | set(getattr(keyword, "softkwlist", [])) ) + +def obfuscate(source: str, obfuscate_ranges, hide_ranges, scan: FileScan) -> str: + ranges = normalize_ranges(obfuscate_ranges) + tree = ast.parse(source) + value_map, attr_map = _build_maps(tree, ranges, scan) + keep_values = _keep_values(scan) + edits: list[tuple[int, int, int, int, str]] = [] tokens = list(tokenize.generate_tokens(io.StringIO(source).readline)) prev_op_dot = False for tok in tokens: srow, scol = tok.start erow, ecol = tok.end - if not _row_in_ranges(srow, ranges): - if tok.type not in ( - tokenize.NL, - tokenize.NEWLINE, - tokenize.INDENT, - tokenize.DEDENT, - ): + if not row_in_ranges(srow, ranges): + if tok.type not in _LAYOUT_TOKENS: prev_op_dot = tok.type == tokenize.OP and tok.string == "." continue @@ -100,6 +113,12 @@ def obfuscate(source: str, obfuscate_ranges, hide_ranges, scan: FileScan) -> str edits.append((srow, scol, erow, ecol, new)) elif tok.type == tokenize.STRING: edits.append((srow, scol, erow, ecol, f'"{_BLANK}"')) + elif tok.type == _FSTRING_MIDDLE: + # The f-string's literal text between `{...}` parts; FSTRING_START/END (the f"/" quote + # delimiters) are left as-is, and any interpolated expression's own tokens (NAME, etc.) + # still flow through the normal token handling above/below like any other reference. + if tok.string: + edits.append((srow, scol, erow, ecol, _BLANK)) elif tok.type == tokenize.NUMBER: edits.append((srow, scol, erow, ecol, _BLANK)) elif tok.type == tokenize.COMMENT: @@ -108,31 +127,39 @@ def obfuscate(source: str, obfuscate_ranges, hide_ranges, scan: FileScan) -> str ) # strip comment text but keep a placeholder (incl. commented-out # configs) so the artifact stays line-aligned without long blank runs - if tok.type not in ( - tokenize.NL, - tokenize.NEWLINE, - tokenize.INDENT, - tokenize.DEDENT, - ): + if tok.type not in _LAYOUT_TOKENS: prev_op_dot = tok.type == tokenize.OP and tok.string == "." - return _apply_hides(_apply_edits(source, edits), _normalize_ranges(hide_ranges)) + return _apply_hides(_apply_edits(source, edits), normalize_ranges(hide_ranges)) # ── build the deterministic rename maps from the AST ───────────────────────────────────── def _build_maps(tree: ast.Module, ranges, scan: FileScan): - keep_attrs: set[str] = set() - mangle_attr_names: set[str] = set() - value_occurrences: list[tuple[tuple[int, int], str]] = [] private_classes = _names_of(tree, ast.ClassDef, ranges) private_funcs = _names_of(tree, ast.FunctionDef, ranges) + mangle_attr_names, keep_attrs, value_occurrences = _classify_nodes( + tree, ranges, scan + ) + attr_map = _assign_attr_placeholders(mangle_attr_names, keep_attrs) + value_map = _assign_value_placeholders( + value_occurrences, scan, private_classes, private_funcs + ) + return value_map, attr_map + +def _classify_nodes(tree: ast.Module, ranges, scan: FileScan): + """One pass over the private nodes: which attribute names must be mangled vs. kept readable + (a public library attr like ``jnp.einsum``), and every Name/arg/keyword/def occurrence that may + need a value placeholder, each tagged with its source position for later ordering.""" + keep_attrs: set[str] = set() + mangle_attr_names: set[str] = set() + value_occurrences: list[tuple[tuple[int, int], str]] = [] for node in ast.walk(tree): - if not _node_in_ranges(node, ranges): + if not node_in_ranges(node, ranges): continue - if isinstance(node, ast.Attribute) and not _is_dunder(node.attr): - root = (_dotted(node.value) or "").split(".")[0] - if root in scan.bindings: + if isinstance(node, ast.Attribute) and not is_dunder(node.attr): + root = (dotted_name(node.value) or "").split(".")[0] + if root in scan.import_bindings: keep_attrs.add( node.attr ) # public library attr (e.g. jnp.einsum) — stays readable @@ -150,21 +177,27 @@ def _build_maps(tree: ast.Module, ranges, scan: FileScan): # the defined name itself, so a class/def signature line is renamed even when the # name is only referenced from hidden (blanked) lines elsewhere value_occurrences.append(((node.lineno, node.col_offset), node.name)) + return mangle_attr_names, keep_attrs, value_occurrences - # attr placeholders, in sorted name order for determinism - attr_map: dict[str, str] = {} - for i, name in enumerate(sorted(mangle_attr_names - keep_attrs)): - attr_map[name] = f"░a{i}" - - # value placeholders, assigned in source order (first occurrence wins) - keep_values = ( - DEFAULT_KEEP - | set(scan.bindings) - | set(scan.visible_defs) - | _KEEP_BUILTINS - | set(keyword.kwlist) - | set(getattr(keyword, "softkwlist", [])) - ) + +def _assign_attr_placeholders( + mangle_attr_names: set[str], keep_attrs: set[str] +) -> dict[str, str]: + """Attribute-name placeholders (``░a0``, ``░a1``, …), in sorted name order for determinism.""" + return { + name: f"░a{i}" for i, name in enumerate(sorted(mangle_attr_names - keep_attrs)) + } + + +def _assign_value_placeholders( + value_occurrences: list[tuple[tuple[int, int], str]], + scan: FileScan, + private_classes: set[str], + private_funcs: set[str], +) -> dict[str, str]: + """Value-name placeholders (``░Cls0``/``░fn0``/``░v0``, …), assigned in source order (first + occurrence wins).""" + keep_values = _keep_values(scan) value_map: dict[str, str] = {} counters = {"cls": 0, "fn": 0, "v": 0} for _pos, name in sorted(value_occurrences): @@ -179,14 +212,14 @@ def _build_maps(tree: ast.Module, ranges, scan: FileScan): else: value_map[name] = f"░v{counters['v']}" counters["v"] += 1 - return value_map, attr_map + return value_map def _names_of(tree, node_type, ranges) -> set[str]: return { n.name for n in ast.walk(tree) - if isinstance(n, node_type) and _node_in_ranges(n, ranges) + if isinstance(n, node_type) and node_in_ranges(n, ranges) } @@ -210,7 +243,7 @@ def _apply_hides(text: str, hide_ranges) -> str: in_block = False # inside a run of consecutive hidden line numbers noted = False # has the explanatory note been added for the current block yet for i, line in enumerate(lines, 1): - if not _row_in_ranges(i, hide_ranges): + if not row_in_ranges(i, hide_ranges): in_block = noted = False # a non-hidden line ends the block continue if not in_block: @@ -225,12 +258,3 @@ def _apply_hides(text: str, hide_ranges) -> str: noted = True lines[i - 1] = f"{indent}{_HIDE}{note}{newline}" return "".join(lines) - - -def _row_in_ranges(row: int, ranges) -> bool: - return any(lo <= row <= hi for lo, hi in ranges) - - -def _node_in_ranges(node: ast.AST, ranges) -> bool: - line = getattr(node, "lineno", None) - return line is not None and _row_in_ranges(line, ranges) diff --git a/packages/syft-restrict/src/syft_restrict/policy.py b/packages/syft-restrict/src/syft_restrict/policy.py index 2d5537c812d..5682d99e9b4 100644 --- a/packages/syft-restrict/src/syft_restrict/policy.py +++ b/packages/syft-restrict/src/syft_restrict/policy.py @@ -1,24 +1,24 @@ -"""Policy: what the hidden region is allowed to call and do. +"""Policy: what the private region is allowed to call and do (see docs/verify.md). -Two channels, mirroring the two verification mechanisms (see research approach-B §3.6.5): +Two channels the author configures per file: -- ``functions`` — dotted paths callable BY NAME (resolved exactly against the import bindings), - e.g. ``jax.*``, ``flax.linen.*``. Checked by glob match, with ``JAX_DENYLIST`` beating the allow. -- ``methods`` — operator *bundles* allowed ON A VALUE, e.g. ``arithmetic``, ``indexing``. These are - language-level operators (``__add__``, ``__getitem__``, …), never named library methods. No named - method may be called on an opaque value at all. +- ``functions`` — dotted paths callable BY NAME (resolved against import bindings), + e.g. ``jax.*``, ``flax.linen.*``. An optional *disallow* list beats the allow. +- ``operators`` — operator *bundles* allowed ON A VALUE, e.g. ``arithmetic``, ``indexing``. + These are language operators (``+``, ``[]``, …), never named library methods on a value. """ from __future__ import annotations import ast import fnmatch +import hashlib from pydantic import BaseModel, Field # ── Operator bundles: bundle name -> the AST node types it enables on a value ────────────── # These are generic, type-agnostic-safe operators (not named-method calls), so the format-string -# escape cannot hide among them (research approach-B §3.6.2). +# escape (`"{0.__class__}".format(x)`) cannot hide among them (docs/verify.md#operator-bundles-on-a-value). OPERATOR_BUNDLES: dict[str, tuple[type[ast.AST], ...]] = { "arithmetic": (ast.BinOp, ast.UnaryOp), "comparison": (ast.Compare, ast.BoolOp), @@ -26,42 +26,10 @@ } # NOTE: there is deliberately no metadata bundle. Reads like `.shape`/`.ndim`/`.dtype` on an opaque # value are named attribute accesses we can't pin to a type, so they're rejected like any other -# attr-on-value and must be routed through a visible wrapper function (research approach-B §3.6.2). +# attr-on-value and must be routed through a public wrapper function (docs/verify.md#operator-bundles-on-a-value). ALL_BUNDLES: frozenset[str] = frozenset(OPERATOR_BUNDLES) -# ── Dangerous JAX / serialization surface — denylist BEATS the allow (approach-B §3.2/§3.3) ── -# Host-callback / IO / FFI / serialization escape hatches that can run host code or touch disk. -JAX_DENYLIST: tuple[str, ...] = ( - "jax.experimental.*", - "jax.debug.*", - "jax.pure_callback", - "*.io_callback", - "*.host_callback", - "*.host_callback.*", - "jax.profiler.*", - "jax.monitoring.*", - "jax.distributed.*", - "jax.dlpack.*", - "jax.ffi", - "jax.ffi.*", - "jax.extend.*", - # array <-> file on disk, even though jax.numpy.* is otherwise allowed - "jax.numpy.save", - "jax.numpy.savez", - "jax.numpy.savez_compressed", - "jax.numpy.load", - "jax.numpy.tofile", - "jax.numpy.fromfile", - "jax.numpy.memmap", - "jax.numpy.savetxt", - "jax.numpy.loadtxt", - "jax.numpy.genfromtxt", - "flax.serialization.*", - "flax.training.checkpoints.*", - "orbax.*", -) - -# ── Builtins that are dynamic-escape / IO hatches and may never be called (approach-B §2.2) ── +# ── Builtins that are dynamic-escape / IO hatches and may never be called (docs/blacklist.md) ── BANNED_NAMES: frozenset[str] = frozenset( { "eval", @@ -80,18 +48,57 @@ "input", "breakpoint", "memoryview", + "type", + "__build_class__", + "print", + "repr", + "str", + "ascii", + "format", + "bytes", } ) -# Decorators allowed above a def/class in the hidden region (approach-B §3.1 #4). +# Builtins safe to call directly by bare name: pure, deterministic, no reflection/IO/dynamic-code +# surface. A bare-name call is default-deny (docs/verify.md#the-full-call-target-rule) -- everything +# NOT on this list must instead resolve to an import binding, a def/class defined in this file, or a +# local traced to one of those (see verifier._check_call / _is_safe_local_source). +SAFE_BUILTIN_CALLS: frozenset[str] = frozenset( + { + "int", + "float", + "bool", + "len", + "range", + "enumerate", + "zip", + "min", + "max", + "sum", + "abs", + "round", + "all", + "any", + "tuple", + "list", + "dict", + "set", + "sorted", + "reversed", + "isinstance", + "super", + } +) + +# Decorators allowed above a def/class in the private region (docs/verify.md#allow_functions--paths-callable-by-name). # NOTE: `property` (and any descriptor) is deliberately absent — it runs code on a bare attribute -# access (`block.w`), the same attribute-access-hook class banned in §3.1 #6, so default-deny +# access (`block.w`), the same attribute-access-hook class banned for dunder defs, so default-deny # rejects it. Pure inference needs only setup/__call__. ALLOWED_DECORATORS: frozenset[str] = frozenset( {"nn.compact", "jax.jit", "jax.named_scope", "flax.linen.compact"} ) -# The only dunder/hook methods a model class may *define* (approach-B §3.1 #6). +# The only dunder/hook methods a model class may *define* (docs/verify.md#allow_functions--paths-callable-by-name). ALLOWED_DUNDER_DEFS: frozenset[str] = frozenset({"__call__", "setup", "__post_init__"}) # Names always preserved verbatim by the obfuscator and never treated as opaque values. @@ -101,42 +108,68 @@ class Policy(BaseModel): - """Parsed allow-lists. ``reserved`` is filled in by the verifier from the file's imports.""" + """Parsed allow-lists. - functions: list[str] = Field(default_factory=list) - methods: set[str] = Field(default_factory=set) - reserved: set[str] = Field(default_factory=set) + ``reserved_names`` is filled by ``verify()`` from the file's import bindings (on a copy of + the policy — the caller's instance is not mutated). + """ + + # allowed functions passed by the user, example: ["jax.*", "flax.linen.*"] + allowed_functions: list[str] = Field(default_factory=list) + + # operator bundles passed by the user, example: ["arithmetic", "indexing", "comparison"] + allowed_operators: set[str] = Field(default_factory=set) + + # optional disallow globs supplied by the user; these beat the allow, example: ["jax.numpy.save"] + disallowed_functions: list[str] = Field(default_factory=list) + + # import aliases reserved against rebinding (set per-file by verify), e.g. {"jnp", "nn"} + reserved_names: set[str] = Field(default_factory=set) @classmethod def parse( cls, allow_functions: list[str] | None = None, - allow_methods: list[str] | None = None, + allow_operators: list[str] | None = None, + disallow_functions: list[str] | None = None, ) -> "Policy": - functions = _clean(allow_functions) - methods = set(_clean(allow_methods)) - unknown = methods - ALL_BUNDLES + allowed_functions = _clean(allow_functions) + allowed_operators = set(_clean(allow_operators)) + disallowed_functions = _clean(disallow_functions) + unknown = allowed_operators - ALL_BUNDLES if unknown: raise ValueError( - f"unknown method bundle(s): {sorted(unknown)}; allowed: {sorted(ALL_BUNDLES)}" + f"unknown operator bundle(s): {sorted(unknown)}; allowed: {sorted(ALL_BUNDLES)}" ) - return cls(functions=functions, methods=methods) + return cls( + allowed_functions=allowed_functions, + allowed_operators=allowed_operators, + disallowed_functions=disallowed_functions, + ) # ── path matching ────────────────────────────────────────────────────────────────── def function_allowed(self, dotted: str) -> bool: - """True iff a fully-qualified dotted path is allowed (and not denylisted).""" - if any(fnmatch.fnmatchcase(dotted, pat) for pat in JAX_DENYLIST): + """True if a fully-qualified dotted path is allowed (and not disallowed). + + An optional user-supplied ``disallowed_functions`` glob beats the allow, so an author can + keep a hard floor over a broad allow like ``jax.*``. Empty (the default) => pure allow-list. + """ + if any(fnmatch.fnmatchcase(dotted, pat) for pat in self.disallowed_functions): return False - return any(_path_matches(dotted, pat) for pat in self.functions) + return any(_path_matches(dotted, pat) for pat in self.allowed_functions) def bundle_enabled(self, name: str) -> bool: - return name in self.methods + return name in self.allowed_operators def policy_id(self) -> str: """A short, stable identifier for the policy (for the certificate).""" - import hashlib - - blob = "|".join(sorted(self.functions)) + "##" + "|".join(sorted(self.methods)) + blob = ( + "|".join(sorted(self.allowed_functions)) + + "##" + + "|".join(sorted(self.allowed_operators)) + + "##" + + "|".join(sorted(self.disallowed_functions)) + ) return hashlib.sha256(blob.encode()).hexdigest()[:16] diff --git a/packages/syft-restrict/src/syft_restrict/runner.py b/packages/syft-restrict/src/syft_restrict/runner.py index 10729b40eae..afee9fcb417 100644 --- a/packages/syft-restrict/src/syft_restrict/runner.py +++ b/packages/syft-restrict/src/syft_restrict/runner.py @@ -8,12 +8,13 @@ from pydantic import BaseModel, Field +from .astutil import normalize_ranges, scan_file from .errors import PolicyViolation from .obfuscator import ( obfuscate as _obfuscate, ) # aliased: `obfuscate` is also a run() kwarg from .policy import Policy -from .verifier import Violation, _normalize_ranges, _scan_file, verify +from .verifier import Violation, verify __all__ = ["run", "RunResult"] @@ -30,7 +31,8 @@ def run( obfuscate=None, hide=None, allow_functions: list[str] | None = None, - allow_methods: list[str] | None = None, + allow_operators: list[str] | None = None, + disallow_functions: list[str] | None = None, out: str | Path | None = None, strict: bool = True, ) -> RunResult: @@ -46,15 +48,18 @@ def run( hide: ``[start, end]`` 1-based inclusive line ranges to *hide* (whole line replaced with a ``■■■■■■■■`` marker, indentation kept). allow_functions: list of dotted-path globs callable by name (e.g. ``["jax.*", "flax.linen.*"]``). - allow_methods: list of operator bundles allowed on a value + allow_operators: list of operator bundles allowed on a value (``["arithmetic", "indexing", "comparison"]``). + disallow_functions: optional list of dotted-path globs that BEAT the allow (e.g. + ``["jax.numpy.save", "jax.experimental.*"]``). A hard floor for authors who allow a + broad glob; empty by default, in which case only ``allow_functions`` applies. out: where to write the obfuscated file (default ``.obfuscated.py`` next to the source). strict: if True (default), raise ``PolicyViolation`` when verification fails; otherwise return a ``RunResult`` with ``ok=False`` and no output written. """ path = Path(path) source = path.read_text() - policy = Policy.parse(allow_functions, allow_methods) + policy = Policy.parse(allow_functions, allow_operators, disallow_functions) obfuscate_ranges = obfuscate or [] hide_ranges = hide or [] @@ -66,7 +71,7 @@ def run( raise PolicyViolation(result.violations) return RunResult(ok=False, violations=result.violations) - scan = _scan_file(ast.parse(source), _normalize_ranges(private)) + scan = scan_file(ast.parse(source), normalize_ranges(private)) obfuscated = _obfuscate(source, obfuscate_ranges, hide_ranges, scan) out_path = Path(out) if out is not None else path.with_suffix(".obfuscated.py") @@ -76,9 +81,9 @@ def run( "source_sha256": hashlib.sha256(source.encode()).hexdigest(), "policy_id": policy.policy_id(), "restrict_version": _version(), - "private_ranges": [list(r) for r in _normalize_ranges(private)], - "obfuscate_ranges": [list(r) for r in _normalize_ranges(obfuscate_ranges)], - "hide_ranges": [list(r) for r in _normalize_ranges(hide_ranges)], + "private_ranges": [list(r) for r in normalize_ranges(private)], + "obfuscate_ranges": [list(r) for r in normalize_ranges(obfuscate_ranges)], + "hide_ranges": [list(r) for r in normalize_ranges(hide_ranges)], "n_calls_checked": result.n_calls_checked, } return RunResult( diff --git a/packages/syft-restrict/src/syft_restrict/verifier.py b/packages/syft-restrict/src/syft_restrict/verifier.py index 5aeff2d3010..78df7b88896 100644 --- a/packages/syft-restrict/src/syft_restrict/verifier.py +++ b/packages/syft-restrict/src/syft_restrict/verifier.py @@ -1,26 +1,43 @@ -"""The static checker (research approach B). +"""The static checker — verifies that the private region only does trusted math. ``verify(source, private, policy)`` parses the file, restricts attention to the *private* line ranges -(the hidden model definition), and walks those nodes default-deny: only explicitly-allowed node types, -calls, operators, and attribute reads pass. It never raises on a policy issue — it returns a -``VerifyResult`` with the violations so callers can inspect them. +(the private model definition), and walks those nodes **default-deny**: a node is allowed only if a +check below explicitly permits it. It never raises on a policy issue — it returns a ``VerifyResult`` +listing the violations, so callers can inspect them. + +The full whitelist and the reasoning behind each rule live in ``docs/verify.md``; the full deny lists +live in ``docs/blacklist.md``. Comments here stay short and cite those docs where the "why" is subtle. """ from __future__ import annotations import ast +from typing import Literal from pydantic import BaseModel, ConfigDict, Field +from .astutil import ( + FileScan, + describe, + dotted_name, + is_dunder, + node_in_ranges, + normalize_ranges, + rooted_in_self, + scan_file, + self_attr_name, +) from .policy import ( ALLOWED_DECORATORS, ALLOWED_DUNDER_DEFS, BANNED_NAMES, OPERATOR_BUNDLES, + SAFE_BUILTIN_CALLS, Policy, ) -# ── allowed AST node types in the hidden region (approach-B §2.1) ──────────────────────────── +# ── node-type allow-list: syntax the private region may use (docs/verify.md#always-on-allow-list) ── +# Anything else is denied by default (docs/blacklist.md). _ALLOWED_NODES: tuple[type[ast.AST], ...] = ( ast.Module, ast.Expr, @@ -64,8 +81,6 @@ ast.Assign, ast.AugAssign, ast.AnnAssign, - ast.JoinedStr, - ast.FormattedValue, # operator/cmpop/boolop/unaryop singletons are leaf nodes under the above; always fine. ast.operator, ast.cmpop, @@ -74,7 +89,9 @@ ast.expr_context, ) -# Banned statement/expr node types (approach-B §2.2): present => violation. +# node-type deny-list: constructs deliberately, permanently denied (docs/blacklist.md) ── +# Listed explicitly (rather than just left off the allow-list) so their violation names them +# clearly, distinct from "node-type" (unreviewed/future syntax). _BANNED_NODES: tuple[type[ast.AST], ...] = ( ast.Import, ast.ImportFrom, @@ -91,14 +108,38 @@ ast.Await, ast.Yield, ast.YieldFrom, + # f-strings: even with no interpolation, just use a plain string; with interpolation, every + # {expr} invokes type(expr).__format__(expr, spec) with no Call node for _check_call to see. + ast.JoinedStr, + ast.FormattedValue, ) +# violation-code registry: every code a check can raise, one line each (docs/blacklist.md) +ViolationCode = Literal[ + "banned-construct", # _enforce (a node type on the permanent deny-list, docs/blacklist.md) + "node-type", # _enforce (node type outside the always-on allow-list) + "dunder-def", # _check_def (defining a magic/hook method) + "class-keyword", # _check_class (metaclass= or other class keyword arg) + "class-base", # _check_class (non-allow-listed base class) + "decorator", # _check_decorators + "reserved-name", # _check_name, _check_arguments_dont_abuse_self_or_cls, _check_reserved_name + "banned-name", # _check_call (banned bare call) / _check_name (any other Load reference) + "call-not-allowed", # _check_call, _check_call_attribute + "call-unresolved", # _check_call (bare-name/value call not traceable to a safe source) + "dunder-attr", # _check_call_attribute, _check_attribute + "attr-on-value", # _check_call_attribute, _check_self_subscript_call, _check_attribute + "method-on-value", # _check_call_attribute + "attr-not-allowed", # _check_attribute + "dunder-name", # _check_name + "operator-disabled", # _require_bundle +] + class Violation(BaseModel): model_config = ConfigDict(frozen=True) line: int - code: str + code: ViolationCode message: str @@ -108,21 +149,12 @@ class VerifyResult(BaseModel): n_calls_checked: int = 0 -class FileScan(BaseModel): - """Names harvested from the whole file, used to classify calls in the hidden region.""" - - bindings: dict[str, str] # alias -> fully-qualified module path (jnp -> jax.numpy) - hidden_defs: set[str] # class/func names defined inside the private region - visible_defs: set[ - str - ] # function names defined in the visible region (the wrappers) - - def verify(source: str, private, policy: Policy) -> VerifyResult: - ranges = _normalize_ranges(private) + ranges = normalize_ranges(private) tree = ast.parse(source) - scan = _scan_file(tree, ranges) - policy.reserved = set(scan.bindings) + scan = scan_file(tree, ranges) + # Copy policy so caller's reserved_names is never mutated across files. + policy = policy.model_copy(update={"reserved_names": set(scan.import_bindings)}) checker = _Checker(policy, scan, ranges) checker.visit(tree) return VerifyResult( @@ -132,29 +164,20 @@ def verify(source: str, private, policy: Policy) -> VerifyResult: ) -# ── file scan ──────────────────────────────────────────────────────────────────────────── -def _scan_file(tree: ast.Module, ranges) -> FileScan: - bindings: dict[str, str] = {} - hidden_defs: set[str] = set() - visible_defs: set[str] = set() - for node in ast.walk(tree): - if isinstance(node, ast.Import): - for alias in node.names: - bindings[alias.asname or alias.name.split(".")[0]] = alias.name - elif isinstance(node, ast.ImportFrom) and node.module: - for alias in node.names: - bindings[alias.asname or alias.name] = f"{node.module}.{alias.name}" - elif isinstance(node, (ast.FunctionDef, ast.ClassDef)): - if _in_ranges(node, ranges): - hidden_defs.add(node.name) - elif isinstance(node, ast.FunctionDef): - visible_defs.add(node.name) - return FileScan( - bindings=bindings, hidden_defs=hidden_defs, visible_defs=visible_defs - ) - - -# ── the checker ────────────────────────────────────────────────────────────────────────── +# ────────────────────────────────────────────────────────────────────────────────────────────── +# One check per node, default-deny. Rationale: docs/verify.md; deny lists: docs/blacklist.md. +# +# dynamic-code / reflection builtins -> _check_call, _check_name +# IO / host-escape statements -> _BANNED_NODES (in _enforce) +# unknown / future syntax -> _ALLOWED_NODES default-deny (in _enforce) +# library call/attr by name -> _check_call_attribute, _check_attribute (resolver + allow/disallow) +# named method / attr on opaque value -> _check_call_attribute, _check_attribute +# f-strings (any shape) -> _BANNED_NODES (in _enforce) +# forged self/cls trust -> _check_arguments_dont_abuse_self_or_cls, _check_name +# aliasing a banned callable -> _check_name (any Load reference, regardless of position) +# unresolved call target (default-deny) -> _check_call, _is_safe_local_source +# class-creation hooks -> _check_class, _check_decorators, _check_def +# ────────────────────────────────────────────────────────────────────────────────────────────── class _Checker: def __init__(self, policy: Policy, scan: FileScan, ranges): self.policy = policy @@ -162,271 +185,817 @@ def __init__(self, policy: Policy, scan: FileScan, ranges): self.ranges = ranges self.violations: list[Violation] = [] self.n_calls = 0 - self._call_funcs: set[int] = ( - set() - ) # Attribute nodes already judged as a call func - def add(self, node: ast.AST, code: str, message: str) -> None: + # nodes already judged as a call's func, so _check_attribute/_check_name skip them + self._checked_call_targets: set[int] = set() + + # enclosing class/def/lambda stack, for _enclosing_class() and self/cls-position checks + self._scope_stack: list[ast.AST] = [] + + # answers "is self. safe to call?" -- see _SelfAttrTrust below + self._self_attr = _SelfAttrTrust(scan, self._resolved_allowed) + + # nodes inside a type annotation (never invoked), exempt from name/container checks + self._annotation_nodes: set[int] = set() + + # per-scope locals provably safe to call. see _track_safe_local / + # _is_safe_local_source. Starts with one frame already present because Module itself never + # pushes a scope in visit() (only ClassDef/FunctionDef/Lambda do) + self._safe_locals_stack: list[dict[str, bool]] = [{}] + + def report(self, node: ast.AST, code: ViolationCode, message: str) -> None: self.violations.append( Violation(line=getattr(node, "lineno", 0), code=code, message=message) ) + # ── tree walk ─────────────────────────────────────────────────────────────────────────── def visit(self, node: ast.AST) -> None: - """Walk the tree; enforce only on nodes inside the private ranges, recurse everywhere.""" - if _in_ranges(node, self.ranges): + """Walk the whole tree; enforce only on nodes inside the private ranges, recurse everywhere.""" + if node_in_ranges(node, self.ranges): self._enforce(node) + else: + # private-defined names are reserved everywhere, including in public + # region. Bare calls trust private_defs by identifier alone, so a + # public-region rebind (helper = evil between private chunks) would + # otherwise reopen the call-target hole. + self._check_private_def_shadow_anywhere(node) + + # push/pop scope stack for ClassDef/FunctionDef/Lambda, so + # _enclosing_class() works + is_scope = isinstance(node, (ast.ClassDef, ast.FunctionDef, ast.Lambda)) + if is_scope: + self._scope_stack.append(node) + self._safe_locals_stack.append({}) + self._mark_annotation_subtrees(node) + + # recurse to children, which may be outside the private ranges for child in ast.iter_child_nodes(node): self.visit(child) - # — per-node enforcement (recursion is handled by visit) — + # pop scope stack after children, so _enclosing_class() sees the right frame + if is_scope: + self._scope_stack.pop() + self._safe_locals_stack.pop() + + def _check_private_def_shadow_anywhere(self, node: ast.AST) -> None: + """Reject rebinding a private-region class/def name even on public lines.""" + + # only check rebinding (Store) and definitions (FunctionDef/ClassDef). + if isinstance(node, ast.Name) and isinstance(node.ctx, ast.Store): + if ( + node.id in self.scan.private_defs + and id(node) not in self.scan.private_def_ids + ): + self.report( + node, + "reserved-name", + f"{node.id!r} is a private-region class/def and may not be rebound", + ) + return + + # methods are not bare-call targets; shared names like `setup`` are fine + if isinstance(node, ast.FunctionDef) and id(node) in self.scan.method_ids: + return + + # only check rebinding (Store) and definitions (FunctionDef/ClassDef). A + if isinstance(node, (ast.FunctionDef, ast.ClassDef)): + if ( + node.name in self.scan.private_defs + and id(node) not in self.scan.private_def_ids + ): + self.report( + node, + "reserved-name", + f"{node.name!r} is a private-region class/def and may not be rebound", + ) + + def _enclosing_class(self) -> ast.ClassDef | None: + """The nearest enclosing ClassDef, skipping any FunctionDef/Lambda frames above it.""" + for frame in reversed(self._scope_stack): + if isinstance(frame, ast.ClassDef): + return frame + return None + + def _mark_annotation_subtrees(self, node: ast.AST) -> None: + # a type annotation (`x: str`, `x: list[str]`, `def f() -> dict[str, + # bytes]`) is never invoked, so it can hold a name that is banned in a + # real reference. + for ann in (getattr(node, "annotation", None), getattr(node, "returns", None)): + if ann is not None: + # Mark the WHOLE subtree, not just its top node: generics nest the + # type names one or more levels down + for descendant in ast.walk(ann): + self._annotation_nodes.add(id(descendant)) + def _enforce(self, node: ast.AST) -> None: + """Run the one check that applies to this node type (recursion is + handled by visit).""" if isinstance(node, _BANNED_NODES): - self.add( + self.report( node, "banned-construct", - f"{type(node).__name__} is not allowed in the hidden region", + f"{type(node).__name__} is not allowed in the private region", ) return if not isinstance(node, _ALLOWED_NODES): - self.add( + self.report( node, "node-type", f"{type(node).__name__} is not on the node-type allow-list", ) return + # --- definitions & classes --- if isinstance(node, ast.FunctionDef): self._check_def(node) + elif isinstance(node, ast.Lambda): + self._check_lambda(node) elif isinstance(node, ast.ClassDef): self._check_class(node) + + # --- calls & attribute access --- elif isinstance(node, ast.Call): self._check_call(node) elif isinstance(node, ast.Attribute): self._check_attribute(node) + + # --- operators (gated by policy bundles) --- elif isinstance(node, (ast.BinOp, ast.UnaryOp)): self._require_bundle(node, "arithmetic") elif isinstance(node, (ast.Compare, ast.BoolOp)): self._require_bundle(node, "comparison") elif isinstance(node, (ast.Subscript, ast.Slice)): self._require_bundle(node, "indexing") + + # --- name binding (assignments, params) --- elif isinstance(node, (ast.Assign, ast.AugAssign, ast.AnnAssign)): - self._check_assign_targets(node) - elif isinstance(node, ast.For): - self._check_reserved_target(node.target) - elif isinstance(node, ast.comprehension): - self._check_reserved_target(node.target) + self._track_safe_local(node) elif isinstance(node, ast.arg): self._check_reserved_name(node, node.arg) - # — defs / classes — + # --- literals & names --- + elif isinstance(node, ast.Name): + self._check_name(node) + + # ── definitions & classes ──────────────────────────────────────────────────────────────── def _check_def(self, node: ast.FunctionDef) -> None: - # Checks a function def: its decorators, and that it defines no non-allow-listed dunder. + """Guards against: + + - defining magic/hook methods (__getattr__, __reduce__, …) that Python + runs automatically without an explicit call + - shadowing a trusted module alias or public wrapper name with a local + def + + """ + + # are the function decorators in the list of allowed decorators? self._check_decorators(node) - if _is_dunder(node.name) and node.name not in ALLOWED_DUNDER_DEFS: - self.add( + + # is the function name not a reserved name? + self._check_reserved_name(node, node.name) + + # if it's a dunder, is it an allowed dunder? + if is_dunder(node.name) and node.name not in ALLOWED_DUNDER_DEFS: + self.report( node, "dunder-def", f"defining magic method {node.name!r} is not allowed", ) + # if the function has cls or self, only allow if it is the first argument for a method + self._check_arguments_dont_abuse_self_or_cls(node.args) + + def _check_lambda(self, node: ast.Lambda) -> None: + """Guards against: + + - a lambda's parameters being named self/cls, forging the self/cls + trust exemption.""" + self._check_arguments_dont_abuse_self_or_cls(node.args) def _check_class(self, node: ast.ClassDef) -> None: - # Checks a class def: its decorators, no class keywords (e.g. metaclass=), and allow-listed bases. + """Guards against: + + - banned base classes + - a class decorator running attacker code when the class is reached + - shadowing a trusted module alias with a local class name. + """ + # check only allowed decorators self._check_decorators(node) + + # is the class name not a reserved name? + self._check_reserved_name(node, node.name) + + # does it use class keyword arguments (like metaclass=)? if node.keywords: - self.add( + self.report( node, "class-keyword", "class keyword arguments (e.g. metaclass=) are not allowed", ) + + # if it has base classes, are they allow-listed? for base in node.bases: - dotted = _dotted(base) - ok = (dotted and self._resolved_allowed(dotted)) or ( - isinstance(base, ast.Name) - and base.id in (self.scan.hidden_defs | {"object"}) - ) - if not ok: - self.add( + if not self._is_base_class_allowed(base): + self.report( base, "class-base", - f"base class {_describe(base)!r} is not allow-listed", + f"base class {describe(base)!r} is not allow-listed", ) + def _is_base_class_allowed(self, base: ast.AST) -> bool: + """Only an allow-listed class may be a base class.""" + path = dotted_name(base) + return bool(path) and self._resolved_allowed(path) + def _check_decorators(self, node) -> None: - # Checks that every decorator on a def/class resolves to one on the decorator allow-list. + """Guards against: + + - a decorator running attacker code when a def/class is reached. + + """ for dec in node.decorator_list: + # resolve the decorator to a dotted name target = dec.func if isinstance(dec, ast.Call) else dec - dotted = _dotted(target) - resolved = self._resolve(dotted) if dotted else None - if not (resolved in ALLOWED_DECORATORS or dotted in ALLOWED_DECORATORS): - self.add( + path = dotted_name(target) + resolved = self._resolve(path) if path else None + + # then check if it's allow-listed + if not (resolved in ALLOWED_DECORATORS or path in ALLOWED_DECORATORS): + self.report( dec, "decorator", - f"decorator {_describe(target)!r} is not allow-listed", + f"decorator {describe(target)!r} is not allow-listed", + ) + + def _check_arguments_dont_abuse_self_or_cls(self, args: ast.arguments) -> None: + """Guards against: + + - a parameter named self/cls anywhere except the first parameter of a + method in a class definition + + """ + # is this def/lambda a method defined directly in a class body? + is_direct_method = bool(self._scope_stack) and isinstance( + self._scope_stack[-1], ast.ClassDef + ) + + # which parameter, if any, is genuinely first (posonly takes precedence + # over plain args) + first = ( + args.posonlyargs[0] + if args.posonlyargs + else (args.args[0] if args.args else None) + ) + + # every parameter this def/lambda declares, including *args/**kwargs + all_params = [*args.posonlyargs, *args.args, *args.kwonlyargs] + if args.vararg is not None: + all_params.append(args.vararg) + if args.kwarg is not None: + all_params.append(args.kwarg) + + # flag self/cls named anywhere except the first position + for a in all_params: + if a.arg in ("self", "cls") and not (is_direct_method and a is first): + self.report( + a, + "reserved-name", + f"{a.arg!r} may only be the first parameter of a method defined directly " + f"in a class body; self/cls attribute access is trusted by identifier alone", ) - # — calls — + # ── calls ──────────────────────────────────────────────────────────────────────────────── def _check_call(self, node: ast.Call) -> None: - # Checks a call: bans dynamic-escape builtins by name and routes attribute-calls for vetting. + """Guards against: calling a non-vetted callable. + + Default-deny semantics. A call target is allowed only if it's provably + one of a small set of safe shapes: + + - a safe builtin + - an allow-listed import + - a def/class defined in this file + - a self. vetted by _SelfAttrTrust + - a local/value traced back to one of those (see _is_safe_local_source) + + Anything else (a parameter, an untraceable local, an unknown name) is + rejected outright: we can't prove it's safe, so per policy we'd rather + force the author to route it through self. or a public wrapper + than risk trusting an opaque callable + + (docs/verify.md#the-full-call-target-rule). + + Attribute-position calls route to a stricter check.""" + + # count total calls for the final result self.n_calls += 1 func = node.func + + # if it's a bare name, check it against the allow-list and the private defs if isinstance(func, ast.Name): + # banned builtins: report once here and mark so _check_name does not double-count if func.id in BANNED_NAMES: - self.add(node, "banned-call", f"call to {func.id!r} is not allowed") - # Otherwise a bare-name call (local var / hidden or visible def / safe builtin) is allowed; - # nothing dangerous can reach a local name given the other rules. + self.report(node, "banned-name", f"call to {func.id!r} is not allowed") + self._checked_call_targets.add(id(func)) + return + if func.id in self.scan.import_bindings: + # resolves through the same import-binding table as a dotted call + if not self._resolved_allowed(func.id): + self.report( + node, + "call-not-allowed", + f"call to {self._resolve(func.id)!r} is not allow-listed", + ) + return + + # a def/class here, a safe builtin, or a local traced to one + if ( + func.id in self.scan.private_defs + or func.id in self.scan.public_defs + or func.id in SAFE_BUILTIN_CALLS + or self._current_safe_locals.get(func.id, False) + ): + return + + self.report( + node, + "call-unresolved", + f"{func.id!r}: could not unambiguously identify what this calls; only an " + f"allow-listed import, a def/class defined in this file, a safe builtin, or a " + f"local traced to one of those may be called", + ) return + + # given x.y.z(...), this checks the Attribute call target x.y if isinstance(func, ast.Attribute): self._check_call_attribute(node, func) return - # func is a Call / Subscript / etc.: calling a *value* (e.g. self.layer[i](...), Block(...)(x)). - # The value's provenance is checked elsewhere; calling it (its __call__) is allowed. + # given self.x[i](...), this checks the self-rooted subscript call + if isinstance(func, ast.Subscript) and rooted_in_self(func): + self._check_self_subscript_call(node, func) + return + + # if we get here, the call is not banned by name or attribute, but must + # still resolve to a safe source + if not self._is_safe_local_source(func): + self.report( + node, + "call-unresolved", + "could not unambiguously identify what this calls; route the callee through " + "self., a local traced to an allow-listed constructor, or a public wrapper " + "instead", + ) + + # Intentionally parallel with _check_self_subscript_call/_check_attribute -- don't merge them. def _check_call_attribute(self, call: ast.Call, func: ast.Attribute) -> None: - # Checks a dotted call: allows self/cls methods, vets library calls, bans methods on opaque values. - self._call_funcs.add( - id(func) - ) # so _check_attribute doesn't re-flag the same node - dotted = _dotted(func) - if dotted is not None: - root = dotted.split(".")[0] + """Guards against: + + - calling a dunder attribute off self/cls (self.__class__(...)) + - calling self.(...) that wasn't inherited or assigned a vetted-safe source (see + _SelfAttrTrust) + - calling a deeper self-rooted chain (self.a.b(...)), not a single self. level + - calling a non-allow-listed library path + - calling a named method on an opaque value (x.reshape(...)) whose type — and thus what + the call does — we can't pin + """ + # add it to the checked set so _check_attribute doesn't double-count it + self._checked_call_targets.add(id(func)) + + # resolve the dotted path to a string, if possible. If not, it's a named + # method on an opaque value (x.reshape(...)) whose type we can't pin, so deny it. + path = dotted_name(func) + if path is not None: + root = path.split(".")[0] if root in ("self", "cls"): - return # self.method(...) — receiver type is the module class, not opaque - if root in self.scan.bindings: - if not self._resolved_allowed(dotted): - self.add( + attr = self_attr_name(func) + # a dunder off self/cls (self.__class__(...)) is always denied + if attr is not None and is_dunder(attr): + self.report( + call, + "dunder-attr", + f"access to dunder attribute {attr!r} is not allowed", + ) + return + # self.(...) — allowed only if is inherited or vetted safe + if attr is not None and self._self_attr.is_safe( + attr, self._enclosing_class() + ): + return # self.(...) — is inherited or was assigned a vetted source + # attr is None: a deeper chain (self.a.b(...)); otherwise attr was assigned unsafely + message = ( + f"self.{attr!r} was assigned a value that isn't an allow-listed constructor " + f"or a locally-defined class; calling it is not allowed" + if attr is not None + else f"{path!r}: only a single self. attribute may be called, " + f"not a deeper attribute chain" + ) + self.report(call, "attr-on-value", message) + return + # a non-self dotted path: must resolve to an allow-listed import + if root in self.scan.import_bindings: + if not self._resolved_allowed(path): + self.report( call, "call-not-allowed", - f"call to {self._resolve(dotted)!r} is not allow-listed", + f"call to {self._resolve(path)!r} is not allow-listed", ) return - # Attribute on an opaque value: this is a NAMED METHOD ON A VALUE — never allowed (§3.6). - self.add( + # not a self/cls chain and not import-rooted: a named method on an opaque value + self.report( call, "method-on-value", f"named method {func.attr!r} called on a value whose type is unknown; " - f"route it through a visible wrapper function instead", + f"route it through a public wrapper function instead", ) - # — attribute reads (not the func of a call) — + def _check_self_subscript_call(self, call: ast.Call, func: ast.Subscript) -> None: + """Guards against: + + - the self.layer[i](x) idiom smuggling a tainted callable. Only allowed + when the callable is inherited or was assigned a safe source + - calling an element of a deeper self-rooted chain (self.a.b[i](x)), not + a single self.[...] level + """ + # only a single self.[...] level is ever eligible; a deeper chain leaves attr None + attr = ( + self_attr_name(func.value) + if isinstance(func.value, ast.Attribute) + else None + ) + # self.[...] — allowed only if is inherited or vetted safe + if attr is not None and self._self_attr.is_safe(attr, self._enclosing_class()): + return + # attr is None: a deeper chain; otherwise was assigned an unsafe value + message = ( + f"self.{attr!r}[...] was assigned a value that isn't a list/tuple of allow-listed " + f"constructors; calling an element of it is not allowed" + if attr is not None + else "only self.[...] may be called this way, not a deeper self-rooted chain" + ) + self.report(call, "attr-on-value", message) + + # ── attribute reads (not the func of a call) ───────────────────────────────────────────── def _check_attribute(self, node: ast.Attribute) -> None: - # Checks an attribute read: bans dunders, vets library refs, bans reads on opaque values. - if id(node) in self._call_funcs: + """Guards against: + + - a dunder attribute read (x.__class__, obj.__dict__) + - a self-rooted chain deeper than a single self./cls. level (self.a.b) + - a non-allow-listed library path (np.dot when only jax.* is allowed) + - an attribute read on an opaque value (x.shape, x.T, x.ndim) + """ + if id(node) in self._checked_call_targets: return # already judged as a call's function position by _check_call_attribute - if _is_dunder(node.attr): - self.add( + # a bare dunder attribute (x.__class__) is always denied + if is_dunder(node.attr): + self.report( node, "dunder-attr", f"access to dunder attribute {node.attr!r} is not allowed", ) return - dotted = _dotted(node) - if dotted is not None: - root = dotted.split(".")[0] + + path = dotted_name(node) + + if path is not None: + root = path.split(".")[0] if root in ("self", "cls"): + # a single self./cls. read is fine; a deeper chain is not + if self_attr_name(node) is not None: + return + self.report( + node, + "attr-on-value", + f"{path!r}: only a single self. attribute may be accessed, " + f"not a deeper attribute chain", + ) return - if root in self.scan.bindings: - if not self._resolved_allowed(dotted): - self.add( + # a non-self dotted path: must resolve to an allow-listed import + if root in self.scan.import_bindings: + if not self._resolved_allowed(path): + self.report( node, "attr-not-allowed", - f"reference to {self._resolve(dotted)!r} is not allow-listed", + f"reference to {self._resolve(path)!r} is not allow-listed", ) return - # Attribute read on an opaque value (including .shape/.ndim/.dtype): we can't pin the - # receiver's type, so it must be routed through a visible wrapper function. - self.add( + + # not self/cls-rooted and not import-rooted: an attribute read on an opaque value + self.report( node, "attr-on-value", f"attribute {node.attr!r} on a value is not allowed; " - f"route it through a visible wrapper function instead", + f"route it through a public wrapper function instead", ) - # — operators — + # ── bare name reads (not the func of a call) ───────────────────────────────────────────── + def _check_name(self, node: ast.Name) -> None: + """Guards against: + + - rebinding self/cls + - rebinding a reserved name + - loading a banned builtin (`open`, `eval`, `exec`, ...) + - loading a bare dunder name (`__class__`, `__dict__`, ...) + """ + if isinstance(node.ctx, ast.Store): + # self/cls itself: always denied, never delegated to _check_reserved_name + if node.id in ("self", "cls"): + self.report( + node, + "reserved-name", + f"{node.id!r} may not be rebound; self/cls attribute access is " + f"trusted by identifier alone, not a verified reference to the real " + f"instance", + ) + return + # any other Store target: is this name reserved by the resolver? + self._check_reserved_name(node, node.id) + return + + # only a reference (Load) can leak/dispatch a name; a Del target is bound elsewhere + if not isinstance(node.ctx, ast.Load): + return + + # skip if already reported as the func of a banned call by _check_call + if id(node) in self._checked_call_targets: + return + + # a reference to a banned builtin, anywhere it can appear, is denied + if node.id in BANNED_NAMES: + # `x: str` / `def f() -> str` — a type annotation, not a reference + if id(node) in self._annotation_nodes: + return + self.report(node, "banned-name", f"reference to {node.id!r} is not allowed") + + # a bare dunder name + elif is_dunder(node.id): + self.report( + node, + "dunder-name", + f"reference to dunder name {node.id!r} is not allowed", + ) + + # ── operators ──────────────────────────────────────────────────────────────────────────── def _require_bundle(self, node: ast.AST, bundle: str) -> None: - # Checks that the operator's bundle (arithmetic/comparison/indexing) is enabled by the policy. + """Guards against: + + - using an operator bundle (arithmetic/comparison/indexing) the policy didn't enable for + this file + """ if not self.policy.bundle_enabled(bundle): ops = "/".join(t.__name__ for t in OPERATOR_BUNDLES[bundle]) - self.add( - node, "bundle-disabled", f"{ops} needs the {bundle!r} method bundle" + self.report( + node, + "operator-disabled", + f"{ops} needs the {bundle!r} operator group in allow_operators", + ) + + # ── assignment / reserved names ────────────────────────────────────────────────────────── + @staticmethod + def _assignment_targets_and_value(node) -> tuple[list[ast.expr], ast.expr | None]: + """Normalize Assign/AnnAssign/AugAssign into a uniform (targets, value) pair.""" + if isinstance(node, ast.Assign): + return node.targets, node.value + if isinstance(node, ast.AnnAssign): + return ([node.target] if node.value is not None else []), node.value + return [node.target], None # AugAssign + + def _track_safe_local(self, node) -> None: + """Track whether each assigned local is a provably-safe call target. + + Any other assignment to a previously-tracked name clears its verdict -- + it no longer traces to that source.""" + + # Module itself never pushes a scope in visit(), so we added a frame at + # init and never pop it. The stack should never be empty. + if not self._safe_locals_stack: + raise RuntimeError( + "safe_locals_stack is empty; visit() should have pushed a frame" ) - # — assignment / reserved names — - def _check_assign_targets(self, node) -> None: - # Checks every target of an assignment for a forbidden rebind of a reserved name. - targets = node.targets if isinstance(node, ast.Assign) else [node.target] + # Track the verdict in the top frame of the stack, which is the current + # scope. + safe = self._safe_locals_stack[-1] + targets, value = self._assignment_targets_and_value(node) + + # check each target in the assignment. If it's a Name, track whether + # it's provably safe to call. If it's not a Name (e.g., a tuple + # unpacking), skip it. If the value is None (AnnAssign without a value), + # clear the verdict. for t in targets: - self._check_reserved_target(t) + if not isinstance(t, ast.Name): + continue + if value is not None and self._is_safe_local_source(value): + safe[t.id] = True + else: + safe.pop(t.id, None) + + def _is_safe_local_source(self, value: ast.AST) -> bool: + """Is a local provably safe to call?""" + attr = self_attr_name(value) + if attr is None and isinstance(value, ast.Subscript): + attr = self_attr_name(value.value) + if attr is not None: + return self._self_attr.is_safe(attr, self._enclosing_class()) + return _all_leaves_safe(value, self._is_safe_local_leaf) + + def _is_safe_local_leaf(self, value: ast.AST) -> bool: + if isinstance(value, ast.Name): + return ( + self._current_safe_locals.get(value.id, False) + or value.id in self.scan.private_defs + or value.id in self.scan.public_defs + or value.id in SAFE_BUILTIN_CALLS + ) + if isinstance(value, ast.Call): + return self._self_attr.is_safe_value(value) + return False - def _check_reserved_target(self, target: ast.AST) -> None: - # Checks each name being bound (Store context) in an assign/for/comprehension target. - for name_node in _iter_names(target): - if isinstance(name_node.ctx, ast.Store): - self._check_reserved_name(name_node, name_node.id) + @property + def _current_safe_locals(self) -> dict[str, bool]: + return self._safe_locals_stack[-1] if self._safe_locals_stack else {} def _check_reserved_name(self, node: ast.AST, name: str) -> None: - # Checks that a bound name doesn't shadow a reserved module alias or a visible wrapper name. - if name in self.policy.reserved: - self.add( + """Guards against: + + - rebinding a trusted module alias (`jnp = evil`), which would poison the path resolver + - rebinding a private-region class/def (`Attn = evil`), which would silently redirect + every bare call that appears to route through the vetted original -- private_defs is + a scope-blind, name-only whole-file scan, so nothing else notices the shadow + - rebinding a public wrapper name (`transpose = evil`), defeating its type guard + - rebinding a safe builtin (`list = evil`) — _check_call trusts a bare call to any of + these four by identifier alone, so shadowing one would silently redirect every call + site that appears to route through it + """ + if name in self.policy.reserved_names: + self.report( node, "reserved-name", f"{name!r} is a reserved module alias and may not be rebound", ) - elif name in self.scan.visible_defs: - self.add( + elif ( + name in self.scan.private_defs and id(node) not in self.scan.private_def_ids + ): + self.report( + node, + "reserved-name", + f"{name!r} is a private-region class/def and may not be rebound", + ) + elif name in self.scan.public_defs: + self.report( + node, + "reserved-name", + f"{name!r} is a public wrapper name and may not be rebound", + ) + elif name in SAFE_BUILTIN_CALLS: + self.report( node, "reserved-name", - f"{name!r} is a visible wrapper name and may not be rebound", + f"{name!r} is a trusted builtin and may not be rebound", ) - # — path resolution — - def _resolve(self, dotted: str) -> str: - root, _, rest = dotted.partition(".") - base = self.scan.bindings.get(root, root) + # ── path resolution ────────────────────────────────────────────────────────────────────── + def _resolve(self, path: str) -> str: + """Rewrite a dotted path's import alias to its fully-qualified form + (`jnp.einsum` -> `jax.numpy.einsum`).""" + root, _, rest = path.partition(".") + base = self.scan.import_bindings.get(root, root) return f"{base}.{rest}" if rest else base - def _resolved_allowed(self, dotted: str) -> bool: - # Checks a dotted path against the policy allow-list after resolving its import alias. - return self.policy.function_allowed(self._resolve(dotted)) - - -# ── helpers ────────────────────────────────────────────────────────────────────────────── -def _normalize_ranges(private) -> list[tuple[int, int]]: - out = [] - for item in private: - lo, hi = item - out.append((int(lo), int(hi))) - return out - - -def _in_ranges(node: ast.AST, ranges) -> bool: - line = getattr(node, "lineno", None) - if line is None: - return False - return any(lo <= line <= hi for lo, hi in ranges) - - -def _dotted(node: ast.AST) -> str | None: - """Return the dotted path for a pure Name/Attribute chain, else None.""" - parts: list[str] = [] - cur = node - while isinstance(cur, ast.Attribute): - parts.append(cur.attr) - cur = cur.value - if isinstance(cur, ast.Name): - parts.append(cur.id) - return ".".join(reversed(parts)) - return None - - -def _describe(node: ast.AST) -> str: - return _dotted(node) or type(node).__name__ - - -def _is_dunder(name: str) -> bool: - return name.startswith("__") - - -def _iter_names(node: ast.AST): - for n in ast.walk(node): - if isinstance(n, ast.Name): - yield n + def _resolved_allowed(self, path: str) -> bool: + """Resolve an import alias, then apply the policy allow-list (disallow + beats allow).""" + return self.policy.function_allowed(self._resolve(path)) + + +# ── self-attribute trust ──────────────────────────────────────────────────────────────────────── +def _iter_self_attrs_in_target(target: ast.AST): + """Yield every self./cls. name found anywhere in a Store-context target tree, + including nested inside tuple/list-unpacking (``self.a, self.b = ...``) or a starred target.""" + attr = self_attr_name(target) + if attr is not None: + yield attr + return + if isinstance(target, (ast.Tuple, ast.List)): + for elt in target.elts: + yield from _iter_self_attrs_in_target(elt) + elif isinstance(target, ast.Starred): + yield from _iter_self_attrs_in_target(target.value) + + +def _all_leaves_safe(value: ast.AST, is_safe_leaf) -> bool: + """Recurse through list/tuple/set literals and list/set/dict comprehensions/generator + expressions down to their leaf elements (the layer idiom, e.g. ``[Block(cfg) for _ in + range(n)]``), checking each leaf with ``is_safe_leaf``. Shared by + ``_SelfAttrTrust.is_safe_value`` and ``_Checker._is_safe_local_source``.""" + if isinstance(value, (ast.List, ast.Tuple, ast.Set)): + return all(_all_leaves_safe(e, is_safe_leaf) for e in value.elts) + if isinstance(value, (ast.ListComp, ast.SetComp, ast.GeneratorExp)): + return _all_leaves_safe(value.elt, is_safe_leaf) + if isinstance(value, ast.DictComp): + return _all_leaves_safe(value.value, is_safe_leaf) + return is_safe_leaf(value) + + +class _SelfAttrTrust: + """Used by ``_check_call_attribute``/``_check_self_subscript_call`` to + determine if ``self.`` (in the given enclosing class) safe to call or + subscript. + + An attribute the class never assigns is presumed inherited from its (already + vetted, since bases are allow-listed) base class, nn.Module's + ``self.param``/``self.variable``. + + Only attributes the class itself assigns are vetted against what was + actually stored there. Results are memoized per class for the lifetime of + one ``_Checker`` (one ``verify()`` call). + """ + + def __init__(self, scan: FileScan, resolved_allowed) -> None: + self._scan = scan + self._resolved_allowed = resolved_allowed + self._cache: dict[int, dict[str, bool]] = {} # id(ClassDef) -> {attr: is-safe} + + def is_safe(self, attr: str, cls_node: ast.ClassDef | None) -> bool: + # if the class is None, we're not in a class body, so self. is not + # allowed + if cls_node is None: + return False + + # check the memoized table for this class; if not present, build it + table = self._cache.get(id(cls_node)) + if table is None: + table = self._build_table(cls_node) + self._cache[id(cls_node)] = table + + # return True if the attribute is safe, False if not, and True if the + # attribute was never assigned in this class (inherited from a vetted + # base) + return table.get(attr, True) + + def _build_table(self, cls_node: ast.ClassDef) -> dict[str, bool]: + """For every self. assigned anywhere in the class, record whether + EVERY value ever stored there is a vetted-safe source. + """ + + # AugAssign always disqualifies (compound-mutating a submodule reference + # isn't a real pattern and is inherently suspect) + # + # So does a self. found nested in a tuple/list-unpack target or + # bound via a for-loop/comprehension target: neither is a real pattern + # either, and unlike a plain `self.x = value` there's no single value + # expression to vet, so we can't tell what ends up in the attribute. + + values: dict[str, list[ast.AST]] = {} + disqualified: set[str] = set() + + for node in ast.walk(cls_node): + # if the node is an assignment, record the value(s) assigned to self. + if isinstance(node, ast.Assign): + for t in node.targets: + attr = self_attr_name(t) + if attr is not None: + values.setdefault(attr, []).append(node.value) + elif isinstance(t, (ast.Tuple, ast.List)): + disqualified.update(_iter_self_attrs_in_target(t)) + + # if the node is an annotated assignment with a value, record the + # value assigned to self. + elif isinstance(node, ast.AnnAssign) and node.value is not None: + attr = self_attr_name(node.target) + if attr is not None: + values.setdefault(attr, []).append(node.value) + # if the node is an AugAssign, disqualify the attribute + elif isinstance(node, ast.AugAssign): + attr = self_attr_name(node.target) + if attr is not None: + disqualified.add(attr) + # if the node is a for-loop or comprehension, disqualify any + # self. in the target + elif isinstance(node, (ast.For, ast.comprehension)): + disqualified.update(_iter_self_attrs_in_target(node.target)) + table = { + attr: attr not in disqualified and all(self.is_safe_value(v) for v in vs) + for attr, vs in values.items() + } + + # any self. that was disqualified but never assigned a value is + # still disqualified + for attr in disqualified: + table.setdefault(attr, False) + + return table + + def is_safe_value(self, value: ast.AST) -> bool: + """A vetted "submodule" source: a call to an allow-listed library constructor or a name defined + in this file (private or public), or a list/tuple/set/comprehension of such (the layer idiom).""" + return _all_leaves_safe(value, self._is_safe_call) + + def _is_safe_call(self, value: ast.AST) -> bool: + if not isinstance(value, ast.Call): + return False + func = value.func + path = dotted_name(func) + if path is not None and path.split(".")[0] in self._scan.import_bindings: + return self._resolved_allowed(path) + return isinstance(func, ast.Name) and ( + func.id in self._scan.private_defs or func.id in self._scan.public_defs + ) diff --git a/packages/syft-restrict/tests/obfuscate/test_obfuscate.py b/packages/syft-restrict/tests/obfuscate/test_obfuscate.py index 2ac6a8f3083..e438481d114 100644 --- a/packages/syft-restrict/tests/obfuscate/test_obfuscate.py +++ b/packages/syft-restrict/tests/obfuscate/test_obfuscate.py @@ -1,13 +1,15 @@ """Tests for the display transform (approach A), exercised through run().""" +import ast import shutil from pathlib import Path -from syft_restrict import run +from syft_restrict import obfuscate, run +from syft_restrict.astutil import scan_file FIXTURES = Path(__file__).parents[1] / "fixtures" ALLOW_FUNCTIONS = ["jax.*", "flax.linen.*"] -ALLOW_METHODS = ["arithmetic", "indexing", "comparison"] +ALLOW_OPERATORS = ["arithmetic", "indexing", "comparison"] def _private_from_config(source: str): @@ -26,7 +28,7 @@ def _obfuscate_fixture(tmp_path: Path): src_path, obfuscate=private, allow_functions=ALLOW_FUNCTIONS, - allow_methods=ALLOW_METHODS, + allow_operators=ALLOW_OPERATORS, ) obf = Path(result.obfuscated_path).read_text() return source, obf, config_line @@ -70,7 +72,7 @@ def test_hide_blanks_whole_body_keeping_indentation(tmp_path): obfuscate=[[def_line, def_line]], hide=[body], allow_functions=ALLOW_FUNCTIONS, - allow_methods=ALLOW_METHODS, + allow_operators=ALLOW_OPERATORS, ) obf_lines = Path(result.obfuscated_path).read_text().splitlines() note = "# hidden/obfuscated lines can only execute restricted python" @@ -88,6 +90,16 @@ def test_hide_blanks_whole_body_keeping_indentation(tmp_path): assert result.certificate["hide_ranges"] == [body] +def test_fstring_literal_text_is_blanked(): + # obfuscate() alone must still blank f-string text, even though verify() + # bans f-strings now. + source = 'def f():\n return f"secret literal text"\n' + scan = scan_file(ast.parse(source), [(1, 2)]) + obf = obfuscate(source, [[1, 2]], [], scan) + assert "secret literal text" not in obf + assert "■" in obf + + def test_obfuscation_is_deterministic(tmp_path): dir_a = tmp_path / "a" dir_b = tmp_path / "b" diff --git a/packages/syft-restrict/tests/test_run.py b/packages/syft-restrict/tests/test_run.py index c784ff94889..9fd935a0132 100644 --- a/packages/syft-restrict/tests/test_run.py +++ b/packages/syft-restrict/tests/test_run.py @@ -9,7 +9,7 @@ FIXTURES = Path(__file__).parent / "fixtures" ALLOW_FUNCTIONS = ["jax.*", "flax.linen.*"] -ALLOW_METHODS = ["arithmetic", "indexing", "comparison"] +ALLOW_OPERATORS = ["arithmetic", "indexing", "comparison"] def _private(source: str): @@ -26,7 +26,7 @@ def test_run_success_writes_obfuscated_and_certificate(tmp_path): src, obfuscate=_private(src.read_text()), allow_functions=ALLOW_FUNCTIONS, - allow_methods=ALLOW_METHODS, + allow_operators=ALLOW_OPERATORS, ) assert result.ok out = Path(result.obfuscated_path) @@ -45,7 +45,7 @@ def test_run_strict_raises_and_writes_nothing(tmp_path): src, obfuscate=[[1, 3]], allow_functions=ALLOW_FUNCTIONS, - allow_methods=ALLOW_METHODS, + allow_operators=ALLOW_OPERATORS, ) assert exc.value.violations assert not (tmp_path / "bad.obfuscated.py").exists() @@ -58,7 +58,7 @@ def test_run_nonstrict_returns_violations(tmp_path): src, obfuscate=[[1, 2]], allow_functions=ALLOW_FUNCTIONS, - allow_methods=ALLOW_METHODS, + allow_operators=ALLOW_OPERATORS, strict=False, ) assert not result.ok diff --git a/packages/syft-restrict/tests/verify/conftest.py b/packages/syft-restrict/tests/verify/conftest.py index 8758cf41520..19514f7f823 100644 --- a/packages/syft-restrict/tests/verify/conftest.py +++ b/packages/syft-restrict/tests/verify/conftest.py @@ -1,27 +1,20 @@ -"""Shared fixtures/helpers for the static-checker tests (research approach B). +"""Shared fixtures/helpers for the static-checker tests. The tests are split by intent: -- ``test_whitelist.py`` — language constructs the hidden region IS allowed to use. -- ``test_blacklist.py`` — constructs/calls/attrs that are rejected (default-deny). -- ``test_whitelisted_lib.py`` — things we *manually* allow: library calls by name and operator bundles. +- ``test_whitelist.py`` — green path: constructs the private region IS allowed to use. +- ``test_disallowed.py`` — default-deny catalog: straightforward rejections. +- ``test_bypasses.py`` — multi-step / subtle escape regressions. +- ``test_whitelisted_lib.py`` — library calls by name and operator bundles (manual allow). +- ``test_ranges.py`` — private-range argument handling, not code policy. """ -from pathlib import Path - import pytest +from syft_restrict import verify -from syft_restrict import Policy, verify -from syft_restrict.verifier import VerifyResult - -FIXTURES = Path(__file__).parents[1] / "fixtures" -REPO_ROOT = Path(__file__).parents[4] - -ALLOW_FUNCTIONS = ["jax.*", "flax.linen.*"] -ALLOW_METHODS = ["arithmetic", "indexing", "comparison"] - - -def make_policy(functions=ALLOW_FUNCTIONS, methods=ALLOW_METHODS): - return Policy.parse(list(functions), list(methods)) +from verify.helpers import ( + make_policy, + normalize_source, +) @pytest.fixture @@ -33,13 +26,10 @@ def policy(): def verify_all(policy): """Verify ``source`` with the whole file marked private, using the standard policy.""" - def _run(source, pol=None): - n = len(source.splitlines()) - return verify(source, [[1, n]], pol or policy) + def _run(source: str | list[str], pol=None, private=None): + source = normalize_source(source) + if private is None: + private = [[1, len(source.splitlines())]] + return verify(source, private, pol or policy) return _run - - -def error_codes(result: VerifyResult): - """The set of violation codes in a VerifyResult (handy for asserts).""" - return {v.code for v in result.violations} diff --git a/packages/syft-restrict/tests/verify/helpers.py b/packages/syft-restrict/tests/verify/helpers.py new file mode 100644 index 00000000000..d7c52fbdbae --- /dev/null +++ b/packages/syft-restrict/tests/verify/helpers.py @@ -0,0 +1,29 @@ +import inspect +from pathlib import Path + +from syft_restrict import Policy, VerifyResult + +FIXTURES = Path(__file__).parents[1] / "fixtures" +REPO_ROOT = Path(__file__).parents[4] + +ALLOW_FUNCTIONS = ["jax.*", "flax.linen.*"] +ALLOW_OPERATORS = ["arithmetic", "indexing", "comparison"] + + +def normalize_source(source: str | list[str]) -> str: + """Normalize test source: join lists, strip common indent, ensure trailing newline.""" + if isinstance(source, list): + source = "\n".join(source) + # Leading newline lets cleandoc strip the common indent of triple-quoted blocks; + # re-add a trailing newline afterward (cleandoc removes it). + source = inspect.cleandoc("\n" + source).strip() + return source + + +def make_policy(functions=ALLOW_FUNCTIONS, operators=ALLOW_OPERATORS, disallow=None): + return Policy.parse(list(functions), list(operators), list(disallow or [])) + + +def get_error_codes(result: VerifyResult): + """The set of violation codes in a VerifyResult (handy for asserts).""" + return {v.code for v in result.violations} diff --git a/packages/syft-restrict/tests/verify/test_bypasses.py b/packages/syft-restrict/tests/verify/test_bypasses.py new file mode 100644 index 00000000000..6ec32834796 --- /dev/null +++ b/packages/syft-restrict/tests/verify/test_bypasses.py @@ -0,0 +1,568 @@ +"""Regression tests for known escape / edge shapes. + +Each documents a specific attack the verifier must reject, and why a narrower +fix could miss it. +""" + +from verify.helpers import get_error_codes, make_policy + + +def _assert_error_code(verify_all, src, expected, private=None, pol=None, strict=False): + err = get_error_codes(verify_all(src, pol=pol, private=private)) + + assert err + # if strict, we fail cases with extra violations + if strict: + assert {expected} == err + else: + assert expected in err + + +# ── banned-name alias surfaces ─────────────────────────────────────────────── + + +def test_private_alias(verify_all): + # Aliasing a banned name to a private name must not let it evade the ban. + _assert_error_code(verify_all, "e = eval", "banned-name") + + +def test_public_alias(verify_all): + # Stash in public, call via untracked name in private — call-unresolved, not + # a free pass because the private call site has no banned identifier. + src = """ + e = eval + e('1/0') + """ + _assert_error_code(verify_all, src, "call-unresolved", private=[[2, 2]]) + + +def test_alias_in_container(verify_all): + _assert_error_code(verify_all, "con = [eval]", "banned-name") + + +def test_container_subscript_store_alias(verify_all): + # Subscript assignment (d["k"] = open) must not smuggle a banned reference + # just because the container name is Load, not Store. + src = """ + d = {} + d['k'] = open + """ + _assert_error_code(verify_all, src, "banned-name") + + +def test_chained_assignment_alias(verify_all): + src = """ + a = b = open + b('/etc/passwd') + """ + _assert_error_code(verify_all, src, "banned-name") + + +def test_tuple_unpack_alias(verify_all): + src = """ + a, b = (1, open) + b('/etc/passwd') + """ + _assert_error_code(verify_all, src, "banned-name") + + +def test_return_value_alias_with_no_local_name(verify_all): + # Banned reference smuggled via `return` and invoked immediately — no alias + # variable near the dangerous call. + src = """ + def leak(): + return open + + def run(x): + return leak()('/etc/passwd') + """ + _assert_error_code(verify_all, src, "banned-name") + + +def test_parameter_passthrough_alias(verify_all): + # Banned callable passed into a generic helper — no suspicious name in the + # helper body itself (the call of `fn` is also call-unresolved). + src = """ + def apply(fn, x): + return fn(x) + + def run(): + return apply(open, '/etc/passwd') + """ + _assert_error_code(verify_all, src, "banned-name") + + +def test_dict_literal_container_alias(verify_all): + src = """ + d = {"o": open} + d["o"]("/etc/passwd") + """ + _assert_error_code(verify_all, src, "banned-name") + + +def test_inline_container_subscript_call(verify_all): + # Subscript-then-call on a bare literal, no assignment step. + _assert_error_code(verify_all, "([eval][0])('1/0')", "banned-name") + + +def test_ifexp_branch_alias(verify_all): + src = """ + c = True + fn = open if c else eval + fn('/etc/passwd') + """ + _assert_error_code(verify_all, src, "banned-name") + + +def test_for_loop_container_alias(verify_all): + src = """ + for fn in [eval]: + fn('1/0') + """ + _assert_error_code(verify_all, src, "banned-name") + + +def test_function_default_alias(verify_all): + src = """ + def run(op=open): + return op('/etc/passwd') + + run() + """ + _assert_error_code(verify_all, src, "banned-name") + + +# ── dynamic import / reflection chains ─────────────────────────────────────── + + +def test_aliased_import_getattr(verify_all): + src = """ + i = __import__ + g = getattr + g(i('os'), 'getcwd')() + """ + _assert_error_code(verify_all, src, "banned-name") + + +def test_aliased_getattr_on_builtins(verify_all): + _assert_error_code(verify_all, "g = getattr", "banned-name") + + +def test_aliased_import_getattr_system(verify_all): + # Full import+invoke chain built entirely in the private region. + src = """ + i = __import__ + g = getattr + g(g(i('os'), 'system'))('id') + """ + _assert_error_code(verify_all, src, "banned-name") + + +def test_vars_alias_on_builtins(verify_all): + src = """ + v = vars + ev = v(__builtins__)['eval'] + ev('1/0') + """ + _assert_error_code(verify_all, src, "banned-name") + + +def test_self_dynamic_import_chain(verify_all): + # Public setup stashes a dynamic-import chain; private __call__ must not run it. + src = """ + class M: + def setup(self): + self.imp = __import__ + self.get = getattr + def __call__(self, x): + os = self.imp('os') + self.get(os, 'system')('id') + """ + _assert_error_code(verify_all, src, "attr-on-value", private=[[5, 7]]) + + +# ── public import resolution / disallow ────────────────────────────────────── + + +def test_bare_name_import_bypass(verify_all): + # Bare name from `from X import f` must use the same allow/deny as a dotted path. + src = """ + from os import system + + def run(): + system('id') + """ + _assert_error_code(verify_all, src, "call-not-allowed", private=[[3, 4]]) + + +def test_bare_name_disallowed_leaf_bypass(verify_all): + src = """ + from jax.numpy import save + + def run(x): + save(x, 'stolen_data.zip') + """ + _assert_error_code( + verify_all, + src, + "call-not-allowed", + private=[[3, 4]], + pol=make_policy(disallow=["jax.numpy.save"]), + ) + + +def test_bare_name_import_alias_bypass(verify_all): + # Renaming (`as persist`) must not bypass a disallow on the underlying leaf. + src = """ + from jax.numpy import save as persist + + def run(x): + persist(x, 'stolen_data.zip') + """ + _assert_error_code( + verify_all, + src, + "call-not-allowed", + private=[[3, 4]], + pol=make_policy(disallow=["jax.numpy.save"]), + ) + + +def test_importlib_import_module_bypass(verify_all): + src = """ + from importlib import import_module + + def run(): + import_module('os') + """ + _assert_error_code(verify_all, src, "call-not-allowed", private=[[3, 4]]) + + +# ── homoglyphs ─────────────────────────────────────────────────────────────── + + +def test_unicode_homoglyphs_should_not_be_allowed(verify_all): + # Public stash + private call through a lookalike name (Cyrillic о р е). + src = """ + ореn = open + f = ореn('/etc/passwd') + """ + _assert_error_code(verify_all, src, "call-unresolved", private=[[2, 3]]) + + +def test_homoglyph_self_stash_and_call(verify_all): + # Public setup stashes open under a homoglyph attr; private __call__ must not invoke it. + src = """ + class M: + def setup(self): + self.ореn = open + def __call__(self, x): + self.ореn('/etc/passwd') + """ + _assert_error_code(verify_all, src, "attr-on-value", private=[[4, 5]]) + + +def test_homoglyph_two_hop_alias(verify_all): + # ASCII stash then homoglyph copy — no banned name textually at the call site. + src = """ + x = open + ореn = x + ореn('/etc/passwd') + """ + _assert_error_code(verify_all, src, "call-unresolved", private=[[3, 4]]) + + +# ── self-attribute trust ───────────────────────────────────────────────────── + + +def test_self_stash_and_call(verify_all): + # Public setup stashes open on self; private __call__ must not get a free pass via self.* + src = """ + class M: + def setup(self): + self.open = open + def __call__(self, x): + self.open('/etc/passwd') + + m = M() + """ + _assert_error_code(verify_all, src, "attr-on-value", private=[[4, 5]]) + + +def test_self_nested_attribute_chain(verify_all): + # Only a single self. level is exempt — self.a.b(...) is not. + # Stash in public setup; deeper chain call is private. + src = """ + class M: + def setup(self): + self.sub = object() + self.sub.evil = open + def __call__(self, x): + self.sub.evil('/etc/passwd') + + m = M() + """ + _assert_error_code(verify_all, src, "attr-on-value", private=[[5, 6]]) + + +def test_self_augassign_alias(verify_all): + # Public setup seeds self.evil; private __call__ mutates and calls it. + src = """ + class M: + def setup(self): + self.evil = print + def __call__(self, x): + self.evil += open + self.evil('/etc/passwd') + + m = M() + """ + _assert_error_code(verify_all, src, "attr-on-value", private=[[4, 6]]) + + +def test_self_layer_subscript_call(verify_all): + # Public setup builds a tainted layer list; private Flax-style call must still be vetted. + src = """ + class M: + def setup(self): + self.layer = [open] + def __call__(self, x): + self.layer[0]('/etc/passwd') + """ + _assert_error_code(verify_all, src, "attr-on-value", private=[[4, 5]]) + + +def test_self_attr_tuple_unpack_must_not_grant_trust(verify_all): + # setup is public/empty; the unpack+call that forges trust is private. + src = """ + class M: + def setup(self): + pass + def __call__(self, x, fn): + self.fn, self.other = fn, 0 + return self.fn(x) + """ + _assert_error_code(verify_all, src, "attr-on-value", private=[[4, 6]]) + + +def test_self_attr_for_loop_target_must_not_grant_trust(verify_all): + # Public setup binds self.leak via for-target; private __call__ must not treat it as inherited. + src = """ + class Model: + def setup(self, cb): + for self.leak in (cb,): + pass + def __call__(self, x, secret): + self.leak(secret) + return x + """ + _assert_error_code(verify_all, src, "attr-on-value", private=[[5, 7]]) + + +def test_self_attr_local_alias_must_not_grant_trust(verify_all): + # Public setup assigns unsafe self.fn; private tmp = self.fn; tmp(x) must not defeat trust. + src = """ + class M: + def setup(self): + self.fn = lambda x: x + 1 + def __call__(self, x): + tmp = self.fn + return tmp(x) + """ + _assert_error_code(verify_all, src, "call-unresolved", private=[[4, 6]]) + + +def test_self_attr_subscript_local_alias_must_not_grant_trust(verify_all): + # Public setup stashes opaque cb; private layer[i] alias must not defeat trust. + src = """ + class M: + def setup(self, cb): + self.layer = [cb] + def __call__(self, x): + tmp = self.layer[0] + return tmp(x) + """ + _assert_error_code(verify_all, src, "call-unresolved", private=[[4, 6]]) + + +def test_self_attr_two_hop_local_alias_must_not_grant_trust(verify_all): + src = """ + class M: + def setup(self): + self.fn = lambda x: x + 1 + def __call__(self, x): + tmp = self.fn + tmp2 = tmp + return tmp2(x) + """ + _assert_error_code(verify_all, src, "call-unresolved", private=[[4, 7]]) + + +def test_self_dunder_call_should_not_be_allowed(verify_all): + # self.__getattribute__(...) must hit the dunder-attr ban on the call path. + src = """ + class M: + def __call__(self, x): + d = self.__getattribute__('__dict__') + return x + """ + _assert_error_code(verify_all, src, "dunder-attr") + + +# ── self/cls lexical trust must not be forgeable ───────────────────────────── + + +def test_self_reassignment_must_not_grant_trust(verify_all): + src = """ + class M: + def __call__(self, x): + self = x + return self.anything_at_all() + """ + _assert_error_code(verify_all, src, "reserved-name") + + +def test_cls_as_local_variable_must_not_grant_trust(verify_all): + src = """ + class M: + def __call__(self, x): + cls = x + return cls.anything_at_all() + """ + _assert_error_code(verify_all, src, "reserved-name") + + +def test_nested_function_self_parameter_must_not_grant_trust(verify_all): + # Public Wrapper.setup stashes system; private M.__call__ must not grant nested + # helper(self) the self. exemption for a foreign object. + src = """ + from os import system + from flax import linen as nn + + class Wrapper(nn.Module): + def setup(self): + self.run = system + + class M(nn.Module): + def __call__(self, x): + w = Wrapper() + def helper(self): + return self.run('id') + return helper(w) + """ + _assert_error_code(verify_all, src, "reserved-name", private=[[9, 13]]) + + +def test_self_only_trusted_as_first_param_of_a_direct_method(verify_all): + src = """ + class M: + def __call__(self, x): + def helper(x, self): + return self.evil() + return helper(1, 2) + """ + _assert_error_code(verify_all, src, "reserved-name") + + lambda_src = """ + class M: + def __call__(self, x): + f = lambda self: self.evil() + return f(x) + """ + _assert_error_code(verify_all, lambda_src, "reserved-name") + + +def test_comprehension_target_rebinding_is_checked(verify_all): + # ast.comprehension has no lineno; targets must still be checked as Store names. + _assert_error_code(verify_all, "y = [cls for cls in [1, 2]]", "reserved-name") + + +def test_class_name_shadowing_reserved_import_alias(verify_all): + # Class name is a Store binding — shadowing jnp must be caught like jnp = evil. + src = """ + import jax.numpy as jnp + + def helper(x, secret_sink): + class jnp: + einsum = secret_sink + return jnp.einsum("ij->i", x) + """ + _assert_error_code(verify_all, src, "reserved-name") + + +def test_def_name_shadowing_public_wrapper(verify_all): + # Def shadowing a public wrapper name must be caught. + src = """ + def transpose(x): + return x + + def helper(x, secret_sink): + def transpose(y): + return secret_sink(y) + return transpose(x) + """ + _assert_error_code(verify_all, src, "reserved-name", private=[[4, 7]]) + + +# ── private_defs trusted by identifier: must not be rebindable ─────────────── + + +def test_private_def_assign_rebind_must_not_grant_call(verify_all): + # Bare calls trust private_defs by name alone. Rebinding helper = evil then + # helper(x) must fail at the rebind. + src = """ + def helper(x): + return x + def f(x, evil): + helper = evil + return helper(x) + """ + _assert_error_code(verify_all, src, "reserved-name") + + +def test_private_def_for_rebind_must_not_grant_call(verify_all): + src = """ + def helper(x): + return x + def f(items): + for helper in items: + helper(1) + """ + _assert_error_code(verify_all, src, "reserved-name") + + +def test_private_def_comprehension_rebind_must_not_grant_call(verify_all): + src = """ + def helper(x): + return x + def f(items): + return [helper for helper in items] + """ + _assert_error_code(verify_all, src, "reserved-name") + + +def test_private_def_nested_def_shadow_must_not_grant_call(verify_all): + src = """ + def helper(x): + return x + def f(x): + def helper(y): + return y + return helper(x) + """ + _assert_error_code(verify_all, src, "reserved-name") + + +def test_private_def_public_region_rebind_must_not_grant_call(verify_all): + # Multi-range hole: private def, public rebind, private call. Public glue + # must not be allowed to shadow a private def name that call sites still trust. + src = """ + def helper(x): + return x + + helper = evil + + def run(x): + return helper(x) + """ + _assert_error_code(verify_all, src, "reserved-name", private=[[1, 2], [7, 8]]) diff --git a/packages/syft-restrict/tests/verify/test_disallowed.py b/packages/syft-restrict/tests/verify/test_disallowed.py index a4b5d5b9cc4..9d0bd04a073 100644 --- a/packages/syft-restrict/tests/verify/test_disallowed.py +++ b/packages/syft-restrict/tests/verify/test_disallowed.py @@ -1,35 +1,88 @@ -"""Things the hidden region is NOT allowed to do — default-deny rejects each one.""" +"""Default-deny tests: things the private region is NOT allowed to do by default. -import pytest +Each test is a straightforward rejection of a policy category (banned-construct, +banned-name, unresolved-call, etc. +""" +import pytest from syft_restrict import verify -from .conftest import error_codes +from verify.helpers import get_error_codes, normalize_source + +# ── banned node types / constructs ─────────────────────────────────────────── @pytest.mark.parametrize( "snippet", [ - "import os\n", - "from os import path\n", - "with ctx() as g:\n pass\n", - "try:\n pass\nfinally:\n pass\n", - "raise ValueError()\n", - "def f():\n global x\n x = 1\n", - "def f():\n y = 1\n def g():\n nonlocal y\n y = 2\n return g\n", - "x = 1\ndel x\n", - "assert x\n", - "async def f():\n return 1\n", - "def f():\n yield 1\n", + "import os", + "from os import path", + """ + with ctx() as g: + pass + """, + """ + try: + pass + finally: + pass + """, + "raise ValueError()", + """ + def f(): + global x + x = 1 + """, + """ + def f(): + y = 1 + def g(): + nonlocal y + y = 2 + return g + """, + """ + x = 1 + del x + """, + "assert x", + """ + async def f(): + return 1 + """, + """ + def f(): + yield 1 + """, + "y = f'no interpolation'", # banned outright, even with no {expr} + "y = f'value={x}'", ], ) def test_banned_constructs(verify_all, snippet): result = verify_all(snippet) - assert "banned-construct" in error_codes(result), [ + assert "banned-construct" in get_error_codes(result), [ (v.code, v.message) for v in result.violations ] +def test_walrus_is_not_on_node_allowlist(verify_all): + """NamedExpr is neither allowed nor explicitly banned -> node-type rejection.""" + assert "node-type" in get_error_codes(verify_all("y = (z := 1)")) + + +def test_match_statement_is_not_on_node_allowlist(verify_all): + """``match``/``case`` is unlisted modern syntax -> default-deny node-type rejection.""" + src = """ + match x: + case 1: + y = 1 + """ + assert "node-type" in get_error_codes(verify_all(src)) + + +# ── banned names (any Load reference) ──────────────────────────────────────── + + @pytest.mark.parametrize( "name", [ @@ -49,96 +102,208 @@ def test_banned_constructs(verify_all, snippet): "input", "breakpoint", "memoryview", + "type", + "__build_class__", + "print", + "repr", + "str", + "ascii", + "format", + "bytes", ], ) -def test_banned_calls(verify_all, name): - result = verify_all(f"y = {name}(z)\n") - assert "banned-call" in error_codes(result), [ +def test_banned_names(verify_all, name): + result = verify_all(f"y = {name}") + assert "banned-name" in get_error_codes(result), [ (v.code, v.message) for v in result.violations ] +def test_banned_call_reports_banned_name_once(verify_all): + """A bare call to a banned builtin must not also emit call-unresolved.""" + result = verify_all("open('/etc/passwd')") + codes = get_error_codes(result) + assert codes == {"banned-name"}, codes + + +def test_comprehension_over_io(verify_all): + """``[v for v in open(f)]`` reintroduces I/O through a denied call in the iterable.""" + assert "banned-name" in get_error_codes(verify_all("y = [v for v in open(f)]")) + + +def test_denied_call_in_passive_position(verify_all): + """Call-checking is position-independent: a denied target in a default arg runs at + def-creation time and is still caught.""" + src = """ + def f(a=eval('1')): + return a + """ + assert "banned-name" in get_error_codes(verify_all(src)) + + +# ── attributes / methods on opaque values ──────────────────────────────────── + + @pytest.mark.parametrize( "snippet", [ - "a = x.reshape(8, -1)\n", - "b = '{0.__class__}'.format(payload)\n", - "c = items.append(1)\n", + "a = x.reshape(8, -1)", + "b = '{0.__class__}'.format(payload)", + "c = items.append(1)", ], ) def test_named_method_on_value(verify_all, snippet): - assert "method-on-value" in error_codes(verify_all(snippet)) + assert "method-on-value" in get_error_codes(verify_all(snippet)) -@pytest.mark.parametrize("snippet", ["a = x.shape\n", "b = x.T\n", "c = x.ndim\n"]) +@pytest.mark.parametrize("snippet", ["a = x.shape", "b = x.T", "c = x.ndim"]) def test_attribute_read_on_value(verify_all, snippet): - assert "attr-on-value" in error_codes(verify_all(snippet)) + assert "attr-on-value" in get_error_codes(verify_all(snippet)) + + +def test_attribute_write_to_foreign_object(verify_all): + """``some_obj.send = data`` — only ``self.`` writes are allowed.""" + assert "attr-on-value" in get_error_codes(verify_all("some_obj.send = data")) -@pytest.mark.parametrize("snippet", ["c = obj.__class__\n", "d = obj.__dict__\n"]) +@pytest.mark.parametrize("snippet", ["c = obj.__class__", "d = obj.__dict__"]) def test_dunder_attribute_read(verify_all, snippet): - assert "dunder-attr" in error_codes(verify_all(snippet)) + assert "dunder-attr" in get_error_codes(verify_all(snippet)) + + +def test_bare_class_dunder_name_is_denied(verify_all): + """``__class__`` as a bare Name (not Attribute) is still a dunder surface.""" + src = """ + class M: + def __call__(self, x): + c = __class__ + return x + """ + assert "dunder-name" in get_error_codes(verify_all(src)) + + +# ── defs / classes / decorators ────────────────────────────────────────────── @pytest.mark.parametrize("dunder", ["__init__", "__getattr__", "__reduce__"]) def test_disallowed_dunder_def(verify_all, dunder): - src = f"class M(object):\n def {dunder}(self):\n return None\n" - assert "dunder-def" in error_codes(verify_all(src)) + src = f""" + class M: + def {dunder}(self): + return None + """ + assert "dunder-def" in get_error_codes(verify_all(src)) @pytest.mark.parametrize( "snippet", [ - "@evil\ndef f():\n return 1\n", - # @property runs code on a bare attribute access — denied like any non-allow-listed decorator - "class B(object):\n @property\n def w(self):\n return 1\n", + """ + @evil + def f(): + return 1 + """, + """ + class B: + @property + def w(self): + return 1 + """, + """ + class B: + @staticmethod + def w(): + return 1 + """, ], ) def test_disallowed_decorator(verify_all, snippet): - assert "decorator" in error_codes(verify_all(snippet)) + assert "decorator" in get_error_codes(verify_all(snippet)) def test_class_keyword_metaclass(verify_all): - src = "class M(object, metaclass=Meta):\n def __call__(self, x):\n return x\n" - assert "class-keyword" in error_codes(verify_all(src)) + src = """ + class M(metaclass=Meta): + def __call__(self, x): + return x + """ + assert "class-keyword" in get_error_codes(verify_all(src)) def test_non_allowlisted_base_class(verify_all): - src = "class M(SomeLib):\n def __call__(self, x):\n return x\n" - assert "class-base" in error_codes(verify_all(src)) + src = """ + class M(SomeLib): + def __call__(self, x): + return x + """ + assert "class-base" in get_error_codes(verify_all(src)) -def test_walrus_is_not_on_node_allowlist(verify_all): - """NamedExpr is neither allowed nor explicitly banned -> node-type rejection.""" - assert "node-type" in error_codes(verify_all("y = (z := 1)\n")) +@pytest.mark.parametrize("base", ["object", "Base"]) +def test_object_and_private_bases_are_denied(verify_all, base): + """``object`` and private-region classes are rejected as bases (metaclass / + __init_subclass__ would run at class-creation time).""" + src = f""" + class Base: + def __call__(self, x): + return x + class Child({base}): + def __call__(self, x): + return x + """ + assert "class-base" in get_error_codes(verify_all(src)) + + +# ── reserved names / trusted identifiers ───────────────────────────────────── def test_reserved_module_alias_cannot_be_rebound(policy): - source = "import jax.numpy as jnp\njnp = make_evil()\n" - # only the rebind line is private (the import is visible glue) + # only the rebind line is private (the import is public glue) + source = normalize_source(""" + import jax.numpy as jnp + jnp = make_evil() + """) result = verify(source, [[2, 2]], policy) - assert "reserved-name" in error_codes(result) + assert "reserved-name" in get_error_codes(result) -def test_attribute_write_to_foreign_object(verify_all): - """`some_obj.send = data` — only `self.` writes are allowed; a foreign attribute - target is an opaque value (exfiltration channel).""" - assert "attr-on-value" in error_codes(verify_all("some_obj.send = data\n")) +def test_safe_builtin_name_cannot_be_rebound(verify_all): + """Bare calls to ``list``/``range``/… are trusted by identifier; shadowing + must be denied.""" + src = """ + def helper(): + list = None + return list + """ + assert "reserved-name" in get_error_codes(verify_all(src)) -def test_comprehension_over_io(verify_all): - """`[v for v in open(f)]` reintroduces I/O through a denied call in the iterable.""" - assert "banned-call" in error_codes(verify_all("y = [v for v in open(f)]\n")) +# ── call-target default-deny ───────────────────────────────────────────────── -def test_denied_call_in_passive_position(verify_all): - """Call-checking is position-independent: a denied target in a default arg runs at - def-creation time and is still caught.""" - src = "def f(a=eval('1')):\n return a\n" - assert "banned-call" in error_codes(verify_all(src)) +def test_call_through_bare_parameter_is_rejected(verify_all): + """A parameter is never a provably safe callee.""" + src = """ + def apply(fn, x): + return fn(x) + """ + assert "call-unresolved" in get_error_codes(verify_all(src)) -def test_match_statement_is_not_on_node_allowlist(verify_all): - """`match`/`case` is unlisted modern syntax -> default-deny node-type rejection.""" - src = "match x:\n case 1:\n y = 1\n" - assert "node-type" in error_codes(verify_all(src)) +def test_call_through_unresolvable_name_is_rejected(verify_all): + """A bare name that isn't an import, def/class, safe builtin, or tracked-safe + local has no provenance.""" + src = """ + def f(x): + return g(x) + """ + assert "call-unresolved" in get_error_codes(verify_all(src)) + + +def test_chained_call_through_unresolved_value_is_rejected(verify_all): + """``d['k'](x)``: callee is a subscript on an arbitrary parameter.""" + src = """ + def helper(d, x): + return d['k'](x) + """ + assert "call-unresolved" in get_error_codes(verify_all(src)) diff --git a/packages/syft-restrict/tests/verify/test_ranges.py b/packages/syft-restrict/tests/verify/test_ranges.py new file mode 100644 index 00000000000..341fe902234 --- /dev/null +++ b/packages/syft-restrict/tests/verify/test_ranges.py @@ -0,0 +1,20 @@ +"""Tests for verify()'s private-range argument itself, not the code inside it.""" + +import pytest +from syft_restrict import verify + +from .conftest import normalize_source + + +def test_inverted_range_must_raise_not_silently_pass(policy): + # A malformed range (end < start) must never be mistaken for "nothing to + # check here" -- it must raise, not silently verify zero nodes and report + # ok=True. + src = normalize_source(""" + import os + def f(x): + os.system("rm -rf /") + return x + """) + with pytest.raises(ValueError): + verify(src, [[4, 2]], policy) diff --git a/packages/syft-restrict/tests/verify/test_whitelist.py b/packages/syft-restrict/tests/verify/test_whitelist.py index 6a98139b10e..662e9d9a669 100644 --- a/packages/syft-restrict/tests/verify/test_whitelist.py +++ b/packages/syft-restrict/tests/verify/test_whitelist.py @@ -1,16 +1,19 @@ -"""Things the hidden region IS allowed to do — these must verify cleanly.""" +"""Whitelisted cases: things the private region is explicitly allowed to do.""" from syft_restrict import verify -from .conftest import FIXTURES, error_codes, make_policy +from verify.helpers import FIXTURES, get_error_codes, make_policy, normalize_source def _ok(result): return result.ok, [f"L{v.line} {v.code}: {v.message}" for v in result.violations] +# ── fixtures / whole-file green path ───────────────────────────────────────── + + def test_compliant_fixture_passes(policy): - """The green-path fixture (model definition) passes with no violations.""" + """The compliant fixture (model definition) passes with no violations.""" source = (FIXTURES / "compliant_model.py").read_text() config_line = next( i for i, ln in enumerate(source.splitlines(), 1) if ln.startswith("CONFIG") @@ -21,119 +24,316 @@ def test_compliant_fixture_passes(policy): assert result.n_calls_checked > 0 -def test_self_method_calls_allowed(verify_all): - """``self.method(...)`` / ``cls.method(...)`` — receiver is the module class, not opaque.""" - src = ( - "class M(object):\n" - " def setup(self):\n" - " self.w = self.param('w')\n" - " def __call__(self, x):\n" - " return self.norm(x)\n" - ) - assert _ok(verify_all(src))[0] +# ── bare names, builtins, local defs ───────────────────────────────────────── def test_bare_name_calls_allowed(verify_all): - """Calling a local var / hidden def / safe builtin by bare name is fine.""" - src = ( - "def helper(n):\n" - " rows = list(range(n))\n" - " vals = tuple(rows)\n" - " total = sum(vals)\n" - " return helper(total)\n" - ) + """Calling a local var / private def / safe builtin by bare name is fine.""" + src = """ + def helper(n): + rows = list(range(n)) + vals = tuple(rows) + total = sum(vals) + return helper(total) + """ assert _ok(verify_all(src))[0] -def test_allowed_decorators(policy): - """Only the allow-listed decorators may sit above a def/class (imports stay visible).""" - header = "from flax import linen as nn\nimport jax\n" - body = ( - "class M(object):\n" - " @nn.compact\n" - " @jax.jit\n" - " def __call__(self, x):\n" - " return x\n" - ) - source = header + body - result = verify(source, [[3, len(source.splitlines())]], policy) +def test_safe_builtin_call_is_allowed(verify_all): + src = """ + def helper(n): + return list(range(n)) + """ + assert _ok(verify_all(src))[0] + + +def test_safe_aliased_builtin_calls(verify_all): + """A local bound to a safe builtin is still callable.""" + src = """ + g = len + g([1, 2, 3]) + """ + assert _ok(verify_all(src))[0] + + +def test_public_call_is_allowed(verify_all): + """A public-region def may be called from the private region.""" + src = """ + def helper(): + return 1 + helper() + """ + # def is public (lines 1–2); call is private (line 3) + result = verify_all(src, private=[[3, 3]]) assert _ok(result)[0] -def test_allowed_dunder_defs(verify_all): - """``__call__``, ``setup`` and ``__post_init__`` are the only definable hooks.""" - src = ( - "class M(object):\n" - " def setup(self):\n" - " return None\n" - " def __post_init__(self):\n" - " return None\n" - " def __call__(self, x):\n" - " return x\n" - ) +def test_public_class_constructor_is_allowed(verify_all): + """A public-region class may be constructed by bare name from the private region.""" + src = """ + class Block: + def __call__(self, x): + return x + def run(x): + b = Block() + return b(x) + """ + # class Block is public (lines 1–3); run is private (lines 4–6) + result = verify_all(src, private=[[4, 6]]) + assert _ok(result)[0] + + +def test_verify_does_not_mutate_caller_policy(policy): + """verify() must not write reserved_names onto the caller's Policy instance.""" + assert policy.reserved_names == set() + src = normalize_source(""" + import jax.numpy as jnp + def f(): + return 1 + """) + verify(src, [[3, 4]], policy) + assert policy.reserved_names == set() + + +def test_private_call_is_allowed(verify_all): + """A private-region def may be called from the private region.""" + src = """ + def helper(): + return 1 + helper() + """ assert _ok(verify_all(src))[0] -def test_class_bases_object_and_hidden_def(verify_all): - """A class may subclass ``object`` or another class defined in the hidden region.""" - src = ( - "class Base(object):\n" - " def __call__(self, x):\n" - " return x\n" - "class Child(Base):\n" - " def __call__(self, x):\n" - " return x\n" - ) +def test_local_bound_to_private_constructor_is_still_callable(verify_all): + """``block = Attn(); block(x)`` — local traced to a class defined here.""" + src = """ + class Attn: + def __call__(self, x): + return x + + def helper(x): + block = Attn() + return block(x) + """ assert _ok(verify_all(src))[0] -def test_control_flow_and_comprehensions(verify_all): - """if/for/while/ternary, comprehensions and a lambda are all on the node allow-list.""" - src = ( - "def f(xs):\n" - " acc = [g(v) for v in xs if v]\n" - " pairs = {k: k for k in xs}\n" - " seen = {v for v in xs}\n" - " h = (lambda z: z)\n" - " out = 0\n" - " for v in xs:\n" - " out = v\n" - " if out:\n" - " break\n" - " while out:\n" - " out = 0\n" - " return out if xs else h(acc)\n" - ) +def test_chained_call_through_private_constructor_is_allowed(verify_all): + """``Block()(x)``: callee is itself a Call to a class defined here.""" + src = """ + class Block: + def __call__(self, x): + return x + + def helper(x): + return Block()(x) + """ assert _ok(verify_all(src))[0] -def test_fstring_is_allowed(verify_all): - """f-strings (JoinedStr/FormattedValue) over bare names are allowed.""" - src = "def f(x):\n return f'value={x}'\n" +# ── self / cls ─────────────────────────────────────────────────────────────── + + +def test_self_method_calls_allowed(verify_all): + """``self.method(...)`` / inherited ``self.param`` — receiver is the module class.""" + src = """ + class M: + def setup(self): + self.w = self.param('w') + def __call__(self, x): + return self.norm(x) + """ assert _ok(verify_all(src))[0] -def test_calling_a_value_is_allowed(verify_all): - """Calling the result of a subscript/call (its ``__call__``) is allowed.""" - src = ( - "class M(object):\n" - " def __call__(self, x):\n" - " layer = self.layers[0]\n" - " return self.layers[0](x) + layer(x)\n" - ) +def test_calling_a_value_through_self_is_allowed(verify_all): + """Calling a trusted self-rooted subscript/local alias is allowed.""" + src = """ + class M: + def __call__(self, x): + layer = self.layers[0] + return self.layers[0](x) + layer(x) + """ + assert _ok(verify_all(src))[0] + + +def test_self_attr_local_alias_of_safe_source_is_still_allowed(verify_all): + src = """ + class M: + def __call__(self, x): + layer = self.layers[0] + return layer(x) + """ + assert _ok(verify_all(src))[0] + + +def test_self_attr_assigned_allowlisted_import_call_is_callable(policy): + src = normalize_source(""" + from flax import linen as nn + + class M(nn.Module): + def setup(self): + self.dense = nn.Dense(8) + def __call__(self, x): + return self.dense(x) + """) + result = verify(src, [[3, 7]], policy) + assert _ok(result)[0] + + +def test_self_attr_list_of_constructors_element_alias_is_callable(verify_all): + """The Gemma-style ``block = self.layer[0]; block(x)`` pattern.""" + src = """ + class Block: + def __call__(self, x): + return x + + class M: + def setup(self): + self.layer = [Block() for _ in range(3)] + def __call__(self, x): + block = self.layer[0] + return block(x) + """ + assert _ok(verify_all(src))[0] + + +def test_self_attr_reassigned_local_alias_is_not_stuck_tainted(verify_all): + """Reassigning a local that once aliased self. clears the taint.""" + src = """ + class M: + def setup(self): + self.fn = lambda x: x + 1 + def __call__(self, x): + tmp = self.fn + tmp = x + return tmp + """ + assert _ok(verify_all(src))[0] + + +# ── decorators / dunder defs / bases ───────────────────────────────────────── + + +def test_allowed_decorators(policy): + """Only the allow-listed decorators may sit above a def/class (imports stay public).""" + src = normalize_source(""" + from flax import linen as nn + import jax + + class M: + @nn.compact + @jax.jit + def __call__(self, x): + return x + """) + result = verify(src, [[4, len(src.splitlines())]], policy) + assert _ok(result)[0] + + +def test_allowed_dunder_defs(verify_all): + """``__call__``, ``setup`` and ``__post_init__`` are the only definable hooks.""" + src = """ + class M: + def setup(self): + return None + def __post_init__(self): + return None + def __call__(self, x): + return x + """ + assert _ok(verify_all(src))[0] + + +def test_allowlisted_import_as_base_class(policy): + """A base class must be an allow-listed import (e.g. ``nn.Module``).""" + src = normalize_source(""" + from flax import linen as nn + class M(nn.Module): + def __call__(self, x): + return x + """) + result = verify(src, [[2, len(src.splitlines())]], make_policy()) + assert _ok(result)[0] + + +# ── control flow / assignment / annotations ────────────────────────────────── + + +def test_control_flow_and_comprehensions(verify_all): + """if/for/while/ternary, comprehensions and a lambda are all on the node allow-list. + + Uses ``abs`` (a safe builtin) rather than an arbitrary bare name for the comprehension's call; + unresolvable bare-name calls are covered in ``test_disallowed.py``. + """ + src = """ + def f(xs): + acc = [abs(v) for v in xs if v] + pairs = {k: k for k in xs} + seen = {v for v in xs} + h = (lambda z: z) + out = 0 + for v in xs: + out = v + if out: + break + while out: + out = 0 + return out if xs else acc + """ assert _ok(verify_all(src))[0] def test_non_reserved_rebind_is_fine(verify_all): """Reassigning ordinary locals (not import aliases / wrapper names) is allowed.""" - src = "def f(x):\n y = x\n y += 1\n z: int = y\n return z\n" + src = """ + def f(x): + y = x + y += 1 + z: int = y + return z + """ result = verify_all(src) assert _ok(result)[0] - assert "reserved-name" not in error_codes(result) + assert "reserved-name" not in get_error_codes(result) def test_arithmetic_only_policy_still_passes_pure_math(verify_all): """A bundle the code doesn't use being disabled doesn't cause spurious failures.""" - pol = make_policy(methods=["arithmetic"]) - src = "def f(a, b):\n return a + b - (-a)\n" + pol = make_policy(operators=["arithmetic"]) + src = """ + def f(a, b): + return a + b - (-a) + """ assert _ok(verify_all(src, pol))[0] + + +def test_nested_generic_annotation_is_not_flagged(verify_all): + """Banned-name identifiers inside generic annotations are not references.""" + src = """ + def f(x: list[str]) -> dict[str, bytes]: + return {} + """ + assert _ok(verify_all(src))[0] + + +def test_future_annotations_import_does_not_change_annotation_handling(policy): + """``from __future__ import annotations`` does not change the AST shape of annotations.""" + src = normalize_source(""" + from __future__ import annotations + + def f(x: str) -> list[bytes]: + return [] + """) + result = verify(src, [[3, 4]], policy) + assert _ok(result)[0] + + +def test_stringized_annotation_is_not_flagged(verify_all): + """Quoted forward-reference annotations are Constants, not Names.""" + src = """ + def f(x: "str") -> "bytes": + return x + """ + assert _ok(verify_all(src))[0] diff --git a/packages/syft-restrict/tests/verify/test_whitelisted_lib.py b/packages/syft-restrict/tests/verify/test_whitelisted_lib.py index 89e0acc499b..b68e2450041 100644 --- a/packages/syft-restrict/tests/verify/test_whitelisted_lib.py +++ b/packages/syft-restrict/tests/verify/test_whitelisted_lib.py @@ -1,95 +1,146 @@ """Things we *manually* allow: library calls resolved by name, and operator bundles. -These need real imports, which are themselves banned inside the hidden region — so the import lines -stay *visible* (glue) and only the usage below them is marked private. +These need real imports, which are themselves banned inside the private region — so the import +lines stay *public* (glue) and only the usage below them is marked private via ``private=``. """ import pytest -from syft_restrict import verify +from verify.helpers import get_error_codes, make_policy -from .conftest import error_codes, make_policy +# ── library calls allowed BY NAME (resolved against the public imports) ───────────── -def _verify(header: str, body: str, policy): - """Keep ``header`` lines visible (imports), mark the ``body`` lines below as private.""" - head_lines = header.splitlines() - source = header + body - lo = len(head_lines) + 1 - hi = len(source.splitlines()) - return verify(source, [[lo, hi]], policy) - - -# ── library calls allowed BY NAME (resolved against the imports) ────────────────────────── @pytest.mark.parametrize( - "header, body", + "src", [ - ("import jax.numpy as jnp\n", "r = jnp.einsum('ij,jk->ik', a, b)\n"), - ("import jax\n", "r = jax.nn.softmax(x)\n"), - ("import jax\n", "r = jax.lax.rsqrt(x)\n"), - ("from flax import linen as nn\n", "layer = nn.Dense(8)\n"), - ("from flax import linen as nn\n", "layer = nn.LayerNorm()\n"), + """ + import jax.numpy as jnp + r = jnp.einsum('ij,jk->ik', a, b) + """, + """ + import jax + r = jax.nn.softmax(x) + """, + """ + import jax + r = jax.lax.rsqrt(x) + """, + """ + from flax import linen as nn + layer = nn.Dense(8) + """, + """ + from flax import linen as nn + layer = nn.LayerNorm() + """, ], ) -def test_allowlisted_library_calls(policy, header, body): - result = _verify(header, body, policy) +def test_allowlisted_library_calls(verify_all, src): + # line 1 = public import; rest = private + result = verify_all(src, private=[[2, 2]]) assert result.ok, [(v.code, v.message) for v in result.violations] -def test_attribute_read_on_allowlisted_namespace(policy): +def test_attribute_read_on_allowlisted_namespace(verify_all): """A constant read off an allow-listed module (``jnp.pi``) resolves and is allowed.""" - result = _verify("import jax.numpy as jnp\n", "r = jnp.pi\n", policy) + src = """ + import jax.numpy as jnp + r = jnp.pi + """ + result = verify_all(src, private=[[2, 2]]) assert result.ok, [(v.code, v.message) for v in result.violations] -# ── operator bundles: enabled => allowed ────────────────────────────────────────────────── +# ── operator bundles: enabled => allowed ───────────────────────────────────── + + @pytest.mark.parametrize( - "body", + "src", [ - "r = a + b - (-a)\n", # arithmetic: BinOp / UnaryOp - "r = (a < b) and (b > a)\n", # comparison: Compare / BoolOp - "r = x[0] + x[1:3]\n", # indexing: Subscript / Slice + "r = a + b - (-a)", # arithmetic: BinOp / UnaryOp + "r = (a < b) and (b > a)", # comparison: Compare / BoolOp + "r = x[0] + x[1:3]", # indexing: Subscript / Slice ], ) -def test_enabled_operator_bundles(verify_all, body): - result = verify_all(body) +def test_enabled_operator_bundles(verify_all, src): + result = verify_all(src) assert result.ok, [(v.code, v.message) for v in result.violations] -# ── operator bundles: disabled => bundle-disabled ───────────────────────────────────────── +# ── operator bundles: disabled => operator-disabled ────────────────────────── + + @pytest.mark.parametrize( - "methods, body", + "operators, src", [ - (["indexing", "comparison"], "r = a + b\n"), # arithmetic off - (["arithmetic", "indexing"], "r = a < b\n"), # comparison off - (["arithmetic", "comparison"], "r = x[0]\n"), # indexing off + (["indexing", "comparison"], "r = a + b"), # arithmetic off + (["arithmetic", "indexing"], "r = a < b"), # comparison off + (["arithmetic", "comparison"], "r = x[0]"), # indexing off ], ) -def test_disabled_operator_bundle_is_rejected(verify_all, methods, body): - result = verify_all(body, make_policy(methods=methods)) - assert "bundle-disabled" in error_codes(result) +def test_disabled_operator_bundle_is_rejected(verify_all, operators, src): + result = verify_all(src, make_policy(operators=operators)) + assert "operator-disabled" in get_error_codes(result) + + +# ── user-supplied disallow beats the allow ─────────────────────────────────── -# ── denylist beats the allow, even under an otherwise-allowed module ────────────────────── @pytest.mark.parametrize( - "header, body", + "src", [ - ("import jax\n", "q = jax.experimental.io_callback(send, x)\n"), - ("import jax.numpy as jnp\n", "q = jnp.save('f', x)\n"), - ("import jax\n", "q = jax.debug.print('x', x)\n"), + """ + import jax + q = jax.experimental.io_callback(send, x) + """, + """ + import jax.numpy as jnp + q = jnp.save('f', x) + """, + """ + import jax + q = jax.debug.print('x', x) + """, ], ) -def test_denylist_beats_allow(policy, header, body): - result = _verify(header, body, policy) - assert "call-not-allowed" in error_codes(result) +def test_user_disallow_beats_allow(verify_all, src): + # Under a broad `jax.*` allow, an explicit disallow still rejects the dangerous leaves. + pol = make_policy(disallow=["jax.experimental.*", "jax.numpy.save", "jax.debug.*"]) + result = verify_all(src, pol, private=[[2, 2]]) + assert "call-not-allowed" in get_error_codes(result) + + +# ── with no disallow, a broad allow permits a formerly-denied leaf ─────────── + + +def test_no_disallow_permits_leaf_under_broad_allow(verify_all): + # Intentional behavior change: safety comes from a *specific* allow-list or an explicit + # disallow list; a bare `jax.*` allow with no disallow now permits `jax.numpy.save`. + src = """ + import jax.numpy as jnp + q = jnp.save('f', x) + """ + result = verify_all(src, private=[[2, 2]]) + assert "call-not-allowed" not in get_error_codes(result) + + +# ── a library that simply isn't on the allow-list ──────────────────────────── -# ── a library that simply isn't on the allow-list ───────────────────────────────────────── -def test_non_allowlisted_library_call(policy): - result = _verify("import numpy as np\n", "r = np.dot(a, b)\n", policy) - assert "call-not-allowed" in error_codes(result) +def test_non_allowlisted_library_call(verify_all): + src = """ + import numpy as np + r = np.dot(a, b) + """ + result = verify_all(src, private=[[2, 2]]) + assert "call-not-allowed" in get_error_codes(result) -def test_non_allowlisted_library_attribute(policy): - result = _verify("import numpy as np\n", "r = np.pi\n", policy) - assert "attr-not-allowed" in error_codes(result) +def test_non_allowlisted_library_attribute(verify_all): + src = """ + import numpy as np + r = np.pi + """ + result = verify_all(src, private=[[2, 2]]) + assert "attr-not-allowed" in get_error_codes(result)