bnb/
codec.rs

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

bnb/
codec.rs
