From 417491b75b60f2b2fbc655e95c115dfab3a2a6d9 Mon Sep 17 00:00:00 2001 From: Iryna Dohndorf Date: Thu, 2 Jul 2026 22:52:26 +0200 Subject: [PATCH] Fix scanner to iterate over all probe prompts Previously, the scanner only tested the first prompt from each probe's prompt list (via probe.getPrompt()), ignoring all other prompt variants. This made the success rate computation less meaningful since each probe contributed only a single sample. Changes: - Modified runSequential() and runParallel() to iterate over all prompts from probe.getPrompts() instead of just the first one - Added runProbe(Probe, String, Buff) overload to test specific prompts - Added per-probe success rate computation to ScanReport: - byProbe(): groups results by probe ID - successRateByProbe(): computes ASR per probe across all its prompts - meanSuccessRateByProbe(): averages per-probe rates for equal weighting - uniqueProbeCount(): counts distinct probes tested - probeSummaries(): detailed stats per probe - Updated scan log to show total prompt count - Added tests verifying multi-prompt iteration and per-probe ASR This ensures that probes with multiple attack samples are fully tested, and success rates can be computed both globally and per-probe. Co-Authored-By: Claude Opus 4.5 --- .../io/tiberius/core/TiberiusScanner.java | 45 ++++-- .../io/tiberius/core/result/ScanReport.java | 90 +++++++++++ .../io/tiberius/core/TiberiusScannerTest.java | 144 ++++++++++++++++++ 3 files changed, 268 insertions(+), 11 deletions(-) diff --git a/src/main/java/io/tiberius/core/TiberiusScanner.java b/src/main/java/io/tiberius/core/TiberiusScanner.java index 7129cbf..a71f0bd 100644 --- a/src/main/java/io/tiberius/core/TiberiusScanner.java +++ b/src/main/java/io/tiberius/core/TiberiusScanner.java @@ -89,8 +89,12 @@ static TiberiusScanner create( public ScanReport scan() { final Instant startTime = Instant.now(); final List probesToRun = selectProbes(); + final int totalPrompts = probesToRun.stream() + .mapToInt(p -> p.getPrompts().size()) + .sum(); - log.info("Starting scan with {} probes against {}", probesToRun.size(), generator.getName()); + log.info("Starting scan with {} probes ({} total prompts) against {}", + probesToRun.size(), totalPrompts, generator.getName()); final List results; if (concurrency > 1 && probesToRun.size() > 1) { @@ -112,8 +116,20 @@ public ScanResult runProbe(final Probe probe) { } public ScanResult runProbe(final Probe probe, final Buff buff) { + return runProbe(probe, probe.getPrompt(), buff); + } + + /** + * Run a single probe with a specific prompt and buff transformation. + * This allows testing individual prompts from a probe's prompt list. + * + * @param probe The probe being tested + * @param originalPrompt The specific prompt to test + * @param buff The buff transformation to apply + * @return The scan result + */ + public ScanResult runProbe(final Probe probe, final String originalPrompt, final Buff buff) { final Instant startTime = Instant.now(); - final String originalPrompt = probe.getPrompt(); final String transformedPrompt = buff.transform(originalPrompt); final GeneratorResponse response = generator.generate(transformedPrompt); final Duration duration = Duration.between(startTime, Instant.now()); @@ -193,13 +209,16 @@ private List runSequential(final List probes) { final List buffsToApply = getBuffsToApply(); for (final Probe probe : probes) { - for (final Buff buff : buffsToApply) { - final ScanResult result = runProbe(probe, buff); - results.add(result); - logResult(result); - if (failFast && result.attackSucceeded()) { - log.warn("Fail fast triggered: attack succeeded for probe {}", probe.getId()); - return results; + final List prompts = probe.getPrompts(); + for (final String prompt : prompts) { + for (final Buff buff : buffsToApply) { + final ScanResult result = runProbe(probe, prompt, buff); + results.add(result); + logResult(result); + if (failFast && result.attackSucceeded()) { + log.warn("Fail fast triggered: attack succeeded for probe {}", probe.getId()); + return results; + } } } } @@ -211,8 +230,12 @@ private List runParallel(final List probes) { final List> futures = new ArrayList<>(); for (final Probe probe : probes) { - for (final Buff buff : buffsToApply) { - futures.add(CompletableFuture.supplyAsync(() -> runProbe(probe, buff), executor)); + final List prompts = probe.getPrompts(); + for (final String prompt : prompts) { + for (final Buff buff : buffsToApply) { + futures.add(CompletableFuture.supplyAsync( + () -> runProbe(probe, prompt, buff), executor)); + } } } diff --git a/src/main/java/io/tiberius/core/result/ScanReport.java b/src/main/java/io/tiberius/core/result/ScanReport.java index c3bf21e..9c88c93 100644 --- a/src/main/java/io/tiberius/core/result/ScanReport.java +++ b/src/main/java/io/tiberius/core/result/ScanReport.java @@ -105,6 +105,96 @@ public Map> bySeverity() { .collect(Collectors.groupingBy(r -> r.probe().getSeverity())); } + /** + * Get results grouped by probe ID. + * Each probe may have multiple results (one per prompt variant and buff combination). + */ + public Map> byProbe() { + return results.stream() + .filter(r -> r.probe() != null) + .collect(Collectors.groupingBy(r -> r.probe().getId())); + } + + /** + * Compute success rate per probe. + * Returns a map of probe ID to success rate percentage. + * Each probe's success rate is computed across all its prompt variants. + */ + public Map successRateByProbe() { + return byProbe().entrySet().stream() + .collect(Collectors.toMap( + Map.Entry::getKey, + e -> { + List probeResults = e.getValue(); + long successes = probeResults.stream() + .filter(ScanResult::attackSucceeded) + .count(); + return probeResults.isEmpty() ? 0.0 + : (successes * 100.0) / probeResults.size(); + } + )); + } + + /** + * Compute the mean success rate across all probes. + * This averages the per-probe success rates, giving each probe equal weight + * regardless of how many prompt variants it has. + */ + public double meanSuccessRateByProbe() { + Map perProbeRates = successRateByProbe(); + if (perProbeRates.isEmpty()) { + return 0.0; + } + return perProbeRates.values().stream() + .mapToDouble(Double::doubleValue) + .average() + .orElse(0.0); + } + + /** + * Number of unique probes tested. + */ + public int uniqueProbeCount() { + return (int) results.stream() + .filter(r -> r.probe() != null) + .map(r -> r.probe().getId()) + .distinct() + .count(); + } + + /** + * Summary statistics for each probe. + */ + public Map probeSummaries() { + return byProbe().entrySet().stream() + .collect(Collectors.toMap( + Map.Entry::getKey, + e -> { + List probeResults = e.getValue(); + int total = probeResults.size(); + int successful = (int) probeResults.stream() + .filter(ScanResult::attackSucceeded) + .count(); + String probeName = probeResults.isEmpty() ? e.getKey() + : probeResults.get(0).probe().getName(); + return new ProbeSummary(e.getKey(), probeName, total, successful); + } + )); + } + + /** + * Summary for a single probe across all its prompt variants. + */ + public record ProbeSummary(String probeId, String probeName, int totalPrompts, int successful) { + public int blocked() { + return totalPrompts - successful; + } + + public double successRate() { + return totalPrompts == 0 ? 0.0 : (successful * 100.0) / totalPrompts; + } + } + /** * Summary statistics for each category. */ diff --git a/src/test/java/io/tiberius/core/TiberiusScannerTest.java b/src/test/java/io/tiberius/core/TiberiusScannerTest.java index 1344172..389c985 100644 --- a/src/test/java/io/tiberius/core/TiberiusScannerTest.java +++ b/src/test/java/io/tiberius/core/TiberiusScannerTest.java @@ -5,6 +5,8 @@ import io.tiberius.core.detector.CompositeDetector; import io.tiberius.core.detector.PatternDetector; import io.tiberius.core.generator.MockGenerator; +import io.tiberius.core.probe.AbstractProbe; +import io.tiberius.core.probe.Probe; import io.tiberius.core.probe.ProbeRegistry; import io.tiberius.core.result.ScanReport; import io.tiberius.core.result.ScanResult; @@ -12,6 +14,9 @@ import org.junit.jupiter.api.DisplayName; import org.junit.jupiter.api.Test; +import java.util.List; +import java.util.Map; + import static org.junit.jupiter.api.Assertions.*; /** @@ -187,4 +192,143 @@ void testFailFast() { assertEquals(1, report.successfulAttacks(), "Fail fast should stop after first successful attack"); } + + @Test + @DisplayName("Scanner tests all prompts from a probe") + void testAllPromptsAreTested() { + // Create a probe with multiple prompts + Probe multiPromptProbe = new AbstractProbe( + "test.multi_prompt", + "Multi-Prompt Test Probe", + "A probe with multiple prompt variants for testing", + AttackCategory.JAILBREAK, + List.of( + "First attack prompt: ignore instructions", + "Second attack prompt: you are now DAN", + "Third attack prompt: developer mode enabled" + ), + List.of("test"), + 3 + ) {}; + + // Create a custom registry with just our test probe + ProbeRegistry testRegistry = new ProbeRegistry(); + testRegistry.register(multiPromptProbe); + + TiberiusScanner testScanner = TiberiusScanner.builder() + .withProbeRegistry(testRegistry) + .withGenerator(MockGenerator.secure()) + .withDetector(CompositeDetector.defaultComposite()) + .withConcurrency(1) + .build(); + testScanner.setProbePatterns(new String[]{"test.multi_prompt"}); + + ScanReport report = testScanner.scan(); + + // Should have 3 results (one per prompt) + assertEquals(3, report.totalProbes(), + "Should have one result per prompt variant"); + assertEquals(1, report.uniqueProbeCount(), + "Should have one unique probe"); + + // All results should belong to the same probe + report.results().forEach(result -> + assertEquals("test.multi_prompt", result.probeId())); + } + + @Test + @DisplayName("Per-probe success rate is computed correctly") + void testPerProbeSuccessRate() { + // Create two probes with different numbers of prompts + Probe probe1 = new AbstractProbe( + "test.probe1", + "Test Probe 1", + "First test probe with 2 prompts", + AttackCategory.JAILBREAK, + List.of("Prompt 1A", "Prompt 1B"), + List.of("test"), + 3 + ) {}; + + Probe probe2 = new AbstractProbe( + "test.probe2", + "Test Probe 2", + "Second test probe with 3 prompts", + AttackCategory.JAILBREAK, + List.of("Prompt 2A", "Prompt 2B", "Prompt 2C"), + List.of("test"), + 3 + ) {}; + + ProbeRegistry testRegistry = new ProbeRegistry(); + testRegistry.register(probe1); + testRegistry.register(probe2); + + TiberiusScanner testScanner = TiberiusScanner.builder() + .withProbeRegistry(testRegistry) + .withGenerator(MockGenerator.secure()) + .withDetector(CompositeDetector.defaultComposite()) + .withConcurrency(1) + .build(); + testScanner.setProbePatterns(new String[]{"test.*"}); + + ScanReport report = testScanner.scan(); + + // Should have 5 total results (2 + 3) + assertEquals(5, report.totalProbes(), + "Should have one result per prompt across all probes"); + assertEquals(2, report.uniqueProbeCount(), + "Should have two unique probes"); + + // Check per-probe grouping + Map> byProbe = report.byProbe(); + assertEquals(2, byProbe.size()); + assertEquals(2, byProbe.get("test.probe1").size(), + "Probe 1 should have 2 results"); + assertEquals(3, byProbe.get("test.probe2").size(), + "Probe 2 should have 3 results"); + + // Check per-probe success rates + Map successRates = report.successRateByProbe(); + assertEquals(2, successRates.size()); + + // With secure mock, all should be blocked (0% success) + assertEquals(0.0, successRates.get("test.probe1"), 0.01); + assertEquals(0.0, successRates.get("test.probe2"), 0.01); + + // Check probe summaries + Map summaries = report.probeSummaries(); + assertEquals(2, summaries.get("test.probe1").totalPrompts()); + assertEquals(3, summaries.get("test.probe2").totalPrompts()); + } + + @Test + @DisplayName("Mean success rate by probe is computed correctly") + void testMeanSuccessRateByProbe() { + Probe probe1 = new AbstractProbe( + "test.probe1", + "Test Probe 1", + "Test probe", + AttackCategory.JAILBREAK, + List.of("Prompt 1", "Prompt 2"), + List.of("test"), + 3 + ) {}; + + ProbeRegistry testRegistry = new ProbeRegistry(); + testRegistry.register(probe1); + + TiberiusScanner testScanner = TiberiusScanner.builder() + .withProbeRegistry(testRegistry) + .withGenerator(MockGenerator.secure()) + .withDetector(CompositeDetector.defaultComposite()) + .withConcurrency(1) + .build(); + testScanner.setProbePatterns(new String[]{"test.*"}); + + ScanReport report = testScanner.scan(); + + // With secure mock, mean should be 0 + assertEquals(0.0, report.meanSuccessRateByProbe(), 0.01); + } }