1#![allow(clippy::unwrap_used)]
8
9use thiserror::Error;
10
11#[derive(Error, Debug)]
13pub enum EncodingError {
14 #[error("IO error: {0}")]
15 Io(#[from] std::io::Error),
16
17 #[error("Field number out of range [1, 32]: {0}")]
18 InvalidFieldNumber(u32),
19
20 #[error("Value exceeds maximum safe integer")]
21 ValueTooLarge,
22
23 #[error("Hash must be exactly 32 bytes, got {0}")]
24 InvalidHashLength(usize),
25
26 #[error("Cannot marshal negative bigint")]
27 NegativeBigInt,
28
29 #[error("Invalid UTF-8 string")]
30 InvalidUtf8,
31}
32
33#[derive(Debug, Clone)]
35pub struct BinaryWriter {
36 buffer: Vec<u8>,
37}
38
39impl BinaryWriter {
40 pub fn new() -> Self {
42 Self { buffer: Vec::new() }
43 }
44
45 pub fn with_capacity(capacity: usize) -> Self {
47 Self {
48 buffer: Vec::with_capacity(capacity),
49 }
50 }
51
52 pub fn into_bytes(self) -> Vec<u8> {
54 self.buffer
55 }
56
57 pub fn bytes(&self) -> &[u8] {
59 &self.buffer
60 }
61
62 pub fn clear(&mut self) {
64 self.buffer.clear();
65 }
66
67 pub fn write_bytes(&mut self, bytes: &[u8]) -> Result<(), EncodingError> {
69 self.buffer.extend_from_slice(bytes);
70 Ok(())
71 }
72
73 pub fn write_field(&mut self, field: u32, value: &[u8]) -> Result<(), EncodingError> {
76 if field < 1 || field > 32 {
77 return Err(EncodingError::InvalidFieldNumber(field));
78 }
79 self.write_uvarint(field as u64)?;
80 self.write_bytes(value)?;
81 Ok(())
82 }
83
84 pub fn write_uvarint(&mut self, mut value: u64) -> Result<(), EncodingError> {
87 while value >= 0x80 {
89 self.buffer.push((value as u8) | 0x80);
90 value >>= 7;
91 }
92 self.buffer.push(value as u8);
93 Ok(())
94 }
95
96 pub fn write_uvarint_field(&mut self, value: u64, field: u32) -> Result<(), EncodingError> {
98 let mut temp_writer = BinaryWriter::new();
99 temp_writer.write_uvarint(value)?;
100 self.write_field(field, temp_writer.bytes())?;
101 Ok(())
102 }
103
104 pub fn write_varint(&mut self, value: i64) -> Result<(), EncodingError> {
107 let unsigned = ((value as u64) << 1) ^ ((value >> 63) as u64);
109 self.write_uvarint(unsigned)
110 }
111
112 pub fn write_varint_field(&mut self, value: i64, field: u32) -> Result<(), EncodingError> {
114 let mut temp_writer = BinaryWriter::new();
115 temp_writer.write_varint(value)?;
116 self.write_field(field, temp_writer.bytes())?;
117 Ok(())
118 }
119
120 pub fn write_big_number(&mut self, value: &num_bigint::BigUint) -> Result<(), EncodingError> {
123 let hex_string = value.to_str_radix(16);
124
125 let padded_hex = if hex_string.len() % 2 == 1 {
127 format!("0{}", hex_string)
128 } else {
129 hex_string
130 };
131
132 let bytes: Result<Vec<u8>, _> = (0..padded_hex.len())
134 .step_by(2)
135 .map(|i| u8::from_str_radix(&padded_hex[i..i + 2], 16))
136 .collect();
137
138 let bytes = bytes.map_err(|_| EncodingError::InvalidUtf8)?;
139 self.write_bytes_with_length(&bytes)?;
140 Ok(())
141 }
142
143 pub fn write_big_number_field(
145 &mut self,
146 value: &num_bigint::BigUint,
147 field: u32,
148 ) -> Result<(), EncodingError> {
149 let mut temp_writer = BinaryWriter::new();
150 temp_writer.write_big_number(value)?;
151 self.write_field(field, temp_writer.bytes())?;
152 Ok(())
153 }
154
155 pub fn write_bool(&mut self, value: bool) -> Result<(), EncodingError> {
158 self.buffer.push(if value { 1 } else { 0 });
159 Ok(())
160 }
161
162 pub fn write_bool_field(&mut self, value: bool, field: u32) -> Result<(), EncodingError> {
164 let mut temp_writer = BinaryWriter::new();
165 temp_writer.write_bool(value)?;
166 self.write_field(field, temp_writer.bytes())?;
167 Ok(())
168 }
169
170 pub fn write_string(&mut self, value: &str) -> Result<(), EncodingError> {
173 let bytes = value.as_bytes();
174 self.write_bytes_with_length(bytes)?;
175 Ok(())
176 }
177
178 pub fn write_string_field(&mut self, value: &str, field: u32) -> Result<(), EncodingError> {
180 let mut temp_writer = BinaryWriter::new();
181 temp_writer.write_string(value)?;
182 self.write_field(field, temp_writer.bytes())?;
183 Ok(())
184 }
185
186 pub fn write_bytes_with_length(&mut self, bytes: &[u8]) -> Result<(), EncodingError> {
189 self.write_uvarint(bytes.len() as u64)?;
190 self.write_bytes(bytes)?;
191 Ok(())
192 }
193
194 pub fn write_bytes_field(&mut self, bytes: &[u8], field: u32) -> Result<(), EncodingError> {
196 let mut temp_writer = BinaryWriter::new();
197 temp_writer.write_bytes_with_length(bytes)?;
198 self.write_field(field, temp_writer.bytes())?;
199 Ok(())
200 }
201
202 pub fn write_hash(&mut self, hash: &[u8; 32]) -> Result<(), EncodingError> {
205 self.write_bytes(hash)?;
206 Ok(())
207 }
208
209 pub fn write_hash_field(&mut self, hash: &[u8; 32], field: u32) -> Result<(), EncodingError> {
211 self.write_field(field, hash)?;
212 Ok(())
213 }
214
215 pub fn write_hash_bytes(&mut self, hash: &[u8]) -> Result<(), EncodingError> {
217 if hash.len() != 32 {
218 return Err(EncodingError::InvalidHashLength(hash.len()));
219 }
220 self.write_bytes(hash)?;
221 Ok(())
222 }
223
224 pub fn write_hash_bytes_field(&mut self, hash: &[u8], field: u32) -> Result<(), EncodingError> {
226 if hash.len() != 32 {
227 return Err(EncodingError::InvalidHashLength(hash.len()));
228 }
229 self.write_field(field, hash)?;
230 Ok(())
231 }
232
233 pub fn write_optional<T, F>(
235 &mut self,
236 value: Option<&T>,
237 _field: u32,
238 writer_fn: F,
239 ) -> Result<(), EncodingError>
240 where
241 T: Clone,
242 F: FnOnce(&mut Self, &T) -> Result<(), EncodingError>,
243 {
244 if let Some(val) = value {
245 writer_fn(self, val)?;
246 }
247 Ok(())
248 }
249
250 pub fn write_array<T, F>(
252 &mut self,
253 items: &[T],
254 _field: u32,
255 writer_fn: F,
256 ) -> Result<(), EncodingError>
257 where
258 F: Fn(&mut Self, &T) -> Result<(), EncodingError>,
259 {
260 for item in items {
261 writer_fn(self, item)?;
262 }
263 Ok(())
264 }
265}
266
267impl Default for BinaryWriter {
268 fn default() -> Self {
269 Self::new()
270 }
271}
272
273impl BinaryWriter {
275 pub fn with_field_number(data: &[u8], field: Option<u32>) -> Result<Vec<u8>, EncodingError> {
277 match field {
278 Some(field_num) => {
279 let mut writer = BinaryWriter::new();
280 writer.write_field(field_num, data)?;
281 Ok(writer.into_bytes())
282 }
283 None => Ok(data.to_vec()),
284 }
285 }
286
287 pub fn encode_uvarint(value: u64) -> Vec<u8> {
289 let mut writer = BinaryWriter::new();
290 writer.write_uvarint(value).unwrap(); writer.into_bytes()
292 }
293
294 pub fn encode_varint(value: i64) -> Vec<u8> {
296 let mut writer = BinaryWriter::new();
297 writer.write_varint(value).unwrap(); writer.into_bytes()
299 }
300
301 pub fn encode_string(value: &str) -> Vec<u8> {
303 let mut writer = BinaryWriter::new();
304 writer.write_string(value).unwrap(); writer.into_bytes()
306 }
307
308 pub fn encode_bytes(bytes: &[u8]) -> Vec<u8> {
310 let mut writer = BinaryWriter::new();
311 writer.write_bytes_with_length(bytes).unwrap(); writer.into_bytes()
313 }
314
315 pub fn encode_bool(value: bool) -> Vec<u8> {
317 vec![if value { 1 } else { 0 }]
318 }
319
320 pub fn encode_hash(hash: &[u8; 32]) -> Vec<u8> {
322 hash.to_vec()
323 }
324}
325
326#[cfg(test)]
327mod tests {
328 use super::*;
329
330 #[test]
331 fn test_uvarint_encoding() {
332 let test_cases = vec![
334 (0u64, vec![0]),
335 (1u64, vec![1]),
336 (127u64, vec![127]),
337 (128u64, vec![128, 1]),
338 (256u64, vec![128, 2]),
339 (16384u64, vec![128, 128, 1]),
340 ];
341
342 for (input, expected) in test_cases {
343 let result = BinaryWriter::encode_uvarint(input);
344 assert_eq!(result, expected, "uvarint({}) failed", input);
345 }
346 }
347
348 #[test]
349 fn test_varint_encoding() {
350 let test_cases = vec![
352 (0i64, vec![0]),
353 (-1i64, vec![1]),
354 (1i64, vec![2]),
355 (-2i64, vec![3]),
356 (2i64, vec![4]),
357 ];
358
359 for (input, expected) in test_cases {
360 let result = BinaryWriter::encode_varint(input);
361 assert_eq!(result, expected, "varint({}) failed", input);
362 }
363 }
364
365 #[test]
366 fn test_string_encoding() {
367 let result = BinaryWriter::encode_string("hello");
368 let expected = vec![5, b'h', b'e', b'l', b'l', b'o'];
370 assert_eq!(result, expected);
371 }
372
373 #[test]
374 fn test_bytes_encoding() {
375 let input = &[1, 2, 3, 4];
376 let result = BinaryWriter::encode_bytes(input);
377 let expected = vec![4, 1, 2, 3, 4];
379 assert_eq!(result, expected);
380 }
381
382 #[test]
383 fn test_bool_encoding() {
384 assert_eq!(BinaryWriter::encode_bool(true), vec![1]);
385 assert_eq!(BinaryWriter::encode_bool(false), vec![0]);
386 }
387
388 #[test]
389 fn test_field_encoding() {
390 let mut writer = BinaryWriter::new();
391 writer.write_field(1, &[42]).unwrap();
392 let expected = vec![1, 42];
394 assert_eq!(writer.bytes(), &expected);
395 }
396
397 #[test]
398 fn test_hash_validation() {
399 let mut writer = BinaryWriter::new();
400
401 let valid_hash = [0u8; 32];
403 assert!(writer.write_hash(&valid_hash).is_ok());
404
405 let invalid_hash = [0u8; 31];
407 assert!(writer.write_hash_bytes(&invalid_hash).is_err());
408 }
409}