tempest_core/encoding/
encoding_raw.rs1use 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 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 #[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]); 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 #[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 #[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 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 let mut buf = BytesMut::new();
197 crate::encoding::varint::encode_varint(&mut buf, 100);
198 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 #[test]
221 fn test_i64_raw_not_order_preserving() {
222 assert!(encode_i64(-1) > encode_i64(0));
226 }
227
228 #[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}