diff --git a/packages/syft-restrict/README.md b/packages/syft-restrict/README.md index b5ebbb557ae..6c50fa05066 100644 --- a/packages/syft-restrict/README.md +++ b/packages/syft-restrict/README.md @@ -57,36 +57,48 @@ import syft_restrict as restrict result = restrict.run( "gemma_inference.py", - allow_functions=["jax.*", "flax.linen.*"], # functions callable BY NAME (path-resolved) + # default-deny: allow ONLY the exact paths this model calls, nothing else in jax/flax + allow_functions=[ + "jax.numpy.einsum", "jax.numpy.mean", "jax.numpy.square", "jax.numpy.arange", + "jax.numpy.sin", "jax.numpy.cos", "jax.numpy.concatenate", "jax.numpy.tril", + "jax.numpy.triu", "jax.numpy.ones", "jax.numpy.where", "jax.numpy.repeat", + "jax.numpy.sqrt", "jax.numpy.transpose", "jax.numpy.array", "jax.numpy.float32", + "jax.numpy.bool_", "jax.lax.rsqrt", "jax.nn.softmax", "jax.nn.gelu", + "flax.linen.Module", + "jax.lax", "jax.nn", # module refs the deep-path calls (jax.lax.rsqrt, jax.nn.softmax) resolve through + ], allow_operators=["arithmetic", "indexing", "comparison"], # operators allowed ON A VALUE ) # On success: writes gemma_inference.obfuscated.py and returns result.certificate. # On a policy violation: raises PolicyViolation naming each offending line. -# If the file has no markers and obfuscate=/hide= are both omitted: raises MarkerError. +# If the file has no `# syft-restrict: ...` markers: raises MarkerError. ``` -`obfuscate=`/`hide=` still work as an explicit escape hatch — pass either one (even an -empty list) and marker scanning is skipped entirely in favor of your own 1-based line -ranges: +> [!IMPORTANT] +> List the **specific** paths your model calls, as above — this is the default-deny posture the +> tool is built for: everything not named is denied. Avoid broad globs like `jax.*`: they pull in +> JAX's own host-callback and disk-IO functions (`jax.numpy.save`, `jax.debug.callback`, +> `jax.experimental.io_callback`, …) that the private region could call to exfiltrate data. If you +> must use a broad glob, pair it with a `disallow_functions` floor — see +> [docs/blacklist.md](docs/blacklist.md#optional-disallow_functions). + +The private region is designated **only** by `# syft-restrict: ...` comment markers in the +source; `run()` resolves them and refuses a file that carries none. This is how +**[examples/gemma_inference.py](examples/gemma_inference.py)** generates its +obfuscated copy — see [examples/generate.py](examples/generate.py). + +Two optional strictness flags, both defaulting to the permissive behavior: ```python 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.*"], + allow_functions=[...], # as above allow_operators=["arithmetic", "indexing", "comparison"], + allow_local_assignments=True, # False: a local aliased to a callable can't be called by name + allow_base_class_attributes=True, # False: a never-assigned self. is not presumed inherited ) ``` -This is how **[examples/gemma_inference.py](examples/gemma_inference.py)** generates -**[examples/gemma_inference.obfuscated.py](examples/gemma_inference.obfuscated.py)** — -see [examples/generate.py](examples/generate.py). - -The same model marked up with comment blocks instead of numeric ranges lives in -**[examples/gemma_inference_marked.py](examples/gemma_inference_marked.py)** / -[examples/generate_marked.py](examples/generate_marked.py). - If syft-restrict was successfully executed on the true original file, the obfuscated file can be safely shared without exposing the private section. Pass `strict=False` to `run` to get a `RunResult` with `.ok` / `.violations` diff --git a/packages/syft-restrict/docs/blacklist.md b/packages/syft-restrict/docs/blacklist.md index 4932df8f563..3de41730dd4 100644 --- a/packages/syft-restrict/docs/blacklist.md +++ b/packages/syft-restrict/docs/blacklist.md @@ -116,8 +116,20 @@ Same idea as calling dunders on a value (`x.__repr__()`), spelled as a bare call - `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. +### Interpreter / site builtins + +Injected by the `site` module. `copyright`/`credits`/`license` write to stdout (same channel as +`print`); `exit`/`quit` shut the interpreter down (`SystemExit`); `help` is an interactive/IO surface. + +- `copyright` +- `credits` +- `license` +- `exit` +- `quit` +- `help` + +> [!WARNING] > `bytes(x)` can dump raw buffer contents. `print`, `copyright`, `credits`, and `license` are stdout +> channels. F-strings are banned separately as constructs (above), including with no `{...}` at all. ```python # all rejected (banned-name) @@ -141,7 +153,7 @@ y = [v for v in open(path)] # passive positions still checked | 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` | +| Dunder attribute | `obj.__class__`, `jnp.fn.__wrapped__(x)` | `dunder-attr` | | Bare dunder name | `c = __class__` | `dunder-name` | | Library attr not allowed | `np.pi` when numpy is not allowed | `attr-not-allowed` | @@ -162,6 +174,11 @@ def apply(fn, 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). +> [!IMPORTANT] > **Dunders are denied in call position too, not just when read.** A dunder segment on an +> allow-listed path — `jnp.einsum.__wrapped__(x)`, `jnp.fn.__globals__[...]` — reaches the +> function-object introspection surface (`__wrapped__`, `__globals__`, `__defaults__`, +> `__closure__`, …) and is rejected as `dunder-attr` **regardless of the allow-list**. + --- ## Classes and definitions @@ -172,7 +189,7 @@ safe are in [verify.md](verify.md#self-and-flax-style-modules). | 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__`. +Allowed hooks only: `setup`, `__call__`, `__post_init__`. Decorators are banned outright as a [forbidden construct](#forbidden-constructs). diff --git a/packages/syft-restrict/docs/verify.md b/packages/syft-restrict/docs/verify.md index bfc46901b01..22025d4860c 100644 --- a/packages/syft-restrict/docs/verify.md +++ b/packages/syft-restrict/docs/verify.md @@ -83,17 +83,16 @@ class Net(nn.Module): The reverse isn't allowed (`obfuscate` can't nest inside `hide`), and neither kind nests inside itself. Any of these raise `MarkerError`: -| Situation | Result | -| --------------------------------------------------------------- | ---------------------------------------------------------------- | -| `start` with no matching `end` (or vice versa) | `MarkerError`, names the line | -| `hide-end`/`obfuscate-end` closing the wrong kind | `MarkerError` (mismatched kind) | -| `obfuscate` nested inside `hide` | `MarkerError` | -| `obfuscate` nested inside `obfuscate` (or `hide` inside `hide`) | `MarkerError` | -| A block with nothing between its start/end | `MarkerError` (empty block) | -| No `# syft-restrict: ...` marker anywhere in the file | `MarkerError`, unless `obfuscate=`/`hide=` are passed explicitly | - -`run()` scans for these markers automatically whenever `obfuscate` and `hide` are both omitted; -passing either explicitly (even `[]`) skips marker scanning entirely and uses your ranges as-is. +| Situation | Result | +| --------------------------------------------------------------- | ----------------------------------------------------------- | +| `start` with no matching `end` (or vice versa) | `MarkerError`, names the line | +| `hide-end`/`obfuscate-end` closing the wrong kind | `MarkerError` (mismatched kind) | +| `obfuscate` nested inside `hide` | `MarkerError` | +| `obfuscate` nested inside `obfuscate` (or `hide` inside `hide`) | `MarkerError` | +| A block with nothing between its start/end | `MarkerError` (empty block) | +| No `# syft-restrict: ...` marker anywhere in the file | `MarkerError` (the private region is designated by markers) | + +`run()` resolves the private region from these markers; a file with none raises `MarkerError`. > [!WARNING] > Private code may **call** public wrappers by name. Public code is trusted, not verified. @@ -101,6 +100,23 @@ passing either explicitly (even `[]`) skips marker scanning entirely and uses yo > A clean `verify()` means the private region cannot escape on its own, not that > the whole file is safe to execute. +### Imports + +Imports happen in the public region but govern what the private region may +reach: + +- **Public imports are unchecked.** Any module may be imported in public code; + syft-restrict never restricts or inspects the import itself. +- **Private imports are banned outright** — an `import` inside the private + region is rejected as `banned-construct`. +- **The private region's _use_ of an import is what's gated.** A private call + like `jnp.einsum(...)` resolves through the public import table and must match + `allow_functions` (see below); otherwise it fails. So the control point is the + private-side call, not the public import. + +Star imports (`from jax import *`) are disallowed everywhere, because they make +it impossible to review the imported names. + --- ## What the verifier does @@ -303,7 +319,7 @@ class Net(nn.Module): # base must be allow-listed (e.g. flax.linen. | ------------------------------------------------------------- | ------------------------------------------------------------------------------- | | Bases that resolve to an allow-listed path (e.g. `nn.Module`) | `object`, random private bases, non-allow-listed libs | | — | Any decorator: `@property`, `@staticmethod`, `@nn.compact`, arbitrary functions | -| Defining `setup`, `__call__` | `__getattr__`, `__reduce__`, `__post_init__`, other magic methods | +| Defining `setup`, `__call__`, `__post_init__` | `__getattr__`, `__reduce__`, other magic methods | | — | `metaclass=` / other class keywords | --- diff --git a/packages/syft-restrict/examples/gemma_inference.certificate.json b/packages/syft-restrict/examples/gemma_inference.certificate.json index 327620e36fe..2c761c2835b 100644 --- a/packages/syft-restrict/examples/gemma_inference.certificate.json +++ b/packages/syft-restrict/examples/gemma_inference.certificate.json @@ -1,86 +1,88 @@ { - "source_sha256": "3b462b5a6490b9dba8a93cf1a06da37601b002ab3a6f7497c01ee6698a671a18", - "policy_id": "b4a557058704eed6", + "source_sha256": "01232c264defbe3fe85b453aefedceac36f7b92ed7a492fd34c790a0cf78f094", + "policy_id": "d84d1e21530fa500", "restrict_version": "0.1.0", "private_ranges": [ - [24, 88], - [91, 91], - [99, 99], - [110, 110], - [122, 122], - [151, 152], - [155, 155], - [159, 160], - [163, 163], - [168, 171], - [178, 178], - [210, 211], - [215, 215], - [221, 225], - [233, 233], - [244, 247], - [250, 250], - [255, 258], - [267, 267], - [92, 93], - [100, 107], - [111, 119], - [123, 129], - [153, 153], - [156, 156], - [161, 161], - [164, 165], - [172, 176], - [179, 207], - [212, 213], - [216, 218], - [226, 231], - [234, 241], - [248, 248], - [251, 252], - [259, 265], - [268, 292] + [24, 91], + [96, 101], + [112, 114], + [126, 128], + [138, 139], + [163, 164], + [169, 170], + [175, 178], + [183, 184], + [190, 195], + [204, 205], + [238, 241], + [247, 248], + [255, 261], + [271, 272], + [284, 289], + [294, 295], + [301, 306], + [317, 318], + [347, 348], + [93, 94], + [103, 110], + [116, 124], + [130, 136], + [166, 167], + [172, 173], + [180, 181], + [186, 188], + [197, 202], + [207, 236], + [243, 245], + [250, 253], + [263, 269], + [274, 282], + [291, 292], + [297, 299], + [308, 315], + [320, 345] ], "obfuscate_ranges": [ - [24, 88], - [91, 91], - [99, 99], - [110, 110], - [122, 122], - [151, 152], - [155, 155], - [159, 160], - [163, 163], - [168, 171], - [178, 178], - [210, 211], - [215, 215], - [221, 225], - [233, 233], - [244, 247], - [250, 250], - [255, 258], - [267, 267] + [24, 91], + [96, 101], + [112, 114], + [126, 128], + [138, 139], + [163, 164], + [169, 170], + [175, 178], + [183, 184], + [190, 195], + [204, 205], + [238, 241], + [247, 248], + [255, 261], + [271, 272], + [284, 289], + [294, 295], + [301, 306], + [317, 318], + [347, 348] ], "hide_ranges": [ - [92, 93], - [100, 107], - [111, 119], - [123, 129], - [153, 153], - [156, 156], - [161, 161], - [164, 165], - [172, 176], - [179, 207], - [212, 213], - [216, 218], - [226, 231], - [234, 241], - [248, 248], - [251, 252], - [259, 265], - [268, 292] + [93, 94], + [103, 110], + [116, 124], + [130, 136], + [166, 167], + [172, 173], + [180, 181], + [186, 188], + [197, 202], + [207, 236], + [243, 245], + [250, 253], + [263, 269], + [274, 282], + [291, 292], + [297, 299], + [308, 315], + [320, 345] ], "n_calls_checked": 77 } diff --git a/packages/syft-restrict/examples/gemma_inference.obfuscated.py b/packages/syft-restrict/examples/gemma_inference.obfuscated.py index ab009defb83..3b268532452 100644 --- a/packages/syft-restrict/examples/gemma_inference.obfuscated.py +++ b/packages/syft-restrict/examples/gemma_inference.obfuscated.py @@ -18,9 +18,9 @@ import sentencepiece as spm from flax import linen as nn - # ── Model config ───────────────────────────────────────────────────────────── # Active model config. Comment/uncomment to switch sizes. +# syft-restrict: obfuscate-start ░v0 = dict( ░v1=■, ░v2=■, @@ -89,14 +89,17 @@ def ░fn0(░v1): + # syft-restrict: hide-start ■■■■■■■■ # hidden/obfuscated lines can only execute restricted python, see restrict docs for more details ■■■■■■■■ + # syft-restrict: hide-end -# ── Standalone helpers ──────────────────────────────────────────────────── +# THIS COMMENT WAS OBFUSCATED def ░fn1(░v14, ░v15, ░v16): + # syft-restrict: hide-start ■■■■■■■■ # hidden/obfuscated lines can only execute restricted python, see restrict docs for more details ■■■■■■■■ ■■■■■■■■ @@ -105,9 +108,11 @@ def ░fn1(░v14, ░v15, ░v16): ■■■■■■■■ ■■■■■■■■ ■■■■■■■■ + # syft-restrict: hide-end def ░fn2(░v17, ░v7): + # syft-restrict: hide-start ■■■■■■■■ # hidden/obfuscated lines can only execute restricted python, see restrict docs for more details ■■■■■■■■ ■■■■■■■■ @@ -117,9 +122,11 @@ def ░fn2(░v17, ░v7): ■■■■■■■■ ■■■■■■■■ ■■■■■■■■ + # syft-restrict: hide-end def ░fn3(░v18, ░v7): + # syft-restrict: hide-start ■■■■■■■■ # hidden/obfuscated lines can only execute restricted python, see restrict docs for more details ■■■■■■■■ ■■■■■■■■ @@ -127,6 +134,10 @@ def ░fn3(░v18, ░v7): ■■■■■■■■ ■■■■■■■■ ■■■■■■■■ + # syft-restrict: hide-end + + +# syft-restrict: obfuscate-end # ── Flax modules ─────────────────────────────────────────────────────────── @@ -148,34 +159,51 @@ def append_to(lst, item): return lst +# syft-restrict: obfuscate-start class ░Cls0(nn.Module): def setup(self): + # syft-restrict: hide-start ■■■■■■■■ # hidden/obfuscated lines can only execute restricted python, see restrict docs for more details + # syft-restrict: hide-end + def __call__(self, ░v19, ░v14): + # syft-restrict: hide-start ■■■■■■■■ # hidden/obfuscated lines can only execute restricted python, see restrict docs for more details + # syft-restrict: hide-end + class ░Cls1(nn.Module): def setup(self): + # syft-restrict: hide-start ■■■■■■■■ # hidden/obfuscated lines can only execute restricted python, see restrict docs for more details + # syft-restrict: hide-end + def __call__(self, ░v14): + # syft-restrict: hide-start ■■■■■■■■ # hidden/obfuscated lines can only execute restricted python, see restrict docs for more details ■■■■■■■■ + # syft-restrict: hide-end + class ░Cls2(nn.Module): ░v20: dict def setup(self): + # syft-restrict: hide-start ■■■■■■■■ # hidden/obfuscated lines can only execute restricted python, see restrict docs for more details ■■■■■■■■ ■■■■■■■■ ■■■■■■■■ ■■■■■■■■ + # syft-restrict: hide-end + def __call__(self, ░v14, ░v15, ░v21, ░v22, ░v23=None): + # syft-restrict: hide-start ■■■■■■■■ # hidden/obfuscated lines can only execute restricted python, see restrict docs for more details ■■■■■■■■ ■■■■■■■■ @@ -206,23 +234,32 @@ def __call__(self, ░v14, ░v15, ░v21, ░v22, ░v23=None): ■■■■■■■■ ■■■■■■■■ + # syft-restrict: hide-end + class ░Cls3(nn.Module): def setup(self): + # syft-restrict: hide-start ■■■■■■■■ # hidden/obfuscated lines can only execute restricted python, see restrict docs for more details ■■■■■■■■ + # syft-restrict: hide-end + def __call__(self, ░v14): + # syft-restrict: hide-start ■■■■■■■■ # hidden/obfuscated lines can only execute restricted python, see restrict docs for more details ■■■■■■■■ ■■■■■■■■ + # syft-restrict: hide-end + class ░Cls4(nn.Module): ░v20: dict ░v22: str = "■" def setup(self): + # syft-restrict: hide-start ■■■■■■■■ # hidden/obfuscated lines can only execute restricted python, see restrict docs for more details ■■■■■■■■ ■■■■■■■■ @@ -230,7 +267,10 @@ def setup(self): ■■■■■■■■ ■■■■■■■■ + # syft-restrict: hide-end + def __call__(self, ░v14, ░v15, ░v21, ░v23=None): + # syft-restrict: hide-start ■■■■■■■■ # hidden/obfuscated lines can only execute restricted python, see restrict docs for more details ■■■■■■■■ ■■■■■■■■ @@ -240,22 +280,31 @@ def __call__(self, ░v14, ░v15, ░v21, ░v23=None): ■■■■■■■■ ■■■■■■■■ + # syft-restrict: hide-end + class ░Cls5(nn.Module): ░v20: dict def setup(self): + # syft-restrict: hide-start ■■■■■■■■ # hidden/obfuscated lines can only execute restricted python, see restrict docs for more details + # syft-restrict: hide-end + def __call__(self, ░v24): + # syft-restrict: hide-start ■■■■■■■■ # hidden/obfuscated lines can only execute restricted python, see restrict docs for more details ■■■■■■■■ + # syft-restrict: hide-end + class ░Cls6(nn.Module): ░v20: dict def setup(self): + # syft-restrict: hide-start ■■■■■■■■ # hidden/obfuscated lines can only execute restricted python, see restrict docs for more details ■■■■■■■■ ■■■■■■■■ @@ -264,7 +313,10 @@ def setup(self): ■■■■■■■■ ■■■■■■■■ + # syft-restrict: hide-end + def __call__(self, ░v25, ░v23=None): + # syft-restrict: hide-start ■■■■■■■■ # hidden/obfuscated lines can only execute restricted python, see restrict docs for more details ■■■■■■■■ ■■■■■■■■ @@ -291,6 +343,11 @@ def __call__(self, ░v25, ░v23=None): ■■■■■■■■ ■■■■■■■■ + # syft-restrict: hide-end + + +# syft-restrict: obfuscate-end + # ── Weight loading ───────────────────────────────────────────────────────── diff --git a/packages/syft-restrict/examples/gemma_inference.py b/packages/syft-restrict/examples/gemma_inference.py index db075c9bc63..934294e3061 100644 --- a/packages/syft-restrict/examples/gemma_inference.py +++ b/packages/syft-restrict/examples/gemma_inference.py @@ -18,9 +18,9 @@ import sentencepiece as spm from flax import linen as nn - # ── Model config ───────────────────────────────────────────────────────────── # Active model config. Comment/uncomment to switch sizes. +# syft-restrict: obfuscate-start CONFIG = dict( num_layers=18, embed_dim=640, @@ -89,14 +89,17 @@ def _attn_types(num_layers): + # syft-restrict: hide-start pattern = ("local",) * 5 + ("global",) return (pattern * ((num_layers + 5) // 6))[:num_layers] + # syft-restrict: hide-end # ── Standalone helpers ──────────────────────────────────────────────────── def apply_rope(x, positions, base_freq): + # syft-restrict: hide-start """Rotary position embeddings (split-half rotation).""" half = shape_of(x)[-1] // 2 freq_exp = (2.0 / shape_of(x)[-1]) * jnp.arange(half, dtype=jnp.float32) @@ -105,9 +108,11 @@ def apply_rope(x, positions, base_freq): sin, cos = jnp.sin(angles), jnp.cos(angles) x1, x2 = x[..., :half], x[..., half:] return jnp.concatenate([x1 * cos - x2 * sin, x2 * cos + x1 * sin], axis=-1) + # syft-restrict: hide-end def make_masks(seq_len, sliding_window): + # syft-restrict: hide-start """Causal masks — local layers also clip to a sliding window.""" causal = jnp.tril(jnp.ones((seq_len, seq_len), dtype=jnp.bool_)) window = jnp.triu( @@ -117,9 +122,11 @@ def make_masks(seq_len, sliding_window): "local": (causal & window)[None, None], "global": causal[None, None], } + # syft-restrict: hide-end def make_decode_masks(pos, sliding_window): + # syft-restrict: hide-start """Masks for single-token decode.""" total_len = pos + 1 positions = jnp.arange(total_len) @@ -127,6 +134,10 @@ def make_decode_masks(pos, sliding_window): "local": (positions >= pos - sliding_window + 1)[None, None, None, :], "global": jnp.ones((1, 1, 1, total_len), dtype=jnp.bool_), } + # syft-restrict: hide-end + + +# syft-restrict: obfuscate-end # ── Flax modules ─────────────────────────────────────────────────────────── @@ -148,34 +159,51 @@ def append_to(lst, item): return lst +# syft-restrict: obfuscate-start class Einsum(nn.Module): def setup(self): + # syft-restrict: hide-start self.w = _get(self, "w") + # syft-restrict: hide-end + def __call__(self, equation, x): + # syft-restrict: hide-start return jnp.einsum(equation, x, self.w) + # syft-restrict: hide-end + class RMSNorm(nn.Module): def setup(self): + # syft-restrict: hide-start self.scale = _get(self, "scale") + # syft-restrict: hide-end + def __call__(self, x): + # syft-restrict: hide-start var = jnp.mean(jnp.square(x), axis=-1, keepdims=True) return x * jax.lax.rsqrt(var + 1e-6) * (1 + self.scale) + # syft-restrict: hide-end + class Attention(nn.Module): cfg: dict def setup(self): + # syft-restrict: hide-start self.q_einsum = Einsum() self.kv_einsum = Einsum() self._query_norm = RMSNorm() self._key_norm = RMSNorm() self.attn_vec_einsum = Einsum() + # syft-restrict: hide-end + def __call__(self, x, positions, mask, attn_type, cache=None): + # syft-restrict: hide-start q = self.q_einsum("bsd,ndh->bsnh", x) kv = self.kv_einsum("bsd,ckdh->cbskh", x) k, v = kv[0], kv[1] @@ -206,23 +234,32 @@ def __call__(self, x, positions, mask, attn_type, cache=None): out = jnp.einsum("bnst,btnh->bsnh", weights, v_exp) return self.attn_vec_einsum("bsnh,nhd->bsd", out), new_cache + # syft-restrict: hide-end + class FeedForward(nn.Module): def setup(self): + # syft-restrict: hide-start self.gating_einsum = Einsum() self.linear = Einsum() + # syft-restrict: hide-end + def __call__(self, x): + # syft-restrict: hide-start gate = self.gating_einsum("bsf,nhf->bsnh", x) h = jax.nn.gelu(gate[:, :, 0, :]) * gate[:, :, 1, :] return self.linear("bsh,hf->bsf", h) + # syft-restrict: hide-end + class Block(nn.Module): cfg: dict attn_type: str = "local" def setup(self): + # syft-restrict: hide-start self.pre_attention_norm = RMSNorm() self.attn = Attention(cfg=self.cfg) self.post_attention_norm = RMSNorm() @@ -230,7 +267,10 @@ def setup(self): self.mlp = FeedForward() self.post_ffw_norm = RMSNorm() + # syft-restrict: hide-end + def __call__(self, x, positions, mask, cache=None): + # syft-restrict: hide-start h = self.pre_attention_norm(x) h, new_cache = self.attn(h, positions, mask, self.attn_type, cache) h = self.post_attention_norm(h) @@ -240,22 +280,31 @@ def __call__(self, x, positions, mask, cache=None): h = self.post_ffw_norm(h) return x + h, new_cache + # syft-restrict: hide-end + class Embedder(nn.Module): cfg: dict def setup(self): + # syft-restrict: hide-start self.input_embedding = _get(self, "input_embedding") + # syft-restrict: hide-end + def __call__(self, token_ids): + # syft-restrict: hide-start table = self.input_embedding return table[token_ids] * jnp.sqrt(float(self.cfg["embed_dim"])), table + # syft-restrict: hide-end + class Transformer(nn.Module): cfg: dict def setup(self): + # syft-restrict: hide-start num_layers = self.cfg["num_layers"] attn_types = _attn_types(num_layers) self.embedder = Embedder(cfg=self.cfg) @@ -264,7 +313,10 @@ def setup(self): ] self.final_norm = RMSNorm() + # syft-restrict: hide-end + def __call__(self, tokens, cache=None): + # syft-restrict: hide-start sliding_window = self.cfg["sliding_window"] num_layers = self.cfg["num_layers"] attn_types = _attn_types(num_layers) @@ -291,6 +343,11 @@ def __call__(self, tokens, cache=None): logits = x @ jnp.transpose(embed_table) return logits, new_cache + # syft-restrict: hide-end + + +# syft-restrict: obfuscate-end + # ── Weight loading ───────────────────────────────────────────────────────── diff --git a/packages/syft-restrict/examples/gemma_inference_marked.certificate.json b/packages/syft-restrict/examples/gemma_inference_marked.certificate.json deleted file mode 100644 index 9940cea5ba0..00000000000 --- a/packages/syft-restrict/examples/gemma_inference_marked.certificate.json +++ /dev/null @@ -1,86 +0,0 @@ -{ - "source_sha256": "ddcf937765dcc67aef9631de54f5f1fce8d8d1fa77c956cd6392904e96eb5a8e", - "policy_id": "b4a557058704eed6", - "restrict_version": "0.1.0", - "private_ranges": [ - [25, 89], - [94, 94], - [106, 106], - [121, 121], - [137, 137], - [170, 171], - [175, 176], - [184, 185], - [189, 190], - [199, 202], - [210, 211], - [247, 248], - [253, 254], - [264, 268], - [277, 278], - [293, 296], - [300, 301], - [310, 313], - [323, 324], - [96, 97], - [108, 115], - [123, 131], - [139, 145], - [173, 173], - [178, 178], - [187, 187], - [192, 193], - [204, 208], - [213, 241], - [250, 251], - [256, 258], - [270, 275], - [280, 287], - [298, 298], - [303, 304], - [315, 321], - [326, 350] - ], - "obfuscate_ranges": [ - [25, 89], - [94, 94], - [106, 106], - [121, 121], - [137, 137], - [170, 171], - [175, 176], - [184, 185], - [189, 190], - [199, 202], - [210, 211], - [247, 248], - [253, 254], - [264, 268], - [277, 278], - [293, 296], - [300, 301], - [310, 313], - [323, 324] - ], - "hide_ranges": [ - [96, 97], - [108, 115], - [123, 131], - [139, 145], - [173, 173], - [178, 178], - [187, 187], - [192, 193], - [204, 208], - [213, 241], - [250, 251], - [256, 258], - [270, 275], - [280, 287], - [298, 298], - [303, 304], - [315, 321], - [326, 350] - ], - "n_calls_checked": 77 -} diff --git a/packages/syft-restrict/examples/gemma_inference_ranges.certificate.json b/packages/syft-restrict/examples/gemma_inference_ranges.certificate.json new file mode 100644 index 00000000000..d0360c8639a --- /dev/null +++ b/packages/syft-restrict/examples/gemma_inference_ranges.certificate.json @@ -0,0 +1,86 @@ +{ + "source_sha256": "3b462b5a6490b9dba8a93cf1a06da37601b002ab3a6f7497c01ee6698a671a18", + "policy_id": "d84d1e21530fa500", + "restrict_version": "0.1.0", + "private_ranges": [ + [24, 88], + [91, 91], + [99, 99], + [110, 110], + [122, 122], + [151, 152], + [155, 155], + [159, 160], + [163, 163], + [168, 171], + [178, 178], + [210, 211], + [215, 215], + [221, 225], + [233, 233], + [244, 247], + [250, 250], + [255, 258], + [267, 267], + [92, 93], + [100, 107], + [111, 119], + [123, 129], + [153, 153], + [156, 156], + [161, 161], + [164, 165], + [172, 176], + [179, 207], + [212, 213], + [216, 218], + [226, 231], + [234, 241], + [248, 248], + [251, 252], + [259, 265], + [268, 292] + ], + "obfuscate_ranges": [ + [24, 88], + [91, 91], + [99, 99], + [110, 110], + [122, 122], + [151, 152], + [155, 155], + [159, 160], + [163, 163], + [168, 171], + [178, 178], + [210, 211], + [215, 215], + [221, 225], + [233, 233], + [244, 247], + [250, 250], + [255, 258], + [267, 267] + ], + "hide_ranges": [ + [92, 93], + [100, 107], + [111, 119], + [123, 129], + [153, 153], + [156, 156], + [161, 161], + [164, 165], + [172, 176], + [179, 207], + [212, 213], + [216, 218], + [226, 231], + [234, 241], + [248, 248], + [251, 252], + [259, 265], + [268, 292] + ], + "n_calls_checked": 77 +} diff --git a/packages/syft-restrict/examples/gemma_inference_marked.obfuscated.py b/packages/syft-restrict/examples/gemma_inference_ranges.obfuscated.py similarity index 87% rename from packages/syft-restrict/examples/gemma_inference_marked.obfuscated.py rename to packages/syft-restrict/examples/gemma_inference_ranges.obfuscated.py index 7d85cf72479..ab009defb83 100644 --- a/packages/syft-restrict/examples/gemma_inference_marked.obfuscated.py +++ b/packages/syft-restrict/examples/gemma_inference_ranges.obfuscated.py @@ -21,7 +21,6 @@ # ── Model config ───────────────────────────────────────────────────────────── # Active model config. Comment/uncomment to switch sizes. -# syft-restrict: obfuscate-start ░v0 = dict( ░v1=■, ░v2=■, @@ -87,24 +86,17 @@ ░v11 = ■ ░v12 = ■ ░v13 = -■ # THIS COMMENT WAS OBFUSCATED -# syft-restrict: obfuscate-end -# syft-restrict: obfuscate-start def ░fn0(░v1): - # syft-restrict: hide-start ■■■■■■■■ # hidden/obfuscated lines can only execute restricted python, see restrict docs for more details ■■■■■■■■ - # syft-restrict: hide-end -# syft-restrict: obfuscate-end # ── Standalone helpers ──────────────────────────────────────────────────── -# syft-restrict: obfuscate-start def ░fn1(░v14, ░v15, ░v16): - # syft-restrict: hide-start ■■■■■■■■ # hidden/obfuscated lines can only execute restricted python, see restrict docs for more details ■■■■■■■■ ■■■■■■■■ @@ -113,13 +105,9 @@ def ░fn1(░v14, ░v15, ░v16): ■■■■■■■■ ■■■■■■■■ ■■■■■■■■ - # syft-restrict: hide-end -# syft-restrict: obfuscate-end -# syft-restrict: obfuscate-start def ░fn2(░v17, ░v7): - # syft-restrict: hide-start ■■■■■■■■ # hidden/obfuscated lines can only execute restricted python, see restrict docs for more details ■■■■■■■■ ■■■■■■■■ @@ -129,13 +117,9 @@ def ░fn2(░v17, ░v7): ■■■■■■■■ ■■■■■■■■ ■■■■■■■■ - # syft-restrict: hide-end -# syft-restrict: obfuscate-end -# syft-restrict: obfuscate-start def ░fn3(░v18, ░v7): - # syft-restrict: hide-start ■■■■■■■■ # hidden/obfuscated lines can only execute restricted python, see restrict docs for more details ■■■■■■■■ ■■■■■■■■ @@ -143,8 +127,6 @@ def ░fn3(░v18, ░v7): ■■■■■■■■ ■■■■■■■■ ■■■■■■■■ - # syft-restrict: hide-end -# syft-restrict: obfuscate-end # ── Flax modules ─────────────────────────────────────────────────────────── @@ -166,50 +148,34 @@ def append_to(lst, item): return lst -# syft-restrict: obfuscate-start class ░Cls0(nn.Module): def setup(self): - # syft-restrict: hide-start ■■■■■■■■ # hidden/obfuscated lines can only execute restricted python, see restrict docs for more details - # syft-restrict: hide-end def __call__(self, ░v19, ░v14): - # syft-restrict: hide-start ■■■■■■■■ # hidden/obfuscated lines can only execute restricted python, see restrict docs for more details - # syft-restrict: hide-end -# syft-restrict: obfuscate-end -# syft-restrict: obfuscate-start class ░Cls1(nn.Module): def setup(self): - # syft-restrict: hide-start ■■■■■■■■ # hidden/obfuscated lines can only execute restricted python, see restrict docs for more details - # syft-restrict: hide-end def __call__(self, ░v14): - # syft-restrict: hide-start ■■■■■■■■ # hidden/obfuscated lines can only execute restricted python, see restrict docs for more details ■■■■■■■■ - # syft-restrict: hide-end -# syft-restrict: obfuscate-end -# syft-restrict: obfuscate-start class ░Cls2(nn.Module): ░v20: dict def setup(self): - # syft-restrict: hide-start ■■■■■■■■ # hidden/obfuscated lines can only execute restricted python, see restrict docs for more details ■■■■■■■■ ■■■■■■■■ ■■■■■■■■ ■■■■■■■■ - # syft-restrict: hide-end def __call__(self, ░v14, ░v15, ░v21, ░v22, ░v23=None): - # syft-restrict: hide-start ■■■■■■■■ # hidden/obfuscated lines can only execute restricted python, see restrict docs for more details ■■■■■■■■ ■■■■■■■■ @@ -239,44 +205,32 @@ def __call__(self, ░v14, ░v15, ░v21, ░v22, ░v23=None): ■■■■■■■■ ■■■■■■■■ - # syft-restrict: hide-end -# syft-restrict: obfuscate-end -# syft-restrict: obfuscate-start class ░Cls3(nn.Module): def setup(self): - # syft-restrict: hide-start ■■■■■■■■ # hidden/obfuscated lines can only execute restricted python, see restrict docs for more details ■■■■■■■■ - # syft-restrict: hide-end def __call__(self, ░v14): - # syft-restrict: hide-start ■■■■■■■■ # hidden/obfuscated lines can only execute restricted python, see restrict docs for more details ■■■■■■■■ ■■■■■■■■ - # syft-restrict: hide-end -# syft-restrict: obfuscate-end -# syft-restrict: obfuscate-start class ░Cls4(nn.Module): ░v20: dict ░v22: str = "■" def setup(self): - # syft-restrict: hide-start ■■■■■■■■ # hidden/obfuscated lines can only execute restricted python, see restrict docs for more details ■■■■■■■■ ■■■■■■■■ ■■■■■■■■ ■■■■■■■■ ■■■■■■■■ - # syft-restrict: hide-end def __call__(self, ░v14, ░v15, ░v21, ░v23=None): - # syft-restrict: hide-start ■■■■■■■■ # hidden/obfuscated lines can only execute restricted python, see restrict docs for more details ■■■■■■■■ ■■■■■■■■ @@ -285,33 +239,23 @@ def __call__(self, ░v14, ░v15, ░v21, ░v23=None): ■■■■■■■■ ■■■■■■■■ ■■■■■■■■ - # syft-restrict: hide-end -# syft-restrict: obfuscate-end -# syft-restrict: obfuscate-start class ░Cls5(nn.Module): ░v20: dict def setup(self): - # syft-restrict: hide-start ■■■■■■■■ # hidden/obfuscated lines can only execute restricted python, see restrict docs for more details - # syft-restrict: hide-end def __call__(self, ░v24): - # syft-restrict: hide-start ■■■■■■■■ # hidden/obfuscated lines can only execute restricted python, see restrict docs for more details ■■■■■■■■ - # syft-restrict: hide-end -# syft-restrict: obfuscate-end -# syft-restrict: obfuscate-start class ░Cls6(nn.Module): ░v20: dict def setup(self): - # syft-restrict: hide-start ■■■■■■■■ # hidden/obfuscated lines can only execute restricted python, see restrict docs for more details ■■■■■■■■ ■■■■■■■■ @@ -319,10 +263,8 @@ def setup(self): ■■■■■■■■ ■■■■■■■■ ■■■■■■■■ - # syft-restrict: hide-end def __call__(self, ░v25, ░v23=None): - # syft-restrict: hide-start ■■■■■■■■ # hidden/obfuscated lines can only execute restricted python, see restrict docs for more details ■■■■■■■■ ■■■■■■■■ @@ -348,8 +290,6 @@ def __call__(self, ░v25, ░v23=None): ■■■■■■■■ ■■■■■■■■ ■■■■■■■■ - # syft-restrict: hide-end -# syft-restrict: obfuscate-end # ── Weight loading ───────────────────────────────────────────────────────── diff --git a/packages/syft-restrict/examples/gemma_inference_marked.py b/packages/syft-restrict/examples/gemma_inference_ranges.py similarity index 86% rename from packages/syft-restrict/examples/gemma_inference_marked.py rename to packages/syft-restrict/examples/gemma_inference_ranges.py index 088cd3b235c..db075c9bc63 100644 --- a/packages/syft-restrict/examples/gemma_inference_marked.py +++ b/packages/syft-restrict/examples/gemma_inference_ranges.py @@ -21,7 +21,6 @@ # ── Model config ───────────────────────────────────────────────────────────── # Active model config. Comment/uncomment to switch sizes. -# syft-restrict: obfuscate-start CONFIG = dict( num_layers=18, embed_dim=640, @@ -87,26 +86,17 @@ LOCAL_ROPE_BASE = 10_000 GLOBAL_ROPE_BASE = 1_000_000 K_MASK = -2.3819763e38 # Google's masking constant (≈ float32 -inf) -# syft-restrict: obfuscate-end -# syft-restrict: obfuscate-start def _attn_types(num_layers): - # syft-restrict: hide-start pattern = ("local",) * 5 + ("global",) return (pattern * ((num_layers + 5) // 6))[:num_layers] - # syft-restrict: hide-end - - -# syft-restrict: obfuscate-end # ── Standalone helpers ──────────────────────────────────────────────────── -# syft-restrict: obfuscate-start def apply_rope(x, positions, base_freq): - # syft-restrict: hide-start """Rotary position embeddings (split-half rotation).""" half = shape_of(x)[-1] // 2 freq_exp = (2.0 / shape_of(x)[-1]) * jnp.arange(half, dtype=jnp.float32) @@ -115,15 +105,9 @@ def apply_rope(x, positions, base_freq): sin, cos = jnp.sin(angles), jnp.cos(angles) x1, x2 = x[..., :half], x[..., half:] return jnp.concatenate([x1 * cos - x2 * sin, x2 * cos + x1 * sin], axis=-1) - # syft-restrict: hide-end - -# syft-restrict: obfuscate-end - -# syft-restrict: obfuscate-start def make_masks(seq_len, sliding_window): - # syft-restrict: hide-start """Causal masks — local layers also clip to a sliding window.""" causal = jnp.tril(jnp.ones((seq_len, seq_len), dtype=jnp.bool_)) window = jnp.triu( @@ -133,15 +117,9 @@ def make_masks(seq_len, sliding_window): "local": (causal & window)[None, None], "global": causal[None, None], } - # syft-restrict: hide-end - -# syft-restrict: obfuscate-end - -# syft-restrict: obfuscate-start def make_decode_masks(pos, sliding_window): - # syft-restrict: hide-start """Masks for single-token decode.""" total_len = pos + 1 positions = jnp.arange(total_len) @@ -149,10 +127,6 @@ def make_decode_masks(pos, sliding_window): "local": (positions >= pos - sliding_window + 1)[None, None, None, :], "global": jnp.ones((1, 1, 1, total_len), dtype=jnp.bool_), } - # syft-restrict: hide-end - - -# syft-restrict: obfuscate-end # ── Flax modules ─────────────────────────────────────────────────────────── @@ -174,59 +148,34 @@ def append_to(lst, item): return lst -# syft-restrict: obfuscate-start class Einsum(nn.Module): def setup(self): - # syft-restrict: hide-start self.w = _get(self, "w") - # syft-restrict: hide-end - def __call__(self, equation, x): - # syft-restrict: hide-start return jnp.einsum(equation, x, self.w) - # syft-restrict: hide-end - -# syft-restrict: obfuscate-end - - -# syft-restrict: obfuscate-start class RMSNorm(nn.Module): def setup(self): - # syft-restrict: hide-start self.scale = _get(self, "scale") - # syft-restrict: hide-end - def __call__(self, x): - # syft-restrict: hide-start var = jnp.mean(jnp.square(x), axis=-1, keepdims=True) return x * jax.lax.rsqrt(var + 1e-6) * (1 + self.scale) - # syft-restrict: hide-end - -# syft-restrict: obfuscate-end - - -# syft-restrict: obfuscate-start class Attention(nn.Module): cfg: dict def setup(self): - # syft-restrict: hide-start self.q_einsum = Einsum() self.kv_einsum = Einsum() self._query_norm = RMSNorm() self._key_norm = RMSNorm() self.attn_vec_einsum = Einsum() - # syft-restrict: hide-end - def __call__(self, x, positions, mask, attn_type, cache=None): - # syft-restrict: hide-start q = self.q_einsum("bsd,ndh->bsnh", x) kv = self.kv_einsum("bsd,ckdh->cbskh", x) k, v = kv[0], kv[1] @@ -257,40 +206,23 @@ def __call__(self, x, positions, mask, attn_type, cache=None): out = jnp.einsum("bnst,btnh->bsnh", weights, v_exp) return self.attn_vec_einsum("bsnh,nhd->bsd", out), new_cache - # syft-restrict: hide-end - -# syft-restrict: obfuscate-end - - -# syft-restrict: obfuscate-start class FeedForward(nn.Module): def setup(self): - # syft-restrict: hide-start self.gating_einsum = Einsum() self.linear = Einsum() - # syft-restrict: hide-end - def __call__(self, x): - # syft-restrict: hide-start gate = self.gating_einsum("bsf,nhf->bsnh", x) h = jax.nn.gelu(gate[:, :, 0, :]) * gate[:, :, 1, :] return self.linear("bsh,hf->bsf", h) - # syft-restrict: hide-end - - -# syft-restrict: obfuscate-end - -# syft-restrict: obfuscate-start class Block(nn.Module): cfg: dict attn_type: str = "local" def setup(self): - # syft-restrict: hide-start self.pre_attention_norm = RMSNorm() self.attn = Attention(cfg=self.cfg) self.post_attention_norm = RMSNorm() @@ -298,10 +230,7 @@ def setup(self): self.mlp = FeedForward() self.post_ffw_norm = RMSNorm() - # syft-restrict: hide-end - def __call__(self, x, positions, mask, cache=None): - # syft-restrict: hide-start h = self.pre_attention_norm(x) h, new_cache = self.attn(h, positions, mask, self.attn_type, cache) h = self.post_attention_norm(h) @@ -311,39 +240,22 @@ def __call__(self, x, positions, mask, cache=None): h = self.post_ffw_norm(h) return x + h, new_cache - # syft-restrict: hide-end - - -# syft-restrict: obfuscate-end - -# syft-restrict: obfuscate-start class Embedder(nn.Module): cfg: dict def setup(self): - # syft-restrict: hide-start self.input_embedding = _get(self, "input_embedding") - # syft-restrict: hide-end - def __call__(self, token_ids): - # syft-restrict: hide-start table = self.input_embedding return table[token_ids] * jnp.sqrt(float(self.cfg["embed_dim"])), table - # syft-restrict: hide-end - - -# syft-restrict: obfuscate-end - -# syft-restrict: obfuscate-start class Transformer(nn.Module): cfg: dict def setup(self): - # syft-restrict: hide-start num_layers = self.cfg["num_layers"] attn_types = _attn_types(num_layers) self.embedder = Embedder(cfg=self.cfg) @@ -352,10 +264,7 @@ def setup(self): ] self.final_norm = RMSNorm() - # syft-restrict: hide-end - def __call__(self, tokens, cache=None): - # syft-restrict: hide-start sliding_window = self.cfg["sliding_window"] num_layers = self.cfg["num_layers"] attn_types = _attn_types(num_layers) @@ -382,11 +291,6 @@ def __call__(self, tokens, cache=None): logits = x @ jnp.transpose(embed_table) return logits, new_cache - # syft-restrict: hide-end - - -# syft-restrict: obfuscate-end - # ── Weight loading ───────────────────────────────────────────────────────── diff --git a/packages/syft-restrict/examples/generate.py b/packages/syft-restrict/examples/generate.py index 0f2d65d0c82..adc1c0ce734 100644 --- a/packages/syft-restrict/examples/generate.py +++ b/packages/syft-restrict/examples/generate.py @@ -2,14 +2,10 @@ Run from the package root: uv run python examples/generate.py -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 -references. `jax.numpy` needs no such entry — it's aliased as `jnp` (a bare name). - -OBFUSCATE keeps structure legible (identifiers renamed, constants blanked); HIDE replaces -whole lines with a ■■■■■■■■ marker. The verified region is their union. +The model carves out its own private region with `# syft-restrict: ...` comment markers, so +run() resolves the region straight from the source — this is the supported UX. (The +generate_ranges.py / gemma_inference_ranges.py variant drives the same model through explicit +hand-counted line ranges instead.) """ import json @@ -21,47 +17,6 @@ result = run( path=EX / "gemma_inference.py", - obfuscate=[ - [24, 88], # CONFIG dict (+ commented size variants) + shared constants - [91, 91], - [99, 99], - [110, 110], - [122, 122], # the 4 standalone helper def lines - [151, 152], - [155, 155], # Einsum: class + setup def | __call__ def - [159, 160], - [163, 163], # RMSNorm - [168, 171], - [178, 178], # Attention (incl. `cfg: dict`) - [210, 211], - [215, 215], # FeedForward - [221, 225], - [233, 233], # Block (incl. `cfg`, `attn_type`) - [244, 247], - [250, 250], # Embedder - [255, 258], - [267, 267], # Transformer - ], - hide=[ - [92, 93], - [100, 107], - [111, 119], - [123, 129], # the 4 standalone helper bodies - [153, 153], - [156, 156], # Einsum setup / __call__ bodies - [161, 161], - [164, 165], # RMSNorm - [172, 176], - [179, 207], # Attention - [212, 213], - [216, 218], # FeedForward - [226, 231], - [234, 241], # Block - [248, 248], - [251, 252], # Embedder - [259, 265], - [268, 292], # Transformer - ], allow_functions=[ "jax.numpy.einsum", "jax.numpy.mean", diff --git a/packages/syft-restrict/examples/generate_marked.py b/packages/syft-restrict/examples/generate_marked.py deleted file mode 100644 index 7e717b056f4..00000000000 --- a/packages/syft-restrict/examples/generate_marked.py +++ /dev/null @@ -1,58 +0,0 @@ -"""Regenerate gemma_inference_marked.obfuscated.py + its certificate. - -Run from the package root: uv run python examples/generate_marked.py - -Same model, same policy, and the same verified region as generate.py / gemma_inference.py -- -the only difference is that gemma_inference_marked.py carves out its own private region with -`# syft-restrict: ...` comments instead of run() taking hand-counted line ranges. Because both -obfuscate and hide are omitted below, run() scans the source for those markers automatically. -""" - -import json -from pathlib import Path - -from syft_restrict import run - -EX = Path(__file__).parent - -result = run( - path=EX / "gemma_inference_marked.py", - allow_functions=[ - "jax.numpy.einsum", - "jax.numpy.mean", - "jax.numpy.square", - "jax.numpy.arange", - "jax.numpy.sin", - "jax.numpy.cos", - "jax.numpy.concatenate", - "jax.numpy.tril", - "jax.numpy.triu", - "jax.numpy.ones", - "jax.numpy.where", - "jax.numpy.repeat", - "jax.numpy.sqrt", - "jax.numpy.transpose", - "jax.numpy.array", - "jax.numpy.float32", - "jax.numpy.bool_", - "jax.lax.rsqrt", - "jax.nn.softmax", - "jax.nn.gelu", - "flax.linen.Module", - # module references required by the deep-path call style (jax.lax.rsqrt, jax.nn.softmax): - "jax.lax", - "jax.nn", - ], - allow_operators=["arithmetic", "indexing", "comparison"], -) - -cert_path = EX / "gemma_inference_marked.certificate.json" -cert_path.write_text(json.dumps(result.certificate, indent=2) + "\n") - -policy_id = result.certificate["policy_id"] -n_calls_checked = result.certificate["n_calls_checked"] -print( - f"""wrote {result.obfuscated_path} -wrote {cert_path} -policy_id={policy_id} n_calls_checked={n_calls_checked}""" -) diff --git a/packages/syft-restrict/examples/generate_ranges.py b/packages/syft-restrict/examples/generate_ranges.py new file mode 100644 index 00000000000..99324f34e90 --- /dev/null +++ b/packages/syft-restrict/examples/generate_ranges.py @@ -0,0 +1,107 @@ +"""Regenerate gemma_inference_ranges.obfuscated.py + its certificate. + +Run from the package root: uv run python examples/generate_ranges.py + +This is the explicit-line-ranges variant of the example. Markers are the supported UX (see +gemma_inference.py / generate.py); this file drives the same model through the internal ``_run`` +entry point with hand-counted ranges instead, exercising that path. + +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 +references. `jax.numpy` needs no such entry — it's aliased as `jnp` (a bare name). + +OBFUSCATE keeps structure legible (identifiers renamed, constants blanked); HIDE replaces +whole lines with a ■■■■■■■■ marker. The verified region is their union. +""" + +import json +from pathlib import Path + +from syft_restrict.runner import _run # explicit numeric ranges -> internal entry point + +EX = Path(__file__).parent + +result = _run( + path=EX / "gemma_inference_ranges.py", + obfuscate=[ + [24, 88], # CONFIG dict (+ commented size variants) + shared constants + [91, 91], + [99, 99], + [110, 110], + [122, 122], # the 4 standalone helper def lines + [151, 152], + [155, 155], # Einsum: class + setup def | __call__ def + [159, 160], + [163, 163], # RMSNorm + [168, 171], + [178, 178], # Attention (incl. `cfg: dict`) + [210, 211], + [215, 215], # FeedForward + [221, 225], + [233, 233], # Block (incl. `cfg`, `attn_type`) + [244, 247], + [250, 250], # Embedder + [255, 258], + [267, 267], # Transformer + ], + hide=[ + [92, 93], + [100, 107], + [111, 119], + [123, 129], # the 4 standalone helper bodies + [153, 153], + [156, 156], # Einsum setup / __call__ bodies + [161, 161], + [164, 165], # RMSNorm + [172, 176], + [179, 207], # Attention + [212, 213], + [216, 218], # FeedForward + [226, 231], + [234, 241], # Block + [248, 248], + [251, 252], # Embedder + [259, 265], + [268, 292], # Transformer + ], + allow_functions=[ + "jax.numpy.einsum", + "jax.numpy.mean", + "jax.numpy.square", + "jax.numpy.arange", + "jax.numpy.sin", + "jax.numpy.cos", + "jax.numpy.concatenate", + "jax.numpy.tril", + "jax.numpy.triu", + "jax.numpy.ones", + "jax.numpy.where", + "jax.numpy.repeat", + "jax.numpy.sqrt", + "jax.numpy.transpose", + "jax.numpy.array", + "jax.numpy.float32", + "jax.numpy.bool_", + "jax.lax.rsqrt", + "jax.nn.softmax", + "jax.nn.gelu", + "flax.linen.Module", + # module references required by the deep-path call style (jax.lax.rsqrt, jax.nn.softmax): + "jax.lax", + "jax.nn", + ], + allow_operators=["arithmetic", "indexing", "comparison"], +) + +cert_path = EX / "gemma_inference_ranges.certificate.json" +cert_path.write_text(json.dumps(result.certificate, indent=2) + "\n") + +policy_id = result.certificate["policy_id"] +n_calls_checked = result.certificate["n_calls_checked"] +print( + f"""wrote {result.obfuscated_path} +wrote {cert_path} +policy_id={policy_id} n_calls_checked={n_calls_checked}""" +) diff --git a/packages/syft-restrict/src/syft_restrict/astutil.py b/packages/syft-restrict/src/syft_restrict/astutil.py index c8e65d0a79c..7923db9f02b 100644 --- a/packages/syft-restrict/src/syft_restrict/astutil.py +++ b/packages/syft-restrict/src/syft_restrict/astutil.py @@ -36,6 +36,20 @@ def node_in_ranges(node: ast.AST, ranges) -> bool: return line is not None and row_in_ranges(line, ranges) +def node_overlaps_ranges(node: ast.AST, ranges) -> bool: + """True if a node's line SPAN intersects any range. + + Unlike ``node_in_ranges`` (which tests only the start line), this also catches a multi-line + node that begins outside the ranges but extends into one -- a statement straddling the + public/private boundary. Nodes with no position never overlap. + """ + start = getattr(node, "lineno", None) + if start is None: + return False + end = getattr(node, "end_lineno", None) or start + return any(start <= hi and lo <= end for lo, hi in 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. @@ -124,8 +138,17 @@ def scan_file(tree: ast.Module, private_ranges) -> FileScan: 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 + # Bind the name Python actually binds at runtime: + # `import jax.numpy as jnp` binds `jnp` -> the jax.numpy module + # `import jax.numpy` binds `jax` -> the jax PACKAGE (not jax.numpy!) -- + # you reach save through `jax.numpy.save`, so the root resolves to itself. + # Binding the root to the full dotted path would mis-resolve `jax.numpy.save` to + # `jax.numpy.numpy.save` and let it slip a disallow floor (see tests/verify). + if alias.asname: + import_bindings[alias.asname] = alias.name + else: + root = alias.name.split(".")[0] + import_bindings[root] = root elif isinstance(node, ast.ImportFrom) and node.module: for alias in node.names: import_bindings[alias.asname or alias.name] = ( diff --git a/packages/syft-restrict/src/syft_restrict/markers.py b/packages/syft-restrict/src/syft_restrict/markers.py index e4647b5d559..081fbe392ef 100644 --- a/packages/syft-restrict/src/syft_restrict/markers.py +++ b/packages/syft-restrict/src/syft_restrict/markers.py @@ -1,7 +1,6 @@ -"""Comment-based markup, an alternative to hand-counted line ranges. +"""Comment-based markup — the way a source file designates its private/hidden regions. -Instead of passing ``obfuscate``/``hide`` as ``[start, end]`` line numbers to ``run()``, a source -file may mark its own private/hidden regions with comments:: +A source file marks its own private/hidden regions with comments:: # syft-restrict: obfuscate-start def attention(x): @@ -15,9 +14,8 @@ def attention(x): renamed), so carving out a stricter sub-region is safe. The reverse is not: obfuscate cannot nest inside hide, and neither kind may nest inside itself. -``parse_markers(source)`` resolves these into the same ``(obfuscate_ranges, hide_ranges)`` shape -``run()`` already accepts. ``run()`` uses this automatically when both ``obfuscate`` and ``hide`` -are omitted. +``parse_markers(source)`` resolves these into ``(obfuscate_ranges, hide_ranges)``. ``run()`` calls +it to locate the private region; a file with no markers raises ``MarkerError``. """ from __future__ import annotations @@ -35,20 +33,54 @@ def attention(x): _MARKER_RE = re.compile(r"^#\s*syft-restrict:\s*(obfuscate|hide)(?:-(start|end))?\s*$") +# Token types that are not code: comments and layout. A line carrying only these (plus a marker +# comment) is a bare marker line; anything else means code shares the line. +_TRIVIA_TOKENS = frozenset( + { + tokenize.COMMENT, + tokenize.NL, + tokenize.NEWLINE, + tokenize.INDENT, + tokenize.DEDENT, + tokenize.ENCODING, + tokenize.ENDMARKER, + } +) + + def _scan_markers(source: str) -> dict[int, tuple[str, str | None]]: """Map each line number carrying a ``# syft-restrict: ...`` comment to its ``(kind, boundary)``. ``boundary`` is ``"start"``, ``"end"``, or ``None`` for a single-line marker. Only real comment tokens count, so a marker-shaped string inside a string literal is never mistaken for a directive. + + A block ``-start``/``-end`` marker must be on a line by itself: code before it would sit on a + boundary line (excluded from every range) and silently escape verification. Such a marker + raises ``MarkerError`` -- use the single-line ``# syft-restrict: obfuscate``/``hide`` form to + mark one line of code. (Code tokens precede the trailing comment in token order, so by the time + a marker comment is seen its line is already known to carry code.) """ markers: dict[int, tuple[str, str | None]] = {} + code_lines: set[int] = set() for tok in tokenize.generate_tokens(io.StringIO(source).readline): + if tok.type not in _TRIVIA_TOKENS: + code_lines.add(tok.start[0]) + continue if tok.type != tokenize.COMMENT: continue match = _MARKER_RE.match(tok.string.strip()) - if match: - markers[tok.start[0]] = (match.group(1), match.group(2)) + if not match: + continue + kind, boundary = match.group(1), match.group(2) + line = tok.start[0] + if boundary is not None and line in code_lines: + raise MarkerError( + f"line {line}: '{kind}-{boundary}' block marker must be on a line by itself; " + f"move the code to its own line, or use a single-line '# syft-restrict: {kind}' " + "marker to mark just this line" + ) + markers[line] = (kind, boundary) return markers @@ -110,8 +142,7 @@ def resolve(self) -> tuple[list[tuple[int, int]], list[tuple[int, int]]]: if not self._lines["obfuscate"] and not self._lines["hide"]: raise MarkerError( "no syft-restrict markers found in source; add `# syft-restrict: obfuscate-start` " - "/ `# syft-restrict: obfuscate-end` (or `hide`) around the private code, or pass " - "obfuscate=/hide= explicitly to run()" + "/ `# syft-restrict: obfuscate-end` (or `hide`) around the private code" ) return _compact(self._lines["obfuscate"]), _compact(self._lines["hide"]) diff --git a/packages/syft-restrict/src/syft_restrict/policy.py b/packages/syft-restrict/src/syft_restrict/policy.py index 10e8cb32583..04395220b4c 100644 --- a/packages/syft-restrict/src/syft_restrict/policy.py +++ b/packages/syft-restrict/src/syft_restrict/policy.py @@ -56,6 +56,13 @@ "ascii", "format", "bytes", + # site-injected builtins: stdout channels, interpreter shutdown, interactive help + "copyright", + "credits", + "license", + "exit", + "quit", + "help", } ) @@ -92,11 +99,11 @@ # 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"}) +ALLOWED_DUNDER_DEFS: frozenset[str] = frozenset({"__call__", "setup", "__post_init__"}) # Names always preserved verbatim by the obfuscator and never treated as opaque values. DEFAULT_KEEP: frozenset[str] = frozenset( - {"self", "cls", "nn", "Module", "setup", "__call__"} + {"self", "cls", "nn", "Module", "setup", "__call__", "__post_init__"} ) @@ -116,6 +123,14 @@ class Policy(BaseModel): # optional disallow globs supplied by the user; these beat the allow, example: ["jax.numpy.save"] disallowed_functions: list[str] = Field(default_factory=list) + # when False, a local assigned a safe callable is no longer trusted as a bare-name call target + # (disables _track_safe_local); the callee must instead be called directly (docs/verify.md). + allow_local_assignments: bool = True + + # when False, a `self.` never assigned in the class body is NOT presumed inherited-safe; + # only attributes the class assigns a vetted source may be called (docs/verify.md). + allow_base_class_attributes: bool = True + # import aliases reserved against rebinding (set per-file by verify), e.g. {"jnp", "nn"} reserved_names: set[str] = Field(default_factory=set) @@ -125,6 +140,8 @@ def parse( allow_functions: list[str] | None = None, allow_operators: list[str] | None = None, disallow_functions: list[str] | None = None, + allow_local_assignments: bool = True, + allow_base_class_attributes: bool = True, ) -> "Policy": allowed_functions = _clean(allow_functions) allowed_operators = set(_clean(allow_operators)) @@ -138,6 +155,8 @@ def parse( allowed_functions=allowed_functions, allowed_operators=allowed_operators, disallowed_functions=disallowed_functions, + allow_local_assignments=allow_local_assignments, + allow_base_class_attributes=allow_base_class_attributes, ) # ── path matching ────────────────────────────────────────────────────────────────── @@ -162,6 +181,8 @@ def policy_id(self) -> str: + "|".join(sorted(self.allowed_operators)) + "##" + "|".join(sorted(self.disallowed_functions)) + + "##" + + f"local={int(self.allow_local_assignments)}|baseattr={int(self.allow_base_class_attributes)}" ) 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 5330916a97c..89b3ea9ea60 100644 --- a/packages/syft-restrict/src/syft_restrict/runner.py +++ b/packages/syft-restrict/src/syft_restrict/runner.py @@ -11,9 +11,7 @@ from .astutil import normalize_ranges, scan_file from .errors import PolicyViolation from .markers import parse_markers -from .obfuscator import ( - obfuscate as _obfuscate, -) # aliased: `obfuscate` is also a run() kwarg +from .obfuscator import obfuscate as _obfuscate from .policy import Policy from .verifier import Violation, verify @@ -29,46 +27,97 @@ class RunResult(BaseModel): def run( path: str | Path, - obfuscate=None, - hide=None, allow_functions: list[str] | None = None, allow_operators: list[str] | None = None, disallow_functions: list[str] | None = None, + allow_local_assignments: bool = True, + allow_base_class_attributes: bool = True, out: str | Path | None = None, strict: bool = True, ) -> RunResult: """Verify the private region, then (on success) write a display copy. - The private region is the *union* of ``obfuscate`` and ``hide`` — both are secret code that runs - in the enclave, so both are verified. They differ only in how the display copy renders them: + The private region is marked in the source with ``# syft-restrict: ...`` comments (see + ``markers.parse_markers`` and docs/verify.md); this is the only supported way to designate it. + Its ``obfuscate`` and ``hide`` sub-regions are resolved from those markers and both verified — + they differ only in how the display copy renders them (identifiers renamed vs. whole lines + blanked). A file with no markers raises ``MarkerError``. Args: - path: the inference source file. - obfuscate: ``[start, end]`` 1-based inclusive line ranges to *obfuscate* (identifiers renamed, - constants blanked, structure preserved). When both ``obfuscate`` and ``hide`` are omitted, - ranges are instead resolved from ``# syft-restrict: ...`` comment markers in the source - (see ``markers.parse_markers``). - hide: ``[start, end]`` 1-based inclusive line ranges to *hide* (whole line replaced with a - ``■■■■■■■■`` marker, indentation kept). + path: the inference source file (must carry ``# syft-restrict: ...`` markers). allow_functions: list of dotted-path globs callable by name (e.g. ``["jax.*", "flax.linen.*"]``). 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. + allow_local_assignments: if True (default), a local aliased to a safe callable may itself be + called by name; if False, callables must be called directly. + allow_base_class_attributes: if True (default), a ``self.`` never assigned in the class + is presumed inherited from the (vetted) base and callable; if False, only assigned attrs are. 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) + # Read the source exactly once and hand it to _run, so marker resolution, verification, and + # obfuscation all operate on identical bytes (no second read that could race / TOCTOU). source = path.read_text() - policy = Policy.parse(allow_functions, allow_operators, disallow_functions) + obfuscate_ranges, hide_ranges = parse_markers(source) + return _run( + path, + obfuscate=obfuscate_ranges, + hide=hide_ranges, + allow_functions=allow_functions, + allow_operators=allow_operators, + disallow_functions=disallow_functions, + allow_local_assignments=allow_local_assignments, + allow_base_class_attributes=allow_base_class_attributes, + out=out, + strict=strict, + source=source, + ) + + +def _run( + path: str | Path, + obfuscate=None, + hide=None, + allow_functions: list[str] | None = None, + allow_operators: list[str] | None = None, + disallow_functions: list[str] | None = None, + allow_local_assignments: bool = True, + allow_base_class_attributes: bool = True, + out: str | Path | None = None, + strict: bool = True, + source: str | None = None, +) -> RunResult: + """Verify and obfuscate using explicit 1-based line ranges (no marker scanning). + + Internal entry point behind ``run()``. Callers that want to bypass the comment-marker UX and + supply their own ranges may use it directly, accepting responsibility for those ranges. + + Args mirror ``run()`` except the private region is given as explicit ranges: + obfuscate: ``[start, end]`` 1-based inclusive line ranges to *obfuscate* (identifiers renamed, + constants blanked, structure preserved). + hide: ``[start, end]`` 1-based inclusive line ranges to *hide* (whole line replaced with a + ``■■■■■■■■`` marker, indentation kept). The verified region is the union of the two. + source: the already-read file contents. When ``run()`` calls this, it passes the exact bytes + it resolved markers against so nothing is read twice; otherwise the file is read here. + """ + path = Path(path) + if source is None: + source = path.read_text() + policy = Policy.parse( + allow_functions, + allow_operators, + disallow_functions, + allow_local_assignments, + allow_base_class_attributes, + ) - if obfuscate is None and hide is None: - obfuscate_ranges, hide_ranges = parse_markers(source) - else: - obfuscate_ranges = obfuscate or [] - hide_ranges = hide or [] + obfuscate_ranges = obfuscate or [] + hide_ranges = hide or [] private = [*obfuscate_ranges, *hide_ranges] # union = the verified region result = verify(source, private, policy) diff --git a/packages/syft-restrict/src/syft_restrict/verifier.py b/packages/syft-restrict/src/syft_restrict/verifier.py index 2cae0933ce7..4ed5a07b560 100644 --- a/packages/syft-restrict/src/syft_restrict/verifier.py +++ b/packages/syft-restrict/src/syft_restrict/verifier.py @@ -22,6 +22,7 @@ dotted_name, is_dunder, node_in_ranges, + node_overlaps_ranges, normalize_ranges, rooted_in_self, scan_file, @@ -113,6 +114,32 @@ ast.FormattedValue, ) +# Compound statements own an indented suite, so a public header + private body is normal marker +# usage (mark the body private, leave the `def`/`for`/`if` line public) -- NOT a boundary straddle. +# Their bodies are separate statements the walk enforces on their own; only *simple* statements and +# expressions that themselves span the boundary count as straddling. Excluded from straddle-enforce. +# `ast.TryStar` (3.11+) and `ast.Match` (3.10+) are resolved via getattr and filtered out when +# absent, so the module still imports on the >=3.10 floor. +_COMPOUND_NODES: tuple[type[ast.AST], ...] = tuple( + t + for t in ( + ast.Module, + ast.FunctionDef, + ast.AsyncFunctionDef, + ast.ClassDef, + ast.For, + ast.AsyncFor, + ast.While, + ast.If, + ast.With, + ast.AsyncWith, + ast.Try, + getattr(ast, "TryStar", None), + getattr(ast, "Match", None), + ) + if t is not None +) + # violation-code registry: every code a check can raise, one line each (docs/blacklist.md) ViolationCode = Literal[ "banned-construct", # _enforce (node type on the permanent deny-list) @@ -131,6 +158,7 @@ "dunder-name", # _check_name "operator-disabled", # _require_bundle "duplicate-method", # _forbid_duplicate_methods + "star-import", # visit (`from ... import *` anywhere in the file) ] @@ -191,8 +219,13 @@ def __init__(self, policy: Policy, scan: FileScan, ranges): # 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) + # answers "is self. safe to call?" -- see _SelfAttrTrust below. allow_base_class_attributes + # controls whether a never-assigned attr is presumed inherited-safe or rejected. + self._self_attr = _SelfAttrTrust( + scan, + self._resolved_allowed, + allow_base_class=policy.allow_base_class_attributes, + ) # nodes inside a type annotation (never invoked), exempt from name/container checks self._annotation_nodes: set[int] = set() @@ -202,6 +235,10 @@ def __init__(self, policy: Policy, scan: FileScan, ranges): # pushes a scope in visit() (only ClassDef/FunctionDef/Lambda do) self._safe_locals_stack: list[dict[str, bool]] = [{}] + # when False, _track_safe_local records nothing, so a local aliased to a safe callable is + # never itself trusted as a bare-name call target -- the callee must be called directly. + self._allow_local_assignments = policy.allow_local_assignments + def report(self, node: ast.AST, code: ViolationCode, message: str) -> None: self.violations.append( Violation(line=getattr(node, "lineno", 0), code=code, message=message) @@ -210,8 +247,30 @@ def report(self, node: ast.AST, code: ViolationCode, message: str) -> None: # ── tree walk ─────────────────────────────────────────────────────────────────────────── def visit(self, node: ast.AST) -> None: """Walk the whole tree; enforce only on nodes inside the private ranges, recurse everywhere.""" + # `from ... import *` is banned everywhere, public region included: it silently pollutes the + # namespace and can shadow a name the private region trusts by spelling alone (a safe + # builtin, an import alias, a wrapper), and it can't be human-reviewed the way an explicit + # import can. Reported here, not via the private-region deny-list, so it fires in public too. + if isinstance(node, ast.ImportFrom) and any(a.name == "*" for a in node.names): + self.report( + node, + "star-import", + "'from ... import *' is not allowed anywhere in the file; import names explicitly", + ) + # Safe to return early: ImportFrom is never a scope node and has no verifiable + # children, so we skip the scope-stack / child-walk bookkeeping below. + return + if node_in_ranges(node, self.ranges): self._enforce(node) + elif not isinstance(node, _COMPOUND_NODES) and node_overlaps_ranges( + node, self.ranges + ): + # A multi-line statement/expression whose start line is public but whose span reaches + # into a private range. Default-deny: resolve the ambiguity by treating the whole node + # as private and verifying it. Compound statements are exempt -- a public header over a + # private body is normal marker usage, and their body statements are enforced by the walk. + self._enforce(node) else: # private-defined names are reserved everywhere, including in public # region. Bare calls trust private_defs by identifier alone, so a @@ -225,6 +284,13 @@ def visit(self, node: ast.AST) -> None: # way, regardless of whether the class statement or either definition is private. self._forbid_duplicate_methods(node) + if isinstance(node, (ast.For, ast.AsyncFor)) and node_overlaps_ranges( + node, self.ranges + ): + # A for-loop rebinds its target in the ENCLOSING scope, to an unvettable element of the + # iterable, so any safe-call verdict the target carries must be cleared. + self._clear_safe_locals(node.target) + # push/pop scope stack for ClassDef/FunctionDef/Lambda, so # _enclosing_class() works is_scope = isinstance(node, (ast.ClassDef, ast.FunctionDef, ast.Lambda)) @@ -349,6 +415,8 @@ def _enforce(self, node: ast.AST) -> None: self._require_bundle(node, "indexing") # --- name binding (assignments, params) --- + # NOTE: for/async-for target clearing is handled in visit() (gated on range OVERLAP), so it + # also fires for a public loop header over a private body -- not here. elif isinstance(node, (ast.Assign, ast.AugAssign, ast.AnnAssign)): self._track_safe_local(node) elif isinstance(node, ast.arg): @@ -604,6 +672,16 @@ def _check_call_attribute(self, call: ast.Call, func: ast.Attribute) -> None: return # a non-self dotted path: must resolve to an allow-listed import if root in self.scan.import_bindings: + # a dunder in call position (allowed.__wrapped__(x)) reaches the function-object + # introspection surface; deny it outright, mirroring the read-position check in + # _check_attribute rather than letting a broad glob (jax.*) admit it by path. + if is_dunder(func.attr): + self.report( + call, + "dunder-attr", + f"access to dunder attribute {func.attr!r} is not allowed", + ) + return if not self._resolved_allowed(path): self.report( call, @@ -776,6 +854,12 @@ def _track_safe_local(self, node) -> None: Any other assignment to a previously-tracked name clears its verdict -- it no longer traces to that source.""" + # Local-alias tracking disabled: never record any local as safe, so a bare-name call to a + # local (block = self.layers[0]; block(x)) falls through to call-unresolved. Callables must + # be called directly (self.layers[0](x)) instead. + if not self._allow_local_assignments: + return + # 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: @@ -788,17 +872,30 @@ def _track_safe_local(self, node) -> None: 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. + # A plain `x = ` target gets the RHS's verdict. Any other target shape -- a + # tuple/list/starred unpack (`b, _ = ...`) -- has no single value to attribute to each + # name, so it CLEARS every name it binds. Clearing (not skipping) matters when the name was + # already tracked safe: `b = self.dense; b, _ = fn, b` must drop b's stale safe verdict, + # since the unpack rebinds b in this scope at runtime. for t in targets: - if not isinstance(t, ast.Name): - continue - if value is not None and self._is_safe_local_source(value): - safe[t.id] = True + if isinstance(t, ast.Name): + if value is not None and self._is_safe_local_source(value): + safe[t.id] = True + else: + safe.pop(t.id, None) else: - safe.pop(t.id, None) + self._clear_safe_locals(t) + + def _clear_safe_locals(self, target: ast.AST) -> None: + """Drop the safe-call verdict for every local name bound by ``target``. + + Used for binding forms that rebind a name in the current scope without a single vettable + value: tuple/list/starred unpack targets and ``for``-loop targets. A comprehension target + is deliberately NOT cleared here -- in Python 3 it has its own scope and does not leak to + the enclosing one, so the outer name's verdict still holds.""" + safe = self._safe_locals_stack[-1] + for name in _iter_bound_local_names(target): + safe.pop(name, None) def _is_safe_local_source(self, value: ast.AST) -> bool: """Is a local provably safe to call?""" @@ -878,6 +975,19 @@ def _resolved_allowed(self, path: str) -> bool: return self.policy.function_allowed(self._resolve(path)) +def _iter_bound_local_names(target: ast.AST): + """Yield every plain local Name id bound by a Store-context target, recursing into + tuple/list-unpack and starred targets (``a, *rest = ...``). Attribute/subscript targets bind + no local name and are skipped.""" + if isinstance(target, ast.Name): + yield target.id + elif isinstance(target, (ast.Tuple, ast.List)): + for elt in target.elts: + yield from _iter_bound_local_names(elt) + elif isinstance(target, ast.Starred): + yield from _iter_bound_local_names(target.value) + + # ── 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, @@ -921,9 +1031,14 @@ class _SelfAttrTrust: one ``_Checker`` (one ``verify()`` call). """ - def __init__(self, scan: FileScan, resolved_allowed) -> None: + def __init__( + self, scan: FileScan, resolved_allowed, allow_base_class: bool = True + ) -> None: self._scan = scan self._resolved_allowed = resolved_allowed + # the verdict for a self. never assigned in the class body: True presumes it is + # inherited from the (already-vetted) base class; False rejects it (docs/verify.md). + self._allow_base_class = allow_base_class self._cache: dict[int, dict[str, bool]] = {} # id(ClassDef) -> {attr: is-safe} def is_safe(self, attr: str, cls_node: ast.ClassDef | None) -> bool: @@ -938,10 +1053,9 @@ def is_safe(self, attr: str, cls_node: ast.ClassDef | None) -> bool: 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) + # a safe/unsafe verdict for an attr the class assigns; for one it never assigns, fall back + # to _allow_base_class (presumed inherited from a vetted base, unless the caller disabled it) + return table.get(attr, self._allow_base_class) def _build_table(self, cls_node: ast.ClassDef) -> dict[str, bool]: """For every self. assigned anywhere in the class, record whether diff --git a/packages/syft-restrict/tests/obfuscate/test_obfuscate.py b/packages/syft-restrict/tests/obfuscate/test_obfuscate.py index e438481d114..d9b826632d8 100644 --- a/packages/syft-restrict/tests/obfuscate/test_obfuscate.py +++ b/packages/syft-restrict/tests/obfuscate/test_obfuscate.py @@ -4,8 +4,9 @@ import shutil from pathlib import Path -from syft_restrict import obfuscate, run +from syft_restrict import obfuscate from syft_restrict.astutil import scan_file +from syft_restrict.runner import _run FIXTURES = Path(__file__).parents[1] / "fixtures" ALLOW_FUNCTIONS = ["jax.*", "flax.linen.*"] @@ -24,7 +25,7 @@ def _obfuscate_fixture(tmp_path: Path): shutil.copy(FIXTURES / "compliant_model.py", src_path) source = src_path.read_text() private, config_line = _private_from_config(source) - result = run( + result = _run( src_path, obfuscate=private, allow_functions=ALLOW_FUNCTIONS, @@ -67,7 +68,7 @@ def test_hide_blanks_whole_body_keeping_indentation(tmp_path): if ln.lstrip().startswith("def scale_pattern") ) body = [def_line + 1, def_line + 2] # the two body lines (4-space indented) - result = run( + result = _run( src_path, obfuscate=[[def_line, def_line]], hide=[body], diff --git a/packages/syft-restrict/tests/test_markers.py b/packages/syft-restrict/tests/test_markers.py index fe1f3249937..54cd8aeef2d 100644 --- a/packages/syft-restrict/tests/test_markers.py +++ b/packages/syft-restrict/tests/test_markers.py @@ -237,3 +237,25 @@ def test_marker_lookalike_comment_prose_is_ignored(): obfuscate, hide = parse_markers(src) assert obfuscate == [(3, 3)] assert hide == [] + + +def test_block_start_marker_sharing_a_line_with_code_raises(): + # A start/end block marker must be on a line by itself; code before it would sit on an + # excluded boundary line and escape verification. The one-line form exists for that case. + src = normalize_source(""" + a = 1 # syft-restrict: obfuscate-start + b = 2 + # syft-restrict: obfuscate-end + """) + with pytest.raises(MarkerError): + parse_markers(src) + + +def test_block_end_marker_sharing_a_line_with_code_raises(): + src = normalize_source(""" + # syft-restrict: obfuscate-start + a = 1 + b = 2 # syft-restrict: obfuscate-end + """) + with pytest.raises(MarkerError): + parse_markers(src) diff --git a/packages/syft-restrict/tests/test_run.py b/packages/syft-restrict/tests/test_run.py index e205d7acf51..c2f2a1fe941 100644 --- a/packages/syft-restrict/tests/test_run.py +++ b/packages/syft-restrict/tests/test_run.py @@ -5,6 +5,7 @@ import pytest from syft_restrict import MarkerError, PolicyViolation, run +from syft_restrict.runner import _run from verify.helpers import normalize_source FIXTURES = Path(__file__).parent / "fixtures" @@ -22,7 +23,7 @@ def _private(source: str): def test_run_success_writes_obfuscated_and_certificate(tmp_path): src = tmp_path / "model.py" shutil.copy(FIXTURES / "compliant_model.py", src) - result = run( + result = _run( src, obfuscate=_private(src.read_text()), allow_functions=ALLOW_FUNCTIONS, @@ -41,7 +42,7 @@ def test_run_strict_raises_and_writes_nothing(tmp_path): src = tmp_path / "bad.py" src.write_text("CONFIG = dict(dim=8)\nimport os\nleak = os.getcwd()\n") with pytest.raises(PolicyViolation) as exc: - run( + _run( src, obfuscate=[[1, 3]], allow_functions=ALLOW_FUNCTIONS, @@ -54,7 +55,7 @@ def test_run_strict_raises_and_writes_nothing(tmp_path): def test_run_nonstrict_returns_violations(tmp_path): src = tmp_path / "bad.py" src.write_text("CONFIG = dict(dim=8)\nleak = x.reshape(1)\n") - result = run( + result = _run( src, obfuscate=[[1, 2]], allow_functions=ALLOW_FUNCTIONS, @@ -91,7 +92,9 @@ def test_run_auto_detects_markers_when_ranges_omitted(tmp_path): ) # the private region's own identifiers were renamed -def test_run_without_markers_or_ranges_raises_marker_error(tmp_path): +def test_run_without_markers_raises_marker_error(tmp_path): + # run() is markers-only: a file with no `# syft-restrict:` markers has no private region to + # resolve, so parse_markers() raises rather than silently verifying nothing. src = tmp_path / "unmarked.py" src.write_text( normalize_source(""" @@ -104,20 +107,20 @@ def f(x): run(src, allow_functions=ALLOW_FUNCTIONS, allow_operators=ALLOW_OPERATORS) -def test_run_explicit_ranges_bypass_marker_scanning(tmp_path): - # A stray unmatched marker would make parse_markers() raise -- explicit obfuscate= must skip - # marker scanning entirely rather than validating markers it isn't going to use. +def test__run_uses_explicit_ranges_and_ignores_markers(tmp_path): + # _run() takes ranges directly and never scans markers, so a stray/unmatched marker in the + # source (which would make run()'s parse_markers raise) is simply ignored. src = tmp_path / "model.py" src.write_text( normalize_source(""" # syft-restrict: obfuscate-start CONFIG = dict(dim=8) - leak = os.getcwd() # unmatched block, no obfuscate-end + x = CONFIG # unmatched block, no obfuscate-end """) ) - result = run( + result = _run( src, - obfuscate=[[1, 2]], + obfuscate=[[1, 3]], allow_functions=ALLOW_FUNCTIONS, allow_operators=ALLOW_OPERATORS, ) diff --git a/packages/syft-restrict/tests/verify/helpers.py b/packages/syft-restrict/tests/verify/helpers.py index d7c52fbdbae..659165c7ded 100644 --- a/packages/syft-restrict/tests/verify/helpers.py +++ b/packages/syft-restrict/tests/verify/helpers.py @@ -20,8 +20,20 @@ def normalize_source(source: str | list[str]) -> str: return source -def make_policy(functions=ALLOW_FUNCTIONS, operators=ALLOW_OPERATORS, disallow=None): - return Policy.parse(list(functions), list(operators), list(disallow or [])) +def make_policy( + functions=ALLOW_FUNCTIONS, + operators=ALLOW_OPERATORS, + disallow=None, + allow_local_assignments=True, + allow_base_class_attributes=True, +): + return Policy.parse( + list(functions), + list(operators), + list(disallow or []), + allow_local_assignments=allow_local_assignments, + allow_base_class_attributes=allow_base_class_attributes, + ) def get_error_codes(result: VerifyResult): diff --git a/packages/syft-restrict/tests/verify/test_bypasses.py b/packages/syft-restrict/tests/verify/test_bypasses.py index cfb87dd1e93..18231b60450 100644 --- a/packages/syft-restrict/tests/verify/test_bypasses.py +++ b/packages/syft-restrict/tests/verify/test_bypasses.py @@ -615,3 +615,154 @@ def __call__(self, x): return x """ _assert_error_code(verify_all, src, "duplicate-method", private=[[8, 9]]) + + +# ── safe-local verdict must be cleared on every rebind form, not just plain assign ─────────── + + +def test_tuple_unpack_rebind_clears_safe_local_verdict(verify_all): + # `b = self.dense` marks b safe; `b, _ = fn, b` rebinds b to an opaque param (tuple-unpack + # targets bind in the current scope, so the rebind leaks). The unpack must clear b's safe + # verdict, so the later b(x) is call-unresolved rather than trusting the stale verdict. + src = """ + class Block: + def __call__(self, x): + return x + class M: + def setup(self): + self.dense = Block() + def steal(self, x, fn): + b = self.dense + b, _ = fn, b + return b(x) + """ + _assert_error_code(verify_all, src, "call-unresolved") + + +def test_for_loop_target_rebind_clears_safe_local_verdict(verify_all): + # `for b in [fn]:` rebinds b to an opaque element (for-loop vars leak to the enclosing scope). + # The for-target must clear b's safe verdict, so b(x) is call-unresolved. + src = """ + class Block: + def __call__(self, x): + return x + class M: + def setup(self): + self.dense = Block() + def steal(self, x, fn): + b = self.dense + for b in [fn]: + return b(x) + return x + """ + _assert_error_code(verify_all, src, "call-unresolved") + + +# ── multi-line statement straddling the public->private boundary ───────────── + + +def test_boundary_straddle_is_verified_as_private(verify_all): + # `jnp.save(` starts on a PUBLIC line but its args are in the private range. Range membership + # is start-line based, so the call was historically skipped. Default-deny: a node whose span + # overlaps the private region is verified as private. Under a tight policy (no jnp.save) the + # straddling call must be caught. + src = """ + import jax.numpy as jnp + def f(arr): + r = jnp.save( + "x.npy", + arr) + return r + """ + result = verify_all( + src, pol=make_policy(functions=["jax.numpy.einsum"]), private=[[4, 5]] + ) + assert "call-not-allowed" in get_error_codes(result) + + +# ── dunder segment in call position on an allow-listed path ────────────────────────────────── + + +def test_dunder_in_call_position_on_allowed_path(verify_all): + # A dunder in DIRECT call position on an allow-listed dotted path must be rejected as + # dunder-attr, not gated only by the allow-list. Under a broad `jax.*` glob, + # `jnp.einsum.__wrapped__` resolves to a `jax.*` path and would otherwise pass, handing + # back the function-object introspection surface (`__wrapped__`, and by the same shape + # `__globals__`/`__defaults__` if reachable this way). Read position was already caught; + # call position must match. + src = """ + import jax.numpy as jnp + def f(x): + return jnp.einsum.__wrapped__(x) + """ + codes = get_error_codes(verify_all(src, private=[[2, 3]])) + assert "dunder-attr" in codes + + +def test_dunder_call_and_read_positions_are_consistent(verify_all): + # The read-then-call form of the same access is caught; the direct-call form must be too. + read_then_call = """ + import jax.numpy as jnp + def f(x): + g = jnp.einsum.__wrapped__ + return g(x) + """ + assert "dunder-attr" in get_error_codes( + verify_all(read_then_call, private=[[2, 4]]) + ) + + +# ── multi-segment `import a.b` must not poison path resolution / defeat the floor ──────────── + + +def test_multisegment_import_does_not_bypass_disallow_floor(verify_all): + # `import jax.numpy` binds the NAME `jax` (to the jax package) at runtime -- not `jax.numpy`. + # A resolver that binds jax -> "jax.numpy" mis-resolves `jax.numpy.save` to + # "jax.numpy.numpy.save", so a disallow floor listing jax.numpy.save misses it while runtime + # still calls the real function. + src = """ + import jax.numpy + def f(x): + return jax.numpy.save("/tmp/leak.npy", x) + """ + codes = get_error_codes( + verify_all(src, pol=make_policy(disallow=["jax.numpy.save"]), private=[[2, 3]]) + ) + assert "call-not-allowed" in codes + + +def test_multi_import_last_wins_does_not_bypass_floor(verify_all): + # `import jax` then `import jax.tree_util` must not clobber the `jax` binding to jax.tree_util. + src = """ + import jax + import jax.tree_util + def f(x): + return jax.pure_callback(x, x, x) + """ + codes = get_error_codes( + verify_all( + src, pol=make_policy(disallow=["jax.pure_callback"]), private=[[3, 4]] + ) + ) + assert "call-not-allowed" in codes + + +def test_public_for_header_over_private_body_clears_safe_local(verify_all): + # A `for` header on a PUBLIC line with a PRIVATE body still rebinds the loop target in the + # enclosing scope. The `for` is a compound node (exempt from straddle-enforce), so the target + # clear must happen in the walk whenever the loop overlaps a private range -- otherwise b keeps + # its stale safe verdict and the private b(x) wrongly trusts it (reintroducing the rebind hole). + src = """ + class Block: + def __call__(self, x): + return x + class M: + def setup(self): + self.dense = Block() + def steal(self, x, fn): + b = self.dense + for b in [fn]: + return b(x) + """ + # line 8 (b = self.dense) private, line 9 (for header) PUBLIC, line 10 (return b(x)) private + _assert_error_code(verify_all, src, "call-unresolved", private=[[8, 8], [10, 10]]) diff --git a/packages/syft-restrict/tests/verify/test_disallowed.py b/packages/syft-restrict/tests/verify/test_disallowed.py index 97e62baff03..3582fff9165 100644 --- a/packages/syft-restrict/tests/verify/test_disallowed.py +++ b/packages/syft-restrict/tests/verify/test_disallowed.py @@ -110,6 +110,14 @@ def test_match_statement_is_not_on_node_allowlist(verify_all): "ascii", "format", "bytes", + # site-injected builtins: stdout channels (copyright/credits/license), + # interpreter shutdown (exit/quit), interactive help (help) + "copyright", + "credits", + "license", + "exit", + "quit", + "help", ], ) def test_banned_names(verify_all, name): @@ -185,9 +193,7 @@ def __call__(self, x): # ── defs / classes / decorators ────────────────────────────────────────────── -@pytest.mark.parametrize( - "dunder", ["__init__", "__post_init__", "__getattr__", "__reduce__"] -) +@pytest.mark.parametrize("dunder", ["__init__", "__getattr__", "__reduce__"]) def test_disallowed_dunder_def(verify_all, dunder): src = f""" class M: @@ -326,3 +332,21 @@ def helper(d, x): return d['k'](x) """ assert "call-unresolved" in get_error_codes(verify_all(src)) + + +# ── star imports (banned anywhere, even in the trusted public region) ───────── + + +def test_star_import_in_public_is_blocked(verify_all): + # `from x import *` can silently shadow a name the private region trusts by spelling (a safe + # builtin, import alias, wrapper) and is unreviewable, so it is banned even in public code. + src = """ + from evil import * + def f(x): + return len(x) + """ + assert "star-import" in get_error_codes(verify_all(src, private=[[2, 3]])) + + +def test_star_import_in_private_is_blocked(verify_all): + assert "star-import" in get_error_codes(verify_all("from os import *")) diff --git a/packages/syft-restrict/tests/verify/test_policy_flags.py b/packages/syft-restrict/tests/verify/test_policy_flags.py new file mode 100644 index 00000000000..dd378b6e67d --- /dev/null +++ b/packages/syft-restrict/tests/verify/test_policy_flags.py @@ -0,0 +1,86 @@ +"""The optional strictness flags: allow_local_assignments / allow_base_class_attributes. + +Both default True (the behavior the rest of the suite assumes). These tests pin what flipping each +to False does: it turns a construct that is normally accepted into a rejection. +""" + +from verify.helpers import get_error_codes, make_policy + +# ── allow_local_assignments ────────────────────────────────────────────────── + +_LOCAL_ALIAS = """ +class M: + def __call__(self, x): + block = self.layers[0] + return block(x) +""" + + +def test_local_alias_call_allowed_by_default(verify_all): + # block = self.layers[0]; block(x) — the alias is tracked as a safe call target. + assert not get_error_codes(verify_all(_LOCAL_ALIAS)) + + +def test_local_alias_call_rejected_when_local_assignments_disabled(verify_all): + # With tracking off, the alias is no longer a trusted callee; the call is unresolved. + codes = get_error_codes( + verify_all(_LOCAL_ALIAS, pol=make_policy(allow_local_assignments=False)) + ) + assert "call-unresolved" in codes + + +def test_direct_self_subscript_call_still_allowed_when_local_assignments_disabled( + verify_all, +): + # The flag only removes the *alias* convenience; calling the value directly still works. + src = """ + class M: + def __call__(self, x): + return self.layers[0](x) + """ + assert not get_error_codes( + verify_all(src, pol=make_policy(allow_local_assignments=False)) + ) + + +# ── allow_base_class_attributes ───────────────────────────────────────────────── + +_UNKNOWN_ATTR = """ +class M: + def __call__(self, x): + return self.norm(x) +""" + + +def test_unknown_self_attr_call_allowed_by_default(verify_all): + # self.norm is never assigned in the class -> presumed inherited from the vetted base. + assert not get_error_codes(verify_all(_UNKNOWN_ATTR)) + + +def test_unknown_self_attr_call_rejected_when_base_class_attributes_disabled( + verify_all, +): + # With the base-class assumption off, a never-assigned self. is not callable. + codes = get_error_codes( + verify_all(_UNKNOWN_ATTR, pol=make_policy(allow_base_class_attributes=False)) + ) + assert "attr-on-value" in codes + + +def test_assigned_self_attr_still_callable_when_base_class_attributes_disabled( + verify_all, +): + # The flag rejects only *unknown* attrs; one the class assigns a vetted source stays callable. + src = """ + class Block: + def __call__(self, x): + return x + class M: + def setup(self): + self.blk = Block() + def __call__(self, x): + return self.blk(x) + """ + assert not get_error_codes( + verify_all(src, pol=make_policy(allow_base_class_attributes=False)) + ) diff --git a/packages/syft-restrict/tests/verify/test_whitelist.py b/packages/syft-restrict/tests/verify/test_whitelist.py index 521334e0169..61168d69dfc 100644 --- a/packages/syft-restrict/tests/verify/test_whitelist.py +++ b/packages/syft-restrict/tests/verify/test_whitelist.py @@ -213,15 +213,36 @@ def __call__(self, x): assert _ok(verify_all(src))[0] +def test_comprehension_target_does_not_taint_outer_local(verify_all): + """A comprehension has its own scope (py3): ``[b for b in ...]`` does NOT rebind the outer + ``b``, so a safe local stays safe and remains callable afterward. Guards against clearing + the verdict for comprehension targets, which would be an over-denial.""" + src = """ + class Block: + def __call__(self, x): + return x + class M: + def setup(self): + self.dense = Block() + def run(self, x): + b = self.dense + [b for b in [x]] + return b(x) + """ + assert _ok(verify_all(src))[0] + + # ── decorators / dunder defs / bases ───────────────────────────────────────── def test_allowed_dunder_defs(verify_all): - """``__call__`` and ``setup`` are the only definable hooks.""" + """``__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 """ @@ -319,3 +340,14 @@ def f(x: "str") -> "bytes": return x """ assert _ok(verify_all(src))[0] + + +def test_fully_private_multiline_call_is_allowed(verify_all): + """A multi-line call entirely within the private region verifies normally.""" + src = """ + def f(xs): + return sum( + xs + ) + """ + assert _ok(verify_all(src))[0]