Handle client disconnects more gracefully.

Give the client an option to disconnect politely by issuing the `bye`
command, but also handle the case where the client just drops the
connection in a better way.

2 files changed, 30 insertions, 8 deletions

diff --git a/README.md b/README.md
index db4cc40..29d1652 100644
--- a/README.md
+++ b/README.md
@@ -28,6 +28,9 @@ You can check if an URL is phishy or not by connecting to the socket and passing
 one URL at the time to the socket. For each URL the server will respond if it's
 phishy or good.
 
+Once the client is done, it can drop the connection gracefully by issuing the `bye`
+command, or simply drop the connection.
+
 An example session can be like this:
 
     % socat - UNIX-CONNECT:/tmp/phisher
@@ -35,6 +38,8 @@ An example session can be like this:
     good url.
     https://safravideos.com/login
     it's a phish!
+    bye
+    bye.
 
 ## How to contribute
 
diff --git a/src/main.rs b/src/main.rs
index ab9f095..db6bb73 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -22,6 +22,7 @@ use std::fs::{ File, remove_file };
 use std::io::{ BufRead, BufReader, Write };
 use std::os::unix::net::{ UnixListener, UnixStream };
 use std::result::Result;
+use std::string::ToString;
 use std::sync::Arc;
 use std::thread;
 use std::time::SystemTime;
@@ -46,7 +47,7 @@ fn main() {
 
         match run_server(tank) {
             Ok(_) => println!("Exiting."),
-            Err(e) => eprintln!("Received error: {}, terminating.", e.to_string())
+            Err(e) => eprintln!("Error: {}, terminating.", e.to_string())
         };
     }
     else {
@@ -64,7 +65,11 @@ fn run_server(tank: Arc<phisher::PhishTank>) -> Result<(), Box<dyn Error>> {
         match stream {
             Ok(stream) => {
                 let t = Arc::clone(&tank);
-                thread::spawn(move || handle_connection(stream, t));
+                thread::spawn(move || {
+                    if let Err(e) = handle_connection(stream, t) {
+                        eprintln!("Client was disconnected: {}", e.to_string());
+                    }
+                });
             }
             Err(e) => {
                 eprintln!("Received error: {}", e.to_string());
@@ -83,19 +88,31 @@ fn run_server(tank: Arc<phisher::PhishTank>) -> Result<(), Box<dyn Error>> {
 // a separate line, and the results are returned in the same order
 // as the URLs were passed in.
 //
-fn handle_connection(mut stream: UnixStream, tank: Arc<phisher::PhishTank>) {
+fn handle_connection(mut stream: UnixStream, tank: Arc<phisher::PhishTank>)
+    -> Result<(), Box<dyn Error>>
+{
     // Clone the stream so we get another ref to the underlying socket,
     // then wrap it in a BufReader, so we can use read_line on it.
-    let mut bufreader = BufReader::new(stream.try_clone().unwrap());
+    let mut bufreader = BufReader::new(stream.try_clone()?);
 
     loop {
         let mut buffer = String::new();
-        bufreader.read_line(&mut buffer).unwrap();
-        if tank.is_phish(buffer.trim_end()) {
-            stream.write(b"it's a phish!\n").unwrap();
+        bufreader.read_line(&mut buffer)?;
+
+        let input = buffer.trim_end();
+
+        if input == "bye" {
+            stream.write(b"bye.")?;
+            break;
+        }
+
+        if tank.is_phish(&input) {
+            stream.write(b"it's a phish!\n")?;
         }
         else {
-            stream.write(b"good url.\n").unwrap();
+            stream.write(b"good url.\n")?;
         }
     }
+
+    Ok(())
 }