diff --git a/tokio-util/src/codec/length_delimited.rs b/tokio-util/src/codec/length_delimited.rs index e114426c2..dec331c5e 100644 --- a/tokio-util/src/codec/length_delimited.rs +++ b/tokio-util/src/codec/length_delimited.rs @@ -603,10 +603,10 @@ impl Decoder for LengthDelimitedCodec { } } -impl Encoder 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 for LengthDelimitedCodec { } // Write the frame to the buffer - dst.extend_from_slice(&data[..]); + dst.extend_from_slice(data); Ok(()) } } +impl Encoder 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() diff --git a/tokio-util/tests/length_delimited.rs b/tokio-util/tests/length_delimited.rs index 2ff60d13e..2995b4e82 100644 --- a/tokio-util/tests/length_delimited.rs +++ b/tokio-util/tests/length_delimited.rs @@ -70,6 +70,26 @@ macro_rules! assert_done { }}; } +fn poll_ready_bytes( + sink: Pin<&mut T>, + cx: &mut Context<'_>, +) -> Poll>::Error>> +where + T: Sink, +{ + Sink::::poll_ready(sink, cx) +} + +fn poll_flush_bytes( + sink: Pin<&mut T>, + cx: &mut Context<'_>, +) -> Poll>::Error>> +where + T: Sink, +{ + Sink::::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()); });