Make scanner gracefully shutdown on Strg-C

1 files changed, 27 insertions, 0 deletions

diff --git a/src/main/java/com/it_jaros/jscanner/App.java b/src/main/java/com/it_jaros/jscanner/App.java
index 27ccfb9..119c42d 100644
--- a/src/main/java/com/it_jaros/jscanner/App.java
+++ b/src/main/java/com/it_jaros/jscanner/App.java
@@ -5,11 +5,16 @@ import java.util.List;
 import java.util.stream.Collectors;
 
 public class App {
+
+    private static Thread shutdownHook;
+
     public static void main(String[] args) {
         CliParser cliParser = new CliParser();
+
         try {
             ScanOptions options = cliParser.parseArgs(args);
             Scanner scanner = new Scanner(options);
+            addShutdownHook(scanner, options.showFilteredPorts());
             printResults(scanner.scan(options), options.showFilteredPorts());
         } catch (IllegalArgumentException e) {
             System.out.printf("Error while parsing arguments: %s", e.getMessage());
@@ -17,6 +22,28 @@ public class App {
             cliParser.usageAndExit();
         } catch (Exception e) {
             System.out.printf("ERROR: %s", e.getMessage());
+        } finally {
+            removeShutdownHook();
+        }
+    }
+
+    private static void addShutdownHook(Scanner scanner, boolean showFilteredPorts) {
+        shutdownHook = new Thread(() -> {
+            System.err.println("\nCtrl-C received.");
+            scanner.cancel();
+        });
+        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
         }
     }