Skip to main content

fraiseql_wire/util/
bytes.rs

1//! Byte manipulation utilities for protocol parsing
2
3use bytes::{Buf, Bytes};
4use std::io;
5
6/// Extension trait for Bytes operations
7pub trait BytesExt {
8    /// Read a null-terminated string
9    ///
10    /// # Errors
11    ///
12    /// Returns `io::Error` with `InvalidData` if no null terminator is found or the bytes are not valid UTF-8.
13    fn read_cstr(&mut self) -> io::Result<String>;
14
15    /// Read a 32-bit big-endian integer
16    ///
17    /// # Errors
18    ///
19    /// Returns `io::Error` with `UnexpectedEof` if fewer than 4 bytes remain.
20    fn read_i32_be(&mut self) -> io::Result<i32>;
21
22    /// Read a 16-bit big-endian integer
23    ///
24    /// # Errors
25    ///
26    /// Returns `io::Error` with `UnexpectedEof` if fewer than 2 bytes remain.
27    fn read_i16_be(&mut self) -> io::Result<i16>;
28}
29
30impl BytesExt for Bytes {
31    fn read_cstr(&mut self) -> io::Result<String> {
32        let null_pos = self
33            .iter()
34            .position(|&b| b == 0)
35            .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "no null terminator"))?;
36
37        let s = String::from_utf8(self.slice(..null_pos).to_vec())
38            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
39
40        self.advance(null_pos + 1);
41        Ok(s)
42    }
43
44    fn read_i32_be(&mut self) -> io::Result<i32> {
45        if self.remaining() < 4 {
46            return Err(io::Error::new(
47                io::ErrorKind::UnexpectedEof,
48                "not enough bytes",
49            ));
50        }
51        Ok(self.get_i32())
52    }
53
54    fn read_i16_be(&mut self) -> io::Result<i16> {
55        if self.remaining() < 2 {
56            return Err(io::Error::new(
57                io::ErrorKind::UnexpectedEof,
58                "not enough bytes",
59            ));
60        }
61        Ok(self.get_i16())
62    }
63}
64
65#[cfg(test)]
66mod tests {
67    #![allow(clippy::unwrap_used)] // Reason: test code, panics are acceptable
68    use super::*;
69
70    #[test]
71    fn test_read_cstr() {
72        let mut data = Bytes::from_static(b"hello\0world\0");
73        assert_eq!(data.read_cstr().unwrap(), "hello");
74        assert_eq!(data.read_cstr().unwrap(), "world");
75    }
76
77    #[test]
78    fn test_read_i32() {
79        let mut data = Bytes::from_static(&[0x00, 0x00, 0x01, 0x00]);
80        assert_eq!(data.read_i32_be().unwrap(), 256);
81    }
82}