pull/144/head
SirCipher 5 years ago
parent 3cf0b83949
commit 2744d1be4f
  1. 70
      src/protocol/mod.rs

@ -18,8 +18,8 @@ use self::message::IncompleteMessage;
use crate::error::{Error, Result};
use crate::extensions::uncompressed::UncompressedExt;
use crate::extensions::WebSocketExtension;
use crate::util::NonBlockingResult;
use crate::protocol::frame::coding::Data;
use crate::util::NonBlockingResult;
pub(crate) const MAX_MESSAGE_SIZE: usize = 64 << 20;
@ -611,16 +611,31 @@ impl<Ext> WebSocketContext<Ext>
let max_frame_size = self.config.max_frame_size.unwrap_or_else(usize::max_value);
if frame.payload().len() > max_frame_size {
let mut chunks = frame.payload().chunks(self.config.max_frame_size.unwrap_or_else(usize::max_value)).peekable();
let data_frame = Frame::message(Vec::from(chunks.next().unwrap()), frame.header().opcode, false);
self.frame.write_frame(stream, data_frame).check_connection_reset(self.state)?;
let mut chunks = frame
.payload()
.chunks(self.config.max_frame_size.unwrap_or_else(usize::max_value))
.peekable();
let data_frame = Frame::message(
Vec::from(chunks.next().unwrap()),
frame.header().opcode,
false,
);
self.frame
.write_frame(stream, data_frame)
.check_connection_reset(self.state)?;
while let Some(chunk) = chunks.next() {
let frame = Frame::message(Vec::from(chunk), OpCode::Data(Data::Continue), chunks.peek().is_none());
let frame = Frame::message(
Vec::from(chunk),
OpCode::Data(Data::Continue),
chunks.peek().is_none(),
);
trace!("Sending frame: {:?}", frame);
self.frame.write_frame(stream, frame).check_connection_reset(self.state)?;
self.frame
.write_frame(stream, frame)
.check_connection_reset(self.state)?;
}
Ok(())
@ -701,10 +716,10 @@ mod tests {
use super::{Message, Role, WebSocket, WebSocketConfig};
use crate::extensions::uncompressed::UncompressedExt;
use crate::protocol::frame::coding::{Data, OpCode};
use crate::protocol::frame::Frame;
use std::io;
use std::io::Cursor;
use crate::protocol::frame::Frame;
use crate::protocol::frame::coding::{OpCode, Data};
struct WriteMoc<Stream>(Stream);
@ -788,16 +803,23 @@ mod tests {
encoder: UncompressedExt::new(Some(max_message_size)),
};
let mut socket = WebSocket::from_raw_socket(Cursor::new(Vec::new()), Role::Client, Some(limit));
let mut socket =
WebSocket::from_raw_socket(Cursor::new(Vec::new()), Role::Client, Some(limit));
socket.write_message(Message::text(input_str)).unwrap();
socket.socket.set_position(0);
let WebSocket { mut socket, mut context } = socket;
let WebSocket {
mut socket,
mut context,
} = socket;
let vec = input_str.chars().collect::<Vec<_>>();
let mut iter = vec.chunks(max_message_size).map(|c| c.iter().collect::<String>())
.into_iter().peekable();
let mut iter = vec
.chunks(max_message_size)
.map(|c| c.iter().collect::<String>())
.into_iter()
.peekable();
let frame_eq = |expected: Frame, actual: Frame| {
assert_eq!(expected.payload(), actual.payload());
@ -806,11 +828,29 @@ mod tests {
};
let expected = Frame::message(iter.next().unwrap().into(), OpCode::Data(Data::Text), false);
frame_eq(expected, context.frame.read_frame(&mut socket, Some(max_message_size)).unwrap().unwrap());
frame_eq(
expected,
context
.frame
.read_frame(&mut socket, Some(max_message_size))
.unwrap()
.unwrap(),
);
while let Some(chars) = iter.next() {
let expected = Frame::message(chars.into(), OpCode::Data(Data::Continue), iter.peek().is_none());
frame_eq(expected, context.frame.read_frame(&mut socket, Some(max_message_size)).unwrap().unwrap());
let expected = Frame::message(
chars.into(),
OpCode::Data(Data::Continue),
iter.peek().is_none(),
);
frame_eq(
expected,
context
.frame
.read_frame(&mut socket, Some(max_message_size))
.unwrap()
.unwrap(),
);
}
}
}

Loading…
Cancel
Save