Made ScanUdpPortTask also non blocking analog to ScanPortTask

2 files changed, 71 insertions, 31 deletions

diff --git a/src/main/java/com/it_jaros/jns/scan/engine/ScanUdpPortTask.java b/src/main/java/com/it_jaros/jns/scan/engine/ScanUdpPortTask.java
index 0c1b9e0..441fd08 100644
--- a/src/main/java/com/it_jaros/jns/scan/engine/ScanUdpPortTask.java
+++ b/src/main/java/com/it_jaros/jns/scan/engine/ScanUdpPortTask.java
@@ -12,6 +12,8 @@ import java.lang.foreign.MemorySegment;
 import java.net.InetAddress;
 import java.util.concurrent.Callable;
 import java.util.concurrent.Semaphore;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.locks.LockSupport;
 
 public class ScanUdpPortTask implements Callable<PortResult> {
 
@@ -23,7 +25,10 @@ public class ScanUdpPortTask implements Callable<PortResult> {
     private final Scan scan;
     private final Semaphore socketLimit;
     private final int port;
-    private final long timeoutInMillis;
+    private final long timeoutInNanos;
+
+    private final long waitInMillis = 100;
+    private final long waitInNanos = TimeUnit.MILLISECONDS.toNanos(waitInMillis);
 
     ScanUdpPortTask(int port, ScanExecutionContext context, ScanHostContext hostContext) {
         this.inetAddress = hostContext.hostAddress();
@@ -32,7 +37,7 @@ public class ScanUdpPortTask implements Callable<PortResult> {
         this.portScanRateLimiter = hostContext.portScanRateLimiter();
         this.scan = context.scan();
         this.socketLimit = context.socketLimit();
-        this.timeoutInMillis = context.scanOptions().timeoutInMillis();
+        this.timeoutInNanos = TimeUnit.MILLISECONDS.toNanos(context.scanOptions().timeoutInMillis());
     }
 
     @Override
@@ -75,31 +80,58 @@ public class ScanUdpPortTask implements Callable<PortResult> {
             final int family = inetAddress.getAddress().length == 16 ? LinuxUdpSocketAPI.AF_INET6 : LinuxUdpSocketAPI.AF_INET;
             int fd = linuxUdpSocketAPI.socket(family);
             try {
-                linuxUdpSocketAPI.setReceiveTimeout(fd, timeoutInMillis);
+                linuxUdpSocketAPI.setNonBlocking(fd);
                 linuxUdpSocketAPI.connect(fd, destination);
                 linuxUdpSocketAPI.sendEmptyDatagram(fd);
 
                 MemorySegment recvBuffer = arena.allocate(RECV_BUFFER_SIZE);
-                try {
-                    long bytesRead = linuxUdpSocketAPI.recv(fd, recvBuffer);
-                    if (bytesRead >= 0) {
-                        return PortState.OPEN;
-                    }
-                } catch (RawLinuxSocketException e) {
-                    int errno = e.errno;
-                    if (errno == LinuxUdpSocketAPI.ECONNREFUSED) {
-                        return PortState.CLOSED;
-                    }
-                    if (errno == LinuxUdpSocketAPI.EAGAIN) {
-                        return PortState.FILTERED;
-                    }
-                    throw new IOException("recv() failed: errno=" + errno);
-                }
+                return recvUntilDeadline(linuxUdpSocketAPI, fd, recvBuffer);
             } finally {
                 linuxUdpSocketAPI.close(fd);
             }
         }
+    }
 
-        return PortState.UNKNOWN;
+    /**
+     * Polls the non-blocking socket until the deadline.
+     * The native recv() call never blocks, so it cannot pin
+     * the virtual thread to its carrier thread; the wait between
+     * polls uses LockSupport.parkNanos, which is safe for virtual threads
+     */
+    private PortState recvUntilDeadline(LinuxUdpSocketAPI socketAPI, int fd, MemorySegment buffer) throws IOException {
+        final long deadlineNanos = System.nanoTime() + timeoutInNanos;
+        while (!cancelledToken.isCancelled()) {
+            final long remainingNanos = deadlineNanos - System.nanoTime();
+            if (remainingNanos <= 0) {
+                break;
+            }
+
+            long bytesRead;
+            try {
+                bytesRead = socketAPI.recv(fd, buffer);
+            } catch (RawLinuxSocketException e) {
+                final int errno = e.errno;
+                if (errno == LinuxUdpSocketAPI.ECONNREFUSED) {
+                    return PortState.CLOSED;
+                }
+                if (errno == LinuxUdpSocketAPI.EAGAIN || errno == LinuxUdpSocketAPI.EWOULDBLOCK) {
+                    // no datagram has arrived yet; keep polling
+                    bytesRead = -1;
+                } else {
+                    throw new IOException("recv() failed: errno=" + errno);
+                }
+            }
+
+            if (bytesRead >= 0) {
+                return PortState.OPEN;
+            }
+
+            LockSupport.parkNanos(Math.min(waitInNanos, remainingNanos));
+        }
+
+        if (cancelledToken.isCancelled()) {
+            return PortState.UNKNOWN;
+        }
+        return PortState.FILTERED;
     }
 }
diff --git a/src/main/java/com/it_jaros/jns/scan/engine/platform/linux/LinuxUdpSocketAPI.java b/src/main/java/com/it_jaros/jns/scan/engine/platform/linux/LinuxUdpSocketAPI.java
index 3f05b67..bfe70d4 100644
--- a/src/main/java/com/it_jaros/jns/scan/engine/platform/linux/LinuxUdpSocketAPI.java
+++ b/src/main/java/com/it_jaros/jns/scan/engine/platform/linux/LinuxUdpSocketAPI.java
@@ -14,10 +14,12 @@ public class LinuxUdpSocketAPI {
     public static final int AF_INET = 2;
     public static final int AF_INET6 = 10;
     public static final int EAGAIN = 11;
+    public static final int EWOULDBLOCK = 11;
     public static final int ECONNREFUSED = 111;
     public static final int SOCK_DGRAM = 2;
-    public static final int SOL_SOCKET = 1;
-    public static final int SO_RCVTIMEO = 20;
+    public static final int F_GETFL = 3;
+    public static final int F_SETFL = 4;
+    public static final int O_NONBLOCK = 2048;
 
     private static final Linker LINKER = Linker.nativeLinker();
     private static final SymbolLookup LOOKUP = LINKER.defaultLookup();
@@ -79,14 +81,22 @@ public class LinuxUdpSocketAPI {
             CAPTURE_ERRNO
     );
 
-    private static final MethodHandle SETSOCKOPT = LINKER.downcallHandle(
-            LOOKUP.findOrThrow("setsockopt"),
+    private static final MethodHandle FCTL_GETFL = LINKER.downcallHandle(
+            LOOKUP.findOrThrow("fcntl"),
             FunctionDescriptor.of(
                     ValueLayout.JAVA_INT,
                     ValueLayout.JAVA_INT,
+                    ValueLayout.JAVA_INT
+            ),
+            CAPTURE_ERRNO
+    );
+
+    private static final MethodHandle FCTL_SETFL = LINKER.downcallHandle(
+            LOOKUP.findOrThrow("fcntl"),
+            FunctionDescriptor.of(
+                    ValueLayout.JAVA_INT,
                     ValueLayout.JAVA_INT,
                     ValueLayout.JAVA_INT,
-                    ValueLayout.ADDRESS,
                     ValueLayout.JAVA_INT
             ),
             CAPTURE_ERRNO
@@ -138,6 +148,11 @@ public class LinuxUdpSocketAPI {
         return this.<Integer>wrapNativeMethodCall(SOCKET, family, SOCK_DGRAM, 0);
     }
 
+    public void setNonBlocking(int fd) throws IOException {
+        int flags = this.<Integer>wrapNativeMethodCall(FCTL_GETFL, fd, F_GETFL);
+        wrapNativeMethodCall(FCTL_SETFL, fd, F_SETFL, flags | O_NONBLOCK);
+    }
+
     public void connect(int fd, MemorySegment destination) throws IOException {
         wrapNativeMethodCall(CONNECT, fd, destination, (int) destination.byteSize());
     }
@@ -151,13 +166,6 @@ public class LinuxUdpSocketAPI {
         return this.<Long>wrapNativeMethodCall(RECV, fd, buffer, buffer.byteSize(), 0);
     }
 
-    public void setReceiveTimeout(int fd, long timeoutInMillis) throws IOException {
-        MemorySegment timeval = arena.allocate(16);  // struct timeval (two longs)
-        timeval.set(ValueLayout.JAVA_LONG, 0L, timeoutInMillis / 1000L);
-        timeval.set(ValueLayout.JAVA_LONG, 8L, (timeoutInMillis % 1000L) * 1000L);
-        wrapNativeMethodCall(SETSOCKOPT, fd, SOL_SOCKET, SO_RCVTIMEO, timeval, 16);
-    }
-
     public void close(int fd) throws IOException {
         wrapNativeMethodCall(CLOSE, fd);
     }