1 files changed, 32 insertions, 13 deletions
diff --git a/src/main/java/com/it_jaros/jscanner/Scanner.java b/src/main/java/com/it_jaros/jscanner/Scanner.java
index 95f0507..7ef6c3e 100644
--- a/src/main/java/com/it_jaros/jscanner/Scanner.java
+++ b/src/main/java/com/it_jaros/jscanner/Scanner.java
@@ -1,8 +1,9 @@
 package com.it_jaros.jscanner;
 
+import java.io.ByteArrayOutputStream;
 import java.io.IOException;
-import java.io.InputStream;
 import java.net.*;
+import java.nio.ByteBuffer;
 import java.nio.channels.SocketChannel;
 import java.nio.charset.StandardCharsets;
 import java.time.Duration;
@@ -375,9 +376,9 @@ public class Scanner implements AutoCloseable {
 
                 if (isConnected) {
                     result.setState(PortState.OPEN);
-                    //if (bannerRecognition) {
-                    //    result.setBanner(getBanner(socket));
-                    //}
+                    if (bannerRecognition) {
+                        result.setBanner(getBanner(socketChannel));
+                    }
                 } else {
                     result.setState(PortState.FILTERED);
                 }
@@ -395,18 +396,36 @@ public class Scanner implements AutoCloseable {
             return result;
         }
 
-        private String getBanner(Socket socket) {
+        private String getBanner(SocketChannel socketChannel) {
+            try {
+                return readBannerFrom(socketChannel);
+            } catch (IOException ignore) {
+                // ignore
+            }
+
+            return null;
+        }
+
+        private String readBannerFrom(SocketChannel socketChannel) throws IOException {
             byte[] buffer = new byte[READ_BUFFER_SIZE];
-            try (InputStream input = socket.getInputStream()) {
-                int bytesRead = input.read(buffer);
-                if (bytesRead <= 0) {
-                    return null;
+            ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
+            long deadlineNanos = System.nanoTime() + TimeUnit.MILLISECONDS.toNanos(timeoutInMillis);
+
+            boolean timedout = false;
+            while (!timedout) {
+                int bytesRead = socketChannel.read(ByteBuffer.wrap(buffer));
+                if (bytesRead == -1) {
+                    break;
+                }
+                byteArrayOutputStream.write(buffer, 0, bytesRead);
+                long remaining = deadlineNanos - System.nanoTime();
+                if (remaining <= 0) {
+                    timedout = true;
+                } else {
+                    LockSupport.parkNanos(Math.min(remaining, TimeUnit.MILLISECONDS.toNanos(1000)));
                 }
-                return new String(buffer, 0, bytesRead, StandardCharsets.UTF_8).trim();
-            } catch (IOException e) {
-                // we ignore this failure
             }
-            return null;
+            return byteArrayOutputStream.toString(StandardCharsets.UTF_8).trim();
         }
     }