mqtt5_protocol/encoding/
variable_int.rs

1//! BeBytes-compatible variable length integer implementation for MQTT
2//!
3//! This module provides a variable length integer type that integrates with `BeBytes` 2.3.0
4//! for automatic serialization/deserialization with size expressions support.
5
6use crate::error::{MqttError, Result};
7use crate::prelude::{format, ToString, Vec};
8use bebytes::BeBytes;
9use bytes::{Buf, BufMut};
10use core::fmt;
11
12/// Maximum value that can be encoded as a variable byte integer (268,435,455)
13pub const VARIABLE_INT_MAX: u32 = 268_435_455;
14
15/// Variable length integer as defined by MQTT specification
16///
17/// Encodes values using 1-4 bytes:
18/// - 0-127: 1 byte
19/// - 128-16,383: 2 bytes
20/// - 16,384-2,097,151: 3 bytes
21/// - 2,097,152-268,435,455: 4 bytes
22#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
23pub struct VariableInt {
24    value: u32,
25}
26
27impl VariableInt {
28    /// Creates a new `VariableInt` from a u32 value
29    ///
30    /// # Errors
31    ///
32    /// Returns `MqttError::ProtocolError` if the value exceeds the maximum
33    pub fn new(value: u32) -> Result<Self> {
34        if value > VARIABLE_INT_MAX {
35            return Err(MqttError::ProtocolError(format!(
36                "Variable integer value {value} exceeds maximum {VARIABLE_INT_MAX}"
37            )));
38        }
39        Ok(Self { value })
40    }
41
42    /// Creates a new `VariableInt` from a u32 value without validation
43    ///
44    /// # Safety
45    ///
46    /// The caller must ensure that value <= `VARIABLE_INT_MAX`
47    #[must_use]
48    pub fn new_unchecked(value: u32) -> Self {
49        debug_assert!(value <= VARIABLE_INT_MAX);
50        Self { value }
51    }
52
53    /// Returns the actual value
54    #[must_use]
55    pub fn value(&self) -> u32 {
56        self.value
57    }
58
59    /// Returns the number of bytes needed to encode this value
60    #[must_use]
61    pub fn encoded_size(&self) -> u32 {
62        match self.value {
63            0..=127 => 1,
64            128..=16_383 => 2,
65            16_384..=2_097_151 => 3,
66            2_097_152..=VARIABLE_INT_MAX => 4,
67            _ => unreachable!("Invalid variable int value"),
68        }
69    }
70
71    /// Encodes this variable integer into the provided buffer
72    ///
73    /// # Errors
74    ///
75    /// Always succeeds for valid `VariableInt` values
76    pub fn encode<B: BufMut>(&self, buf: &mut B) -> Result<()> {
77        let mut val = self.value;
78        loop {
79            let mut byte = (val % 128) as u8;
80            val /= 128;
81            if val > 0 {
82                byte |= 0x80; // Set continuation bit
83            }
84            buf.put_u8(byte);
85            if val == 0 {
86                break;
87            }
88        }
89        Ok(())
90    }
91
92    /// Decodes a variable integer from the buffer
93    ///
94    /// # Errors
95    ///
96    /// Returns an error if:
97    /// - The buffer doesn't contain enough bytes
98    /// - The encoded value exceeds the maximum
99    /// - More than 4 bytes are used (protocol violation)
100    pub fn decode<B: Buf>(buf: &mut B) -> Result<Self> {
101        let mut value = 0u32;
102        let mut multiplier = 1u32;
103        let mut byte_count = 0;
104
105        loop {
106            if !buf.has_remaining() {
107                return Err(MqttError::MalformedPacket(
108                    "Insufficient bytes for variable integer".to_string(),
109                ));
110            }
111
112            byte_count += 1;
113            if byte_count > 4 {
114                return Err(MqttError::MalformedPacket(
115                    "Variable integer exceeds 4 bytes".to_string(),
116                ));
117            }
118
119            let byte = buf.get_u8();
120            value += u32::from(byte & 0x7F) * multiplier;
121
122            if (byte & 0x80) == 0 {
123                break;
124            }
125
126            multiplier *= 128;
127            if multiplier > 128 * 128 * 128 {
128                return Err(MqttError::MalformedPacket(
129                    "Variable integer overflow".to_string(),
130                ));
131            }
132        }
133
134        if value > VARIABLE_INT_MAX {
135            return Err(MqttError::MalformedPacket(format!(
136                "Variable integer value {value} exceeds maximum"
137            )));
138        }
139
140        Ok(Self { value })
141    }
142}
143
144impl fmt::Display for VariableInt {
145    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
146        write!(f, "{}", self.value)
147    }
148}
149
150impl From<VariableInt> for u32 {
151    fn from(v: VariableInt) -> Self {
152        v.value
153    }
154}
155
156impl TryFrom<u32> for VariableInt {
157    type Error = MqttError;
158
159    fn try_from(value: u32) -> Result<Self> {
160        Self::new(value)
161    }
162}
163
164impl TryFrom<usize> for VariableInt {
165    type Error = MqttError;
166
167    fn try_from(value: usize) -> Result<Self> {
168        let value = u32::try_from(value).map_err(|_| {
169            MqttError::ProtocolError("Value too large for variable integer".to_string())
170        })?;
171        Self::new(value)
172    }
173}
174
175// BeBytes implementation
176impl BeBytes for VariableInt {
177    fn field_size() -> usize {
178        // Variable length, so we return 0 as it's determined at runtime
179        0
180    }
181
182    fn to_be_bytes(&self) -> Vec<u8> {
183        let mut buf = Vec::new();
184        let _ = self.encode(&mut buf); // Encoding can't fail for valid VariableInt
185        buf
186    }
187
188    fn try_from_be_bytes(
189        bytes: &[u8],
190    ) -> core::result::Result<(Self, usize), bebytes::BeBytesError> {
191        use bytes::Bytes;
192        let mut buf = Bytes::copy_from_slice(bytes);
193        let start_len = buf.len();
194
195        match Self::decode(&mut buf) {
196            Ok(var_int) => {
197                let consumed = start_len - buf.len();
198                Ok((var_int, consumed))
199            }
200            Err(_) => Err(bebytes::BeBytesError::InsufficientData {
201                expected: 1, // At least 1 byte expected for variable int
202                actual: bytes.len(),
203            }),
204        }
205    }
206
207    fn to_le_bytes(&self) -> Vec<u8> {
208        // MQTT uses big-endian encoding only
209        self.to_be_bytes()
210    }
211
212    fn try_from_le_bytes(
213        bytes: &[u8],
214    ) -> core::result::Result<(Self, usize), bebytes::BeBytesError> {
215        // MQTT uses big-endian encoding only
216        Self::try_from_be_bytes(bytes)
217    }
218}
219
220/// Encodes a u32 value as a variable byte integer (compatibility function)
221///
222/// This function provides compatibility with the old `variable_byte` module API.
223/// Prefer using `VariableInt::new(value)?.encode(buf)` for new code.
224///
225/// # Errors
226///
227/// Returns `MqttError::ProtocolError` if the value exceeds the maximum
228pub fn encode_variable_int<B: BufMut>(buf: &mut B, value: u32) -> Result<()> {
229    VariableInt::new(value)?.encode(buf)
230}
231
232/// Decodes a variable byte integer from the buffer (compatibility function)
233///
234/// This function provides compatibility with the old `variable_byte` module API.
235/// Prefer using `VariableInt::decode(buf)?.value()` for new code.
236///
237/// # Errors
238///
239/// Returns an error if decoding fails
240pub fn decode_variable_int<B: Buf>(buf: &mut B) -> Result<u32> {
241    Ok(VariableInt::decode(buf)?.value())
242}
243
244/// Calculates the number of bytes needed to encode a value (compatibility function)
245///
246/// This function provides compatibility with the old `variable_byte` module API.
247/// Prefer using `VariableInt::new(value)?.encoded_size()` for new code.
248#[must_use]
249pub fn variable_int_len(value: u32) -> usize {
250    VariableInt::new_unchecked(value.min(VARIABLE_INT_MAX)).encoded_size() as usize
251}
252
253/// Alias for `variable_int_len` for consistency (compatibility function)
254#[must_use]
255pub fn encoded_variable_int_len(value: u32) -> usize {
256    variable_int_len(value)
257}
258
259/// Maximum value that can be encoded as a variable byte integer (compatibility constant)
260pub const VARIABLE_BYTE_INT_MAX: u32 = VARIABLE_INT_MAX;
261
262#[cfg(test)]
263mod tests {
264    use super::*;
265    use bytes::BytesMut;
266
267    #[test]
268    fn test_new_valid_values() {
269        assert!(VariableInt::new(0).is_ok());
270        assert!(VariableInt::new(127).is_ok());
271        assert!(VariableInt::new(128).is_ok());
272        assert!(VariableInt::new(16_383).is_ok());
273        assert!(VariableInt::new(16_384).is_ok());
274        assert!(VariableInt::new(2_097_151).is_ok());
275        assert!(VariableInt::new(2_097_152).is_ok());
276        assert!(VariableInt::new(VARIABLE_INT_MAX).is_ok());
277    }
278
279    #[test]
280    fn test_new_invalid_values() {
281        assert!(VariableInt::new(VARIABLE_INT_MAX + 1).is_err());
282        assert!(VariableInt::new(u32::MAX).is_err());
283    }
284
285    #[test]
286    fn test_encoded_size() {
287        assert_eq!(VariableInt::new_unchecked(0).encoded_size(), 1);
288        assert_eq!(VariableInt::new_unchecked(127).encoded_size(), 1);
289        assert_eq!(VariableInt::new_unchecked(128).encoded_size(), 2);
290        assert_eq!(VariableInt::new_unchecked(16_383).encoded_size(), 2);
291        assert_eq!(VariableInt::new_unchecked(16_384).encoded_size(), 3);
292        assert_eq!(VariableInt::new_unchecked(2_097_151).encoded_size(), 3);
293        assert_eq!(VariableInt::new_unchecked(2_097_152).encoded_size(), 4);
294        assert_eq!(
295            VariableInt::new_unchecked(VARIABLE_INT_MAX).encoded_size(),
296            4
297        );
298    }
299
300    #[test]
301    fn test_encode_decode_single_byte() {
302        let mut buf = BytesMut::new();
303
304        for value in [0, 1, 64, 127] {
305            buf.clear();
306            let var_int = VariableInt::new(value).unwrap();
307            var_int.encode(&mut buf).unwrap();
308            assert_eq!(buf.len(), 1);
309
310            let decoded = VariableInt::decode(&mut buf).unwrap();
311            assert_eq!(decoded.value(), value);
312        }
313    }
314
315    #[test]
316    fn test_encode_decode_two_bytes() {
317        let mut buf = BytesMut::new();
318
319        for value in [128, 129, 321, 16_383] {
320            buf.clear();
321            let var_int = VariableInt::new(value).unwrap();
322            var_int.encode(&mut buf).unwrap();
323            assert_eq!(buf.len(), 2);
324
325            let decoded = VariableInt::decode(&mut buf).unwrap();
326            assert_eq!(decoded.value(), value);
327        }
328    }
329
330    #[test]
331    fn test_encode_decode_three_bytes() {
332        let mut buf = BytesMut::new();
333
334        for value in [16_384, 65_535, 2_097_151] {
335            buf.clear();
336            let var_int = VariableInt::new(value).unwrap();
337            var_int.encode(&mut buf).unwrap();
338            assert_eq!(buf.len(), 3);
339
340            let decoded = VariableInt::decode(&mut buf).unwrap();
341            assert_eq!(decoded.value(), value);
342        }
343    }
344
345    #[test]
346    fn test_encode_decode_four_bytes() {
347        let mut buf = BytesMut::new();
348
349        for value in [2_097_152, 10_000_000, VARIABLE_INT_MAX] {
350            buf.clear();
351            let var_int = VariableInt::new(value).unwrap();
352            var_int.encode(&mut buf).unwrap();
353            assert_eq!(buf.len(), 4);
354
355            let decoded = VariableInt::decode(&mut buf).unwrap();
356            assert_eq!(decoded.value(), value);
357        }
358    }
359
360    #[test]
361    fn test_mqtt_spec_examples() {
362        let mut buf = BytesMut::new();
363
364        // Example from MQTT spec: 64 decimal = 0x40
365        let var_int = VariableInt::new(64).unwrap();
366        var_int.encode(&mut buf).unwrap();
367        assert_eq!(buf[0], 0x40);
368
369        // Example: 321 decimal = 0xC1 0x02
370        buf.clear();
371        let var_int = VariableInt::new(321).unwrap();
372        var_int.encode(&mut buf).unwrap();
373        assert_eq!(buf[0], 0xC1);
374        assert_eq!(buf[1], 0x02);
375    }
376
377    #[test]
378    fn test_bebytes_integration() {
379        use bebytes::BeBytes;
380
381        // Test BeBytes to_be_bytes
382        let var_int = VariableInt::new(321).unwrap();
383        let bytes = var_int.to_be_bytes();
384        assert_eq!(bytes.len(), 2);
385        assert_eq!(bytes[0], 0xC1);
386        assert_eq!(bytes[1], 0x02);
387
388        // Test BeBytes try_from_be_bytes
389        let (decoded, consumed) = VariableInt::try_from_be_bytes(&bytes).unwrap();
390        assert_eq!(decoded.value(), 321);
391        assert_eq!(consumed, 2);
392
393        // Test field_size (static method)
394        assert_eq!(VariableInt::field_size(), 0); // Variable length
395    }
396
397    #[test]
398    fn test_decode_insufficient_bytes() {
399        let mut buf = BytesMut::new();
400        buf.put_u8(0x80); // Continuation bit set but no following byte
401
402        let result = VariableInt::decode(&mut buf);
403        assert!(result.is_err());
404    }
405
406    #[test]
407    fn test_decode_too_many_bytes() {
408        let mut buf = BytesMut::new();
409        // 5 bytes with continuation bits
410        buf.put_u8(0x80);
411        buf.put_u8(0x80);
412        buf.put_u8(0x80);
413        buf.put_u8(0x80);
414        buf.put_u8(0x01);
415
416        let result = VariableInt::decode(&mut buf);
417        assert!(result.is_err());
418    }
419
420    #[test]
421    fn test_conversions() {
422        let var_int = VariableInt::new(123).unwrap();
423
424        // Test From<VariableInt> for u32
425        let value: u32 = var_int.into();
426        assert_eq!(value, 123);
427
428        // Test TryFrom<u32> for VariableInt
429        let var_int2 = VariableInt::try_from(456u32).unwrap();
430        assert_eq!(var_int2.value(), 456);
431
432        // Test TryFrom<usize> for VariableInt
433        let var_int3 = VariableInt::try_from(789usize).unwrap();
434        assert_eq!(var_int3.value(), 789);
435
436        // Test invalid conversions
437        assert!(VariableInt::try_from(VARIABLE_INT_MAX + 1).is_err());
438    }
439
440    #[test]
441    fn test_display() {
442        let var_int = VariableInt::new(12345).unwrap();
443        assert_eq!(format!("{var_int}"), "12345");
444    }
445
446    #[cfg(test)]
447    mod property_tests {
448        use super::*;
449        use proptest::prelude::*;
450
451        proptest! {
452            #[test]
453            fn prop_round_trip(value in 0u32..=VARIABLE_INT_MAX) {
454                let mut buf = BytesMut::new();
455
456                let var_int = VariableInt::new(value).unwrap();
457                var_int.encode(&mut buf).unwrap();
458                let decoded = VariableInt::decode(&mut buf).unwrap();
459
460                prop_assert_eq!(decoded.value(), value);
461            }
462
463            #[test]
464            fn prop_encoded_size_matches_actual(value in 0u32..=VARIABLE_INT_MAX) {
465                let mut buf = BytesMut::new();
466
467                let var_int = VariableInt::new(value).unwrap();
468                let predicted_size = var_int.encoded_size();
469                var_int.encode(&mut buf).unwrap();
470                let actual_size = u32::try_from(buf.len()).expect("buffer size should fit in u32");
471
472                prop_assert_eq!(predicted_size, actual_size);
473            }
474
475            #[test]
476            fn prop_bebytes_round_trip(value in 0u32..=VARIABLE_INT_MAX) {
477                use bebytes::BeBytes;
478
479                let var_int = VariableInt::new(value).unwrap();
480                let bytes = var_int.to_be_bytes();
481                prop_assert_eq!(bytes.len(), var_int.encoded_size() as usize);
482
483                let (decoded, consumed) = VariableInt::try_from_be_bytes(&bytes).unwrap();
484                prop_assert_eq!(decoded.value(), value);
485                prop_assert_eq!(consumed, bytes.len());
486            }
487
488            #[test]
489            fn prop_invalid_values_rejected(value in (VARIABLE_INT_MAX + 1)..=u32::MAX) {
490                let result = VariableInt::new(value);
491                prop_assert!(result.is_err());
492            }
493
494            #[test]
495            fn prop_size_boundaries(value in 0u32..=VARIABLE_INT_MAX) {
496                let var_int = VariableInt::new(value).unwrap();
497                let size = var_int.encoded_size();
498
499                match value {
500                    0..=127 => prop_assert_eq!(size, 1),
501                    128..=16_383 => prop_assert_eq!(size, 2),
502                    16_384..=2_097_151 => prop_assert_eq!(size, 3),
503                    2_097_152..=VARIABLE_INT_MAX => prop_assert_eq!(size, 4),
504                    _ => unreachable!(),
505                }
506            }
507        }
508    }
509}