feat: Further syft-restrict improvements and red-team hardening

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.<attr> 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`

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.<name>(...)`                    | 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.<name>` / `cls.<name>` 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).
 

packages/syft-restrict/docs/verify.md

@@ -83,24 +83,40 @@ 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.
 >
 > 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                                             |
 
 ---

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
 }