Clean up socket on termination.

2 files changed, 10 insertions, 1 deletions

diff --git a/Cargo.toml b/Cargo.toml
index 181ffcf..b37a2b1 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -7,5 +7,6 @@ edition = "2018"
 [dependencies]
 chrono = { version = "0.4.6", features = ["serde"] }
 clap = "2.33"
+ctrlc = { version = "3.1.2", features = ["termination"] }
 serde = { version = "1.0.91", features = ["derive"] }
 serde_json = "1.0.39"
diff --git a/src/main.rs b/src/main.rs
index db6bb73..36c5cf9 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -21,6 +21,7 @@ use std::error::Error;
 use std::fs::{ File, remove_file };
 use std::io::{ BufRead, BufReader, Write };
 use std::os::unix::net::{ UnixListener, UnixStream };
+use std::process::exit;
 use std::result::Result;
 use std::string::ToString;
 use std::sync::Arc;
@@ -28,6 +29,13 @@ use std::thread;
 use std::time::SystemTime;
 
 fn main() {
+    // Set signal handler to clean up on termination.
+    ctrlc::set_handler(move || {
+        // We don't care about errors in remove_file.
+        remove_file("/tmp/phisher").unwrap_or(());
+        exit(0);
+    }).expect("Couldn't set signal handler!");
+
     let args = clap_app!(app =>
         (name: crate_name!())
         (version: crate_version!())
@@ -78,7 +86,7 @@ fn run_server(tank: Arc<phisher::PhishTank>) -> Result<(), Box<dyn Error>> {
         }
     }
 
-    Ok(remove_file("/tmp/phisher")?)
+    Ok(())
 }
 
 // Handle a connected client.