imks/engine/packet.rs

use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u8)]
pub enum PacketType {
    Open = 0,
    Close = 1,
    Ping = 2,
    Pong = 3,
    Message = 4,
    Upgrade = 5,
    Noop = 6,
}

impl TryFrom<u8> for PacketType {
    type Error = PacketError;

    fn try_from(value: u8) -> Result<Self, Self::Error> {
        match value {
            0 => Ok(Self::Open),
            1 => Ok(Self::Close),
            2 => Ok(Self::Ping),
            3 => Ok(Self::Pong),
            4 => Ok(Self::Message),
            5 => Ok(Self::Upgrade),
            6 => Ok(Self::Noop),
            _ => Err(PacketError::InvalidType(value)),
        }
    }
}

impl TryFrom<char> for PacketType {
    type Error = PacketError;

    fn try_from(value: char) -> Result<Self, Self::Error> {
        match value {
            '0' => Ok(Self::Open),
            '1' => Ok(Self::Close),
            '2' => Ok(Self::Ping),
            '3' => Ok(Self::Pong),
            '4' => Ok(Self::Message),
            '5' => Ok(Self::Upgrade),
            '6' => Ok(Self::Noop),
            _ => Err(PacketError::InvalidTypeChar(value)),
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PacketData {
    Text(String),
    Binary(Vec<u8>),
    Empty,
}

#[derive(Debug, Clone)]
pub struct Packet {
    pub packet_type: PacketType,
    pub data: PacketData,
}

impl Packet {
    pub fn open(handshake: &HandshakeData) -> Self {
        let data = serde_json::to_string(handshake)
            .unwrap_or_else(|e| {
                tracing::error!("Failed to serialize handshake data: {}", e);
                "{}".to_string()
            });
        Self {
            packet_type: PacketType::Open,
            data: PacketData::Text(data),
        }
    }

    pub fn close() -> Self {
        Self {
            packet_type: PacketType::Close,
            data: PacketData::Empty,
        }
    }

    pub fn ping(data: impl Into<String>) -> Self {
        Self {
            packet_type: PacketType::Ping,
            data: PacketData::Text(data.into()),
        }
    }

    pub fn pong(data: impl Into<String>) -> Self {
        Self {
            packet_type: PacketType::Pong,
            data: PacketData::Text(data.into()),
        }
    }

    pub fn message_text(data: impl Into<String>) -> Self {
        Self {
            packet_type: PacketType::Message,
            data: PacketData::Text(data.into()),
        }
    }

    pub fn message_binary(data: Vec<u8>) -> Self {
        Self {
            packet_type: PacketType::Message,
            data: PacketData::Binary(data),
        }
    }

    pub fn upgrade() -> Self {
        Self {
            packet_type: PacketType::Upgrade,
            data: PacketData::Empty,
        }
    }

    pub fn noop() -> Self {
        Self {
            packet_type: PacketType::Noop,
            data: PacketData::Empty,
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HandshakeData {
    pub sid: String,
    pub upgrades: Vec<String>,
    #[serde(rename = "pingInterval")]
    pub ping_interval: u64,
    #[serde(rename = "pingTimeout")]
    pub ping_timeout: u64,
    #[serde(rename = "maxPayload")]
    pub max_payload: usize,
}

#[derive(Debug, thiserror::Error)]
pub enum PacketError {
    #[error("invalid packet type: {0}")]
    InvalidType(u8),
    #[error("invalid packet type char: {0}")]
    InvalidTypeChar(char),
    #[error("empty packet")]
    Empty,
    #[error("invalid base64: {0}")]
    InvalidBase64(#[from] base64::DecodeError),
    #[error("invalid utf8: {0}")]
    InvalidUtf8(#[from] std::string::FromUtf8Error),
    #[error("serialization error: {0}")]
    Serialization(String),
}