add proper error handling
This commit is contained in:
parent
b7669181a1
commit
809d24a70d
3 changed files with 67 additions and 37 deletions
|
|
@ -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<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();
|
||||
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<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();
|
||||
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<Self> {
|
||||
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<String> {
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
|
@ -1,2 +1,2 @@
|
|||
mod client;
|
||||
pub use client::RconClient;
|
||||
pub use client::{RconClient, RconError, RconResult};
|
||||
|
|
@ -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(())
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue