1 files changed, 55 insertions, 38 deletions

diff --git a/src/main/java/com/it_jaros/jscanner/App.java b/src/main/java/com/it_jaros/jscanner/App.java
index 9693af1..cd82140 100644
--- a/src/main/java/com/it_jaros/jscanner/App.java
+++ b/src/main/java/com/it_jaros/jscanner/App.java
@@ -1,12 +1,18 @@
 package com.it_jaros.jscanner;
 
+import java.time.Duration;
 import java.util.BitSet;
 import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import java.util.stream.Collectors;
 
 public class App {
 
     private static Thread shutdownHook;
+    private static final CountDownLatch shutdownLatch = new CountDownLatch(1);
+    private static final int shutdownTimeoutInMillis = 10000;
 
     public static void main(String[] args) {
         try {
@@ -21,10 +27,10 @@ public class App {
         if (options == null) return;
         try (Scanner scanner = new Scanner(options)) {
             Scan scan = Scan.create(options);
-            addShutdownHook(scan, scanner, options);
-            scanner.runScan(scan);
-            printScanResults(scan, options);
+            addShutdownHook(scanner, options);
+            printScanResults(scanner.runScan(scan), scan.getPeakConcurrentConnects(), options.showFilteredPorts());
         } finally {
+            shutdownLatch.countDown();
             removeShutdownHook();
         }
     }
@@ -42,42 +48,9 @@ public class App {
         return null;
     }
 
-    /**
-     * Allows us to handle signal int events
-     * Prints intermediate result
-     *
-     * @param scan
-     * @param scanner
-     * @param scanOptions
-     */
-    private static void addShutdownHook(Scan scan, Scanner scanner, ScanOptions scanOptions) {
-        shutdownHook = new Thread(() -> {
-            System.err.println("\nCtrl-C received.");
-            scanner.cancel();
-            printScanResults(scan, scanOptions);
-        });
-        Runtime.getRuntime().addShutdownHook(shutdownHook);
-    }
-
-    private static void removeShutdownHook() {
-        if (shutdownHook == null) {
-            return;
-        }
-
-        try {
-            Runtime.getRuntime().removeShutdownHook(shutdownHook);
-        } catch (IllegalStateException ignored) {
-            // JVM is already shutting down
-        }
-    }
-
-    public static void printScanResults(Scan scan, ScanOptions scanOptions) {
-        printScanResults(scan.getResults(), scan.getPeakConcurrentConnects(), scanOptions);
-    }
-
-    private static void printScanResults(List<ScanResult> result, int max, ScanOptions options) {
+    private static void printScanResults(List<ScanResult> result, int max, boolean showFilteredPorts) {
         printStats(max);
-        result.forEach((ScanResult r) -> printPortsResults(r, options.showFilteredPorts()));
+        result.forEach((ScanResult r) -> printPortsResults(r, showFilteredPorts));
     }
 
     private static void printStats(int max) {
@@ -117,4 +90,48 @@ public class App {
             return port + "/" + serviceType.name().toLowerCase();
         }).collect(Collectors.joining(","));
     }
+
+    /**
+     * Notice user that the signal was received
+     * Start shutdown of scanner
+     *
+     * @param scanner
+     */
+    private static void addShutdownHook(Scanner scanner, ScanOptions options) {
+        shutdownHook = new Thread(() -> {
+            System.err.printf("%nCtrl-C received. Shutting down... Waiting for timeout of %d seconds%n", shutdownTimeoutInMillis / 1000);
+            shutdown(scanner);
+        });
+        Runtime.getRuntime().addShutdownHook(shutdownHook);
+    }
+
+    /**
+     * Try to gracefully shutdown scanner
+     * and wait
+     *
+     * @param scanner
+     */
+    private static void shutdown(Scanner scanner) {
+        try {
+            scanner.cancel();
+            scanner.awaitTermination(Duration.ofMillis(shutdownTimeoutInMillis));
+            shutdownLatch.await(shutdownTimeoutInMillis, TimeUnit.MILLISECONDS);
+        } catch (InterruptedException e) {
+            Thread.currentThread().interrupt();
+        } catch (TimeoutException e) {
+            // Time ran out
+        }
+    }
+
+    private static void removeShutdownHook() {
+        if (shutdownHook == null) {
+            return;
+        }
+
+        try {
+            Runtime.getRuntime().removeShutdownHook(shutdownHook);
+        } catch (IllegalStateException ignored) {
+            // JVM is already shutting down
+        }
+    }
 }