Skip to main content

bnb/
wirelen.rs

1//! [`WireLen<T>`] — a length/count field that auto-derives by default but can be pinned.
2//!
3//! A wire length or count (a UDP length, a DNS `qdcount`, an `rdlength`) normally equals the
4//! real size of what it describes, but a **dual-use** codec must also let you write a value
5//! that *disagrees* (a forged/malformed frame). `WireLen<T>` is that field: it is either
6//!
7//! - [`Auto`](WireLen::Auto) — "derive me at encode" (the [`Default`]), or
8//! - [`Set`](WireLen::Set) — "write exactly this value" (an override, or a decoded value).
9//!
10//! On **decode** it always yields `Set(decoded)`, so `decode → encode` stays byte-identical
11//! (a forged length survives a round-trip). On **encode**, an `Auto` is resolved from a
12//! declared target via the `#[bw(auto_len = count(x) | bytes(x))]` field directive (same-struct
13//! target) or the `#[bin(auto_len(path = count(field), …))]` struct directive (cross-struct);
14//! a `Set` is written as-is. So plain `to_bytes()` is correct by default, yet
15//! `WireLen::set(n)` deliberately deviates. Derivation is checked ([`CountPrefix`]) — an
16//! oversized length is a [`BitError`], never a silent truncation.
17
18use crate::bitstream::{BitDecode, BitEncode, BitError, CountPrefix, FixedBitLen, Sink, Source};
19use crate::field::Bits;
20
21/// A wire length/count field: [`Auto`](Self::Auto)-derived by default, or an explicit
22/// [`Set`](Self::Set) override. See the [module docs](self).
23#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Default)]
24pub enum WireLen<T> {
25    /// Derive this length from its declared target when the message is encoded.
26    #[default]
27    Auto,
28    /// Write exactly this value — a deliberate override, or the value read on decode.
29    Set(T),
30}
31
32impl<T> WireLen<T> {
33    /// An auto-deriving length (the default). Equivalent to [`WireLen::Auto`].
34    #[must_use]
35    pub fn auto() -> Self {
36        WireLen::Auto
37    }
38
39    /// A pinned, explicit length. Equivalent to [`WireLen::Set`] — the dual-use override.
40    #[must_use]
41    pub fn set(value: T) -> Self {
42        WireLen::Set(value)
43    }
44
45    /// Whether this length auto-derives (has no explicit value yet).
46    #[must_use]
47    pub fn is_auto(&self) -> bool {
48        matches!(self, WireLen::Auto)
49    }
50
51    /// The explicit value, or `None` if it auto-derives. A decoded `WireLen` is always
52    /// `Some`; a freshly-built `Auto` is `None` until resolved at encode.
53    #[must_use]
54    pub fn get(&self) -> Option<&T> {
55        match self {
56            WireLen::Set(v) => Some(v),
57            WireLen::Auto => None,
58        }
59    }
60}
61
62impl<T: CountPrefix> WireLen<T> {
63    /// The value as a `usize` element count (for `#[br(count = …)]` on the decode side,
64    /// where a `WireLen` is always [`Set`](Self::Set)). An unresolved [`Auto`](Self::Auto)
65    /// reads as `0` (it never occurs on decode).
66    #[must_use]
67    pub fn to_count(&self) -> usize {
68        match self {
69            WireLen::Set(v) => v.to_count(),
70            WireLen::Auto => 0,
71        }
72    }
73
74    /// Resolve an [`Auto`](Self::Auto) to the checked count of `len` (leaving a
75    /// [`Set`](Self::Set) untouched) — the encode-time derivation the macro emits.
76    ///
77    /// # Errors
78    /// [`BitError`] if `len` exceeds `T`'s range (a checked conversion, never a truncation).
79    pub fn resolve_count(&self, len: usize) -> Result<Self, BitError> {
80        match self {
81            WireLen::Set(v) => Ok(WireLen::Set(*v)),
82            WireLen::Auto => Ok(WireLen::Set(
83                T::try_from_count(len).map_err(BitError::from)?,
84            )),
85        }
86    }
87}
88
89impl<T: Bits> BitDecode for WireLen<T> {
90    #[inline]
91    fn bit_decode<S: Source>(r: &mut S) -> Result<Self, BitError> {
92        Ok(WireLen::Set(r.read::<T>()?))
93    }
94}
95
96impl<T: Bits> BitEncode for WireLen<T> {
97    #[inline]
98    fn bit_encode<K: Sink>(&self, w: &mut K) -> Result<(), BitError> {
99        match self {
100            WireLen::Set(v) => w.write(*v),
101            WireLen::Auto => Err(BitError::convert(
102                "unresolved `WireLen::Auto`: this field needs a \
103                 `#[bw(auto_len = count(f)|bytes(f))]` (or `#[bin(auto_len(...))]`) directive \
104                 to derive its value, or set it explicitly with `WireLen::set(n)`"
105                    .into(),
106                w.bit_pos(),
107            )),
108        }
109    }
110}
111
112/// A `WireLen<T>` occupies exactly `T`'s width, so it stays a fixed-width field (a DNS
113/// header with `WireLen<u16>` counts is still 12 bytes).
114impl<T: Bits> FixedBitLen for WireLen<T> {
115    const BIT_LEN: u32 = <T as Bits>::BITS;
116}
117
118#[cfg(test)]
119mod unit {
120    use super::*;
121    use crate::bitstream::{BitReader, BitWriter};
122
123    #[test]
124    fn default_is_auto() {
125        assert_eq!(WireLen::<u16>::default(), WireLen::Auto);
126        assert!(WireLen::<u16>::auto().is_auto());
127        assert!(!WireLen::set(5u16).is_auto());
128    }
129
130    #[test]
131    fn decode_yields_set() {
132        let mut r = BitReader::new(&[0x12, 0x34]);
133        let v = WireLen::<u16>::bit_decode(&mut r).unwrap();
134        assert_eq!(v, WireLen::Set(0x1234));
135        assert_eq!(v.get(), Some(&0x1234));
136        assert_eq!(v.to_count(), 0x1234);
137    }
138
139    #[test]
140    fn encode_set_writes_the_value_auto_errors() {
141        let mut w = BitWriter::new();
142        WireLen::set(0xABCDu16).bit_encode(&mut w).unwrap();
143        assert_eq!(w.into_bytes(), [0xAB, 0xCD]);
144
145        let mut w = BitWriter::new();
146        assert!(WireLen::<u16>::auto().bit_encode(&mut w).is_err());
147    }
148
149    #[test]
150    fn resolve_count_fills_auto_leaves_set_and_checks_overflow() {
151        assert_eq!(
152            WireLen::<u16>::auto().resolve_count(7).unwrap(),
153            WireLen::Set(7)
154        );
155        assert_eq!(
156            WireLen::set(9u16).resolve_count(7).unwrap(),
157            WireLen::Set(9)
158        ); // Set wins
159        // 300 doesn't fit a u8 → checked error, not a truncation to 44.
160        assert!(WireLen::<u8>::auto().resolve_count(300).is_err());
161    }
162
163    #[test]
164    fn fixed_bit_len_matches_the_inner_width() {
165        assert_eq!(<WireLen<u16> as FixedBitLen>::BIT_LEN, 16);
166        assert_eq!(<WireLen<u32> as FixedBitLen>::BIT_LEN, 32);
167    }
168}