Minor Refactoring

4 files changed, 17 insertions, 16 deletions

diff --git a/src/main/java/com/it_jaros/jns/scan/engine/ProducerThread.java b/src/main/java/com/it_jaros/jns/scan/engine/ProducerThread.java
index 013f3cc..f02e379 100644
--- a/src/main/java/com/it_jaros/jns/scan/engine/ProducerThread.java
+++ b/src/main/java/com/it_jaros/jns/scan/engine/ProducerThread.java
@@ -70,7 +70,7 @@ public class ProducerThread {
         }
     }
 
-    private <OUTPUT, INPUT> void submitNext(ProducerContext<INPUT, OUTPUT> context) throws InterruptedException {
+    private <INPUT, OUTPUT> void submitNext(ProducerContext<INPUT, OUTPUT> context) throws InterruptedException {
         context.activeWorkers().acquire();
 
         // just in case something
diff --git a/src/main/java/com/it_jaros/jns/scan/engine/ScanHostContext.java b/src/main/java/com/it_jaros/jns/scan/engine/ScanHostContext.java
index a4a6777..0653448 100644
--- a/src/main/java/com/it_jaros/jns/scan/engine/ScanHostContext.java
+++ b/src/main/java/com/it_jaros/jns/scan/engine/ScanHostContext.java
@@ -1,5 +1,8 @@
 package com.it_jaros.jns.scan.engine;
 
+import java.net.InetAddress;
+
 public record ScanHostContext(
+        InetAddress hostAddress,
         PortScanRateLimiter portScanRateLimiter
 ) {}
diff --git a/src/main/java/com/it_jaros/jns/scan/engine/ScanHostTask.java b/src/main/java/com/it_jaros/jns/scan/engine/ScanHostTask.java
index b3b062b..02834ff 100644
--- a/src/main/java/com/it_jaros/jns/scan/engine/ScanHostTask.java
+++ b/src/main/java/com/it_jaros/jns/scan/engine/ScanHostTask.java
@@ -14,7 +14,6 @@ import java.util.concurrent.TimeUnit;
 
 public class ScanHostTask implements Callable<ScanResult> {
 
-    private InetAddress hostAddress;
     private final CancelledToken cancelledToken;
     private final PortScanResultAccumulator accumulator;
     private final Scan scan;
@@ -22,7 +21,7 @@ public class ScanHostTask implements Callable<ScanResult> {
     private final String host; // input parameter
     private final boolean disableOnlineCheck;
     private final int maxWorkersPerHost;
-    private final ScanHostContext hostContext;
+    private ScanHostContext hostContext;
 
     public ScanHostTask(String host, ScanExecutionContext context) {
         this.host = host;
@@ -32,14 +31,6 @@ public class ScanHostTask implements Callable<ScanResult> {
         this.disableOnlineCheck = context.scanOptions().disableOnlineCheck();
         this.maxWorkersPerHost = context.scanOptions().maxWorkersPerHost();
         this.scan = context.scan();
-
-        long delayInNanos = TimeUnit.MILLISECONDS.toNanos(Math.max(0, context.scanOptions().delayInMillis()));
-        this.hostContext = new ScanHostContext(
-                new PortScanRateLimiter(
-                        context.cancelledToken(),
-                        delayInNanos
-                )
-        );
     }
 
     @Override
@@ -49,7 +40,14 @@ public class ScanHostTask implements Callable<ScanResult> {
             if (context.cancelledToken().isCancelled()) {
                 return ScanResult.empty(host);
             }
-            this.hostAddress = InetAddress.getByName(host);
+            long delayInNanos = TimeUnit.MILLISECONDS.toNanos(Math.max(0, context.scanOptions().delayInMillis()));
+            this.hostContext = new ScanHostContext(
+                    InetAddress.getByName(host),
+                    new PortScanRateLimiter(
+                            context.cancelledToken(),
+                            delayInNanos
+                    )
+            );
             return scanHostPorts();
         } catch (Exception e) {
             throw new ScanHostException(host, e);
@@ -72,7 +70,7 @@ public class ScanHostTask implements Callable<ScanResult> {
         ProducerState<PortResult> state = new ProducerThread(context).startProducer(
                 portRange.iterator(),
                 maxWorkersPerHost,
-                port -> new ScanPortTask(hostAddress, port, context, hostContext)
+                port -> new ScanPortTask(port, context, hostContext)
         );
 
         runConsumerLoop(state);
@@ -81,7 +79,7 @@ public class ScanHostTask implements Callable<ScanResult> {
 
     private boolean checkHostOnline() {
         try {
-            PortResult result = new ScanPortTask(hostAddress, 7, context, hostContext).call();
+            PortResult result = new ScanPortTask(7, context, hostContext).call();
             Exception exception = result.getException();
             if (exception == null && !result.getState().equals(PortState.FILTERED)) {
                 return true;
diff --git a/src/main/java/com/it_jaros/jns/scan/engine/ScanPortTask.java b/src/main/java/com/it_jaros/jns/scan/engine/ScanPortTask.java
index 1e709d5..e9244c2 100644
--- a/src/main/java/com/it_jaros/jns/scan/engine/ScanPortTask.java
+++ b/src/main/java/com/it_jaros/jns/scan/engine/ScanPortTask.java
@@ -29,8 +29,8 @@ public class ScanPortTask implements Callable<PortResult> {
     private final int port;
     private final long timeoutInNanos;
 
-    ScanPortTask(InetAddress inetAddress, int port, ScanExecutionContext context, ScanHostContext hostContext) {
-        this.inetAddress = inetAddress;
+    ScanPortTask(int port, ScanExecutionContext context, ScanHostContext hostContext) {
+        this.inetAddress = hostContext.hostAddress();
         this.port = port;
         this.bannerRecognition = context.scanOptions().bannerRecognition();
         this.cancelledToken = context.cancelledToken();