Merge pull request #23 from tls-attacker/handle-interrupts

Handle interrupts instead of ignoring them
tls-attacker · Nov 16, 2023 · 4cad651 · 4cad651
2 parents 905d1d0 + cd8bbe2
commit 4cad651
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 32 deletions.
diff --git a/src/main/java/de/rub/nds/scanner/core/execution/ScanJobExecutor.java b/src/main/java/de/rub/nds/scanner/core/execution/ScanJobExecutor.java
@@ -12,7 +12,7 @@
 
 public abstract class ScanJobExecutor<ReportT extends ScanReport> {
 
-    public abstract void execute(ReportT report);
+    public abstract void execute(ReportT report) throws InterruptedException;
 
     public abstract void shutdown();
 }
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
@@ -149,11 +149,20 @@ public ReportT scan() {
         // Scan Execution
         LOGGER.debug("Starting scan execution");
         ScanJob<ReportT, ProbeT, AfterProbeT, StateT> scanJob = new ScanJob<>(probeList, afterList);
-        ThreadedScanJobExecutor<ReportT, ProbeT, AfterProbeT, StateT> executor =
+        try (ThreadedScanJobExecutor<ReportT, ProbeT, AfterProbeT, StateT> scanJobExecutor =
                 new ThreadedScanJobExecutor<>(
-                        executorConfig, scanJob, executorConfig.getParallelProbes(), "");
-        report.setScanStartTime(System.currentTimeMillis());
-        executor.execute(report);
+                        executorConfig,
+                        scanJob,
+                        executorConfig.getParallelProbes(),
+                        "ScannerProbeExecutor")) {
+            report.setScanStartTime(System.currentTimeMillis());
+            scanJobExecutor.execute(report);
+        } catch (InterruptedException e) {
+            LOGGER.warn("Scan execution interrupted");
+            report.setScanEndTime(System.currentTimeMillis());
+            Thread.currentThread().interrupt();
+            return report;
+        }
         LOGGER.debug("Scan execution complete");
 
         // Rating
@@ -181,7 +190,6 @@ public ReportT scan() {
         LOGGER.debug(
                 "Scan finished, took a total of {} milliseconds",
                 report.getScanEndTime() - report.getScanStartTime());
-        executor.shutdown();
 
         // Serialize report to file
         if (executorConfig.isWriteReportToFile()) {

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
@@ -20,7 +20,6 @@
 import java.util.List;
 import java.util.Observable;
 import java.util.Observer;
-import java.util.concurrent.CancellationException;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Future;
 import java.util.concurrent.Semaphore;
@@ -33,7 +32,7 @@ public class ThreadedScanJobExecutor<
                 ProbeT extends ScannerProbe<ReportT, StateT>,
                 AfterProbeT extends AfterProbe<ReportT>,
                 StateT>
-        extends ScanJobExecutor<ReportT> implements Observer {
+        extends ScanJobExecutor<ReportT> implements Observer, AutoCloseable {
 
     private static final Logger LOGGER = LogManager.getLogger();
 
@@ -75,7 +74,7 @@ public ThreadedScanJobExecutor(
     }
 
     @Override
-    public void execute(ReportT report) {
+    public void execute(ReportT report) throws InterruptedException {
         notScheduledTasks = new ArrayList<>(scanJob.getProbeList());
         report.addObserver(this);
 
@@ -101,31 +100,24 @@ private void checkForExecutableProbes(ReportT report) {
         update(report, null);
     }
 
-    private void executeProbesTillNoneCanBeExecuted(ReportT report) {
+    private void executeProbesTillNoneCanBeExecuted(ReportT report) throws InterruptedException {
         boolean probesQueued = true;
         while (probesQueued) {
             // handle all Finished Results
             List<Future<ScannerProbe<ReportT, StateT>>> finishedFutures = new LinkedList<>();
             for (Future<ScannerProbe<ReportT, StateT>> result : futureResults) {
                 if (result.isDone()) {
+                    ScannerProbe<ReportT, StateT> probeResult = null;
                     try {
-                        ScannerProbe<ReportT, StateT> probeResult = result.get();
-                        LOGGER.info(probeResult.getType().getName() + " probe executed");
-                        finishedFutures.add(result);
-                        report.markProbeAsExecuted(result.get());
-                        probeResult.merge(report);
-                    } catch (InterruptedException | ExecutionException ex) {
-                        LOGGER.error(
-                                "Encountered an exception before we could merge the result. Killing the task.",
-                                ex);
-                        result.cancel(true);
-                        finishedFutures.add(result);
-                    } catch (CancellationException ex) {
-                        LOGGER.info(
-                                "Could not retrieve a task because it was cancelled after {} milliseconds",
-                                config.getProbeTimeout());
-                        finishedFutures.add(result);
+                        probeResult = result.get();
+                        LOGGER.info("{} probe executed", probeResult.getType().getName());
+                    } catch (ExecutionException e) {
+                        LOGGER.error("Some probe execution failed", e);
+                        throw new RuntimeException(e);
                     }
+                    finishedFutures.add(result);
+                    report.markProbeAsExecuted(probeResult);
+                    probeResult.merge(report);
                 }
             }
             futureResults.removeAll(finishedFutures);
@@ -135,12 +127,8 @@ private void executeProbesTillNoneCanBeExecuted(ReportT report) {
                 // nothing can be executed anymore
                 probesQueued = false;
             } else {
-                try {
-                    // wait for at least one probe to finish executing before checking again
-                    semaphore.acquire();
-                } catch (Exception e) {
-                    LOGGER.info("Interrupted while waiting for probe execution");
-                }
+                // wait for at least one probe to finish executing before checking again
+                semaphore.acquire();
             }
         }
     }
@@ -188,6 +176,11 @@ private void executeAfterProbes(ReportT report) {
         LOGGER.debug("Finished analysis");
     }
 
+    @Override
+    public void close() {
+        shutdown();
+    }
+
     @Override
     public void shutdown() {
         executor.shutdown();