rust-websocket-lite
rust-websocket-lite copied to clipboard
Add size limit to server Message Codec
This PR adds an optional size limit to the MessageCodec
, which allows server to prevent clients from sending large messages that could cause the server to run out of memory.
Thanks! Could you add a test? I think you can make one based on the reserves_buffer
test: use the client codec to encode a message, then try to decode it with a server codec under a very small limit. The logic in the reserves_buffer
test is needed since the size limit is applied when reserving extra space in the server's receive buffer: in other words, encoding then decoding the same buffer skips the check, since the buffer is already large enough to receive the incoming data.
Can you run cargo fmt
? I can see that it would make some changes to the newly added lines.
I think this test triggers the size limit logic correctly:
#[test]
fn applies_message_size_limit() {
let message = Message::binary(
std::iter::repeat(b"abcdefghijklmnopqrstuvwxyz")
.take(100)
.flat_map(|b| b.iter().copied())
.collect::<Vec<u8>>(),
);
let mut data = BytesMut::new();
MessageCodec::client()
.encode(&message, &mut data)
.expect("didn't expect MessageCodec::encode to return an error");
let mut src = &data[..];
assert_eq!(src.len(), 2608);
let mut decoder = MessageCodec::server_with_message_size_limit(20);
let mut decoder_buf = BytesMut::new();
let err = loop {
match decoder.decode(&mut decoder_buf) {
Ok(Some(_)) => panic!("didn't expect decoder to succeed"),
Ok(None) => {
let n = decoder_buf.remaining_mut().min(src.len()).min(16);
assert!(n > 0, "expected decoder to reserve at least one byte");
decoder_buf.put_slice(&src[..n]);
src = &src[n..];
}
Err(err) => break err,
}
};
assert_eq!(err.to_string(), "frame exceeds size limit: 2608 bytes (a30)");
}