gitdataai/lib/socketio/packet.rs

use serde_json::{Value, json};

use crate::error::{Result, SocketIoError};

#[derive(
    Clone, Copy, Debug, Eq, PartialEq, serde::Deserialize, serde::Serialize,
)]
pub enum PacketType {
    Connect,
    Disconnect,
    Event,
    Ack,
    ConnectError,
    BinaryEvent,
    BinaryAck,
}

#[derive(Clone, Debug, PartialEq)]
pub struct Packet {
    pub packet_type: PacketType,
    pub namespace: String,
    pub id: Option<u64>,
    pub data: Option<Value>,
    pub attachments: Vec<Vec<u8>>,
    pub expected_attachments: usize,
}

#[derive(Clone, Debug)]
pub struct EventPayload {
    pub event: String,
    pub args: Vec<Value>,
    pub binary: Vec<Vec<u8>>,
    pub ack_id: Option<u64>,
    pub ack: Option<crate::socket::AckSender>,
}

impl Packet {
    pub fn connect(namespace: impl Into<String>, data: Option<Value>) -> Self {
        Self::new(PacketType::Connect, namespace, None, data)
    }

    pub fn event(
        namespace: impl Into<String>,
        event: &str,
        args: Vec<Value>,
    ) -> Self {
        let mut data = Vec::with_capacity(args.len() + 1);
        data.push(Value::String(event.to_owned()));
        data.extend(args);
        Self::new(PacketType::Event, namespace, None, Some(Value::Array(data)))
    }

    pub fn ack(
        namespace: impl Into<String>,
        id: u64,
        args: Vec<Value>,
    ) -> Self {
        Self::new(
            PacketType::Ack,
            namespace,
            Some(id),
            Some(Value::Array(args)),
        )
    }

    pub fn connect_error(
        namespace: impl Into<String>,
        message: impl Into<String>,
    ) -> Self {
        Self::new(
            PacketType::ConnectError,
            namespace,
            None,
            Some(json!({ "message": message.into() })),
        )
    }

    pub fn new(
        packet_type: PacketType,
        namespace: impl Into<String>,
        id: Option<u64>,
        data: Option<Value>,
    ) -> Self {
        Self {
            packet_type,
            namespace: namespace.into(),
            id,
            data,
            attachments: Vec::new(),
            expected_attachments: 0,
        }
    }

    pub fn with_binary(mut self, attachments: Vec<Vec<u8>>) -> Self {
        self.expected_attachments = attachments.len();
        if self.expected_attachments > 0
            && let Some(Value::Array(values)) = &mut self.data
        {
            for num in 0..self.expected_attachments {
                values.push(json!({ "_placeholder": true, "num": num }));
            }
        }
        self.attachments = attachments;
        self.packet_type = match self.packet_type {
            PacketType::Ack => PacketType::BinaryAck,
            PacketType::Event => PacketType::BinaryEvent,
            other => other,
        };
        self
    }

    pub fn encode(&self) -> String {
        let mut out = String::new();
        out.push(packet_type_digit(self.packet_type));

        if matches!(
            self.packet_type,
            PacketType::BinaryEvent | PacketType::BinaryAck
        ) {
            out.push_str(&self.expected_attachments.to_string());
            out.push('-');
        }

        if self.namespace != "/" {
            out.push_str(&self.namespace);
            out.push(',');
        }

        if let Some(id) = self.id {
            out.push_str(&id.to_string());
        }

        if let Some(data) = &self.data {
            out.push_str(&data.to_string());
        }

        out
    }

    pub fn decode(input: &str) -> Result<Self> {
        let mut chars = input.char_indices();
        let (_, first) = chars.next().ok_or_else(|| {
            SocketIoError::InvalidPacket("empty packet".to_owned())
        })?;
        let packet_type = packet_type_from_digit(first)?;
        let mut index = first.len_utf8();
        let mut expected_attachments = 0;

        if matches!(
            packet_type,
            PacketType::BinaryEvent | PacketType::BinaryAck
        ) {
            let rest = &input[index..];
            let dash = rest.find('-').ok_or_else(|| {
                SocketIoError::InvalidPacket(
                    "binary packet missing attachment count".to_owned(),
                )
            })?;
            expected_attachments = rest[..dash].parse().map_err(|_| {
                SocketIoError::InvalidPacket(
                    "invalid attachment count".to_owned(),
                )
            })?;
            index += dash + 1;
        }

        let namespace = if input[index..].starts_with('/') {
            let rest = &input[index..];
            if let Some(comma) = rest.find(',') {
                index += comma + 1;
                rest[..comma].to_owned()
            } else {
                index = input.len();
                rest.to_owned()
            }
        } else {
            "/".to_owned()
        };

        let id_start = index;
        while let Some(ch) = input[index..].chars().next() {
            if ch.is_ascii_digit() {
                index += ch.len_utf8();
            } else {
                break;
            }
        }
        let id = if index > id_start {
            Some(input[id_start..index].parse().map_err(|_| {
                SocketIoError::InvalidPacket(
                    "invalid acknowledgment id".to_owned(),
                )
            })?)
        } else {
            None
        };

        let data = if index < input.len() {
            Some(serde_json::from_str(&input[index..])?)
        } else {
            None
        };

        Ok(Self {
            packet_type,
            namespace,
            id,
            data,
            attachments: Vec::new(),
            expected_attachments,
        })
    }

    pub(crate) fn into_event_payload(
        self,
        ack: Option<crate::socket::AckSender>,
    ) -> Result<EventPayload> {
        let values = match self.data {
            Some(Value::Array(values)) if !values.is_empty() => values,
            _ => {
                return Err(SocketIoError::InvalidPacket(
                    "event payload must be a non-empty array".to_owned(),
                ));
            }
        };
        let mut values = values.into_iter();
        let event = values
            .next()
            .and_then(|value| value.as_str().map(ToOwned::to_owned))
            .ok_or_else(|| {
                SocketIoError::InvalidPacket(
                    "event name must be a string".to_owned(),
                )
            })?;

        let args = values.filter(|value| !is_placeholder(value)).collect();

        Ok(EventPayload {
            event,
            args,
            binary: self.attachments,
            ack_id: self.id,
            ack,
        })
    }
}

fn is_placeholder(value: &Value) -> bool {
    value
        .as_object()
        .and_then(|object| object.get("_placeholder"))
        .and_then(Value::as_bool)
        .unwrap_or(false)
}

fn packet_type_digit(packet_type: PacketType) -> char {
    match packet_type {
        PacketType::Connect => '0',
        PacketType::Disconnect => '1',
        PacketType::Event => '2',
        PacketType::Ack => '3',
        PacketType::ConnectError => '4',
        PacketType::BinaryEvent => '5',
        PacketType::BinaryAck => '6',
    }
}

fn packet_type_from_digit(value: char) -> Result<PacketType> {
    match value {
        '0' => Ok(PacketType::Connect),
        '1' => Ok(PacketType::Disconnect),
        '2' => Ok(PacketType::Event),
        '3' => Ok(PacketType::Ack),
        '4' => Ok(PacketType::ConnectError),
        '5' => Ok(PacketType::BinaryEvent),
        '6' => Ok(PacketType::BinaryAck),
        _ => Err(SocketIoError::InvalidPacket(format!(
            "unknown socket packet type {value}"
        ))),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn event_packet_round_trips() {
        let packet =
            Packet::event("/", "message", vec![json!({ "body": "hello" })]);
        assert_eq!(packet.encode(), "2[\"message\",{\"body\":\"hello\"}]");

        let decoded = Packet::decode(&packet.encode()).unwrap();
        assert_eq!(decoded.packet_type, PacketType::Event);
        assert_eq!(decoded.namespace, "/");
        assert_eq!(decoded.data, packet.data);
    }

    #[test]
    fn event_packet_accepts_namespace_and_ack_id() {
        let decoded =
            Packet::decode("2/admin,17[\"save\",{\"ok\":true}]").unwrap();
        assert_eq!(decoded.packet_type, PacketType::Event);
        assert_eq!(decoded.namespace, "/admin");
        assert_eq!(decoded.id, Some(17));
    }

    #[test]
    fn binary_packet_accepts_attachment_count() {
        let decoded = Packet::decode(
            "51-/admin,13[\"file\",{\"_placeholder\":true,\"num\":0}]",
        )
        .unwrap();
        assert_eq!(decoded.packet_type, PacketType::BinaryEvent);
        assert_eq!(decoded.expected_attachments, 1);
        assert_eq!(decoded.namespace, "/admin");
        assert_eq!(decoded.id, Some(13));
    }

    #[test]
    fn binary_emit_adds_placeholders() {
        let packet = Packet::event("/", "file", vec![json!("meta")])
            .with_binary(vec![vec![1, 2]]);
        assert_eq!(
            packet.encode(),
            "51-[\"file\",\"meta\",{\"_placeholder\":true,\"num\":0}]"
        );
    }

    #[test]
    fn decode_empty_packet_returns_error() {
        assert!(Packet::decode("").is_err());
    }

    #[test]
    fn decode_unknown_type_returns_error() {
        assert!(Packet::decode("9[]").is_err());
    }

    #[test]
    fn default_namespace_omits_prefix_in_encode() {
        let packet = Packet::event("/", "ping", vec![]);
        assert_eq!(packet.encode(), "2[\"ping\"]");
    }

    #[test]
    fn custom_namespace_includes_trailing_comma() {
        let packet = Packet::event("/chat", "msg", vec![json!("hi")]);
        let encoded = packet.encode();
        assert!(encoded.starts_with("2/chat,"));
        let decoded = Packet::decode(&encoded).unwrap();
        assert_eq!(decoded.namespace, "/chat");
    }
}