use std::{
convert::{AsRef, From, Into, TryFrom},
fmt,
result::Result as StdResult,
str,
};
use super::frame::{CloseFrame, Frame};
use crate::error::{CapacityError, Error, Result};
mod string_collect {
use utf8::DecodeError;
use crate::error::{Error, Result};
#[derive(Debug)]
pub struct StringCollector {
data: String,
incomplete: Option<utf8::Incomplete>,
}
impl StringCollector {
pub fn new() -> Self {
StringCollector { data: String::new(), incomplete: None }
}
pub fn len(&self) -> usize {
self.data
.len()
.saturating_add(self.incomplete.map(|i| i.buffer_len as usize).unwrap_or(0))
}
pub fn extend<T: AsRef<[u8]>>(&mut self, tail: T) -> Result<()> {
let mut input: &[u8] = tail.as_ref();
if let Some(mut incomplete) = self.incomplete.take() {
if let Some((result, rest)) = incomplete.try_complete(input) {
input = rest;
if let Ok(text) = result {
self.data.push_str(text);
} else {
return Err(Error::Utf8);
}
} else {
input = &[];
self.incomplete = Some(incomplete);
}
}
if !input.is_empty() {
match utf8::decode(input) {
Ok(text) => {
self.data.push_str(text);
Ok(())
}
Err(DecodeError::Incomplete { valid_prefix, incomplete_suffix }) => {
self.data.push_str(valid_prefix);
self.incomplete = Some(incomplete_suffix);
Ok(())
}
Err(DecodeError::Invalid { valid_prefix, .. }) => {
self.data.push_str(valid_prefix);
Err(Error::Utf8)
}
}
} else {
Ok(())
}
}
pub fn into_string(self) -> Result<String> {
if self.incomplete.is_some() {
Err(Error::Utf8)
} else {
Ok(self.data)
}
}
}
}
use self::string_collect::StringCollector;
#[derive(Debug)]
pub struct IncompleteMessage {
collector: IncompleteMessageCollector,
}
#[derive(Debug)]
enum IncompleteMessageCollector {
Text(StringCollector),
Binary(Vec<u8>),
}
impl IncompleteMessage {
pub fn new(message_type: IncompleteMessageType) -> Self {
IncompleteMessage {
collector: match message_type {
IncompleteMessageType::Binary => IncompleteMessageCollector::Binary(Vec::new()),
IncompleteMessageType::Text => {
IncompleteMessageCollector::Text(StringCollector::new())
}
},
}
}
pub fn len(&self) -> usize {
match self.collector {
IncompleteMessageCollector::Text(ref t) => t.len(),
IncompleteMessageCollector::Binary(ref b) => b.len(),
}
}
pub fn extend<T: AsRef<[u8]>>(&mut self, tail: T, size_limit: Option<usize>) -> Result<()> {
let max_size = size_limit.unwrap_or_else(usize::max_value);
let my_size = self.len();
let portion_size = tail.as_ref().len();
if my_size > max_size || portion_size > max_size - my_size {
return Err(Error::Capacity(CapacityError::MessageTooLong {
size: my_size + portion_size,
max_size,
}));
}
match self.collector {
IncompleteMessageCollector::Binary(ref mut v) => {
v.extend(tail.as_ref());
Ok(())
}
IncompleteMessageCollector::Text(ref mut t) => t.extend(tail),
}
}
pub fn complete(self) -> Result<Message> {
match self.collector {
IncompleteMessageCollector::Binary(v) => Ok(Message::Binary(v)),
IncompleteMessageCollector::Text(t) => {
let text = t.into_string()?;
Ok(Message::Text(text))
}
}
}
}
pub enum IncompleteMessageType {
Text,
Binary,
}
#[derive(Debug, Eq, PartialEq, Clone)]
pub enum Message {
Text(String),
Binary(Vec<u8>),
Ping(Vec<u8>),
Pong(Vec<u8>),
Close(Option<CloseFrame<'static>>),
Frame(Frame),
}
impl Message {
pub fn text<S>(string: S) -> Message
where
S: Into<String>,
{
Message::Text(string.into())
}
pub fn binary<B>(bin: B) -> Message
where
B: Into<Vec<u8>>,
{
Message::Binary(bin.into())
}
pub fn is_text(&self) -> bool {
matches!(*self, Message::Text(_))
}
pub fn is_binary(&self) -> bool {
matches!(*self, Message::Binary(_))
}
pub fn is_ping(&self) -> bool {
matches!(*self, Message::Ping(_))
}
pub fn is_pong(&self) -> bool {
matches!(*self, Message::Pong(_))
}
pub fn is_close(&self) -> bool {
matches!(*self, Message::Close(_))
}
pub fn len(&self) -> usize {
match *self {
Message::Text(ref string) => string.len(),
Message::Binary(ref data) | Message::Ping(ref data) | Message::Pong(ref data) => {
data.len()
}
Message::Close(ref data) => data.as_ref().map(|d| d.reason.len()).unwrap_or(0),
Message::Frame(ref frame) => frame.len(),
}
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn into_data(self) -> Vec<u8> {
match self {
Message::Text(string) => string.into_bytes(),
Message::Binary(data) | Message::Ping(data) | Message::Pong(data) => data,
Message::Close(None) => Vec::new(),
Message::Close(Some(frame)) => frame.reason.into_owned().into_bytes(),
Message::Frame(frame) => frame.into_data(),
}
}
pub fn into_text(self) -> Result<String> {
match self {
Message::Text(string) => Ok(string),
Message::Binary(data) | Message::Ping(data) | Message::Pong(data) => {
Ok(String::from_utf8(data)?)
}
Message::Close(None) => Ok(String::new()),
Message::Close(Some(frame)) => Ok(frame.reason.into_owned()),
Message::Frame(frame) => Ok(frame.into_string()?),
}
}
pub fn to_text(&self) -> Result<&str> {
match *self {
Message::Text(ref string) => Ok(string),
Message::Binary(ref data) | Message::Ping(ref data) | Message::Pong(ref data) => {
Ok(str::from_utf8(data)?)
}
Message::Close(None) => Ok(""),
Message::Close(Some(ref frame)) => Ok(&frame.reason),
Message::Frame(ref frame) => Ok(frame.to_text()?),
}
}
}
impl From<String> for Message {
fn from(string: String) -> Self {
Message::text(string)
}
}
impl<'s> From<&'s str> for Message {
fn from(string: &'s str) -> Self {
Message::text(string)
}
}
impl<'b> From<&'b [u8]> for Message {
fn from(data: &'b [u8]) -> Self {
Message::binary(data)
}
}
impl From<Vec<u8>> for Message {
fn from(data: Vec<u8>) -> Self {
Message::binary(data)
}
}
impl From<Message> for Vec<u8> {
fn from(message: Message) -> Self {
message.into_data()
}
}
impl TryFrom<Message> for String {
type Error = Error;
fn try_from(value: Message) -> StdResult<Self, Self::Error> {
value.into_text()
}
}
impl fmt::Display for Message {
fn fmt(&self, f: &mut fmt::Formatter) -> StdResult<(), fmt::Error> {
if let Ok(string) = self.to_text() {
write!(f, "{string}")
} else {
write!(f, "Binary Data<length={}>", self.len())
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn display() {
let t = Message::text("test".to_owned());
assert_eq!(t.to_string(), "test".to_owned());
let bin = Message::binary(vec![0, 1, 3, 4, 241]);
assert_eq!(bin.to_string(), "Binary Data<length=5>".to_owned());
}
#[test]
fn binary_convert() {
let bin = [6u8, 7, 8, 9, 10, 241];
let msg = Message::from(&bin[..]);
assert!(msg.is_binary());
assert!(msg.into_text().is_err());
}
#[test]
fn binary_convert_vec() {
let bin = vec![6u8, 7, 8, 9, 10, 241];
let msg = Message::from(bin);
assert!(msg.is_binary());
assert!(msg.into_text().is_err());
}
#[test]
fn binary_convert_into_vec() {
let bin = vec![6u8, 7, 8, 9, 10, 241];
let bin_copy = bin.clone();
let msg = Message::from(bin);
let serialized: Vec<u8> = msg.into();
assert_eq!(bin_copy, serialized);
}
#[test]
fn text_convert() {
let s = "kiwotsukete";
let msg = Message::from(s);
assert!(msg.is_text());
}
}