use bytes::Bytes;
use serde::{de::DeserializeOwned, Serialize};
use crate::error::WsError;
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub enum Frame {
Text(String),
Binary(Bytes),
}
impl Frame {
pub fn binary(data: impl Into<Bytes>) -> Self {
Frame::Binary(data.into())
}
pub fn len(&self) -> usize {
match self {
Frame::Text(s) => s.len(),
Frame::Binary(b) => b.len(),
}
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn is_text(&self) -> bool {
matches!(self, Frame::Text(_))
}
pub fn is_binary(&self) -> bool {
matches!(self, Frame::Binary(_))
}
pub fn as_bytes(&self) -> &[u8] {
match self {
Frame::Text(s) => s.as_bytes(),
Frame::Binary(b) => b.as_ref(),
}
}
pub fn into_bytes(self) -> Bytes {
match self {
Frame::Text(s) => Bytes::from(s),
Frame::Binary(b) => b,
}
}
pub fn into_text(self) -> Result<String, WsError> {
match self {
Frame::Text(s) => Ok(s),
Frame::Binary(_) => Err(WsError::InvalidMessage),
}
}
}
impl From<String> for Frame {
fn from(s: String) -> Self {
Frame::Text(s)
}
}
impl From<&str> for Frame {
fn from(s: &str) -> Self {
Frame::Text(s.to_string())
}
}
impl From<Bytes> for Frame {
fn from(b: Bytes) -> Self {
Frame::Binary(b)
}
}
impl From<Vec<u8>> for Frame {
fn from(v: Vec<u8>) -> Self {
Frame::Binary(Bytes::from(v))
}
}
impl From<&[u8]> for Frame {
fn from(v: &[u8]) -> Self {
Frame::Binary(Bytes::copy_from_slice(v))
}
}
pub trait Codec: Sized {
fn encode(&self) -> String
where
Self: Serialize,
{
serde_json::to_string(self).unwrap_or_default()
}
fn encode_result(&self) -> Result<String, WsError>
where
Self: Serialize,
{
serde_json::to_string(self).map_err(|_| WsError::InvalidMessage)
}
fn decode(s: &str) -> Result<Self, WsError>
where
Self: DeserializeOwned,
{
serde_json::from_str(s).map_err(|_| WsError::InvalidMessage)
}
}
impl<T> Codec for T where T: Serialize + DeserializeOwned {}
pub struct JsonCodec;
impl JsonCodec {
pub fn encode<T: Serialize>(value: &T) -> Result<String, WsError> {
serde_json::to_string(value).map_err(|_| WsError::InvalidMessage)
}
pub fn decode<T: DeserializeOwned>(s: &str) -> Result<T, WsError> {
serde_json::from_str(s).map_err(|_| WsError::InvalidMessage)
}
pub fn encode_pretty<T: Serialize>(value: &T) -> Result<String, WsError> {
serde_json::to_string_pretty(value).map_err(|_| WsError::InvalidMessage)
}
}
#[allow(
clippy::unwrap_used,
clippy::expect_used,
clippy::indexing_slicing,
clippy::panic
)]
#[cfg(test)]
mod tests {
use super::*;
use serde::{Deserialize, Serialize};
#[derive(Serialize, Deserialize, Debug, PartialEq, Clone)]
struct Payload {
id: u32,
msg: String,
}
#[test]
fn codec_roundtrip_trait() {
let p = Payload {
id: 42,
msg: "hello".into(),
};
let s = p.encode();
let d = Payload::decode(&s).unwrap();
assert_eq!(p, d);
}
#[test]
fn codec_roundtrip_result() {
let p = Payload {
id: 1,
msg: "world".into(),
};
let s = p.encode_result().unwrap();
let d = Payload::decode(&s).unwrap();
assert_eq!(p, d);
}
#[test]
fn codec_invalid_decode() {
let err = Payload::decode("not json").unwrap_err();
assert_eq!(err, WsError::InvalidMessage);
}
#[test]
fn json_codec_helpers() {
let p = Payload {
id: 7,
msg: "test".into(),
};
let s = JsonCodec::encode(&p).unwrap();
let d: Payload = JsonCodec::decode(&s).unwrap();
assert_eq!(p, d);
let pretty = JsonCodec::encode_pretty(&p).unwrap();
assert!(pretty.contains('\n'));
}
#[test]
fn encode_vec() {
let v = vec![1, 2, 3];
let s = v.encode();
let d: Vec<i32> = Vec::decode(&s).unwrap();
assert_eq!(v, d);
}
#[test]
fn frame_constructors_and_kinds() {
assert!(Frame::from("hello".to_string()).is_text());
assert!(Frame::from("hello").is_text());
assert!(Frame::from(Bytes::from_static(b"\x00\x01")).is_binary());
assert!(Frame::from(vec![1u8, 2]).is_binary());
assert!(Frame::from(&[1u8, 2][..]).is_binary());
assert!(Frame::binary(Vec::new()).is_binary());
}
#[test]
fn frame_len_and_bytes() {
let t = Frame::from("héllo"); assert_eq!(t.len(), 6);
assert_eq!(t.as_bytes(), "héllo".as_bytes());
assert!(!t.is_empty());
let b = Frame::from(vec![0u8; 4]);
assert_eq!(b.len(), 4);
assert!(Frame::from(String::new()).is_empty());
}
#[test]
fn frame_into_bytes_and_text() {
let t = Frame::from("abc".to_string());
assert_eq!(&t.into_bytes()[..], b"abc");
let b = Frame::from(vec![9u8]);
assert_eq!(b.into_bytes()[0], 9);
let t2 = Frame::from("text".to_string()).into_text();
assert_eq!(t2.unwrap(), "text");
let b2 = Frame::from(vec![0xffu8, 0xfe]).into_text();
assert_eq!(b2.unwrap_err(), WsError::InvalidMessage);
}
#[test]
fn frame_zero_copy_moves() {
let s = String::from("owned");
let f = Frame::from(s);
assert_eq!(f, Frame::Text("owned".to_string()));
let v = vec![1u8; 8];
let f2 = Frame::from(v);
assert_eq!(f2.len(), 8);
}
#[test]
fn frame_clone_eq() {
let f = Frame::from("dup".to_string());
assert_eq!(f.clone(), f);
assert_ne!(
Frame::from("a".to_string()),
Frame::from(Bytes::from_static(b"a"))
);
}
}