Repository navigation
FEAT: Rebuild saved registry entries in dependency order (part 3) - #3908
varunj-msft wants to merge 67 commits into
Conversation
…API input External construction now rejects a parameter whose type includes Path unless the calling service passes the name in allowed_path_params after making the value safe. The converter service allows its upload parameters once they are stored, and the target service allows the server-set upload folder and the GitHubCopilotTarget working directory. The scorer-only check is removed, and the "Valid parameters" list leaves out path parameters callers cannot set.
The target catalog now applies the path rule to required parameters too, so a target that needs a disallowed path is left out instead of listed without the field, matching the scorer catalog. allowed_path_params is typed frozenset[str] so a bare string cannot turn the allow-list into a substring match.
The target catalog still lists a target whose required path parameter is its server-set upload folder when the server has one, and leaves it out otherwise.
# Conflicts: # tests/unit/backend/test_converter_service.py
…navailable storage
Withdrawing my approval for now so I can review this PR again.
Richard Lundeen (richlundeen)
left a comment
There was a problem hiding this comment.
Main point: keep the app up when a saved entry can't be restored, and simplify around that
I asked for "stay not ready if any saved entry fails" in the design review. Having seen what it costs across the stack, I'd like to change course.
The rule needs a whole recovery system: a failure list for every user, an explanation for each reason, routes that answer while the app is down, and Retry with partial-state tracking and a restart fallback (#3965 alone is about 1,050 lines), plus deleting while down in #4071. It also turns routine upgrades into outages. A renamed constructor parameter or a class moved to another module fails every saved entry that uses it, and #4026 notes another case: a target whose class only a scan's launch scripts register can't come back after a restart.
The goal stays the same: a run must never quietly use something other than what it expects. I think we can keep that without stopping the app:
- Restore builds what it can, the runtime becomes ready, and the failures are listed in the runtime status and the banner.
- A saved entry that failed keeps its name. Looking it up, or starting a run that uses it, fails with the restore reason instead of finding nothing or something else.
- When
adversarial_chatorobjective_scorer_chatis missing, the backend raises instead of building a plainOpenAIChatTargetin_get_default_chat_target. That fallback is a problem today anyway, for example after a typo in a startup script. - You found that runs picking techniques by all, tag or default would have to stop anyway, since a broken file can hide its tags. Agreed: while any saved technique is broken, refuse exactly those launches. Launches that name their techniques still run.
- Fixing and deleting use the normal routes, because the app is up.
For the stack: #3891, #3892 and this PR stay as they are, apart from the review comments. #3909 would publish even when something fails and add points 2 and 4. #3965 would shrink to the failure list and the banner, with Retry optional. #4071's routes would no longer need to answer while the app is down. Point 3 can be its own small fix on main, and sections 1 and 4.6 of the design change to match.
The same simplification applies inside this PR. The class docstring already sets the trust model: "Stored files are otherwise trusted the way configuration is." With that as the rule, _safe_message can become URL redaction plus a length cap, and building can go through the same function the API's create uses instead of a partial copy that has already drifted. Details are in the inline comments.
Inline comments marked Blocking need to be addressed before merge; Non-blocking ones are judgment calls.
| except KeyError as error: | ||
| message = f"'{recipe.class_name}' is not a registered {self._singular(ref.kind)} class." | ||
| raise _RecipeRejectedError(reason=RestoreFailureReason.UNKNOWN_CLASS, message=message) from error | ||
| if cls.__module__ != recipe.class_module: |
There was a problem hiding this comment.
Blocking: don't fail an entry because its class moved to another module.
Classes are looked up by registered name, and the API builds by name alone, so this makes restore stricter than the API it mirrors: any refactor that moves a class fails every saved entry of it. Suggested fix: accept the registered class when both the saved module and its current module are inside pyrit (a move within PyRIT), and fail only when a class from outside PyRIT now holds the name. Keep class_module in the file as information.
There was a problem hiding this comment.
This move already happened once; compare with pyrit. including the dot
Agree, and PyRIT already made such a move. #2161 moved every converter module from pyrit.prompt_converter to pyrit.converter and kept class names such as Base64Converter. With this check, every saved converter would fail after that upgrade, although the lookup by name on line 203 finds the correct class. For the suggested fix, treat a module as inside PyRIT only when it is pyrit or starts with pyrit.. A plain startswith("pyrit") also accepts a module such as pyrit_evil.x.
| if stored.ref.kind is RegistryRecipeKind.TARGETS: | ||
| # As in the target API, a target may let callers set some path parameters, such as a working folder. | ||
| allowed_path_params = registry.get_class(recipe.class_name).caller_path_parameters | ||
| registry.create_named_instance( |
There was a problem hiding this comment.
Blocking: build through the API's own create function.
Calling create_named_instance directly skips steps the target API applies in target_service.create_target_async: the auth-mode translation (get_auth_mode_parameters), the Copilot token rule and the admin-only checks. Converters with upload parameters also fail here as a vague build_error, because allowed_path_params is empty for them. Editing (part 21) already rebuilds through the API's path. Suggested fix: move each kind's create logic into one function that the API, edit, restore and Retry all call, with restore passing a server context. If that has to wait, correct the class docstring (lines 53-54), which says a restored entry "passes the same checks", and list what restore skips and why.
There was a problem hiding this comment.
Build through the service create methods; store auth_mode in the recipe first
Agree. The service create methods (create_target_async, create_converter_async, create_scorer_async, TechniqueService._create) already are the one function per kind. Restore can't call them only because it is in pyrit.registry, which imports nothing from pyrit.backend, and its only caller (runtime_lifecycle.py in #3909 and #3965) is in the backend. Moving the build step into pyrit/backend/services and calling those methods with a server context is smaller than moving four services' rules into the library. It also removes the upload-folder refusal (lines 213-217), the copy of TechniqueService._create (lines 433-450) and the copy of the reserved technique names (line 65) that repeats backend/models/techniques.py line 96, and create_target_async already releases a target when registering fails (thread on line 423). One format decision comes first: ComponentRecipe has no request-level auth_mode, and create_target_async raises when params['auth_mode'] differs from request.auth_mode (default api_key). So an identity target saved like the test fixture, with params: {'auth_mode': 'identity'}, fails on replay. While nothing writes recipes, store auth_mode and the params before translation. Two tests in test_registry_recipe_restore.py pin the current path and must change: line 713 restores a GitHubCopilotTarget without the token that create_target_async requires, and line 332 mocks create_named_instance and checks its arguments. Use test-registered classes there and check the result.
| for kind, storage in self._storages.items(): | ||
| try: | ||
| listings.append(storage.list_recipes()) | ||
| except Exception as error: # Any listing error, not only storage ones, must become a failure. |
There was a problem hiding this comment.
Blocking: catch storage errors only.
Every exception becomes unreadable_storage and is described with describe_storage_error, so a code bug shows up as "The saved targets couldn't be listed: KeyError" and its traceback is lost. Suggested fix: catch OSError and AzureError here, and log anything else with its traceback (or let it raise).
| ) | ||
| except _RecipeRejectedError as rejection: | ||
| return self._failure(stored=stored, reason=rejection.reason, message=str(rejection)) | ||
| except Exception as error: # Planning one entry must never stop the others from being checked. |
There was a problem hiding this comment.
Blocking: don't blame the file for our own bugs.
An unexpected exception while planning is reported as invalid_recipe, which tells the owner their file is wrong, and the traceback is dropped. Log it with its traceback and report it as a build error.
| ) | ||
| except _RecipeRejectedError as rejection: | ||
| return self._failure(stored=stored, reason=rejection.reason, message=str(rejection)) | ||
| except Exception as error: # A constructor can raise anything; one bad entry must not stop the rest. |
There was a problem hiding this comment.
Blocking: log the exception.
Constructor failures keep only the cleaned-up text, and nothing in this module logs, so the full error and traceback never reach anyone, although the design promises admins the full error. Log each failure once with its traceback, with SAS tokens redacted, or carry the exception on the failure so the caller can log it.
| name=name, attack_type=recipe.type, technique_tags=recipe.tags, strict_json=True, **args | ||
| ) | ||
| try: | ||
| registry.instances.register_runtime(factory) |
There was a problem hiding this comment.
Blocking (small): not every ValueError here is a name clash.
register_runtime also raises when scenario-pool validation fails, which this reports as name_clash. Check name and tag availability before registering, or catch only the clash cases, and report anything else as invalid.
There was a problem hiding this comment.
The component path has the same mapping, and reserved names get two reasons
Agree, and two more cases share this cause. A technique named value fails here as name_clash with "maps to enum member name 'value', which conflicts with an inherited ScenarioTechnique attribute". That error comes from enum generation, not from another entry. Lines 221-223 also map every validate_name_available error to name_clash, and that includes reserved names: a saved converter named preview is reported as name_clash, but a technique named types is invalid_recipe (line 241). Use one reason for a reserved or unusable name in every kind, and keep name_clash for a name that another entry holds.
| listings, failures = list(saved.listings), list(saved.failures) | ||
| ignored = [file for listing in listings for file in listing.ignored] | ||
| failures.extend(failure for listing in listings for failure in listing.failures) | ||
| if len(listings) < len(self._storages): |
There was a problem hiding this comment.
Blocking (small): this means "a folder couldn't be listed". if saved.failures: says that directly and doesn't depend on how many listings a caller passes in.
| references.extend(SavedEntryRef(kind=kind, name=name) for name in names if isinstance(name, str)) | ||
| return references | ||
|
|
||
| def _build_in_order( |
There was a problem hiding this comment.
Non-blocking: consider graphlib.TopologicalSorter.
It hands out entries in dependency order and keeps going past a loop after prepare() raises CycleError, so _build_in_order and _ready could shrink. The loop-versus-waiting labels in _unresolved would still need custom code.
| params=recipe.params, | ||
| external_input=True, | ||
| allowed_path_params=allowed_path_params, | ||
| ) |
There was a problem hiding this comment.
Non-blocking: if registering fails after create_named_instance has built the object, the object isn't released; the target API calls release_component_async on that path. It's unlikely at startup, but more likely once restore or Retry runs while the app is up. Worth handling in create_named_instance.
| f"'{recipe.class_name}' uploads files from a folder the server chooses, which restore can't set yet." | ||
| ) | ||
| raise _RecipeRejectedError(reason=RestoreFailureReason.UNSUPPORTED_CLASS, message=message) | ||
| parameters = registry.get_class_metadata(cls).parameters |
There was a problem hiding this comment.
Non-blocking (minor): planning reads class metadata with get_class_metadata for every recipe instead of the registry's cached get_registered_class_metadata, and _get_registry (line 452) re-imports its modules on every call.
FileDocumentStorage now raises DocumentStorageUnavailableError when the folder or container can't be used, and DocumentUnreadableError when one document exists but can't be read. The message names the location, with any SAS token hidden, and a short reason. One backend handler maps both to 503 and logs the location and reason, with the traceback at debug level, so the per-route storage helpers are gone. A single read, a conditional save and a delete now follow the listing's rules: a folder, pipe or broken link named like a document is unreadable rather than missing, and on Blob only a deleted blob counts as missing. Blob sign-in, permission, network and missing-container errors stop the listing instead of marking every file unreadable. Each unreadable report keeps the real file name and the original error, and the listing's documents are read-only. Initializer delete removes a script that can't be read or decoded, and registering one no longer turns a storage write failure into a 400. Tests patch module-level aliases instead of os.
List recipes escapes the OS text it reports, such as a file name that isn't UTF-8, and a recipe whose JSON holds an unpaired surrogate escape is reported as invalid, so the models no longer need a JSON-only serializer. Recipes may start with a byte-order mark. A document without schema_version reads as version 1, and each file is checked against the format version of the model that reads it. Validation summaries shorten each location part and the whole message, since a location can repeat a stored key. Recipes and StoredRecipe are no longer frozen, since their params and tags never were, and StoredRecipe checks that its recipe type matches its registry. The listing's recipes are read-only, and it keeps the error behind each recipe that couldn't be read, for the log.
Listing catches only storage errors, so a bug is raised instead of being reported as unreadable storage, and a bug while planning is a build error rather than an invalid recipe. Each failure is logged once, with the error behind it and URL credentials removed. The report shows the error's text with each URL's user name, query and fragment removed, any SAS signature hidden, and the length capped. A class that moved within PyRIT is restored; one now registered from outside PyRIT still fails. A technique whose name or tag clashes is a name clash, and anything else registration refuses is an invalid recipe. A recipe that sets a server file or folder is refused before it is built. Class metadata comes from the registry's cache. The class docstring lists the steps the API's create runs that restore skips.
| assert failures["moved"].reason is RestoreFailureReason.MODULE_MISMATCH | ||
| assert _OPENAI_MODULE in failures["moved"].message | ||
| assert failures["unknown"].reason is RestoreFailureReason.UNKNOWN_CLASS | ||
| assert module not in sys.modules |
There was a problem hiding this comment.
[P3] Make the never-imports test fail when restore imports a module
This test names a module that does not exist, then checks that the module is not in sys.modules. A failed import never adds a module to sys.modules, so this check cannot fail. A restore that calls importlib.import_module(recipe.class_module) inside try/except ImportError and then looks the class up by name passes every assertion here. The class docstring promises that a file never chooses what gets imported, and the class-move fix in registry_recipe_restore.py (line 207) changes code near this check. Write a small module file under tmp_path, add tmp_path to sys.path with patch.object, name that module in class_module, and assert that the module is not in sys.modules after the restore. Use patch.dict(sys.modules) to keep the test isolated.
| assert "uploader" not in TargetRegistry.get_registry_singleton().instances | ||
|
|
||
| @pytest.mark.parametrize("auth_mode", [[], {}, ["identity"]]) | ||
| def test_reports_malformed_values_without_stopping_the_restore(self, tmp_path: Path, auth_mode: object) -> None: |
There was a problem hiding this comment.
[P3] Keep one test for a build error that does not stop the restore
This test (3 cases) and test_reports_constructor_errors_and_keeps_going (line 828) each save an OpenAIChatTarget with a value that cannot be coerced, next to a good entry. All four cases fail at the same point: create_named_instance raises ValueError, and restore reports build_error. test_error_text_that_is_not_a_valid_url_fails_only_that_entry (line 439) takes the same path and checks the same results and more. This test also checks only the names, not the reason, so it still passes if the failure moves to planning. Keep one of these tests, assert its reason, and delete the others.
Description
This is one of the changes for saving Registry entries across restarts. It builds on #3892. Until #3891 and #3892 merge, the diff here also shows their changes; this PR's own code is
pyrit/registry/registry_recipe_restore.pyandRegistryRestoreReport, with their tests.Problem. Nothing turns saved recipes back into registry entries.
Change.
RegistryRecipeRestorer:read()reads every folder without building anything, so a caller on an event loop can run it in a worker thread;restore()then builds on the calling thread, as the target API does, because a constructor can need the running loop.restore():__module__, and looks a technique's attack type up in the attack registry. It never imports a module a file names.Parameter.referenceand the technique fields. A name resolves to a saved entry first (even a broken one), then to a startup entry. Missing and failed dependencies and loops are reported with the entry each one waits on.create_named_instance(..., external_input=True), so a server path is refused exactly as for API input (FIX: Reject path parameters by default when building components from API input #3594): a target may set only the path parameters its class lets callers set (caller_path_parameters, such asGitHubCopilotTarget'sworking_directory), and converters and scorers none. Techniques go throughcreate_factory(..., strict_json=True)andregister_runtime, as the technique API does. A test checks that technique recipes and the create request accept and reject the same documents and have the same defaults.upload_directory_parameter), which restore can't set yet.?or#, where a URL's query or fragment starts, and says the rest is hidden. A message with an@, as in a URL's user name and password, is hidden entirely. Look-alike characters such as a full-width@count as the real ones, and the text that is shown keeps its original characters.RegistryRestoreReport: restored entries, failures (kind, name, reason, message, version, dependency) and ignored files.Not in this PR.
Tests and Documentation
tests/unit/registry/test_registry_recipe_restore.py(89 tests): order across the four kinds (including converters in either of a technique's lists), loops, missing and failed dependencies, clashes with startup entries,read()building nothing andrestore()using what it returned without reading again, a startup entry found by its exact name, stored credentials, server paths refused as the API refuses them (a target's declared caller path is accepted, and only for targets), class lookup without imports, unknown classes and attack types, technique recipes matching the create request, and failure messages (queries, fragments, user info, the exact message for every form of?,#and@, original characters kept, long messages checked one character at a time).pytest tests/unit/registry tests/unit/models -n 4: 3,935 passed, 5 skipped, 4 xfailedpre-commit run --from-ref origin/main --to-ref HEAD: passed