sioc 0.2.0

//! Socket.IO v4 packet types and wire encoding.

use crate::error::PacketError;
use bytes::Bytes;
use bytestring::ByteString;
use miette::Diagnostic;
use serde::Deserialize;
use serde_json::{Map, Value};
use thiserror::Error;
use tokio::sync::{mpsc, oneshot};

/// A value tagged with a Socket.IO namespace path (e.g. `"/chat"`).
#[derive(Debug)]
pub struct Ns<T>(pub ByteString, pub T);

/// Server payload confirming a successful namespace connection.
#[derive(Debug, Deserialize)]
pub struct Connect {
    /// Server-assigned session ID.
    pub sid: ByteString,
    #[serde(flatten)]
    pub extra: Map<String, Value>,
}

/// Server payload for a rejected namespace connection.
#[derive(Debug, Error, Diagnostic, Deserialize)]
#[error("{message}")]
#[diagnostic(
    code(sioc::connect_error),
    help(
        "Server rejected the namespace connection. Verify your auth payload and server middleware."
    ),
    url("https://socket.io/docs/v4/socket-io-protocol/#connection-to-a-namespace")
)]
pub struct ConnectError {
    pub message: ByteString,
    #[serde(flatten)]
    pub extra: Map<String, Value>,
}

/// Type-erased inbound event after binary reassembly.
#[derive(Debug, Clone)]
pub struct DynEvent {
    pub payload: ByteString,
    pub id: Option<u64>,
    pub attachments: Option<Vec<Bytes>>,
}

impl std::fmt::Display for DynEvent {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let mut map = f.debug_map();
        map.entry(&"payload", &format_args!("{}", self.payload));
        if let Some(id) = self.id {
            map.entry(&"id", &id);
        }
        if let Some(attachments) = &self.attachments {
            map.entry(&"count", &attachments.len());
        }
        map.finish()
    }
}

impl DynEvent {
    pub fn new<T>(payload: T, id: Option<u64>) -> Self
    where
        T: Into<ByteString>,
    {
        Self {
            payload: payload.into(),
            id,
            attachments: None,
        }
    }

    pub fn with_attachments(mut self, attachments: Vec<Bytes>) -> Self {
        self.attachments = Some(attachments);
        self
    }
}

/// Type-erased inbound acknowledgement after binary reassembly.
///
/// Convert to a typed `sioc::Ack` via `TryFrom<DynAck>`.
#[derive(Debug, Clone)]
pub struct DynAck {
    pub payload: ByteString,
    pub attachments: Option<Vec<Bytes>>,
}

impl std::fmt::Display for DynAck {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let mut map = f.debug_map();
        map.entry(&"payload", &format_args!("{}", self.payload));
        if let Some(attachments) = &self.attachments {
            map.entry(&"count", &attachments.len());
        }
        map.finish()
    }
}

impl DynAck {
    pub fn new<T>(payload: T) -> Self
    where
        T: Into<ByteString>,
    {
        Self {
            payload: payload.into(),
            attachments: None,
        }
    }

    pub fn with_attachments(mut self, attachments: Vec<Bytes>) -> Self {
        self.attachments = Some(attachments);
        self
    }
}

/// A fully decoded inbound packet.
#[derive(Debug)]
pub enum Signal<E = DynEvent> {
    /// The server confirmed the namespace connection.
    Connect(Connect),
    /// The namespace was disconnected. Does not close the receiver; a reconnect delivers a new [`Signal::Connect`].
    Disconnect,
    /// The server rejected a namespace connection attempt. Does not close the receiver.
    ConnectError(ConnectError),
    /// An application-level event (possibly with binary attachments).
    Event(E),
}

impl<E> std::fmt::Display for Signal<E>
where
    E: std::fmt::Display,
{
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Connect(c) => f
                .debug_tuple("Connect")
                .field(&format_args!("{}", c.sid))
                .finish(),
            Self::Disconnect => f.write_str("Disconnect"),
            Self::ConnectError(e) => f
                .debug_tuple("ConnectError")
                .field(&format_args!("{}", e))
                .finish(),
            Self::Event(e) => f
                .debug_tuple("Event")
                .field(&format_args!("{}", e))
                .finish(),
        }
    }
}

impl<E> Signal<E> {
    /// Returns the inner event if this is [`Event`](Signal::Event), otherwise `None`.
    pub fn take_event(self) -> Option<E> {
        match self {
            Self::Event(e) => Some(e),
            _ => None,
        }
    }

    /// Applies `f` to the inner event if this is [`Event`](Signal::Event), otherwise returns `None`.
    pub fn and_then<F, T>(self, f: F) -> Option<T>
    where
        F: FnOnce(E) -> Option<T>,
    {
        match self {
            Self::Event(e) => f(e),
            _ => None,
        }
    }

    /// Applies `f` to the event in [`Event`](Signal::Event), passing other variants through unchanged.
    pub fn map<F, U>(self, f: F) -> Signal<U>
    where
        F: FnOnce(E) -> U,
    {
        match self {
            Self::Connect(c) => Signal::Connect(c),
            Self::Disconnect => Signal::Disconnect,
            Self::ConnectError(e) => Signal::ConnectError(e),
            Self::Event(e) => Signal::Event(f(e)),
        }
    }
}

/// An outbound packet to be encoded and sent to the server.
#[derive(Debug)]
pub enum Directive {
    /// Opens a namespace; `data` is an optional authentication payload.
    Connect {
        tx: mpsc::Sender<Signal>,
        payload: ByteString,
    },
    /// Closes the namespace.
    Disconnect,
    /// Emits an event; if `tx` is set, an ack ID is assigned and the response routed to it.
    Event {
        payload: ByteString,
        tx: Option<oneshot::Sender<DynAck>>,
        attachments: Option<Vec<Bytes>>,
    },
    /// Acknowledges a previously received event.
    Ack {
        payload: ByteString,
        id: u64,
        attachments: Option<Vec<Bytes>>,
    },
    /// Sent by `SocketSenderInner::drop`; the manager warns if the namespace is still live.
    Dropped,
}

/// A wire-level packet decoded from a single text frame.
///
/// Binary variants carry an attachment count; the socket router collects
/// the follow-up binary frames and reassembles them into a [`Signal`].
#[derive(Debug)]
pub enum Packet {
    /// Type `0`: namespace connection confirmed.
    Connect(ByteString),
    /// Type `1`: namespace disconnection.
    Disconnect,
    /// Type `2`: event (text or binary).
    Event {
        payload: ByteString,
        id: Option<u64>,
    },
    /// Type `3`: acknowledgement (text or binary).
    Ack { payload: ByteString, id: u64 },
    /// Type `4`: namespace connection rejected.
    ConnectError(ByteString),

    /// Type `5`: binary event with `count` follow-up binary frames.
    BinaryEvent {
        payload: ByteString,
        id: Option<u64>,
        count: usize,
    },

    /// Type `6`: binary acknowledgement with `count` follow-up binary frames.
    BinaryAck {
        payload: ByteString,
        id: u64,
        count: usize,
    },
}

impl Packet {
    /// Returns a conservative upper bound on the serialised text-frame byte length.
    pub fn size_hint(&self, ns: &str) -> usize {
        match self {
            Self::Connect(payload) => hint_packet_size(ns, false, false, Some(payload)),
            Self::Disconnect => hint_packet_size(ns, false, false, None),
            Self::Event { payload, id } => hint_packet_size(ns, false, id.is_some(), Some(payload)),
            Self::Ack { payload, .. } => hint_packet_size(ns, false, true, Some(payload)),
            Self::ConnectError(payload) => hint_packet_size(ns, false, false, Some(payload)),
            Self::BinaryEvent { payload, id, .. } => {
                hint_packet_size(ns, true, id.is_some(), Some(payload))
            }
            Self::BinaryAck { payload, .. } => hint_packet_size(ns, true, true, Some(payload)),
        }
    }

    pub fn encode(&self, ns: &str) -> String {
        let mut buffer = String::with_capacity(self.size_hint(ns));

        match self {
            Self::Connect(bytes) => write_packet(&mut buffer, b'0', None, ns, None, Some(bytes)),

            Self::Disconnect => write_packet(&mut buffer, b'1', None, ns, None, None),
            Self::Event { payload, id } => {
                write_packet(&mut buffer, b'2', None, ns, *id, Some(payload))
            }
            Self::Ack { payload, id } => {
                write_packet(&mut buffer, b'3', None, ns, Some(*id), Some(payload))
            }
            Self::ConnectError(payload) => {
                write_packet(&mut buffer, b'4', None, ns, None, Some(payload))
            }
            Self::BinaryEvent { payload, id, count } => {
                write_packet(&mut buffer, b'5', Some(*count), ns, *id, Some(payload))
            }
            Self::BinaryAck { payload, id, count } => write_packet(
                &mut buffer,
                b'6',
                Some(*count),
                ns,
                Some(*id),
                Some(payload),
            ),
        }

        buffer
    }
}

impl std::fmt::Display for Packet {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Connect(payload) => f
                .debug_tuple("Connect")
                .field(&format_args!("{}", payload))
                .finish(),
            Self::Disconnect => f.write_str("Disconnect"),
            Self::Event { payload, id } => {
                let mut s = f.debug_struct("Event");
                s.field("payload", &format_args!("{}", payload));
                if let Some(id) = id {
                    s.field("id", id);
                }
                s.finish()
            }
            Self::Ack { payload, id } => f
                .debug_struct("Ack")
                .field("payload", &format_args!("{}", payload))
                .field("id", id)
                .finish(),
            Self::ConnectError(payload) => f
                .debug_tuple("ConnectError")
                .field(&format_args!("{}", payload))
                .finish(),
            Self::BinaryEvent { payload, id, count } => {
                let mut s = f.debug_struct("BinaryEvent");
                s.field("payload", &format_args!("{}", payload));
                if let Some(id) = id {
                    s.field("id", id);
                }
                s.field("count", count);
                s.finish()
            }
            Self::BinaryAck { payload, id, count } => f
                .debug_struct("BinaryAck")
                .field("payload", &format_args!("{}", payload))
                .field("id", id)
                .field("count", count)
                .finish(),
        }
    }
}

impl Ns<Packet> {
    pub fn size_hint(&self) -> usize {
        self.1.size_hint(&self.0)
    }

    pub fn encode(&self) -> String {
        self.1.encode(&self.0)
    }
}

impl TryFrom<ByteString> for Ns<Packet> {
    type Error = PacketError;

    fn try_from(bytes: ByteString) -> Result<Self, PacketError> {
        let mut chars = bytes.chars();

        let id = chars.next().ok_or(PacketError::Empty)?;

        let bytes = bytes.slice_ref(chars.as_str());

        let packet = match id {
            '0' => {
                let (ns, payload) = split_namespace(bytes)?;
                Ns(ns, Packet::Connect(payload))
            }
            '1' => {
                let (ns, _) = split_namespace(bytes)?;
                Ns(ns, Packet::Disconnect)
            }
            '2' => {
                let (count, bytes) = split_attachments(bytes)?;
                if let Some(count) = count {
                    return Err(PacketError::UnexpectedAttachments { count });
                }
                let (ns, bytes) = split_namespace(bytes)?;
                let (id, payload) = split_id(bytes)?;
                Ns(ns, Packet::Event { payload, id })
            }
            '3' => {
                let (ns, bytes) = split_namespace(bytes)?;
                let (id, payload) = split_id(bytes)?;
                let id = id.ok_or(PacketError::MissingAckId)?;
                Ns(ns, Packet::Ack { payload, id })
            }
            '4' => {
                let (ns, payload) = split_namespace(bytes)?;
                Ns(ns, Packet::ConnectError(payload))
            }
            '5' => {
                let (count, bytes) = split_attachments(bytes)?;
                let count = count.ok_or(PacketError::MissingAttachmentCount)?;
                let (ns, bytes) = split_namespace(bytes)?;
                let (id, payload) = split_id(bytes)?;
                Ns(ns, Packet::BinaryEvent { payload, id, count })
            }
            '6' => {
                let (count, bytes) = split_attachments(bytes)?;
                let count = count.ok_or(PacketError::MissingAttachmentCount)?;
                let (ns, bytes) = split_namespace(bytes)?;
                let (id, payload) = split_id(bytes)?;
                let id = id.ok_or(PacketError::MissingAckId)?;
                Ns(ns, Packet::BinaryAck { payload, id, count })
            }
            id => return Err(PacketError::InvalidId { id }),
        };

        Ok(packet)
    }
}

const U64_MAX_LEN: usize = 20; // max decimal digits in a u64

const fn ack_size_hint() -> usize {
    U64_MAX_LEN
}

const fn binary_size_hint() -> usize {
    U64_MAX_LEN + 1
}

/// Returns the encoded byte length of a namespace field (`0` for the default `"/"`).
fn namespace_size(ns: &str) -> usize {
    if ns == "/" { 0 } else { ns.len() + 1 }
}

pub fn hint_packet_size(ns: &str, binary: bool, ack: bool, payload: Option<&str>) -> usize {
    let mut n = 1 + namespace_size(ns);
    if ack {
        n += ack_size_hint();
    }
    if binary {
        n += binary_size_hint();
    }
    if let Some(payload) = payload {
        n += payload.len();
    }
    n
}

/// Writes `<count>-` into `buffer`.
fn write_attachments(buffer: &mut String, count: usize) {
    buffer.push_str(&count.to_string());
    buffer.push('-');
}

/// Writes the namespace followed by `,`; skipped for the default `"/"`.
fn write_namespace(buffer: &mut String, ns: &str) {
    if ns != "/" {
        buffer.push_str(ns);
        buffer.push(',');
    }
}

/// Writes a decimal ack ID.
fn write_id(buffer: &mut String, id: u64) {
    buffer.push_str(&id.to_string());
}

fn write_payload(buffer: &mut String, payload: &str) {
    buffer.push_str(payload);
}

pub fn write_packet(
    buffer: &mut String,
    type_id: u8,
    count: Option<usize>,
    ns: &str,
    id: Option<u64>,
    payload: Option<&str>,
) {
    buffer.push(type_id as char);
    if let Some(count) = count {
        write_attachments(buffer, count);
    }
    write_namespace(buffer, ns);
    if let Some(id) = id {
        write_id(buffer, id);
    }
    if let Some(payload) = payload {
        write_payload(buffer, payload);
    }
}

/// Consumes a `<n>-` attachment count prefix.
///
/// Only matches when the data starts with one or more ASCII digits immediately
/// followed by `-`. This avoids misinterpreting a namespace that contains a
/// hyphen (e.g. `/admin-ns,["data"]`) as an attachment count.
///
/// Returns `(Some(count), rest)` when the prefix is present, `(None, bytes)`
/// when absent.
pub fn split_attachments(bytes: ByteString) -> Result<(Option<usize>, ByteString), PacketError> {
    let pair = match bytes.char_indices().find(|(_, c)| !c.is_ascii_digit()) {
        Some((i, '-')) => {
            let count = bytes[..i]
                .parse()
                .map_err(PacketError::InvalidAttachmentCount)?;

            let rest = bytes.slice_ref(&bytes[i + 1..]);

            (Some(count), rest)
        }
        _ => (None, bytes),
    };

    Ok(pair)
}

/// Consumes a `/ns,` prefix and returns `(namespace, rest)`.
///
/// Returns `"/"` for the default namespace.
pub fn split_namespace(bytes: ByteString) -> Result<(ByteString, ByteString), PacketError> {
    match bytes.chars().next() {
        Some('/') => match bytes.split_once(',') {
            Some((ns, payload)) => Ok((bytes.slice_ref(ns), bytes.slice_ref(payload))),
            None => Err(PacketError::MissingNamespaceDelimiter),
        },
        _ => Ok((ByteString::from_static("/"), bytes)),
    }
}

/// Consumes a run of leading ASCII digits as a `u64`.
///
/// Returns `(Some(id), rest)` when digits are present and fit in `u64`,
/// `(None, bytes)` when no leading digits are found or the input is empty.
pub fn split_id(bytes: ByteString) -> Result<(Option<u64>, ByteString), PacketError> {
    let pair = match bytes.char_indices().find(|(_, c)| !c.is_ascii_digit()) {
        Some((i, _)) if i > 0 => {
            let id = bytes[..i].parse().map_err(PacketError::InvalidAckId)?;
            let rest = bytes.slice_ref(&bytes[i..]);
            (Some(id), rest)
        }
        _ => (None, bytes),
    };

    Ok(pair)
}