From f860fe3689a53a47c5c7d9d8339aca5c9c452733 Mon Sep 17 00:00:00 2001
From: Tobias Reisinger <tobias@msrg.cc>
Date: Wed, 27 Mar 2024 15:37:52 +0100
Subject: [PATCH] Improve error handling for events and in general

---
 src/command_utils/events.rs |  40 +++++++++++
 src/command_utils/mod.rs    |   1 +
 src/commands.rs             |  14 ++--
 src/main.rs                 | 134 ++++++++++++++----------------------
 src/wrappers.rs             |  25 +++----
 ts-control                  |   4 +-
 6 files changed, 114 insertions(+), 104 deletions(-)
 create mode 100644 src/command_utils/events.rs
 create mode 100644 src/command_utils/mod.rs

diff --git a/src/command_utils/events.rs b/src/command_utils/events.rs
new file mode 100644
index 0000000..b4a8773
--- /dev/null
+++ b/src/command_utils/events.rs
@@ -0,0 +1,40 @@
+use telnet::Telnet;
+use crate::commands;
+use crate::models::{Event, EventType};
+use crate::response::Response;
+
+pub fn register_events(connection: &mut Telnet, events: Vec<EventType>) -> Result<(), String> {
+    for event in events {
+        if commands::clientnotifyregister(connection, 1, event).is_err() {
+            return Err(String::from("Failed to register event listener."));
+        }
+    }
+    Ok(())
+}
+
+pub fn handle_event_response(connection: &mut Telnet, response: Response) {
+    if let Response::Event(event_type, params) = response {
+
+        let event = Event::new(connection, event_type, params);
+        match serde_json::to_string(&event) {
+            Ok(json) => println!("{}", json),
+            Err(err) => {
+                // TODO: Handle serialization error
+                eprintln!("Serialization error: {}", err);
+            }
+        }
+
+    }
+}
+
+pub fn loop_response_reader(connection: &mut Telnet) {
+    loop {
+        match commands::read_response(connection, true, String::new()) {
+            Ok(response) => handle_event_response(connection, response),
+            Err(_) => {
+                // print error?
+                return;
+            }
+        }
+    }
+}
\ No newline at end of file
diff --git a/src/command_utils/mod.rs b/src/command_utils/mod.rs
new file mode 100644
index 0000000..6226263
--- /dev/null
+++ b/src/command_utils/mod.rs
@@ -0,0 +1 @@
+pub mod events;
\ No newline at end of file
diff --git a/src/commands.rs b/src/commands.rs
index 204e455..f5fdea2 100644
--- a/src/commands.rs
+++ b/src/commands.rs
@@ -1,5 +1,6 @@
+use std::time::Duration;
 use crate::utils::SendTextMessageTarget;
-use telnet::Event::Data;
+use telnet::Event::{Data, TimedOut};
 use telnet::Telnet;
 use crate::models::EventType;
 
@@ -14,18 +15,21 @@ fn to_single_response(resp: Response) -> Response {
 }
 
 fn read_part(connection: &mut Telnet) -> Result<String, String> {
-    match connection.read() {
+    match connection.read_timeout(Duration::new(5, 0)) {
         Ok(event) => {
             match event {
                 Data(bytes) => Ok(String::from_utf8(bytes.to_vec())
-                    .map_err(|_| "Teamspeak returned a badly formatted models.")?),
+                    .map_err(|_| "Teamspeak returned an invalid response.")?),
+                TimedOut => {
+                    Ok(String::new())
+                },
                 _ => {
                     Err(String::from("Received unknown event from Teamspeak."))
                 }
             }
         }
-        Err(_) => {
-            Err(String::from("Failed to read from Teamspeak."))
+        Err(err) => {
+            Err(format!("Failed to read from Teamspeak: {}", err))
         }
     }
 }
diff --git a/src/main.rs b/src/main.rs
index 6af36ac..24d8645 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,11 +1,10 @@
-use std::process::exit;
+use std::io::{Error, ErrorKind};
 
 use telnet::Telnet;
 
 use crate::cli::Commands;
-use crate::models::{Channel, Event};
+use crate::models::Channel;
 use crate::models::Client;
-use crate::response::Response;
 
 mod wrappers;
 mod commands;
@@ -14,41 +13,42 @@ mod cli;
 mod utils;
 mod models;
 mod response;
+mod command_utils;
 
-fn channel_or_exit(channel_res: Result<Option<Channel>, String>) -> Channel {
-    channel_res.unwrap_or_else(|err| {
-        println!("Failed to find channel: {}", err);
-        exit(1);
-    })
-        .unwrap_or_else(|| {
-            println!("Failed to find channel.");
-            exit(1);
-        })
+fn channel_or_error(channel_res: Result<Option<Channel>, String>) -> Result<Channel, Error> {
+    channel_res.map_err(|err| make_action_error("find channel", err))?
+        .ok_or_else(|| make_action_error("find channel", String::from("Not Found.")))
 }
 
-fn client_or_exit(client_res: Result<Option<Client>, String>) -> Client {
-    client_res.unwrap_or_else(|err| {
-        println!("Failed to find client: {}", err);
-        exit(1);
-    })
-        .unwrap_or_else(|| {
-            println!("Failed to find client.");
-            exit(1);
-        })
+fn client_or_error(client_res: Result<Option<Client>, String>) -> Result<Client, Error> {
+    client_res.map_err(|err| make_action_error("find client", err))?
+        .ok_or_else(|| make_action_error("find client", String::from("Not Found.")))
 }
 
-fn main() {
+fn connect() -> Result<Telnet, Error> {
+    let mut connection = Telnet::connect(("127.0.0.1", 25639), 512 * 1024)
+        .map_err(|_| Error::new(ErrorKind::AddrNotAvailable, String::from("Failed to connect to Teamspeak.")))?;
+
+    wrappers::skip_welcome(&mut connection)
+        .map_err(to_other_error)?;
+    wrappers::login(&mut connection)
+        .map_err(|msg| Error::new(ErrorKind::PermissionDenied, msg))?;
+
+    Ok(connection)
+}
+
+fn to_other_error(msg: String) -> Error {
+    Error::new(ErrorKind::Other, msg)
+}
+
+fn make_action_error(action: &str, msg: String) -> Error {
+    Error::new(ErrorKind::Other, format!("Failed to {}: {}", action, msg))
+}
+
+fn main() -> Result<(), Error> {
     let cli = cli::init();
 
-    let connection = Telnet::connect(("127.0.0.1", 25639), 512 * 1024);
-    if connection.is_err() {
-        println!("Failed to connect to Teamspeak.");
-        exit(1);
-    }
-    let mut connection = connection.unwrap();
-
-    wrappers::skip_welcome(&mut connection);
-    wrappers::login(&mut connection);
+    let mut connection = connect()?;
 
     // You can check for the existence of subcommands, and if found use their
     // matches just as you would the top level cmd
@@ -61,8 +61,7 @@ fn main() {
                     }
                 }
                 Err(msg) => {
-                    println!("Failed to get channels: {}", msg);
-                    exit(1);
+                    return Err(make_action_error("get channels", msg));
                 }
             }
         }
@@ -75,71 +74,62 @@ fn main() {
                     }
                 }
                 Err(msg) => {
-                    println!("Failed to get clients: {}", msg);
-                    exit(1);
+                    return Err(make_action_error("get clients", msg));
                 }
             }
         }
 
         Commands::Fetch(args) => {
             if args.want_client() && args.want_channel() {
-                println!("Fetching both clients and channels is not supported.");
-                exit(1);
+                return Err(Error::new(ErrorKind::InvalidInput, "Fetching both clients and channels is not supported."));
             }
             if !args.want_client() && !args.want_channel() {
-                println!("No clients or channels specified.");
-                exit(1);
+                return Err(Error::new(ErrorKind::InvalidInput, "No clients or channels specified."));
             }
 
             if args.want_client() {
-                let client = client_or_exit(args.client(&mut connection));
+                let client = client_or_error(args.client(&mut connection))?;
 
                 match wrappers::fetch_client(&mut connection, &[client]) {
                     Ok(_) => println!("Successfully fetched client."),
                     Err(msg) => {
-                        println!("Failed to fetch client: {}", msg);
-                        exit(1);
+                        return Err(make_action_error("fetch client", msg));
                     }
                 }
             }
 
             if args.want_channel() {
-                let channel = channel_or_exit(args.channel(&mut connection));
+                let channel = channel_or_error(args.channel(&mut connection))?;
 
                 match wrappers::fetch_channel(&mut connection, channel) {
                     Ok(_) => println!("Successfully fetched channel."),
                     Err(msg) => {
-                        println!("Failed to fetch channel: {}", msg);
-                        exit(1);
+                        return Err(make_action_error("fetch channel", msg));
                     }
                 }
             }
         }
 
         Commands::Message(args) => {
-            let target = args.target(&mut connection).unwrap_or_else(|err| {
-                println!("Failed to get message target: {}", err);
-                exit(1);
-            });
+            let target = args.target(&mut connection)
+                .map_err(|msg| make_action_error("message target", msg))?;
 
             match wrappers::send_text_message(&mut connection, target, args.message) {
                 Ok(_) => println!("Successfully sent message."),
                 Err(msg) => {
-                    println!("Failed to send message: {}", msg);
-                    exit(1);
+                    return Err(make_action_error("send message", msg));
                 }
             }
         }
 
         Commands::Move(args) => {
-            let channel = channel_or_exit(args.channel(&mut connection));
-            let client = client_or_exit(args.client(&mut connection));
+            let channel = channel_or_error(args.channel(&mut connection))?;
+            let client = client_or_error(args.client(&mut connection))?;
 
             match wrappers::move_client(&mut connection, &channel, &[client]) {
                 Ok(resp) => println!("Successfully moved client: {}", resp),
                 Err(msg) => {
-                    println!("Failed to move client: {}", msg);
-                    exit(1);
+                    return Err(make_action_error("move client", msg));
                 }
             }
         }
@@ -148,40 +138,22 @@ fn main() {
             match wrappers::update_client(&mut connection, args.to_parameter_list()) {
                 Ok(_) => println!("Successfully updated client."),
                 Err(msg) => {
-                    println!("Failed to update client: {}", msg);
-                    exit(1);
+                    return Err(make_action_error("update client", msg));
                 }
             }
         }
 
         Commands::Events(args) => {
-            for event in args.event {
-                if commands::clientnotifyregister(&mut connection, 1, event).is_err() {
-                    println!("Failed to register event listener.");
-                    exit(1);
-                }
-            }
-
             loop {
-                match commands::read_response(&mut connection, true, String::new()) {
-                    Ok(response) => {
-                        if let Response::Event(event_type, params) = response {
+                command_utils::events::register_events(&mut connection, args.event.clone())
+                    .map_err(to_other_error)?;
+                command_utils::events::loop_response_reader(&mut connection);
 
-                            let event = Event::new(&mut connection, event_type, params);
-                            match serde_json::to_string(&event) {
-                                Ok(json) => println!("{}", json),
-                                Err(_) => {
-                                    // TODO: Handle error
-                                }
-                            }
-
-                        }
-                    }
-                    Err(_) => {
-                        // TODO: Handle error
-                    }
-                }
+                // loop_response_reader failed. Let's try to reconnect after 1 second.
+                std::thread::sleep(std::time::Duration::from_secs(1));
+                connection = connect()?;
             }
         }
     }
+    Ok(())
 }
\ No newline at end of file
diff --git a/src/wrappers.rs b/src/wrappers.rs
index 9c50680..fd9f64f 100644
--- a/src/wrappers.rs
+++ b/src/wrappers.rs
@@ -1,4 +1,3 @@
-use std::process::exit;
 use std::time::Duration;
 
 use telnet::Event::TimedOut;
@@ -11,34 +10,28 @@ use crate::models::Client;
 use crate::response::Response;
 use crate::utils::SendTextMessageTarget;
 
-pub fn skip_welcome(connection: &mut Telnet) {
+pub fn skip_welcome(connection: &mut Telnet) -> Result<(), String> {
     loop {
         let event_result = connection.read_timeout(Duration::from_millis(100));
         match event_result {
             Ok(event) => {
                 if let TimedOut = event {
-                    break;
+                    return Ok(());
                 }
             }
-            Err(_) => println!("Failed to read from Teamspeak."),
+            Err(_) => return Err(String::from("Failed to read from Teamspeak.")),
         }
     }
 }
 
-pub fn login(connection: &mut Telnet) {
+pub fn login(connection: &mut Telnet) -> Result<(), String> {
     // read api key from environment variable
-    let apikey = std::env::var("TS3_CLIENT_API_KEY").unwrap_or_else(|_| {
-        println!("No API key found in environment variable TS3_CLIENT_API_KEY.");
-        exit(1);
-    });
+    let apikey = std::env::var("TS3_CLIENT_API_KEY")
+        .map_err(|_| String::from("No API key found in environment variable TS3_CLIENT_API_KEY."))?;
 
-    match commands::login(connection, &apikey) {
-        Ok(_) => {}
-        Err(msg) => {
-            println!("Failed to authenticate with Teamspeak: {}", msg);
-            exit(1);
-        }
-    }
+    commands::login(connection, &apikey)
+        .map(|_| ())
+        .map_err(|err| format!("Failed to authenticate with Teamspeak: {}", err))
 }
 
 pub fn get_channels(connection: &mut Telnet, spacers: bool) -> Result<Vec<Channel>, String> {
diff --git a/ts-control b/ts-control
index 5f26108..462cb7c 100755
--- a/ts-control
+++ b/ts-control
@@ -1,4 +1,4 @@
-#!/usr/bin/env sh
+#!/usr/bin/env bash
 
 actions="quick
 move
@@ -101,5 +101,5 @@ case $action in
 	"events-move")
 		teamspeak-query-lib events NotifyClientMoved NotifyClientEnterView \
 			| jq -r --unbuffered '.client.client_nickname + " joined " + .channel.channel_name // "the server"' \
-			| xargs -I{} notify-send "TS3 movement" "{}"
+			| tee >(xargs -I{} notify-send "TS3 movement" "{}")
 esac