Skip to main content

tempest_core/encoding/
encoding_raw.rs

1use std::string::FromUtf8Error;
2
3use bytes::{Buf, BufMut, Bytes, BytesMut};
4
5use crate::encoding::varint::{decode_varint, encode_varint};
6
7pub trait BufPutRawExt {
8    fn put_i64_raw(&mut self, i: i64);
9    fn put_bool_raw(&mut self, b: bool);
10    fn put_str_raw(&mut self, s: &str);
11}
12
13impl BufPutRawExt for BytesMut {
14    fn put_i64_raw(&mut self, i: i64) {
15        self.put_i64(i);
16    }
17
18    fn put_bool_raw(&mut self, b: bool) {
19        self.put_u8(b as u8);
20    }
21
22    fn put_str_raw(&mut self, s: &str) {
23        encode_varint(self, s.len());
24        self.put_slice(s.as_bytes());
25    }
26}
27
28#[derive(Debug, Display, Error, From)]
29pub enum RawDecodeError {
30    UnexpectedEof,
31    DecodeVarintError,
32    FromUtf8Error(FromUtf8Error),
33}
34
35pub trait BufGetRawExt {
36    fn get_i64_raw(&mut self) -> Result<i64, RawDecodeError>;
37    fn get_bool_raw(&mut self) -> Result<bool, RawDecodeError>;
38    fn get_str_raw(&mut self) -> Result<String, RawDecodeError>;
39}
40
41impl BufGetRawExt for Bytes {
42    fn get_i64_raw(&mut self) -> Result<i64, RawDecodeError> {
43        if self.len() < 8 {
44            return Err(RawDecodeError::UnexpectedEof);
45        }
46        Ok(self.get_i64())
47    }
48
49    fn get_bool_raw(&mut self) -> Result<bool, RawDecodeError> {
50        if self.is_empty() {
51            return Err(RawDecodeError::UnexpectedEof);
52        }
53        Ok(self.get_u8() != 0)
54    }
55
56    fn get_str_raw(&mut self) -> Result<String, RawDecodeError> {
57        let (len, bytes_read) =
58            decode_varint(self).ok_or_else(|| RawDecodeError::DecodeVarintError)?;
59        self.advance(bytes_read);
60        if self.len() < len {
61            return Err(RawDecodeError::UnexpectedEof);
62        }
63        let bytes = self.split_to(len);
64
65        String::from_utf8(bytes.to_vec()).map_err(RawDecodeError::FromUtf8Error)
66    }
67}
68
69#[cfg(test)]
70mod tests {
71    use bytes::{Bytes, BytesMut};
72
73    use super::*;
74
75    // -- helpers --
76
77    fn encode_i64(i: i64) -> Bytes {
78        let mut buf = BytesMut::new();
79        buf.put_i64_raw(i);
80        buf.freeze()
81    }
82
83    fn encode_bool(b: bool) -> Bytes {
84        let mut buf = BytesMut::new();
85        buf.put_bool_raw(b);
86        buf.freeze()
87    }
88
89    fn encode_str(s: &str) -> Bytes {
90        let mut buf = BytesMut::new();
91        buf.put_str_raw(s);
92        buf.freeze()
93    }
94
95    // -- i64 roundtrip --
96
97    #[test]
98    fn test_i64_roundtrip() {
99        for val in [0i64, 1, -1, i64::MIN, i64::MAX, -1000, 1000, 42] {
100            let mut bytes = encode_i64(val);
101            assert_eq!(
102                bytes.get_i64_raw().unwrap(),
103                val,
104                "roundtrip failed for {}",
105                val
106            );
107        }
108    }
109
110    #[test]
111    fn test_i64_eof() {
112        let mut bytes = Bytes::from_static(&[0x00, 0x00]); // only 2 bytes, need 8
113        assert!(matches!(
114            bytes.get_i64_raw(),
115            Err(RawDecodeError::UnexpectedEof)
116        ));
117    }
118
119    #[test]
120    fn test_i64_exact_size() {
121        let mut bytes = encode_i64(42);
122        assert_eq!(bytes.len(), 8);
123        bytes.get_i64_raw().unwrap();
124        assert!(bytes.is_empty());
125    }
126
127    // -- bool roundtrip --
128
129    #[test]
130    fn test_bool_roundtrip() {
131        for val in [true, false] {
132            let mut bytes = encode_bool(val);
133            assert_eq!(bytes.get_bool_raw().unwrap(), val);
134        }
135    }
136
137    #[test]
138    fn test_bool_eof() {
139        let mut bytes = Bytes::new();
140        assert!(matches!(
141            bytes.get_bool_raw(),
142            Err(RawDecodeError::UnexpectedEof)
143        ));
144    }
145
146    #[test]
147    fn test_bool_exact_size() {
148        let mut bytes = encode_bool(true);
149        assert_eq!(bytes.len(), 1);
150        bytes.get_bool_raw().unwrap();
151        assert!(bytes.is_empty());
152    }
153
154    // -- string roundtrip --
155
156    #[test]
157    fn test_str_roundtrip() {
158        for val in ["", "hello", "hello world", "unicode: ??"] {
159            let mut bytes = encode_str(val);
160            assert_eq!(
161                bytes.get_str_raw().unwrap(),
162                val,
163                "roundtrip failed for {:?}",
164                val
165            );
166        }
167    }
168
169    #[test]
170    fn test_str_with_null_bytes() {
171        // raw encoding is not null-escaped - null bytes are stored as-is
172        let s = "hel\x00lo";
173        let mut bytes = encode_str(s);
174        assert_eq!(bytes.get_str_raw().unwrap(), s);
175    }
176
177    #[test]
178    fn test_str_empty() {
179        let mut bytes = encode_str("");
180        assert_eq!(bytes.get_str_raw().unwrap(), "");
181        assert!(bytes.is_empty());
182    }
183
184    #[test]
185    fn test_str_eof_in_length() {
186        let mut bytes = Bytes::new();
187        assert!(matches!(
188            bytes.get_str_raw(),
189            Err(RawDecodeError::DecodeVarintError)
190        ));
191    }
192
193    #[test]
194    fn test_str_eof_in_body() {
195        // write a length varint claiming 100 bytes but provide none
196        let mut buf = BytesMut::new();
197        crate::encoding::varint::encode_varint(&mut buf, 100);
198        // no actual string bytes follow
199        let mut bytes = buf.freeze();
200        assert!(matches!(
201            bytes.get_str_raw(),
202            Err(RawDecodeError::UnexpectedEof)
203        ));
204    }
205
206    #[test]
207    fn test_str_advances_cursor_correctly() {
208        let mut buf = BytesMut::new();
209        buf.put_str_raw("foo");
210        buf.put_str_raw("bar");
211        let mut bytes = buf.freeze();
212
213        assert_eq!(bytes.get_str_raw().unwrap(), "foo");
214        assert_eq!(bytes.get_str_raw().unwrap(), "bar");
215        assert!(bytes.is_empty());
216    }
217
218    // -- note: raw i64 is NOT order-preserving --
219
220    #[test]
221    fn test_i64_raw_not_order_preserving() {
222        // raw uses put_i64 (two's complement BE), NOT sign-bit-flipped
223        // so -1 encodes as 0xFFFF... which sorts AFTER 0 - opposite of numeric order
224        // this confirms raw != lexical
225        assert!(encode_i64(-1) > encode_i64(0));
226    }
227
228    // -- mixed sequence --
229
230    #[test]
231    fn test_mixed_sequence() {
232        let mut buf = BytesMut::new();
233        buf.put_i64_raw(99);
234        buf.put_bool_raw(false);
235        buf.put_str_raw("tempest");
236        buf.put_i64_raw(-1);
237        buf.put_str_raw("");
238
239        let mut bytes = buf.freeze();
240        assert_eq!(bytes.get_i64_raw().unwrap(), 99);
241        assert_eq!(bytes.get_bool_raw().unwrap(), false);
242        assert_eq!(bytes.get_str_raw().unwrap(), "tempest");
243        assert_eq!(bytes.get_i64_raw().unwrap(), -1);
244        assert_eq!(bytes.get_str_raw().unwrap(), "");
245        assert!(bytes.is_empty());
246    }
247}