armature_websocket/
message.rs

1//! WebSocket message types.
2
3use bytes::Bytes;
4use serde::{Deserialize, Serialize};
5
6/// Message type enumeration.
7#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
8#[serde(rename_all = "lowercase")]
9pub enum MessageType {
10    /// Text message
11    Text,
12    /// Binary message
13    Binary,
14    /// Ping message
15    Ping,
16    /// Pong message
17    Pong,
18    /// Close message
19    Close,
20}
21
22/// A WebSocket message.
23#[derive(Debug, Clone)]
24pub struct Message {
25    /// The message type
26    pub message_type: MessageType,
27    /// The message payload
28    pub payload: Bytes,
29}
30
31impl Message {
32    /// Create a new text message.
33    pub fn text<S: Into<String>>(text: S) -> Self {
34        Self {
35            message_type: MessageType::Text,
36            payload: Bytes::from(text.into()),
37        }
38    }
39
40    /// Create a new binary message.
41    pub fn binary<B: Into<Bytes>>(data: B) -> Self {
42        Self {
43            message_type: MessageType::Binary,
44            payload: data.into(),
45        }
46    }
47
48    /// Create a new ping message.
49    pub fn ping<B: Into<Bytes>>(data: B) -> Self {
50        Self {
51            message_type: MessageType::Ping,
52            payload: data.into(),
53        }
54    }
55
56    /// Create a new pong message.
57    pub fn pong<B: Into<Bytes>>(data: B) -> Self {
58        Self {
59            message_type: MessageType::Pong,
60            payload: data.into(),
61        }
62    }
63
64    /// Create a close message.
65    pub fn close() -> Self {
66        Self {
67            message_type: MessageType::Close,
68            payload: Bytes::new(),
69        }
70    }
71
72    /// Create a JSON message from a serializable value.
73    pub fn json<T: Serialize>(value: &T) -> Result<Self, serde_json::Error> {
74        let json = serde_json::to_string(value)?;
75        Ok(Self::text(json))
76    }
77
78    /// Parse the message payload as JSON.
79    pub fn parse_json<'a, T: Deserialize<'a>>(&'a self) -> Result<T, serde_json::Error> {
80        serde_json::from_slice(&self.payload)
81    }
82
83    /// Get the message payload as a string.
84    pub fn as_text(&self) -> Option<&str> {
85        if self.message_type == MessageType::Text {
86            std::str::from_utf8(&self.payload).ok()
87        } else {
88            None
89        }
90    }
91
92    /// Get the message payload as bytes.
93    pub fn as_bytes(&self) -> &[u8] {
94        &self.payload
95    }
96
97    /// Check if this is a text message.
98    pub fn is_text(&self) -> bool {
99        self.message_type == MessageType::Text
100    }
101
102    /// Check if this is a binary message.
103    pub fn is_binary(&self) -> bool {
104        self.message_type == MessageType::Binary
105    }
106
107    /// Check if this is a ping message.
108    pub fn is_ping(&self) -> bool {
109        self.message_type == MessageType::Ping
110    }
111
112    /// Check if this is a pong message.
113    pub fn is_pong(&self) -> bool {
114        self.message_type == MessageType::Pong
115    }
116
117    /// Check if this is a close message.
118    pub fn is_close(&self) -> bool {
119        self.message_type == MessageType::Close
120    }
121}
122
123impl From<tungstenite::Message> for Message {
124    fn from(msg: tungstenite::Message) -> Self {
125        match msg {
126            tungstenite::Message::Text(text) => Self::text(text.to_string()),
127            tungstenite::Message::Binary(data) => Self::binary(Bytes::from(data)),
128            tungstenite::Message::Ping(data) => Self::ping(Bytes::from(data)),
129            tungstenite::Message::Pong(data) => Self::pong(Bytes::from(data)),
130            tungstenite::Message::Close(_) => Self::close(),
131            tungstenite::Message::Frame(_) => Self::binary(Bytes::new()),
132        }
133    }
134}
135
136impl From<Message> for tungstenite::Message {
137    fn from(msg: Message) -> Self {
138        match msg.message_type {
139            MessageType::Text => {
140                tungstenite::Message::Text(String::from_utf8_lossy(&msg.payload).into_owned().into())
141            }
142            MessageType::Binary => tungstenite::Message::Binary(msg.payload.to_vec().into()),
143            MessageType::Ping => tungstenite::Message::Ping(msg.payload.to_vec().into()),
144            MessageType::Pong => tungstenite::Message::Pong(msg.payload.to_vec().into()),
145            MessageType::Close => tungstenite::Message::Close(None),
146        }
147    }
148}
149