wire_codec/
codec.rs

1//! Adapters that run a [`Framer`] over a growable [`bytes::BytesMut`] buffer.
2//!
3//! Available with the `bytes` feature. [`FrameCodec`] wraps any framer and
4//! exposes two inherent methods, [`FrameCodec::decode_frame`] and
5//! [`FrameCodec::encode_frame`], that work on `BytesMut` directly. The
6//! `tokio-util` and `asynchronous-codec` features add implementations of
7//! those crates' `Decoder` and `Encoder` traits on top, so a framer plugs
8//! straight into `tokio_util::codec::Framed` or `asynchronous_codec::Framed`:
9//!
10//! ```ignore
11//! use tokio_util::codec::Framed;
12//! use wire_codec::codec::FrameCodec;
13//! use wire_codec::framing::{Endian, LengthPrefixed, LengthWidth};
14//!
15//! let codec = FrameCodec::new(LengthPrefixed::bounded(LengthWidth::U32, Endian::Big));
16//! let mut framed = Framed::new(tcp_stream, codec);
17//! // framed: Stream<Item = io::Result<BytesMut>> + Sink<&[u8]>
18//! ```
19//!
20//! Decoded payloads are split off the receive buffer without copying when
21//! the framer returns a payload that lies inside the frame (all built-in
22//! framers do). The decoder keeps the resumable-scan state described on
23//! [`Framer::next_frame_from`], so a slow peer cannot force quadratic work.
24//!
25//! The trait implementations use [`std::io::Error`] as their error type,
26//! like tokio-util's own codecs; the original [`Error`] is available through
27//! `io::Error::get_ref` and `downcast_ref`.
28
29use bytes::{Buf, BytesMut};
30
31use crate::buf::WriteBuf;
32use crate::error::{Error, Result};
33use crate::framing::{Framer, StreamDecoder};
34
35/// A [`Framer`] driven over a `BytesMut` buffer.
36///
37/// # Example
38///
39/// ```
40/// use bytes::BytesMut;
41/// use wire_codec::codec::FrameCodec;
42/// use wire_codec::framing::VarintPrefixed;
43///
44/// let mut codec = FrameCodec::new(VarintPrefixed::new());
45/// let mut wire = BytesMut::new();
46/// codec.encode_frame(b"hello", &mut wire).unwrap();
47/// codec.encode_frame(b"world", &mut wire).unwrap();
48///
49/// assert_eq!(&codec.decode_frame(&mut wire).unwrap().unwrap()[..], b"hello");
50/// assert_eq!(&codec.decode_frame(&mut wire).unwrap().unwrap()[..], b"world");
51/// assert_eq!(codec.decode_frame(&mut wire).unwrap(), None);
52/// assert!(wire.is_empty());
53/// ```
54#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)]
55pub struct FrameCodec<F> {
56    decoder: StreamDecoder<F>,
57}
58
59impl<F: Framer> FrameCodec<F> {
60    /// Wrap `framer`.
61    #[inline]
62    #[must_use]
63    pub const fn new(framer: F) -> Self {
64        Self {
65            decoder: StreamDecoder::new(framer),
66        }
67    }
68
69    /// Borrow the wrapped framer.
70    #[inline]
71    #[must_use]
72    pub const fn framer(&self) -> &F {
73        self.decoder.framer()
74    }
75
76    /// Consume the codec and return the wrapped framer.
77    #[inline]
78    #[must_use]
79    pub fn into_inner(self) -> F {
80        self.decoder.into_inner()
81    }
82
83    /// Forget the partial-scan progress kept between
84    /// [`FrameCodec::decode_frame`] calls.
85    ///
86    /// The codec assumes the receive buffer only grows at the back between
87    /// calls, apart from the frames it splits off itself. Call `reset` when
88    /// you clear, truncate, or replace the buffer (for example to
89    /// resynchronise, after `Framed::from_parts` with a fresh read buffer, or
90    /// when reusing the codec for a new connection). Progress is also cleared
91    /// automatically whenever decoding returns an error.
92    ///
93    /// # Example
94    ///
95    /// ```
96    /// use bytes::BytesMut;
97    /// use wire_codec::codec::FrameCodec;
98    /// use wire_codec::framing::Delimited;
99    ///
100    /// let mut codec = FrameCodec::new(Delimited::bounded(b"\n").unwrap());
101    /// let mut src = BytesMut::from(&b"partial line"[..]);
102    /// assert_eq!(codec.decode_frame(&mut src).unwrap(), None);
103    ///
104    /// // Switch to a different stream: drop the buffer and the progress.
105    /// src.clear();
106    /// codec.reset();
107    /// src.extend_from_slice(b"a\nb\n");
108    /// assert_eq!(&codec.decode_frame(&mut src).unwrap().unwrap()[..], b"a");
109    /// ```
110    #[inline]
111    pub fn reset(&mut self) {
112        self.decoder.reset();
113    }
114
115    /// Remove the next complete frame from the front of `src` and return its
116    /// payload, or `Ok(None)` if `src` does not hold a full frame yet.
117    ///
118    /// # Errors
119    ///
120    /// Propagates the framer's errors. Returns [`Error::InvalidEncoding`] if
121    /// a custom framer violates the [`Frame`][`crate::framing::Frame`]
122    /// contract (zero bytes consumed, or more than the input). `src` is left
123    /// untouched on error, and the scan progress is cleared so the caller
124    /// may discard the buffer and resynchronise.
125    pub fn decode_frame(&mut self, src: &mut BytesMut) -> Result<Option<BytesMut>> {
126        let (consumed, payload_at) = match self.decoder.decode(src)? {
127            None => return Ok(None),
128            Some(frame) => {
129                let consumed = frame.consumed();
130                if consumed == 0 || consumed > src.len() {
131                    return Err(Error::InvalidEncoding);
132                }
133                let payload = frame.payload();
134                let start = payload.as_ptr().addr().wrapping_sub(src.as_ptr().addr());
135                if start <= consumed && payload.len() <= consumed - start {
136                    (consumed, Ok((start, payload.len())))
137                } else {
138                    // The payload does not live inside the frame (a custom
139                    // framer may return a static slice): copy it out.
140                    (consumed, Err(BytesMut::from(payload)))
141                }
142            }
143        };
144        match payload_at {
145            Ok((start, len)) => {
146                let mut frame = src.split_to(consumed);
147                frame.advance(start);
148                frame.truncate(len);
149                Ok(Some(frame))
150            }
151            Err(copied) => {
152                src.advance(consumed);
153                Ok(Some(copied))
154            }
155        }
156    }
157
158    /// Like [`FrameCodec::decode_frame`], for the end of the stream: once no
159    /// more input will arrive, leftover bytes that do not form a frame are an
160    /// error.
161    ///
162    /// # Errors
163    ///
164    /// Everything [`FrameCodec::decode_frame`] returns, plus
165    /// [`Error::UnexpectedEof`] when `src` holds a partial frame.
166    pub fn decode_frame_eof(&mut self, src: &mut BytesMut) -> Result<Option<BytesMut>> {
167        match self.decode_frame(src)? {
168            Some(frame) => Ok(Some(frame)),
169            None if src.is_empty() => Ok(None),
170            None => Err(Error::UnexpectedEof),
171        }
172    }
173
174    /// Append `payload` as one frame to `dst`, growing it as needed.
175    ///
176    /// Space is reserved in one step using [`Framer::max_frame_len`]. For a
177    /// custom framer that does not implement it, the codec retries with a
178    /// doubling scratch region, up to four times the payload size plus
179    /// 4 KiB.
180    ///
181    /// # Errors
182    ///
183    /// Propagates the framer's errors, or returns [`Error::BufferFull`] if a
184    /// custom framer needs more room than the retry limit. `dst` is left
185    /// untouched on error.
186    pub fn encode_frame(&mut self, payload: &[u8], dst: &mut BytesMut) -> Result<()> {
187        let start = dst.len();
188        let framer = self.decoder.framer();
189        if let Some(max) = framer.max_frame_len(payload.len()) {
190            return write_into(framer, payload, dst, start, max);
191        }
192        let limit = payload.len().saturating_mul(4).saturating_add(4096);
193        let mut room = payload.len().saturating_add(64);
194        loop {
195            match write_into(framer, payload, dst, start, room) {
196                Err(Error::BufferFull) if room < limit => room = room.saturating_mul(2).min(limit),
197                result => return result,
198            }
199        }
200    }
201}
202
203/// Write one frame into `room` freshly zeroed bytes at the end of `dst`,
204/// then trim `dst` to what was written (or back to `start` on error).
205fn write_into<F: Framer>(
206    framer: &F,
207    payload: &[u8],
208    dst: &mut BytesMut,
209    start: usize,
210    room: usize,
211) -> Result<()> {
212    let end = start.checked_add(room).ok_or(Error::BufferFull)?;
213    dst.resize(end, 0);
214    let mut out = WriteBuf::new(&mut dst[start..]);
215    let result = framer.write_frame(payload, &mut out);
216    let written = out.position();
217    match result {
218        Ok(()) => {
219            dst.truncate(start + written);
220            Ok(())
221        }
222        Err(e) => {
223            dst.truncate(start);
224            Err(e)
225        }
226    }
227}
228
229#[cfg(feature = "tokio-util")]
230mod tokio_impls {
231    use std::io;
232
233    use bytes::{Bytes, BytesMut};
234    use tokio_util::codec::{Decoder, Encoder};
235
236    use super::FrameCodec;
237    use crate::framing::Framer;
238
239    #[cfg_attr(docsrs, doc(cfg(feature = "tokio-util")))]
240    impl<F: Framer> Decoder for FrameCodec<F> {
241        type Item = BytesMut;
242        type Error = io::Error;
243
244        #[inline]
245        fn decode(&mut self, src: &mut BytesMut) -> io::Result<Option<BytesMut>> {
246            self.decode_frame(src).map_err(io::Error::from)
247        }
248
249        #[inline]
250        fn decode_eof(&mut self, src: &mut BytesMut) -> io::Result<Option<BytesMut>> {
251            self.decode_frame_eof(src).map_err(io::Error::from)
252        }
253    }
254
255    #[cfg_attr(docsrs, doc(cfg(feature = "tokio-util")))]
256    impl<F: Framer> Encoder<&[u8]> for FrameCodec<F> {
257        type Error = io::Error;
258
259        #[inline]
260        fn encode(&mut self, item: &[u8], dst: &mut BytesMut) -> io::Result<()> {
261            self.encode_frame(item, dst).map_err(io::Error::from)
262        }
263    }
264
265    #[cfg_attr(docsrs, doc(cfg(feature = "tokio-util")))]
266    impl<F: Framer> Encoder<Bytes> for FrameCodec<F> {
267        type Error = io::Error;
268
269        #[inline]
270        fn encode(&mut self, item: Bytes, dst: &mut BytesMut) -> io::Result<()> {
271            self.encode_frame(&item, dst).map_err(io::Error::from)
272        }
273    }
274
275    #[cfg_attr(docsrs, doc(cfg(feature = "tokio-util")))]
276    impl<F: Framer> Encoder<BytesMut> for FrameCodec<F> {
277        type Error = io::Error;
278
279        #[inline]
280        fn encode(&mut self, item: BytesMut, dst: &mut BytesMut) -> io::Result<()> {
281            self.encode_frame(&item, dst).map_err(io::Error::from)
282        }
283    }
284}
285
286#[cfg(feature = "asynchronous-codec")]
287mod futures_impls {
288    use std::io;
289
290    use asynchronous_codec::{Decoder, Encoder};
291    use bytes::BytesMut;
292
293    use super::FrameCodec;
294    use crate::framing::Framer;
295
296    #[cfg_attr(docsrs, doc(cfg(feature = "asynchronous-codec")))]
297    impl<F: Framer> Decoder for FrameCodec<F> {
298        type Item = BytesMut;
299        type Error = io::Error;
300
301        #[inline]
302        fn decode(&mut self, src: &mut BytesMut) -> io::Result<Option<BytesMut>> {
303            self.decode_frame(src).map_err(io::Error::from)
304        }
305
306        #[inline]
307        fn decode_eof(&mut self, src: &mut BytesMut) -> io::Result<Option<BytesMut>> {
308            self.decode_frame_eof(src).map_err(io::Error::from)
309        }
310    }
311
312    #[cfg_attr(docsrs, doc(cfg(feature = "asynchronous-codec")))]
313    impl<F: Framer> Encoder for FrameCodec<F> {
314        type Item<'a> = &'a [u8];
315        type Error = io::Error;
316
317        #[inline]
318        fn encode(&mut self, item: &[u8], dst: &mut BytesMut) -> io::Result<()> {
319            self.encode_frame(item, dst).map_err(io::Error::from)
320        }
321    }
322}
323
324#[cfg(test)]
325mod tests {
326    use super::*;
327    use crate::framing::{Delimited, Endian, Frame, LengthPrefixed, LengthWidth};
328
329    /// Custom framer without `max_frame_len` whose payload is a static
330    /// slice, to exercise the fallback paths.
331    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
332    struct Fixed {
333        overhead: usize,
334        consumed_override: Option<usize>,
335    }
336
337    impl Framer for Fixed {
338        fn next_frame<'a>(&self, input: &'a [u8]) -> Result<Option<Frame<'a>>> {
339            if let Some(consumed) = self.consumed_override {
340                return Ok(Some(Frame::new(b"", consumed)));
341            }
342            Ok(input.get(..2).map(|_| Frame::new(b"st", 2)))
343        }
344        fn write_frame(&self, payload: &[u8], out: &mut WriteBuf<'_>) -> Result<()> {
345            if out.remaining() < payload.len() + self.overhead {
346                return Err(Error::BufferFull);
347            }
348            for _ in 0..self.overhead {
349                out.write_u8(0xAA)?;
350            }
351            out.write_bytes(payload)
352        }
353    }
354
355    #[test]
356    fn zero_copy_split_for_builtin_framers() {
357        let mut codec = FrameCodec::new(LengthPrefixed::bounded(LengthWidth::U16, Endian::Big));
358        let mut wire = BytesMut::new();
359        codec.encode_frame(b"abc", &mut wire).unwrap();
360        assert_eq!(&wire[..], &[0, 3, b'a', b'b', b'c']);
361        wire.extend_from_slice(&[0, 2, b'x']);
362        let first = codec.decode_frame(&mut wire).unwrap().unwrap();
363        assert_eq!(&first[..], b"abc");
364        assert_eq!(codec.decode_frame(&mut wire).unwrap(), None);
365        assert_eq!(codec.decode_frame_eof(&mut wire), Err(Error::UnexpectedEof));
366        wire.extend_from_slice(b"y");
367        assert_eq!(
368            &codec.decode_frame_eof(&mut wire).unwrap().unwrap()[..],
369            b"xy"
370        );
371        assert_eq!(codec.decode_frame_eof(&mut wire), Ok(None));
372    }
373
374    #[test]
375    fn byte_at_a_time_delimited() {
376        let mut codec = FrameCodec::new(Delimited::bounded(b"\r\n").unwrap());
377        let mut wire = BytesMut::new();
378        let mut lines = Vec::new();
379        for &byte in b"GET /\r\nHost: x\r\n\r\n" {
380            wire.extend_from_slice(&[byte]);
381            while let Some(line) = codec.decode_frame(&mut wire).unwrap() {
382                lines.push(line.to_vec());
383            }
384        }
385        assert_eq!(lines, [b"GET /".to_vec(), b"Host: x".to_vec(), Vec::new()]);
386        assert_eq!(codec.framer().delimiter(), b"\r\n");
387    }
388
389    #[test]
390    fn resynchronising_after_an_error_does_not_merge_lines() {
391        // Regression test for WC2-001, the verifier's scenario: an over-long
392        // line, then the caller clears the buffer and feeds fresh lines.
393        let framer = Delimited::new(b"\n").unwrap().with_max_payload(8);
394        let mut codec = FrameCodec::new(framer);
395        let mut src = BytesMut::from(&b"aaaaaaa"[..]);
396        assert_eq!(codec.decode_frame(&mut src).unwrap(), None);
397        src.extend_from_slice(b"aaaaa");
398        assert!(codec.decode_frame(&mut src).is_err());
399        src.clear();
400        src.extend_from_slice(b"hi\nthere\n");
401        assert_eq!(&codec.decode_frame(&mut src).unwrap().unwrap()[..], b"hi");
402        assert_eq!(
403            &codec.decode_frame(&mut src).unwrap().unwrap()[..],
404            b"there"
405        );
406        assert!(src.is_empty());
407    }
408
409    #[test]
410    fn reset_after_swapping_the_buffer() {
411        let mut codec = FrameCodec::new(Delimited::bounded(b"\n").unwrap());
412        let mut old = BytesMut::from(&b"0123456789"[..]);
413        assert_eq!(codec.decode_frame(&mut old).unwrap(), None);
414        codec.reset();
415        let mut fresh = BytesMut::from(&b"x\nyyyyyy\n"[..]);
416        assert_eq!(&codec.decode_frame(&mut fresh).unwrap().unwrap()[..], b"x");
417        assert_eq!(
418            &codec.decode_frame(&mut fresh).unwrap().unwrap()[..],
419            b"yyyyyy"
420        );
421    }
422
423    #[test]
424    fn errors_leave_buffers_untouched() {
425        let mut codec = FrameCodec::new(Delimited::new(b"\n").unwrap().with_max_payload(3));
426        let mut dst = BytesMut::from(&b"keep"[..]);
427        assert_eq!(
428            codec.encode_frame(b"a\nb", &mut dst),
429            Err(Error::DelimiterInPayload)
430        );
431        assert_eq!(
432            codec.encode_frame(b"abcd", &mut dst),
433            Err(Error::FrameTooLarge { len: 4, limit: 3 })
434        );
435        assert_eq!(&dst[..], b"keep");
436        let mut src = BytesMut::from(&b"abcde"[..]);
437        assert!(codec.decode_frame(&mut src).is_err());
438        assert_eq!(&src[..], b"abcde");
439        let framer = codec.into_inner();
440        assert_eq!(framer.max_payload(), 3);
441    }
442
443    #[test]
444    fn custom_framer_fallbacks() {
445        let mut codec = FrameCodec::new(Fixed {
446            overhead: 100,
447            consumed_override: None,
448        });
449        let mut dst = BytesMut::new();
450        // No max_frame_len: the scratch region grows until the frame fits.
451        codec.encode_frame(b"p", &mut dst).unwrap();
452        assert_eq!(dst.len(), 101);
453        // Payload outside the input: copied instead of split.
454        let mut src = BytesMut::from(&b"zzz"[..]);
455        assert_eq!(&codec.decode_frame(&mut src).unwrap().unwrap()[..], b"st");
456        assert_eq!(&src[..], b"z");
457
458        // A framer that needs more than the retry limit gets BufferFull.
459        let mut codec = FrameCodec::new(Fixed {
460            overhead: 10_000,
461            consumed_override: None,
462        });
463        let mut dst = BytesMut::new();
464        assert_eq!(codec.encode_frame(b"p", &mut dst), Err(Error::BufferFull));
465        assert!(dst.is_empty());
466    }
467
468    #[test]
469    fn contract_violations_are_reported() {
470        for consumed in [0, 99] {
471            let mut codec = FrameCodec::new(Fixed {
472                overhead: 0,
473                consumed_override: Some(consumed),
474            });
475            let mut src = BytesMut::from(&b"abc"[..]);
476            assert_eq!(codec.decode_frame(&mut src), Err(Error::InvalidEncoding));
477            assert_eq!(src.len(), 3);
478        }
479    }
480}
wire_codec/codec.rs

wire_codec/
codec.rs
