rusty_common/websocket/

message.rs

//! WebSocket message types
//!
//! Provides a unified message type that wraps yawc's FrameView.

use smartstring::alias::String;
use std::fmt;
use yawc::close::CloseCode;
use yawc::frame::{FrameView, OpCode};

/// WebSocket message type
///
/// Note: This enum holds owned data (String, Vec<u8>) which prevents true zero-copy
/// operations with memory pools. For high-performance scenarios requiring zero-copy,
/// consider using the lower-level yawc::FrameView directly or implementing a custom
/// message type that holds buffer handles instead of owned data.
#[derive(Clone)]
pub enum Message {
    /// Text message
    Text(String),

    /// Binary message
    Binary(Vec<u8>),

    /// Ping message
    Ping(Vec<u8>),

    /// Pong message
    Pong(Vec<u8>),

    /// Close message
    Close(Option<(u16, String)>),

    /// Raw frame (for compatibility)
    Frame(FrameView),
}

impl Message {
    /// Create a text message
    pub fn text<S: Into<String>>(text: S) -> Self {
        Message::Text(text.into())
    }

    /// Create a binary message
    pub fn binary<V: Into<Vec<u8>>>(data: V) -> Self {
        Message::Binary(data.into())
    }

    /// Create a ping message
    pub fn ping<V: Into<Vec<u8>>>(data: V) -> Self {
        Message::Ping(data.into())
    }

    /// Create a pong message
    pub fn pong<V: Into<Vec<u8>>>(data: V) -> Self {
        Message::Pong(data.into())
    }

    /// Create a close message
    #[must_use]
    pub fn close(code: u16, reason: &str) -> Self {
        Message::Close(Some((code, reason.into())))
    }

    /// Check if this is a text message
    pub const fn is_text(&self) -> bool {
        matches!(self, Message::Text(_))
    }

    /// Check if this is a binary message
    pub const fn is_binary(&self) -> bool {
        matches!(self, Message::Binary(_))
    }

    /// Check if this is a ping message
    pub const fn is_ping(&self) -> bool {
        matches!(self, Message::Ping(_))
    }

    /// Check if this is a pong message
    pub const fn is_pong(&self) -> bool {
        matches!(self, Message::Pong(_))
    }

    /// Check if this is a close message
    pub const fn is_close(&self) -> bool {
        matches!(self, Message::Close(_))
    }

    /// Check if this is a frame message
    pub const fn is_frame(&self) -> bool {
        matches!(self, Message::Frame(_))
    }

    /// Get the text content if this is a text message
    pub fn as_text(&self) -> Option<&str> {
        match self {
            Message::Text(text) => Some(text.as_str()),
            _ => None,
        }
    }

    /// Get the binary content if this is a binary message
    pub fn as_binary(&self) -> Option<&[u8]> {
        match self {
            Message::Binary(data) => Some(data),
            _ => None,
        }
    }

    /// Convert to yawc FrameView
    pub fn to_frame_view(self) -> FrameView {
        match self {
            Message::Text(text) => FrameView::text(text.to_string()),
            Message::Binary(data) => FrameView::binary(data),
            Message::Ping(data) => FrameView::ping(data),
            Message::Pong(data) => FrameView::pong(data),
            Message::Close(Some((code, reason))) => {
                let close_code = CloseCode::Iana(code);
                FrameView::close(close_code, reason.as_str())
            }
            Message::Close(None) => FrameView::close(CloseCode::Iana(1000), ""),
            Message::Frame(frame) => frame,
        }
    }

    /// Create from yawc FrameView
    #[must_use]
    pub fn from_frame_view(frame: FrameView) -> Self {
        match frame.opcode {
            OpCode::Text => {
                let text = std::str::from_utf8(&frame.payload)
                    .unwrap_or_default()
                    .into();
                Message::Text(text)
            }
            OpCode::Binary => Message::Binary(frame.payload.to_vec()),
            OpCode::Ping => Message::Ping(frame.payload.to_vec()),
            OpCode::Pong => Message::Pong(frame.payload.to_vec()),
            OpCode::Close => {
                if frame.payload.len() >= 2 {
                    let code = u16::from_be_bytes([frame.payload[0], frame.payload[1]]);
                    let reason = if frame.payload.len() > 2 {
                        std::str::from_utf8(&frame.payload[2..])
                            .unwrap_or_default()
                            .into()
                    } else {
                        String::new()
                    };
                    Message::Close(Some((code, reason)))
                } else {
                    Message::Close(None)
                }
            }
            _ => Message::Binary(vec![]), // Unknown opcodes
        }
    }
}

/// Message type classification
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MessageType {
    /// Data message (text or binary)
    Data,

    /// Control message (ping, pong, close)
    Control,
}

impl Message {
    /// Get the message type
    pub const fn message_type(&self) -> MessageType {
        match self {
            Message::Text(_) | Message::Binary(_) => MessageType::Data,
            Message::Ping(_) | Message::Pong(_) | Message::Close(_) | Message::Frame(_) => {
                MessageType::Control
            }
        }
    }
}

impl fmt::Debug for Message {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Message::Text(text) => f.debug_tuple("Text").field(text).finish(),
            Message::Binary(data) => f
                .debug_tuple("Binary")
                .field(&format!("{} bytes", data.len()))
                .finish(),
            Message::Ping(data) => f
                .debug_tuple("Ping")
                .field(&format!("{} bytes", data.len()))
                .finish(),
            Message::Pong(data) => f
                .debug_tuple("Pong")
                .field(&format!("{} bytes", data.len()))
                .finish(),
            Message::Close(frame) => f.debug_tuple("Close").field(frame).finish(),
            Message::Frame(_) => f.debug_struct("Frame").finish_non_exhaustive(),
        }
    }
}

impl fmt::Display for Message {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Message::Text(text) => write!(f, "Text message: {text}"),
            Message::Binary(data) => write!(f, "Binary message: {} bytes", data.len()),
            Message::Ping(data) => write!(f, "Ping message: {} bytes", data.len()),
            Message::Pong(data) => write!(f, "Pong message: {} bytes", data.len()),
            Message::Close(Some((code, reason))) => {
                write!(f, "Close message: {code} - {reason}")
            }
            Message::Close(None) => write!(f, "Close message"),
            Message::Frame(_) => write!(f, "Frame message"),
        }
    }
}