use std::{
fmt::{self, Display},
str::from_utf8,
};
use anyhow::{anyhow, Result};
pub mod error;
use error::{PacketBuildError, PacketDeserialiseError, PacketTypeError};
#[derive(Debug, Copy, Clone)]
pub enum PacketType {
Client(ClientPacketType),
Server(ServerPacketType),
}
impl From<PacketType> for i32 {
fn from(value: PacketType) -> Self {
match value {
PacketType::Client(msg) => msg.into(),
PacketType::Server(msg) => msg.into(),
}
}
}
#[derive(Debug, Copy, Clone)]
pub enum ClientPacketType {
ExecCommand,
Auth,
}
impl From<ClientPacketType> for PacketType {
fn from(value: ClientPacketType) -> Self {
Self::Client(value)
}
}
impl From<ClientPacketType> for i32 {
fn from(value: ClientPacketType) -> Self {
match value {
ClientPacketType::ExecCommand => 2,
ClientPacketType::Auth => 3,
}
}
}
impl TryFrom<i32> for ClientPacketType {
type Error = PacketTypeError;
fn try_from(value: i32) -> Result<Self, Self::Error> {
match value {
2 => Ok(Self::ExecCommand),
3 => Ok(Self::Auth),
v => Err(PacketTypeError(v)),
}
}
}
#[derive(Debug, Copy, Clone)]
pub enum ServerPacketType {
ResponseValue,
AuthResponse,
}
impl From<ServerPacketType> for PacketType {
fn from(value: ServerPacketType) -> Self {
Self::Server(value)
}
}
impl From<ServerPacketType> for i32 {
fn from(value: ServerPacketType) -> Self {
match value {
ServerPacketType::ResponseValue => 0,
ServerPacketType::AuthResponse => 2,
}
}
}
impl TryFrom<i32> for ServerPacketType {
type Error = PacketTypeError;
fn try_from(value: i32) -> Result<Self, Self::Error> {
match value {
0 => Ok(Self::ResponseValue),
2 => Ok(Self::AuthResponse),
v => Err(PacketTypeError(v)),
}
}
}
#[derive(Debug)]
pub struct Packet {
size: i32,
id: i32,
msg_type: PacketType,
body: String,
terminator: [u8; 2],
}
impl Packet {
pub fn new(
id: Option<i32>,
msg_type: PacketType,
body: &str,
) -> Result<Self, PacketBuildError> {
if body.len() > 4096 - 10 {
return Err(PacketBuildError::BodyTooBig(body.len()));
}
Ok(Self {
size: body.len() as i32 + 10,
id: id.unwrap_or(0),
msg_type,
body: body.into(),
terminator: [0; 2],
})
}
pub fn serialize(&self) -> Vec<u8> {
let mut bytes = Vec::with_capacity(self.size as usize);
let msg_type: &i32 = &self.msg_type.into();
bytes.extend_from_slice(&self.size.to_le_bytes());
bytes.extend_from_slice(&self.id.to_le_bytes());
bytes.extend_from_slice(&msg_type.to_le_bytes());
bytes.extend_from_slice(self.body.as_bytes());
bytes.extend_from_slice(&self.terminator);
bytes
}
pub fn deserialize_client_msg(bytes: Vec<u8>) -> Result<Self> {
Self::deserialize::<ClientPacketType>(bytes)
}
pub fn deserialize_server_msg(bytes: Vec<u8>) -> Result<Self> {
Self::deserialize::<ServerPacketType>(bytes)
}
fn deserialize<T: Into<PacketType> + TryFrom<i32>>(bytes: Vec<u8>) -> Result<Self>
where
<T as std::convert::TryFrom<i32>>::Error: std::error::Error,
<T as std::convert::TryFrom<i32>>::Error: std::marker::Send,
<T as std::convert::TryFrom<i32>>::Error: std::marker::Sync,
<T as std::convert::TryFrom<i32>>::Error: 'static,
{
if bytes.len() < 10 + 4 {
return Err(anyhow!(PacketDeserialiseError::InsuficientSize(
bytes.len() as i32
)));
}
let size = i32::from_le_bytes(bytes[0..4].try_into()?);
if bytes.len() < (size + 4) as usize {
return Err(anyhow!(PacketDeserialiseError::WrongSize {
expected: size + 4,
got: bytes.len() as i32,
}));
}
let id = i32::from_le_bytes(bytes[4..8].try_into()?);
let msg_type: T = i32::from_le_bytes(bytes[8..12].try_into()?).try_into()?;
let body = from_utf8(&bytes[12..12 + (size - 10) as usize])?.to_string();
Ok(Self {
size,
id,
msg_type: msg_type.into(),
body,
terminator: [0; 2],
})
}
}
impl Display for Packet {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let msg = match self.msg_type {
PacketType::Server(ServerPacketType::AuthResponse) => format!(
"Authentication has {}",
if self.id != -1 { "SUCCEDED" } else { "FAILED" }
),
PacketType::Server(ServerPacketType::ResponseValue) => self.body.to_string(),
PacketType::Client(ClientPacketType::Auth) => {
format!("Authentication with password '{}'", self.body)
}
PacketType::Client(ClientPacketType::ExecCommand) => format!("Command: {}", self.body),
};
write!(f, "{}", msg)
}
}