1 files changed, 28 insertions, 8 deletions

diff --git a/src/main/java/com/it_jaros/jscanner/Scanner.java b/src/main/java/com/it_jaros/jscanner/Scanner.java
index 2c0a36c..4d951f5 100644
--- a/src/main/java/com/it_jaros/jscanner/Scanner.java
+++ b/src/main/java/com/it_jaros/jscanner/Scanner.java
@@ -17,10 +17,13 @@ import java.util.concurrent.locks.LockSupport;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
-public class Scanner {
+public class Scanner implements AutoCloseable {
 
     private static final int READ_BUFFER_SIZE = 1024;
 
+    private ArrayList<ScanResult> scanResult;
+    private volatile boolean cancelled = false;
+
     private final Counter counter = new Counter();
     private final ProgressBar progressBar = new ProgressBar();
     private final Semaphore socketLimit;
@@ -75,7 +78,7 @@ public class Scanner {
     }
 
     public List<ScanResult> scanHosts(Stream<String> hosts, String ports) throws IOException {
-        List<ScanResult> results = new ArrayList<>();
+        this.scanResult = new ArrayList<>();
         progressBar.start(quiet);
         try (ExecutorService executor = Executors.newVirtualThreadPerTaskExecutor()) {
             CompletionService<ScanResult> completionService = new ExecutorCompletionService<>(executor);
@@ -83,9 +86,9 @@ public class Scanner {
 
             int hostsToProcess = hostsInQueue.size();
             int activeHostWorkers = 0;
-            while (hostsToProcess > 0) {
+            while (hostsToProcess > 0 && !cancelled) {
                 // process queue
-                while (!hostsInQueue.isEmpty() && activeHostWorkers <= maxHostsLimit) {
+                while (!hostsInQueue.isEmpty() && activeHostWorkers <= maxHostsLimit && !cancelled) {
                     activeHostWorkers++;
                     final String host = hostsInQueue.poll();
                     completionService.submit(() -> scanHost(host, ports));
@@ -98,13 +101,17 @@ public class Scanner {
                     }
                     activeHostWorkers--;
                     hostsToProcess--;
-                    results.add(finishedHost.get());
+                    this.scanResult.add(finishedHost.get());
                 } catch (InterruptedException e) {
                     Thread.currentThread().interrupt();
                 } catch (ExecutionException e) {
                     System.out.printf("%s -> %s%n", e.getClass().getSimpleName(), e.getMessage());
                 }
             }
+
+            if (cancelled) {
+                executor.shutdown();
+            }
         }
 
         progressBar.stop(quiet);
@@ -113,7 +120,7 @@ public class Scanner {
         System.out.println("Stats:");
         System.out.println("Peak concurrent connects: " + counter.max());
         System.out.println("--------------------");
-        return results;
+        return scanResult;
     }
 
     private ScanResult scanHost(String host, String ports) {
@@ -138,8 +145,8 @@ public class Scanner {
             List<Throwable> errors = new ArrayList<>();
             int activePerHostWorkers = 0;
             int portsToProcess = portRange.getTotal();
-            while (portsToProcess > 0) {
-                while (portRange.hasNext() && activePerHostWorkers <= maxWorkersPerHost) {
+            while (portsToProcess > 0 && !cancelled) {
+                while (!cancelled && portRange.hasNext() && activePerHostWorkers <= maxWorkersPerHost) {
                     activePerHostWorkers++;
                     final int currentPort = portRange.next();
                     completionService.submit(() -> {
@@ -174,6 +181,10 @@ public class Scanner {
                 }
             }
 
+            if (!cancelled) {
+                executor.shutdown();
+            }
+
             if (!errors.isEmpty()) {
                 System.out.printf("Errors happened during scan of host %s%nErrors:%s -> %s", host, errors.size(), errors);
             }
@@ -240,4 +251,13 @@ public class Scanner {
         }
         return null;
     }
+
+    public void cancel() {
+        cancelled = true;
+    }
+
+    @Override
+    public void close() throws Exception {
+        cancel();
+    }
 }