rust-websocket-lite icon indicating copy to clipboard operation
rust-websocket-lite copied to clipboard

Add size limit to server Message Codec

Open the10thWiz opened this issue 3 years ago • 2 comments

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.

the10thWiz avatar May 14 '21 22:05 the10thWiz

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.

1tgr avatar May 15 '21 10:05 1tgr

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)");
    }

1tgr avatar May 15 '21 10:05 1tgr