diff --git a/src/main/java/de/rub/nds/scanner/core/execution/ProbeProgressCallback.java b/src/main/java/de/rub/nds/scanner/core/execution/ProbeProgressCallback.java new file mode 100644 index 0000000..ef9f9be --- /dev/null +++ b/src/main/java/de/rub/nds/scanner/core/execution/ProbeProgressCallback.java @@ -0,0 +1,52 @@ +/* + * Scanner Core - A Modular Framework for Probe Definition, Execution, and Result Analysis. + * + * Copyright 2017-2023 Ruhr University Bochum, Paderborn University, Technology Innovation Institute, and Hackmanit GmbH + * + * Licensed under Apache License, Version 2.0 + * http://www.apache.org/licenses/LICENSE-2.0.txt + */ +package de.rub.nds.scanner.core.execution; + +import de.rub.nds.scanner.core.probe.ScannerProbe; +import de.rub.nds.scanner.core.report.ScanReport; + +/** + * Callback interface for receiving probe execution progress updates. This interface allows external + * components to be notified when individual probes complete during a scan, enabling real-time + * progress monitoring and streaming of partial results. + * + * @param the type of scan report + * @param the type of state object used by probes + */ +@FunctionalInterface +public interface ProbeProgressCallback { + + /** + * Called when a probe has completed execution and merged its results into the report. + * + * @param probe the probe that completed execution + * @param report the scan report with the probe's results merged in + * @param completedProbes the number of probes that have completed so far + * @param totalProbes the total number of probes scheduled for this scan + */ + void onProbeCompleted( + ScannerProbe probe, + ReportT report, + int completedProbes, + int totalProbes); + + /** + * Creates a no-op callback that does nothing when probes complete. Useful as a default when no + * progress tracking is needed. + * + * @param the type of scan report + * @param the type of state object + * @return a callback that performs no operations + */ + static ProbeProgressCallback noOp() { + return (probe, report, completedProbes, totalProbes) -> { + // No operation + }; + } +} diff --git a/src/main/java/de/rub/nds/scanner/core/execution/Scanner.java b/src/main/java/de/rub/nds/scanner/core/execution/Scanner.java index 3222639..1aa896c 100644 --- a/src/main/java/de/rub/nds/scanner/core/execution/Scanner.java +++ b/src/main/java/de/rub/nds/scanner/core/execution/Scanner.java @@ -49,6 +49,9 @@ public abstract class Scanner< private final List afterList; private final boolean fillProbeListsAtScanStart; + // Optional callback for probe progress updates + private ProbeProgressCallback progressCallback = ProbeProgressCallback.noOp(); + /** * Creates a new scanner instance. * @@ -131,6 +134,18 @@ protected List getGuidelines() { return List.of(); } + /** + * Sets the progress callback to be invoked when probes complete during scanning. This allows + * external components to receive real-time updates about scan progress and partial results. + * + * @param progressCallback the callback to invoke on probe completion, or null to disable + * callbacks + */ + public void setProgressCallback(ProbeProgressCallback progressCallback) { + this.progressCallback = + progressCallback != null ? progressCallback : ProbeProgressCallback.noOp(); + } + /** * Performs the scan. It will take care of all the necessary steps to perform a scan, including * filling the probe list by calling {@link #fillProbeLists}, checking the scan prerequisites by @@ -166,6 +181,8 @@ public ReportT scan() { scanJob, executorConfig.getParallelProbes(), "ScannerProbeExecutor " + report.getRemoteName())) { + // Set the progress callback on the executor + scanJobExecutor.setProgressCallback(progressCallback); ProgressSpinner.startSpinnerTask("Executing:"); report.setScanStartTime(System.currentTimeMillis()); scanJobExecutor.execute(report); diff --git a/src/main/java/de/rub/nds/scanner/core/execution/ThreadedScanJobExecutor.java b/src/main/java/de/rub/nds/scanner/core/execution/ThreadedScanJobExecutor.java index a27e2fa..17abfc5 100644 --- a/src/main/java/de/rub/nds/scanner/core/execution/ThreadedScanJobExecutor.java +++ b/src/main/java/de/rub/nds/scanner/core/execution/ThreadedScanJobExecutor.java @@ -60,6 +60,9 @@ public class ThreadedScanJobExecutor< private volatile int probeCount; private final AtomicInteger finishedProbes = new AtomicInteger(0); + // Callback for probe progress updates (optional) + private ProbeProgressCallback progressCallback = ProbeProgressCallback.noOp(); + /** * Creates a new ThreadedScanJobExecutor with a custom thread pool. * @@ -103,6 +106,18 @@ public ThreadedScanJobExecutor( this.futureResults = new LinkedList<>(); } + /** + * Sets the progress callback to be invoked when probes complete. This allows external + * components to receive real-time updates about scan progress. + * + * @param progressCallback the callback to invoke on probe completion, or null to disable + * callbacks + */ + public void setProgressCallback(ProbeProgressCallback progressCallback) { + this.progressCallback = + progressCallback != null ? progressCallback : ProbeProgressCallback.noOp(); + } + /** * Executes the scan job by running probes concurrently and populating the report with results. * This method manages probe dependencies and ensures probes are executed in the correct order. @@ -146,7 +161,8 @@ private void executeProbesTillNoneCanBeExecuted(ReportT report) throws Interrupt try { probeResult = result.get(); LOGGER.info( - "[{}/{}] {} probe executed", + "[{}] [{}/{}] {} probe executed", + report.getRemoteName(), String.format("%2d", currentFinishedProbes), String.format("%2d", probeCount), probeResult.getType().getName()); @@ -157,6 +173,14 @@ private void executeProbesTillNoneCanBeExecuted(ReportT report) throws Interrupt finishedFutures.add(result); probeResult.merge(report); report.markProbeAsExecuted(probeResult); + + // Notify progress callback + try { + progressCallback.onProbeCompleted( + probeResult, report, currentFinishedProbes, probeCount); + } catch (Exception e) { + LOGGER.warn("Progress callback threw exception, continuing scan", e); + } } } futureResults.removeAll(finishedFutures); diff --git a/src/test/java/de/rub/nds/scanner/core/execution/ScannerTest.java b/src/test/java/de/rub/nds/scanner/core/execution/ScannerTest.java index 8d4908b..60c3622 100644 --- a/src/test/java/de/rub/nds/scanner/core/execution/ScannerTest.java +++ b/src/test/java/de/rub/nds/scanner/core/execution/ScannerTest.java @@ -431,4 +431,79 @@ public void testInterruptedScan() throws InterruptedException { assertTrue(testThread.isInterrupted() || !testThread.isAlive()); } + + @Test + public void testSetProgressCallbackWithValidCallback() { + try (TestScanner scanner = new TestScanner(executorConfig)) { + ProbeProgressCallback callback = + (probe, report, completedProbes, totalProbes) -> { + // Custom callback logic + }; + + assertDoesNotThrow(() -> scanner.setProgressCallback(callback)); + } + } + + @Test + public void testSetProgressCallbackWithNull() { + try (TestScanner scanner = new TestScanner(executorConfig)) { + // Should not throw when setting null callback + assertDoesNotThrow(() -> scanner.setProgressCallback(null)); + + // Scanner should still work with null callback + TestReport report = scanner.scan(); + assertNotNull(report); + } + } + + @Test + public void testProgressCallbackIsInvoked() { + List probeList = new ArrayList<>(); + probeList.add(new TestProbe(new TestProbeType("probe1"))); + probeList.add(new TestProbe(new TestProbeType("probe2"))); + + List afterList = new ArrayList<>(); + + final int[] callbackInvocations = {0}; + ProbeProgressCallback callback = + (probe, report, completedProbes, totalProbes) -> { + callbackInvocations[0]++; + assertTrue(completedProbes <= totalProbes); + assertTrue(completedProbes > 0); + }; + + try (TestScanner scanner = new TestScanner(executorConfig, probeList, afterList)) { + scanner.setProgressCallback(callback); + scanner.scan(); + } + + // Callback should have been invoked for each probe + assertTrue(callbackInvocations[0] > 0, "Progress callback was not invoked"); + } + + @Test + public void testProgressCallbackReportsCorrectProbeCount() { + List probeList = new ArrayList<>(); + int expectedProbeCount = 3; + for (int i = 0; i < expectedProbeCount; i++) { + probeList.add(new TestProbe(new TestProbeType("probe" + i))); + } + + List afterList = new ArrayList<>(); + + final List completedCounts = new ArrayList<>(); + ProbeProgressCallback callback = + (probe, report, completedProbes, totalProbes) -> { + assertEquals(expectedProbeCount, totalProbes); + completedCounts.add(completedProbes); + }; + + try (TestScanner scanner = new TestScanner(executorConfig, probeList, afterList)) { + scanner.setProgressCallback(callback); + scanner.scan(); + } + + // Verify callback was invoked for each probe + assertEquals(expectedProbeCount, completedCounts.size()); + } } diff --git a/src/test/java/de/rub/nds/scanner/core/execution/ThreadedScanJobExecutorTest.java b/src/test/java/de/rub/nds/scanner/core/execution/ThreadedScanJobExecutorTest.java index 59d5a05..0a21045 100644 --- a/src/test/java/de/rub/nds/scanner/core/execution/ThreadedScanJobExecutorTest.java +++ b/src/test/java/de/rub/nds/scanner/core/execution/ThreadedScanJobExecutorTest.java @@ -491,4 +491,156 @@ public void testMultipleExtractedValueContainersOfSameType() throws InterruptedE assertEquals(4, container.getExtractedValueList().size()); } } + + @Test + public void testSetProgressCallbackInvokesAfterEachProbe() throws InterruptedException { + List probeList = new ArrayList<>(); + int expectedProbes = 3; + for (int i = 0; i < expectedProbes; i++) { + probeList.add(new TestProbe(new TestProbeType("probe" + i))); + } + + List afterList = new ArrayList<>(); + ScanJob scanJob = + new ScanJob<>(probeList, afterList); + + try (ThreadedScanJobExecutor executor = + new ThreadedScanJobExecutor<>(executorConfig, scanJob, 1, "Test")) { + + final List completedCounts = new ArrayList<>(); + final List totalCounts = new ArrayList<>(); + + ProbeProgressCallback callback = + (probe, report, completedProbes, totalProbes) -> { + completedCounts.add(completedProbes); + totalCounts.add(totalProbes); + }; + + executor.setProgressCallback(callback); + + TestReport report = new TestReport(); + executor.execute(report); + + // Callback should have been invoked for each probe + assertEquals(expectedProbes, completedCounts.size()); + + // Each callback should report the correct total + for (int total : totalCounts) { + assertEquals(expectedProbes, total); + } + + // Verify completed counts are in ascending order + for (int i = 0; i < completedCounts.size(); i++) { + assertTrue(completedCounts.get(i) > 0); + assertTrue(completedCounts.get(i) <= expectedProbes); + } + } + } + + @Test + public void testSetProgressCallbackWithNull() throws InterruptedException { + List probeList = List.of(new TestProbe(new TestProbeType("probe1"))); + List afterList = new ArrayList<>(); + + ScanJob scanJob = + new ScanJob<>(probeList, afterList); + + try (ThreadedScanJobExecutor executor = + new ThreadedScanJobExecutor<>(executorConfig, scanJob, 1, "Test")) { + + // Should not throw when setting null callback + assertDoesNotThrow(() -> executor.setProgressCallback(null)); + + // Executor should still work with null callback + TestReport report = new TestReport(); + executor.execute(report); + + assertEquals(1, report.getExecutedProbeTypes().size()); + } + } + + @Test + public void testProgressCallbackExceptionDoesNotStopExecution() throws InterruptedException { + TestProbe probe1 = new TestProbe(new TestProbeType("probe1")); + TestProbe probe2 = new TestProbe(new TestProbeType("probe2")); + List probeList = new ArrayList<>(); + probeList.add(probe1); + probeList.add(probe2); + + List afterList = new ArrayList<>(); + ScanJob scanJob = + new ScanJob<>(probeList, afterList); + + try (ThreadedScanJobExecutor executor = + new ThreadedScanJobExecutor<>(executorConfig, scanJob, 1, "Test")) { + + final int[] callbackInvocations = {0}; + + ProbeProgressCallback callback = + (probe, report, completedProbes, totalProbes) -> { + callbackInvocations[0]++; + throw new RuntimeException("Test exception in callback"); + }; + + executor.setProgressCallback(callback); + + TestReport report = new TestReport(); + // Should not throw despite callback exceptions + assertDoesNotThrow(() -> executor.execute(report)); + + // All probes should still execute and be marked as executed (not failed) + assertEquals(2, report.getExecutedProbeTypes().size()); + assertTrue( + report.getExecutedProbeTypes().contains(probe1.getType()), + "probe1 should be marked as executed"); + assertTrue( + report.getExecutedProbeTypes().contains(probe2.getType()), + "probe2 should be marked as executed"); + + // Verify probes actually ran (not just marked) + assertTrue(probe1.wasExecuted(), "probe1 should have been executed"); + assertTrue(probe2.wasExecuted(), "probe2 should have been executed"); + + // Verify no probes were marked as unexecuted/failed + assertTrue( + report.getUnexecutedProbeTypes().isEmpty(), + "No probes should be marked as unexecuted"); + + // Callback should have been invoked for each probe despite throwing + assertEquals(2, callbackInvocations[0]); + } + } + + @Test + public void testProgressCallbackReceivesCorrectProbeReference() throws InterruptedException { + TestProbe probe1 = new TestProbe(new TestProbeType("probe1")); + TestProbe probe2 = new TestProbe(new TestProbeType("probe2")); + + List probeList = Arrays.asList(probe1, probe2); + List afterList = new ArrayList<>(); + + ScanJob scanJob = + new ScanJob<>(probeList, afterList); + + try (ThreadedScanJobExecutor executor = + new ThreadedScanJobExecutor<>(executorConfig, scanJob, 1, "Test")) { + + final List probeNames = new ArrayList<>(); + + ProbeProgressCallback callback = + (probe, report, completedProbes, totalProbes) -> { + probeNames.add(probe.getType().getName()); + }; + + executor.setProgressCallback(callback); + + TestReport report = new TestReport(); + executor.execute(report); + + // Verify callback received both probes + assertEquals(2, probeNames.size()); + assertTrue(probeNames.contains("probe1")); + assertTrue(probeNames.contains("probe2")); + } + } }