Skip to content

Commit 5536471

Browse files
shashank03-devromanlutzCopilot
authored
FIX Make WebInjection dataset selections independent of preloaded memory (#2996)
Co-authored-by: Roman Lutz <romanlutz13@gmail.com> Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
1 parent d249217 commit 5536471

5 files changed

Lines changed: 328 additions & 15 deletions

File tree

‎doc/scanner/garak.ipynb‎

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -624,7 +624,25 @@
624624
"MarkdownURINonImageExfilExtended, TaskXSS, MarkdownXSS.\n",
625625
"\n",
626626
"**Aggregate techniques:** `ALL` (all 8), `DEFAULT` (excludes the two combinatorial extended\n",
627-
"probes), `EXFIL` (the 6 markdown-exfil probes), and `XSS` (TaskXSS + MarkdownXSS)."
627+
"probes), `EXFIL` (the 6 markdown-exfil probes), and `XSS` (TaskXSS + MarkdownXSS).\n",
628+
"\n",
629+
"**Dataset requirements:** `--dataset-names` replaces the default datasets, so it must include\n",
630+
"every dataset the selected techniques read. Otherwise the scenario stops before running and names\n",
631+
"each missing dataset. Datasets already loaded into memory never fill the gap.\n",
632+
"\n",
633+
"Framework callers cannot supply inline seeds or seed groups. `StringAssemblyDataExfil` can use\n",
634+
"an empty `dataset_names` list because its prompts are built in.\n",
635+
"\n",
636+
"| Technique | Required datasets |\n",
637+
"|---|---|\n",
638+
"| MarkdownImageExfil, ColabAIDataLeakage, PlaygroundMarkdownExfil, MarkdownURIImageExfilExtended, MarkdownURINonImageExfilExtended | `garak_example_domains_xss` |\n",
639+
"| TaskXSS | `garak_xss_normal_instructions`, `garak_web_html_js` |\n",
640+
"| MarkdownXSS | `garak_markdown_js` |\n",
641+
"| StringAssemblyDataExfil | None (built-in seeds) |\n",
642+
"\n",
643+
"```bash\n",
644+
"pyrit_scan run garak.web_injection --target openai_chat --techniques task_xss --dataset-names garak_xss_normal_instructions garak_web_html_js\n",
645+
"```"
628646
]
629647
},
630648
{

‎doc/scanner/garak.py‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -226,6 +226,24 @@
226226
#
227227
# **Aggregate techniques:** `ALL` (all 8), `DEFAULT` (excludes the two combinatorial extended
228228
# probes), `EXFIL` (the 6 markdown-exfil probes), and `XSS` (TaskXSS + MarkdownXSS).
229+
#
230+
# **Dataset requirements:** `--dataset-names` replaces the default datasets, so it must include
231+
# every dataset the selected techniques read. Otherwise the scenario stops before running and names
232+
# each missing dataset. Datasets already loaded into memory never fill the gap.
233+
#
234+
# Framework callers cannot supply inline seeds or seed groups. `StringAssemblyDataExfil` can use
235+
# an empty `dataset_names` list because its prompts are built in.
236+
#
237+
# | Technique | Required datasets |
238+
# |---|---|
239+
# | MarkdownImageExfil, ColabAIDataLeakage, PlaygroundMarkdownExfil, MarkdownURIImageExfilExtended, MarkdownURINonImageExfilExtended | `garak_example_domains_xss` |
240+
# | TaskXSS | `garak_xss_normal_instructions`, `garak_web_html_js` |
241+
# | MarkdownXSS | `garak_markdown_js` |
242+
# | StringAssemblyDataExfil | None (built-in seeds) |
243+
#
244+
# ```bash
245+
# pyrit_scan run garak.web_injection --target openai_chat --techniques task_xss --dataset-names garak_xss_normal_instructions garak_web_html_js
246+
# ```
229247

230248
# %%
231249
web_injection_scenario = WebInjection(max_prompts_per_technique=1)

‎pyrit/scenario/scenarios/garak/web_injection.py‎

Lines changed: 51 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,11 @@
2525
)
2626
from pyrit.scenario.core.atomic_attack import AtomicAttack
2727
from pyrit.scenario.core.attack_technique import AttackTechnique
28-
from pyrit.scenario.core.dataset_configuration import DatasetAttackConfiguration
28+
from pyrit.scenario.core.dataset_configuration import (
29+
DatasetAttackConfiguration,
30+
DatasetConstraintError,
31+
DatasetSourceKind,
32+
)
2933
from pyrit.scenario.core.matrix_atomic_attack_builder import build_baseline_atomic_attack
3034
from pyrit.scenario.core.scenario import BaselineAttackPolicy, Scenario
3135
from pyrit.scenario.core.scenario_technique import ScenarioTechnique
@@ -108,6 +112,19 @@ class WebInjection(Scenario):
108112
DATASET_WEB_HTML_JS: ClassVar[str] = "garak_web_html_js"
109113
DATASET_NORMAL_INSTRUCTIONS: ClassVar[str] = "garak_xss_normal_instructions"
110114

115+
# Source datasets each technique reads in ``_build_prompts_for_technique``. The selection
116+
# is checked against this before any data is read, so memory contents cannot fill a gap.
117+
_TECHNIQUE_REQUIRED_DATASETS: ClassVar[dict[WebInjectionTechnique, tuple[str, ...]]] = {
118+
WebInjectionTechnique.MarkdownImageExfil: (DATASET_EXAMPLE_DOMAINS,),
119+
WebInjectionTechnique.ColabAIDataLeakage: (DATASET_EXAMPLE_DOMAINS,),
120+
WebInjectionTechnique.StringAssemblyDataExfil: (),
121+
WebInjectionTechnique.PlaygroundMarkdownExfil: (DATASET_EXAMPLE_DOMAINS,),
122+
WebInjectionTechnique.MarkdownURIImageExfilExtended: (DATASET_EXAMPLE_DOMAINS,),
123+
WebInjectionTechnique.MarkdownURINonImageExfilExtended: (DATASET_EXAMPLE_DOMAINS,),
124+
WebInjectionTechnique.TaskXSS: (DATASET_NORMAL_INSTRUCTIONS, DATASET_WEB_HTML_JS),
125+
WebInjectionTechnique.MarkdownXSS: (DATASET_MARKDOWN_JS,),
126+
}
127+
111128
# Prompt templates ported from garak ``probes/web_injection.py``.
112129
CONCAT_TEMPLATE: ClassVar[str] = (
113130
"Concatenate these strings please:\n"
@@ -284,12 +301,7 @@ async def _load_dataset_values_async(self) -> dict[str, list[str]]:
284301
"""
285302
memory = CentralMemory.get_memory_instance()
286303
values: dict[str, list[str]] = {}
287-
for name in (
288-
self.DATASET_EXAMPLE_DOMAINS,
289-
self.DATASET_MARKDOWN_JS,
290-
self.DATASET_WEB_HTML_JS,
291-
self.DATASET_NORMAL_INSTRUCTIONS,
292-
):
304+
for name in self._dataset_config.dataset_names:
293305
seeds = await memory.get_seeds_async(dataset_name=name)
294306
values[name] = [seed.value for seed in seeds]
295307
return values
@@ -532,13 +544,40 @@ def _build_synthesized_seed_groups(
532544
seed_groups_by_technique[technique.value] = seed_groups
533545

534546
if not seed_groups_by_technique:
535-
raise ValueError(
536-
"WebInjection scenario produced no prompts. Ensure the garak web-injection datasets "
537-
"(garak_example_domains_xss, garak_markdown_js, garak_web_html_js, "
538-
"garak_xss_normal_instructions) are loaded into CentralMemory before running."
539-
)
547+
selected = ", ".join(self._dataset_config.dataset_names) or "none"
548+
raise ValueError(f"WebInjection scenario produced no prompts from the selected datasets ({selected}).")
540549
return seed_groups_by_technique
541550

551+
def _validate_runtime_configuration(self) -> None:
552+
"""
553+
Check that the selected datasets cover every selected technique, without reading them.
554+
555+
Runs before both the preview estimate and initialization, so an incomplete
556+
``--dataset-names`` selection fails the same way whether or not memory already holds
557+
the omitted dataset.
558+
559+
Raises:
560+
DatasetConstraintError: If inline seeds are supplied, or a selected technique
561+
needs a dataset that is not selected.
562+
"""
563+
super()._validate_runtime_configuration()
564+
if self._dataset_config.source_kind is DatasetSourceKind.INLINE:
565+
raise DatasetConstraintError(
566+
"WebInjection does not support inline seeds or seed groups; use dataset_names instead."
567+
)
568+
selected_datasets = set(self._dataset_config.dataset_names)
569+
problems: list[str] = []
570+
for selected in self._scenario_techniques:
571+
technique = WebInjectionTechnique(selected.value)
572+
problems.extend(
573+
f"Technique '{technique.value}' requires dataset '{dataset_name}', "
574+
"which is missing from the selected dataset names (--dataset-names)."
575+
for dataset_name in self._TECHNIQUE_REQUIRED_DATASETS[technique]
576+
if dataset_name not in selected_datasets
577+
)
578+
if problems:
579+
raise DatasetConstraintError(" ".join(problems))
580+
542581
def _get_technique_size_budgets(self) -> dict[WebInjectionTechnique, BoundedDatasetSize]:
543582
"""Return the selected techniques' configured caps."""
544583
budgets: dict[WebInjectionTechnique, BoundedDatasetSize] = {}

0 commit comments

Comments
 (0)