Use SocketChannel for better timeout control

1 files changed, 29 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 9f5ac52..5277327 100644
--- a/src/main/java/com/it_jaros/jscanner/Scanner.java
+++ b/src/main/java/com/it_jaros/jscanner/Scanner.java
@@ -3,6 +3,9 @@ package com.it_jaros.jscanner;
 import java.io.IOException;
 import java.io.InputStream;
 import java.net.*;
+import java.nio.channels.SelectionKey;
+import java.nio.channels.Selector;
+import java.nio.channels.SocketChannel;
 import java.nio.charset.StandardCharsets;
 import java.time.Duration;
 import java.util.*;
@@ -347,11 +350,32 @@ public class Scanner implements AutoCloseable {
 
             socketLimit.acquire();
             scan.socketStart();
-            try (Socket socket = new Socket()) {
-                socket.connect(new InetSocketAddress(host, port), timeoutInMillis);
-                result.setState(PortState.OPEN);
-                if (bannerRecognition) {
-                    result.setBanner(getBanner(socket));
+            try(
+                    SocketChannel socketChannel = SocketChannel.open();
+                    Selector isOpen = Selector.open()
+            ) {
+                socketChannel.configureBlocking(false);
+                SelectionKey isOpenKey = socketChannel.register(isOpen, SelectionKey.OP_CONNECT);
+                socketChannel.connect(new InetSocketAddress(host, port));
+                long deadlineNanos = System.nanoTime() + TimeUnit.MILLISECONDS.toNanos(timeoutInMillis);
+                long waitInMillis = Math.min(timeoutInMillis, 5);
+                while (isOpenKey.isValid() && !isOpenKey.isConnectable()) {
+                    if (isOpen.select(waitInMillis) == 1) {
+                        if (socketChannel.finishConnect()) {
+                            result.setState(PortState.OPEN);
+                            //if (bannerRecognition) {
+                            //    result.setBanner(getBanner(socket));
+                            //}
+                        } else {
+                            result.setState(PortState.FILTERED);
+                        }
+                        break;
+                    }
+
+                    if (deadlineNanos - System.nanoTime() < 0) {
+                        result.setState(PortState.FILTERED);
+                        break;
+                    }
                 }
             } catch (SocketTimeoutException ignored) {
                 result.setState(PortState.FILTERED);