osirose-new/packet-service/src/packet.rs

use crate::metrics::PACKETS_SENT;
use crate::packet_type::PacketType;
use bincode::{Decode, Encode};
use std::error::Error;
use tokio::io::{AsyncWriteExt, WriteHalf};
use tokio::net::TcpStream;
use tracing::{debug, error};

#[derive(Debug)]
pub struct Packet {
    pub packet_size: u16,
    pub packet_type: PacketType,
    pub packet_crc: u16,
    pub payload: Vec<u8>,
}

pub trait PacketPayload {
    fn encode(&self) -> Result<Vec<u8>, Box<dyn std::error::Error + Send + Sync>>
    where
        Self: Encode,
    {
        let config = bincode::config::standard().with_fixed_int_encoding();
        Ok(bincode::encode_to_vec(self, config)?)
    }

    fn decode(data: &[u8]) -> Result<Self, Box<dyn std::error::Error + Send + Sync>>
    where
        Self: Sized,
        Self: Decode<()>,
    {
        let config = bincode::config::standard().with_fixed_int_encoding();
        Ok(bincode::decode_from_slice(data, config)?.0)
    }
}

impl Packet {
    pub fn new<T: PacketPayload + bincode::Encode>(
        packet_type: PacketType,
        payload: &T,
    ) -> Result<Self, Box<dyn Error + Send + Sync>> {
        let encoded_payload = <T as PacketPayload>::encode(payload)?;
        let packet_size = (6 + encoded_payload.len()) as u16;
        // let packet_crc = crc::crc16::checksum_x25(&encoded_payload);
        let packet_crc = 0;

        Ok(Self {
            packet_size,
            packet_type,
            packet_crc,
            payload: encoded_payload,
        })
    }

    pub fn from_raw(data: &[u8]) -> Result<Self, Box<dyn Error + Send + Sync>> {
        if data.len() < 6 {
            error!("Invalid packet: Size too small {}", data.len());
            return Err("Invalid packet: Too short".into());
        }

        // Extract packet fields
        let packet_size = u16::from_le_bytes(data[0..2].try_into()?);
        let raw_packet_type = u16::from_le_bytes(data[2..4].try_into()?);
        let packet_crc = u16::from_le_bytes(data[4..6].try_into()?);

        // Validate the size
        if packet_size as usize > data.len() {
            error!(
                "Invalid packet {:#X}: Size mismatch, expected: {}, received: {}",
                raw_packet_type,
                packet_size,
                data.len()
            );
            return Err("Invalid packet: Size mismatch".into());
        }

        debug!(
            "size: {:#X}, raw_type: {:#X}, crc: {:#X}",
            packet_size, raw_packet_type, packet_crc
        );

        // Convert raw packet type into `PacketType` enum
        let packet_type = PacketType::try_from(raw_packet_type)?;

        // Extract the payload
        let payload = data[6..packet_size as usize].to_vec();

        Ok(Self {
            packet_size,
            packet_type,
            packet_crc,
            payload,
        })
    }

    /// Serialize the `Packet` back to raw bytes
    pub fn to_raw(&self) -> Vec<u8> {
        let mut raw = Vec::with_capacity(self.packet_size as usize);

        // Serialize fields
        raw.extend(&self.packet_size.to_le_bytes());
        raw.extend(&(self.packet_type as u16).to_le_bytes());
        raw.extend(&self.packet_crc.to_le_bytes());
        raw.extend(&self.payload);

        raw
    }

    pub fn parse<T: PacketPayload + bincode::Decode<()>>(&self) -> Result<T, Box<dyn Error + Send + Sync>> {
        <T as PacketPayload>::decode(&self.payload)
    }
}

pub async fn send_packet(stream: &mut WriteHalf<TcpStream>, packet: &Packet) -> Result<(), std::io::Error> {
    let data = packet.to_raw();
    debug!("Sending '{:#X}' bytes of data. {:?}", data.len(), data);
    PACKETS_SENT.inc();
    stream.write_all(&data).await?;
    Ok(())
}