From 5731358d3c1a4603d10be84557542d48dcc068a1 Mon Sep 17 00:00:00 2001
From: Harald Eilertsen <haraldei@anduin.net>
Date: Thu, 9 May 2019 14:39:04 +0200
Subject: 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.
---
 src/main.rs | 33 +++++++++++++++++++++++++--------
 1 file changed, 25 insertions(+), 8 deletions(-)

(limited to 'src/main.rs')

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(())
 }
-- 
cgit v1.2.3