Changed aquireUni.. to aquiere so that it can be interrupted in case of sudden shutdown

1 files changed, 9 insertions, 5 deletions

diff --git a/src/main/java/com/it_jaros/jscanner/Scanner.java b/src/main/java/com/it_jaros/jscanner/Scanner.java
index eaf2cc9..7e52775 100644
--- a/src/main/java/com/it_jaros/jscanner/Scanner.java
+++ b/src/main/java/com/it_jaros/jscanner/Scanner.java
@@ -76,12 +76,16 @@ public class Scanner implements AutoCloseable {
             try {
                 while (queue.hasNext() && !cancelled) {
                     activeHostWorkers.acquire();
+                    // just in case while waiting something changed
+                    if (cancelled) {
+                        break;
+                    }
 
                     // here we are filling the completion service
                     // host <-> virtual thread
                     final String host = queue.next();
                     inPipeline.incrementAndGet();
-                    completionService.submit(() -> scanHostPorts(host, scan.getPorts(), scan));
+                    completionService.submit(() -> scanHostPorts(host, scan));
                 }
             } catch (InterruptedException e) {
                 Thread.currentThread().interrupt();
@@ -125,8 +129,8 @@ public class Scanner implements AutoCloseable {
         return scan.getResults();
     }
 
-    private ScanResult scanHostPorts(final String host, final String ports, final Scan scan) {
-        PortRange portRange = new PortRange(ports);
+    private ScanResult scanHostPorts(final String host, final Scan scan) {
+        PortRange portRange = new PortRange(scan.getPorts());
         BitSet openPorts = new BitSet(PortRange.MAX_PORT);
         BitSet filteredPorts = new BitSet(PortRange.MAX_PORT);
         HashMap<Integer, ServiceType> serviceTypes = new HashMap<>();
@@ -233,13 +237,13 @@ public class Scanner implements AutoCloseable {
         }
     }
 
-    private PortResult checkPort(String host, int port, Scan scan) {
+    private PortResult checkPort(String host, int port, Scan scan) throws InterruptedException {
         PortResult result = new PortResult();
         result.setPort(port);
         result.setState(PortState.UNKNOWN);
 
         // don't allow more sockets then specified
-        socketLimit.acquireUninterruptibly();
+        socketLimit.acquire();
         // count how many ports are concurrently checked
         scan.incSocketCounter();
         try (Socket socket = new Socket()) {