Added timeout to cli options

1 files changed, 32 insertions, 22 deletions

diff --git a/src/main/java/com/it_jaros/networkScanner/Scanner.java b/src/main/java/com/it_jaros/networkScanner/Scanner.java
index b3f7767..eb3af29 100644
--- a/src/main/java/com/it_jaros/networkScanner/Scanner.java
+++ b/src/main/java/com/it_jaros/networkScanner/Scanner.java
@@ -4,6 +4,7 @@ import java.io.IOException;
 import java.net.ConnectException;
 import java.net.InetSocketAddress;
 import java.net.Socket;
+import java.net.SocketTimeoutException;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Queue;
@@ -17,14 +18,18 @@ import java.util.concurrent.locks.LockSupport;
 
 public class Scanner {
                 
-    private Counter counter = new Counter();
-    private Semaphore socketCap;
+    private final Counter counter = new Counter();
+    private final Semaphore socketCap;
+    private final int timeoutInMillis;
     
-    public Scanner(int socketLimit) {
+    public Scanner(int socketLimit, int timeoutInMillis) {
+        this.timeoutInMillis = timeoutInMillis;
         socketCap = new Semaphore(socketLimit);
     }
 
-    public Scanner() {}
+    public Scanner() {
+        this(1000, 1000);
+    }
 
     public List<ScanResult> scanPorts(ScanOptions target) {
         return scanPorts(target.targets(), target.delayInMillis());
@@ -48,26 +53,31 @@ public class Scanner {
         for (int port = 1; port < 65536; port++) {
             final int currentPort = port;
             executor.submit(() -> {
-                if (delayInNanos > 0) {
-                    long mySlot = nextStartNanos.getAndAdd(delayInNanos);
-                    long wait = mySlot - System.nanoTime();
-                    if (wait > 0)
-                        LockSupport.parkNanos(wait);
-                }
-                socketCap.acquireUninterruptibly();
-                counter.inc();
-                try (Socket socket = new Socket()) {
-                    socket.connect(new InetSocketAddress(target, currentPort), 10000);
-                    openPorts.add(currentPort);
-                } catch (ConnectException ignored) {
-                } catch (IOException e) {
-                    System.out.println(e.getClass().getSimpleName() + ": " + e.getMessage());
-                } finally {
-                    counter.dec();
-                    socketCap.release();
-                }
+                scanPortsAsync(target, delayInNanos, nextStartNanos, openPorts, currentPort);
             });
         }
         return new ScanResult(target, openPorts);
     }
+
+    private void scanPortsAsync(String target, long delayInNanos, AtomicLong nextStartNanos, Queue<Integer> openPorts,
+            int currentPort) {
+        if (delayInNanos > 0) {
+            long mySlot = nextStartNanos.getAndAdd(delayInNanos);
+            long wait = mySlot - System.nanoTime();
+            if (wait > 0)
+                LockSupport.parkNanos(wait);
+        }
+        socketCap.acquireUninterruptibly();
+        counter.inc();
+        try (Socket socket = new Socket()) {
+            socket.connect(new InetSocketAddress(target, currentPort), timeoutInMillis);
+            openPorts.add(currentPort);
+        } catch (ConnectException | SocketTimeoutException ignored) {
+        } catch (IOException e) {
+            System.out.println(e.getClass().getSimpleName() + ": " + e.getMessage());
+        } finally {
+            counter.dec();
+            socketCap.release();
+        }
+    }
 }