diff --git a/src/client.rs b/src/client.rs index 40c8c56..83354e3 100644 --- a/src/client.rs +++ b/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] { - (*self as i32).to_le_bytes() + fn to_le_bytes(self) -> [u8; 4] { + (self as i32).to_le_bytes() } } @@ -26,39 +26,66 @@ impl TryFrom for PacketType { } } +#[derive(Debug)] +pub enum RconError { + Io(io::Error), + AuthFailed, + InvalidResponse(String) +} +pub type RconResult = Result; + +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 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 { let mut buf = [0u8; 4]; - self.stream.read_exact(&mut buf).unwrap(); - i32::from_le_bytes(buf) + self.stream.read_exact(&mut buf)?; + Ok(i32::from_le_bytes(buf)) } - fn recv_packet(&mut self) -> (i32, PacketType, String) { - let length = self.read_i32_le(); + fn recv_packet(&mut self) -> RconResult<(i32, PacketType, String)> { + 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 ptype = PacketType::try_from(self.read_i32_le()).unwrap(); + let id = self.read_i32_le()?; + 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 { 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(); - id + self.stream.write_all(&packet)?; + Ok(id) } - fn authenticate(&mut self, password: &str) { - self.send_packet(PacketType::Login, password); + fn authenticate(&mut self, password: &str) -> RconResult<()> { + 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); } + + Ok(()) } - pub fn connect(addr: impl ToSocketAddrs, password: &str) -> Self { - let stream = TcpStream::connect(addr).unwrap(); + pub fn connect(addr: impl ToSocketAddrs, password: &str) -> RconResult { + let stream = TcpStream::connect(addr)?; let mut client = RconClient { stream, next_id: 1 }; - client.authenticate(password); - client + client.authenticate(password)?; + Ok(client) } - pub fn command(&mut self, cmd: &str) -> String { - let id = self.send_packet(PacketType::Command, cmd); + pub fn command(&mut self, cmd: &str) -> RconResult { + 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) } } \ No newline at end of file diff --git a/src/lib.rs b/src/lib.rs index bb722f5..79ffb8d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,2 +1,2 @@ mod client; -pub use client::RconClient; \ No newline at end of file +pub use client::{RconClient, RconError, RconResult}; \ No newline at end of file diff --git a/src/main.rs b/src/main.rs index 915b952..737c72e 100644 --- a/src/main.rs +++ b/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 resp = client.command(&args.cmd); + let mut client = RconClient::connect(args.addr, args.password.as_str())?; + let resp = client.command(&args.cmd)?; println!("{resp}"); + Ok(()) } \ No newline at end of file