add proper error handling

src/client.rs

@@ -1,4 +1,4 @@
-use std::{io::{Read, Write}, net::{TcpStream, ToSocketAddrs}};
+use std::{fmt::Display, io::{self, Read, Write}, net::{TcpStream, ToSocketAddrs}};
 
 #[derive(Clone, Copy, PartialEq)]
 enum PacketType {
@@ -8,8 +8,8 @@ enum PacketType {
 }
 
 impl PacketType {
-    fn to_le_bytes(&self) -> [u8; 4] {
+    fn to_le_bytes(self) -> [u8; 4] {
-        (*self as i32).to_le_bytes()
+        (self as i32).to_le_bytes()
     }
 }
 
@@ -26,39 +26,66 @@ impl TryFrom<i32> for PacketType {
     }
 }
 
+#[derive(Debug)]
+pub enum RconError {
+    Io(io::Error),
+    AuthFailed,
+    InvalidResponse(String)
+}
+pub type RconResult<T> = Result<T, RconError>;
+
+impl Display for RconError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            RconError::Io(e) => write!(f, "io error: {e}"),
+            RconError::AuthFailed => write!(f, "auth failed"),
+            RconError::InvalidResponse(e) => write!(f, "invalid response from server: {e}")
+        }
+    }
+}
+
+impl From<io::Error> for RconError {
+    fn from(value: io::Error) -> Self {
+        Self::Io(value)
+    }
+}
+
 pub struct RconClient {
     stream: TcpStream,
     next_id: i32
 }
 
 impl RconClient {
-    fn read_i32_le(&mut self) -> i32 {
+    fn read_i32_le(&mut self) -> RconResult<i32> {
         let mut buf = [0u8; 4];
-        self.stream.read_exact(&mut buf).unwrap();
+        self.stream.read_exact(&mut buf)?;
-        i32::from_le_bytes(buf)
+        Ok(i32::from_le_bytes(buf))
     }
 
-    fn recv_packet(&mut self) -> (i32, PacketType, String) {
+    fn recv_packet(&mut self) -> RconResult<(i32, PacketType, String)> {
-        let length = self.read_i32_le();
+        let length = self.read_i32_le()?;
         if length < 10 {
-            panic!("response length < 10");
+            return Err(RconError::InvalidResponse("response length < 10".to_string()))
         }
 
-        let id = self.read_i32_le();
+        let id = self.read_i32_le()?;
-        let ptype = PacketType::try_from(self.read_i32_le()).unwrap();
+        let ptype = match PacketType::try_from(self.read_i32_le()?) {
+            Ok(ptype) => ptype,
+            Err(()) => return Err(RconError::InvalidResponse("unknown packet type".to_string()))
+        };
 
         let payload_len = (length - 10) as usize;
         let mut payload_bytes = vec![0u8; payload_len];
-        self.stream.read_exact(&mut payload_bytes).unwrap();
+        self.stream.read_exact(&mut payload_bytes)?;
 
         let mut nulls = [0u8; 2];
-        self.stream.read_exact(&mut nulls).unwrap();
+        self.stream.read_exact(&mut nulls)?;
 
         let payload = String::from_utf8_lossy(&payload_bytes).into_owned();
-        (id, ptype, payload)
+        Ok((id, ptype, payload))
     }
 
-    fn send_packet(&mut self, ptype: PacketType, payload: &str) -> i32 {
+    fn send_packet(&mut self, ptype: PacketType, payload: &str) -> RconResult<i32> {
         let id = self.next_id;
         self.next_id = self.next_id.wrapping_add(1);
 
@@ -69,42 +96,44 @@ impl RconClient {
         packet.extend_from_slice(&length.to_le_bytes());
         packet.extend_from_slice(&id.to_le_bytes());
         packet.extend_from_slice(&ptype.to_le_bytes());
-        packet.extend_from_slice(&payload_bytes);
+        packet.extend_from_slice(payload_bytes);
         packet.push(0); // terminate payload
         packet.push(0); // padding
 
-        self.stream.write_all(&packet).unwrap();
+        self.stream.write_all(&packet)?;
-        id
+        Ok(id)
     }
 
-    fn authenticate(&mut self, password: &str) {
+    fn authenticate(&mut self, password: &str) -> RconResult<()> {
-        self.send_packet(PacketType::Login, password);
+        self.send_packet(PacketType::Login, password)?;
 
-        let (auth_id, _, _) = self.recv_packet();
+        let (auth_id, _, _) = self.recv_packet()?;
 
         if auth_id == -1 {
-            panic!("auth failed");
+            return Err(RconError::AuthFailed);
-        }
         }
 
-    pub fn connect(addr: impl ToSocketAddrs, password: &str) -> Self {
+        Ok(())
-        let stream = TcpStream::connect(addr).unwrap();
+    }
+    
+    pub fn connect(addr: impl ToSocketAddrs, password: &str) -> RconResult<Self> {
+        let stream = TcpStream::connect(addr)?;
         let mut client = RconClient { stream, next_id: 1 };
-        client.authenticate(password);
+        client.authenticate(password)?;
-        client
+        Ok(client)
     }
 
-    pub fn command(&mut self, cmd: &str) -> String {
+    pub fn command(&mut self, cmd: &str) -> RconResult<String> {
-        let id = self.send_packet(PacketType::Command, cmd);
+        let id = self.send_packet(PacketType::Command, cmd)?;
 
-        let (resp_id, resp_type, payload) = self.recv_packet();
+        let (resp_id, resp_type, payload) = self.recv_packet()?;
         if resp_type != PacketType::Response {
-            panic!("invalid response")
+            return Err(RconError::InvalidResponse("received packet not of response type".to_string()))
         }
         if resp_id != id {
-            panic!("invalid response")
+            return Err(RconError::InvalidResponse("mismatched packet id".to_string()))
         }
 
-        payload
+        Ok(payload)
     }
 }

src/lib.rs

@@ -1,2 +1,2 @@
 mod client;
-pub use client::RconClient;
+pub use client::{RconClient, RconError, RconResult};

src/main.rs

@@ -1,6 +1,6 @@
 use clap::Parser;
 
-use mc_rcon::RconClient;
+use mc_rcon::{RconClient, RconResult};
 
 #[derive(Parser, Debug)]
 struct Args {
@@ -14,10 +14,11 @@ struct Args {
     cmd: String
 }
 
-fn main() {
+fn main() -> RconResult<()> {
     let args = Args::parse();
 
-    let mut client = RconClient::connect(args.addr, args.password.as_str());
+    let mut client = RconClient::connect(args.addr, args.password.as_str())?;
-    let resp = client.command(&args.cmd);
+    let resp = client.command(&args.cmd)?;
     println!("{resp}");
+    Ok(())
 }