Clean up socket on termination.

1 files changed, 9 insertions, 1 deletions

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.