tokio-rs/tokio
1
0
Fork 0
mirror of https://github.com/tokio-rs/tokio.git synced 2026-08-23 00:00:10 +02:00
Code Issues Packages Projects Releases Wiki Activity

codec: support byte slices in LengthDelimitedCodec (#8355)

Browse Source
This commit is contained in:
Dhruv Vaishnav
2026-08-10 12:18:22 +02:00
committed by GitHub
parent 8b13642a1f
commit 43e6ef50e4
2 changed files with 84 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tokio-util/src/codec/length_delimited.rs
+11 -3
View File
@@ -603,10 +603,10 @@ impl Decoder for LengthDelimitedCodec {
}
}
impl Encoder<Bytes> for LengthDelimitedCodec {
impl Encoder<&[u8]> for LengthDelimitedCodec {
type Error = io::Error;
fn encode(&mut self, data: Bytes, dst: &mut BytesMut) -> Result<(), io::Error> {
fn encode(&mut self, data: &[u8], dst: &mut BytesMut) -> Result<(), io::Error> {
let n = data.len();
if n > self.builder.max_frame_len {
@@ -641,12 +641,20 @@ impl Encoder<Bytes> for LengthDelimitedCodec {
}
// Write the frame to the buffer
dst.extend_from_slice(&data[..]);
dst.extend_from_slice(data);
Ok(())
}
}
impl Encoder<Bytes> for LengthDelimitedCodec {
type Error = io::Error;
fn encode(&mut self, data: Bytes, dst: &mut BytesMut) -> Result<(), io::Error> {
Encoder::<&[u8]>::encode(self, data.as_ref(), dst)
}
}
impl Default for LengthDelimitedCodec {
fn default() -> Self {
Self::new()
tokio-util/tests/length_delimited.rs
+73 -31
View File
@@ -70,6 +70,26 @@ macro_rules! assert_done {
}};
}
fn poll_ready_bytes<T>(
sink: Pin<&mut T>,
cx: &mut Context<'_>,
) -> Poll<Result<(), <T as Sink<Bytes>>::Error>>
where
T: Sink<Bytes>,
{
Sink::<Bytes>::poll_ready(sink, cx)
}
fn poll_flush_bytes<T>(
sink: Pin<&mut T>,
cx: &mut Context<'_>,
) -> Poll<Result<(), <T as Sink<Bytes>>::Error>>
where
T: Sink<Bytes>,
{
Sink::<Bytes>::poll_flush(sink, cx)
}
#[test]
fn read_empty_io_yields_nothing() {
let io = Box::pin(FramedRead::new(mock!(), LengthDelimitedCodec::new()));
@@ -423,9 +443,9 @@ fn write_single_frame_length_adjusted() {
pin_mut!(io);
task::spawn(()).enter(|cx, _| {
assert_ready_ok!(io.as_mut().poll_ready(cx));
assert_ready_ok!(poll_ready_bytes(io.as_mut(), cx));
assert_ok!(io.as_mut().start_send(Bytes::from("abcdefghi")));
assert_ready_ok!(io.as_mut().poll_flush(cx));
assert_ready_ok!(poll_flush_bytes(io.as_mut(), cx));
assert!(io.get_ref().calls.is_empty());
});
}
@@ -436,7 +456,7 @@ fn write_nothing_yields_nothing() {
pin_mut!(io);
task::spawn(()).enter(|cx, _| {
assert_ready_ok!(io.poll_flush(cx));
assert_ready_ok!(poll_flush_bytes(io, cx));
});
}
@@ -453,9 +473,31 @@ fn write_single_frame_one_packet() {
pin_mut!(io);
task::spawn(()).enter(|cx, _| {
assert_ready_ok!(io.as_mut().poll_ready(cx));
assert_ready_ok!(poll_ready_bytes(io.as_mut(), cx));
assert_ok!(io.as_mut().start_send(Bytes::from("abcdefghi")));
assert_ready_ok!(io.as_mut().poll_flush(cx));
assert_ready_ok!(poll_flush_bytes(io.as_mut(), cx));
assert!(io.get_ref().calls.is_empty());
});
}
#[test]
fn write_single_frame_from_slice() {
let io = FramedWrite::new(
mock! {
data(b"\x00\x00\x00\x09"),
data(b"abcdefghi"),
flush(),
},
LengthDelimitedCodec::new(),
);
pin_mut!(io);
task::spawn(()).enter(|cx, _| {
let data: &[u8] = b"abcdefghi";
assert_ready_ok!(Sink::<&[u8]>::poll_ready(io.as_mut(), cx));
assert_ok!(io.as_mut().start_send(data));
assert_ready_ok!(Sink::<&[u8]>::poll_flush(io.as_mut(), cx));
assert!(io.get_ref().calls.is_empty());
});
}
@@ -477,16 +519,16 @@ fn write_single_multi_frame_one_packet() {
pin_mut!(io);
task::spawn(()).enter(|cx, _| {
assert_ready_ok!(io.as_mut().poll_ready(cx));
assert_ready_ok!(poll_ready_bytes(io.as_mut(), cx));
assert_ok!(io.as_mut().start_send(Bytes::from("abcdefghi")));
assert_ready_ok!(io.as_mut().poll_ready(cx));
assert_ready_ok!(poll_ready_bytes(io.as_mut(), cx));
assert_ok!(io.as_mut().start_send(Bytes::from("123")));
assert_ready_ok!(io.as_mut().poll_ready(cx));
assert_ready_ok!(poll_ready_bytes(io.as_mut(), cx));
assert_ok!(io.as_mut().start_send(Bytes::from("hello world")));
assert_ready_ok!(io.as_mut().poll_flush(cx));
assert_ready_ok!(poll_flush_bytes(io.as_mut(), cx));
assert!(io.get_ref().calls.is_empty());
});
}
@@ -510,20 +552,20 @@ fn write_single_multi_frame_multi_packet() {
pin_mut!(io);
task::spawn(()).enter(|cx, _| {
assert_ready_ok!(io.as_mut().poll_ready(cx));
assert_ready_ok!(poll_ready_bytes(io.as_mut(), cx));
assert_ok!(io.as_mut().start_send(Bytes::from("abcdefghi")));
assert_ready_ok!(io.as_mut().poll_flush(cx));
assert_ready_ok!(poll_flush_bytes(io.as_mut(), cx));
assert_ready_ok!(io.as_mut().poll_ready(cx));
assert_ready_ok!(poll_ready_bytes(io.as_mut(), cx));
assert_ok!(io.as_mut().start_send(Bytes::from("123")));
assert_ready_ok!(io.as_mut().poll_flush(cx));
assert_ready_ok!(poll_flush_bytes(io.as_mut(), cx));
assert_ready_ok!(io.as_mut().poll_ready(cx));
assert_ready_ok!(poll_ready_bytes(io.as_mut(), cx));
assert_ok!(io.as_mut().start_send(Bytes::from("hello world")));
assert_ready_ok!(io.as_mut().poll_flush(cx));
assert_ready_ok!(poll_flush_bytes(io.as_mut(), cx));
assert!(io.get_ref().calls.is_empty());
});
}
@@ -544,12 +586,12 @@ fn write_single_frame_would_block() {
pin_mut!(io);
task::spawn(()).enter(|cx, _| {
assert_ready_ok!(io.as_mut().poll_ready(cx));
assert_ready_ok!(poll_ready_bytes(io.as_mut(), cx));
assert_ok!(io.as_mut().start_send(Bytes::from("abcdefghi")));
assert_pending!(io.as_mut().poll_flush(cx));
assert_pending!(io.as_mut().poll_flush(cx));
assert_ready_ok!(io.as_mut().poll_flush(cx));
assert_pending!(poll_flush_bytes(io.as_mut(), cx));
assert_pending!(poll_flush_bytes(io.as_mut(), cx));
assert_ready_ok!(poll_flush_bytes(io.as_mut(), cx));
assert!(io.get_ref().calls.is_empty());
});
@@ -567,10 +609,10 @@ fn write_single_frame_little_endian() {
pin_mut!(io);
task::spawn(()).enter(|cx, _| {
assert_ready_ok!(io.as_mut().poll_ready(cx));
assert_ready_ok!(poll_ready_bytes(io.as_mut(), cx));
assert_ok!(io.as_mut().start_send(Bytes::from("abcdefghi")));
assert_ready_ok!(io.as_mut().poll_flush(cx));
assert_ready_ok!(poll_flush_bytes(io.as_mut(), cx));
assert!(io.get_ref().calls.is_empty());
});
}
@@ -587,10 +629,10 @@ fn write_single_frame_with_short_length_field() {
pin_mut!(io);
task::spawn(()).enter(|cx, _| {
assert_ready_ok!(io.as_mut().poll_ready(cx));
assert_ready_ok!(poll_ready_bytes(io.as_mut(), cx));
assert_ok!(io.as_mut().start_send(Bytes::from("abcdefghi")));
assert_ready_ok!(io.as_mut().poll_flush(cx));
assert_ready_ok!(poll_flush_bytes(io.as_mut(), cx));
assert!(io.get_ref().calls.is_empty());
});
@@ -604,7 +646,7 @@ fn write_max_frame_len() {
pin_mut!(io);
task::spawn(()).enter(|cx, _| {
assert_ready_ok!(io.as_mut().poll_ready(cx));
assert_ready_ok!(poll_ready_bytes(io.as_mut(), cx));
assert_err!(io.as_mut().start_send(Bytes::from("abcdef")));
assert!(io.get_ref().calls.is_empty());
@@ -621,10 +663,10 @@ fn write_update_max_frame_len_at_rest() {
pin_mut!(io);
task::spawn(()).enter(|cx, _| {
assert_ready_ok!(io.as_mut().poll_ready(cx));
assert_ready_ok!(poll_ready_bytes(io.as_mut(), cx));
assert_ok!(io.as_mut().start_send(Bytes::from("abcdef")));
assert_ready_ok!(io.as_mut().poll_flush(cx));
assert_ready_ok!(poll_flush_bytes(io.as_mut(), cx));
io.encoder_mut().set_max_frame_length(5);
@@ -646,14 +688,14 @@ fn write_update_max_frame_len_in_flight() {
pin_mut!(io);
task::spawn(()).enter(|cx, _| {
assert_ready_ok!(io.as_mut().poll_ready(cx));
assert_ready_ok!(poll_ready_bytes(io.as_mut(), cx));
assert_ok!(io.as_mut().start_send(Bytes::from("abcdef")));
assert_pending!(io.as_mut().poll_flush(cx));
assert_pending!(poll_flush_bytes(io.as_mut(), cx));
io.encoder_mut().set_max_frame_length(5);
assert_ready_ok!(io.as_mut().poll_flush(cx));
assert_ready_ok!(poll_flush_bytes(io.as_mut(), cx));
assert_err!(io.as_mut().start_send(Bytes::from("abcdef")));
assert!(io.get_ref().calls.is_empty());
@@ -666,10 +708,10 @@ fn write_zero() {
pin_mut!(io);
task::spawn(()).enter(|cx, _| {
assert_ready_ok!(io.as_mut().poll_ready(cx));
assert_ready_ok!(poll_ready_bytes(io.as_mut(), cx));
assert_ok!(io.as_mut().start_send(Bytes::from("abcdef")));
assert_ready_err!(io.as_mut().poll_flush(cx));
assert_ready_err!(poll_flush_bytes(io.as_mut(), cx));
assert!(io.get_ref().calls.is_empty());
});