bnb/
framing.rs

1//! Async framing via [`tokio_util::codec`] — the `tokio` feature.
2//!
3//! [`BinCodec<T>`] implements [`Decoder`](tokio_util::codec::Decoder) +
4//! [`Encoder`](tokio_util::codec::Encoder) for **any** `#[bin]` message (anything that is
5//! [`BitDecode`]/[`BitEncode`]), so wrapping a tokio stream is one line:
6//!
7//! ```ignore
8//! use bnb::BinCodec;
9//! use tokio_util::codec::Framed;
10//! use futures_util::{SinkExt, StreamExt};
11//!
12//! let mut conn = Framed::new(tcp_stream, BinCodec::<MyMsg>::new());
13//! conn.send(MyMsg::Ping).await?;          // it's a `Sink<MyMsg>`
14//! let reply = conn.next().await.unwrap()?; // …and a `Stream<Item = MyMsg>`
15//! ```
16//!
17//! The codec relies on each message being **self-delimiting** (its `#[bin]` structure or a
18//! `magic`/length prefix bounds it) — `decode` reads exactly one message and returns `None`
19//! when only a partial frame has arrived, so `Framed` reads more and retries.
20//!
21//! **Byte alignment over `Framed`.** `Framed`'s buffer is byte-granular, so `BinCodec`
22//! advances by whole bytes and starts each `decode` at a byte boundary. A message whose
23//! encoded length is **not** a whole number of bytes (a sub-byte `#[bin]` frame) therefore
24//! cannot be packed back-to-back on a `Framed` byte *stream* — use a byte-aligned message, or
25//! carry the bit cursor yourself with [`BitBuf`](crate::BitBuf)/[`MessageStream`](crate::MessageStream).
26//! (Each `UdpFramed` datagram is framed independently, so sub-byte messages are fine there.)
27//!
28//! The same `BinCodec` drives **datagrams**, too: `tokio_util::udp::UdpFramed::new(udp_socket,
29//! BinCodec::<T>::new())` is a `Stream<Item = (T, SocketAddr)>` + `Sink<(T, SocketAddr)>`. So
30//! one codec covers async streams (`Framed`, TCP) and async datagrams (`UdpFramed`, UDP) — the
31//! async mirror of the sync `MessageStream` / `MessageDatagram` split. See `examples/tokio_udp`.
32
33use crate::{BitDecode, BitEncode, BitReader, BitWriter, ErrorKind};
34use bytes::{Buf, BytesMut};
35use core::marker::PhantomData;
36use std::io;
37use tokio_util::codec::{Decoder, Encoder};
38
39/// A [`tokio_util::codec`] codec that frames any bnb `#[bin]` message `T`. Construct with
40/// [`BinCodec::new`] and hand it to [`Framed`](tokio_util::codec::Framed).
41pub struct BinCodec<T>(PhantomData<T>);
42
43impl<T> BinCodec<T> {
44    /// A codec for messages of type `T`.
45    #[must_use]
46    pub fn new() -> Self {
47        Self(PhantomData)
48    }
49}
50
51impl<T> Default for BinCodec<T> {
52    fn default() -> Self {
53        Self::new()
54    }
55}
56
57// `Decoder` also bounds `T: BitEncode` so it can read in the message's declared
58// [`LAYOUT`](BitEncode::LAYOUT) — a `#[bin(little)]`/`bits = lsb` message must decode in
59// the same order the paired `Encoder` writes it, or the round-trip silently corrupts. (A
60// `BinCodec` used over `Framed`/`UdpFramed` always implements both directions anyway.)
61impl<T: BitDecode + BitEncode> Decoder for BinCodec<T> {
62    type Item = T;
63    type Error = io::Error;
64
65    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<T>, io::Error> {
66        if src.is_empty() {
67            return Ok(None);
68        }
69        let mut reader = BitReader::with_layout(&src[..], <T as BitEncode>::LAYOUT);
70        match <T as BitDecode>::bit_decode(&mut reader) {
71            Ok(item) => {
72                // Consume exactly what this message used; leave the rest for the next call.
73                let consumed = reader.bit_pos() / 8;
74                src.advance(consumed);
75                Ok(Some(item))
76            }
77            // Only a partial frame is buffered — ask `Framed` to read more (don't consume).
78            Err(e)
79                if matches!(
80                    e.kind,
81                    ErrorKind::UnexpectedEof { .. } | ErrorKind::Incomplete { .. }
82                ) =>
83            {
84                Ok(None)
85            }
86            // A genuine framing error.
87            Err(e) => Err(io::Error::new(io::ErrorKind::InvalidData, e.to_string())),
88        }
89    }
90}
91
92impl<T: BitEncode> Encoder<T> for BinCodec<T> {
93    type Error = io::Error;
94
95    fn encode(&mut self, item: T, dst: &mut BytesMut) -> Result<(), io::Error> {
96        let mut w = BitWriter::with_layout(<T as BitEncode>::LAYOUT);
97        item.bit_encode(&mut w)
98            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e.to_string()))?;
99        dst.extend_from_slice(&w.into_bytes());
100        Ok(())
101    }
102}
103
104#[cfg(test)]
105mod component {
106    //! Component tests: the `BinCodec` Decoder/Encoder framing logic over a `BytesMut` (one
107    //! message per `decode`, partial-frame `None`, exact consumption, error mapping).
108    use super::BinCodec;
109    use bnb::bin;
110    use bytes::BytesMut;
111    use tokio_util::codec::{Decoder, Encoder};
112
113    /// A fixed 4-byte message — its length is implicit in its `#[bin]` structure.
114    #[bin(big)]
115    #[derive(Debug, Clone, PartialEq, Eq)]
116    struct Msg {
117        a: u16,
118        b: u16,
119    }
120
121    /// A magic-prefixed message — wrong bytes make `decode` a hard error, not "read more".
122    #[bin(big, magic = 0xCAFEu16)]
123    #[derive(Debug, Clone, PartialEq, Eq)]
124    struct Magic {
125        v: u8,
126    }
127
128    /// A **little-endian** message — the Decoder must read in the same order the Encoder
129    /// writes, or the round-trip silently corrupts.
130    #[bin(little)]
131    #[derive(Debug, Clone, PartialEq, Eq)]
132    struct LeMsg {
133        a: u16,
134        b: u32,
135    }
136
137    #[test]
138    fn little_endian_message_round_trips() {
139        let mut codec = BinCodec::<LeMsg>::new();
140        let mut buf = BytesMut::new();
141        let m = LeMsg {
142            a: 0x1234,
143            b: 0xAABB_CCDD,
144        };
145        codec.encode(m.clone(), &mut buf).unwrap();
146        // Little-endian on the wire.
147        assert_eq!(&buf[..], &[0x34, 0x12, 0xDD, 0xCC, 0xBB, 0xAA]);
148        // And the Decoder reads it back in the same order.
149        assert_eq!(codec.decode(&mut buf).unwrap(), Some(m));
150        assert!(buf.is_empty());
151    }
152
153    #[test]
154    fn encoder_writes_exact_message_bytes() {
155        let mut codec = BinCodec::<Msg>::new();
156        let mut buf = BytesMut::new();
157        codec
158            .encode(
159                Msg {
160                    a: 0x0102,
161                    b: 0x0304,
162                },
163                &mut buf,
164            )
165            .unwrap();
166        assert_eq!(&buf[..], &[0x01, 0x02, 0x03, 0x04]);
167    }
168
169    #[test]
170    fn encode_then_decode_round_trips_and_drains() {
171        let mut codec = BinCodec::<Msg>::new();
172        let mut buf = BytesMut::new();
173        let m = Msg {
174            a: 0xAABB,
175            b: 0xCCDD,
176        };
177        codec.encode(m.clone(), &mut buf).unwrap();
178        assert_eq!(codec.decode(&mut buf).unwrap(), Some(m));
179        assert!(buf.is_empty(), "decode consumed exactly the one frame");
180    }
181
182    #[test]
183    fn decode_empty_buffer_is_none() {
184        let mut codec = BinCodec::<Msg>::new();
185        let mut buf = BytesMut::new();
186        assert_eq!(codec.decode(&mut buf).unwrap(), None);
187    }
188
189    #[test]
190    fn decode_partial_frame_is_none_and_keeps_bytes() {
191        let mut codec = BinCodec::<Msg>::new();
192        let mut buf = BytesMut::from(&[0x01, 0x02][..]); // only 2 of the 4 bytes
193        assert_eq!(codec.decode(&mut buf).unwrap(), None);
194        assert_eq!(
195            &buf[..],
196            &[0x01, 0x02],
197            "a partial frame is left for the next read"
198        );
199    }
200
201    #[test]
202    fn decode_consumes_one_message_and_leaves_the_tail() {
203        let mut codec = BinCodec::<Msg>::new();
204        let mut buf = BytesMut::from(&[0x01, 0x02, 0x03, 0x04, 0xEE, 0xFF][..]);
205        assert_eq!(
206            codec.decode(&mut buf).unwrap(),
207            Some(Msg {
208                a: 0x0102,
209                b: 0x0304
210            })
211        );
212        assert_eq!(
213            &buf[..],
214            &[0xEE, 0xFF],
215            "trailing bytes remain for the next frame"
216        );
217    }
218
219    #[test]
220    fn decode_walks_back_to_back_messages() {
221        let mut codec = BinCodec::<Msg>::new();
222        let mut buf = BytesMut::from(&[0x00, 0x01, 0x00, 0x02, 0x00, 0x03, 0x00, 0x04][..]);
223        assert_eq!(codec.decode(&mut buf).unwrap(), Some(Msg { a: 1, b: 2 }));
224        assert_eq!(codec.decode(&mut buf).unwrap(), Some(Msg { a: 3, b: 4 }));
225        assert_eq!(codec.decode(&mut buf).unwrap(), None);
226    }
227
228    #[test]
229    fn decode_bad_magic_is_an_invalid_data_error() {
230        let mut codec = BinCodec::<Magic>::new();
231        let mut buf = BytesMut::from(&[0x00, 0x00, 0x07][..]); // full frame, wrong magic
232        let err = codec.decode(&mut buf).unwrap_err();
233        assert_eq!(err.kind(), std::io::ErrorKind::InvalidData);
234    }
235
236    #[test]
237    fn default_constructs_a_codec() {
238        let _c: BinCodec<Msg> = BinCodec::default();
239    }
240}
bnb/framing.rs

bnb/
framing.rs
