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    fn read_cstr(&mut self) -> io::Result<String>;
10
11    /// Read a 32-bit big-endian integer
12    fn read_i32_be(&mut self) -> io::Result<i32>;
13
14    /// Read a 16-bit big-endian integer
15    fn read_i16_be(&mut self) -> io::Result<i16>;
16}
17
18impl BytesExt for Bytes {
19    fn read_cstr(&mut self) -> io::Result<String> {
20        let null_pos = self
21            .iter()
22            .position(|&b| b == 0)
23            .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "no null terminator"))?;
24
25        let s = String::from_utf8(self.slice(..null_pos).to_vec())
26            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
27
28        self.advance(null_pos + 1);
29        Ok(s)
30    }
31
32    fn read_i32_be(&mut self) -> io::Result<i32> {
33        if self.remaining() < 4 {
34            return Err(io::Error::new(
35                io::ErrorKind::UnexpectedEof,
36                "not enough bytes",
37            ));
38        }
39        Ok(self.get_i32())
40    }
41
42    fn read_i16_be(&mut self) -> io::Result<i16> {
43        if self.remaining() < 2 {
44            return Err(io::Error::new(
45                io::ErrorKind::UnexpectedEof,
46                "not enough bytes",
47            ));
48        }
49        Ok(self.get_i16())
50    }
51}
52
53#[cfg(test)]
54mod tests {
55    use super::*;
56
57    #[test]
58    fn test_read_cstr() {
59        let mut data = Bytes::from_static(b"hello\0world\0");
60        assert_eq!(data.read_cstr().unwrap(), "hello");
61        assert_eq!(data.read_cstr().unwrap(), "world");
62    }
63
64    #[test]
65    fn test_read_i32() {
66        let mut data = Bytes::from_static(&[0x00, 0x00, 0x01, 0x00]);
67        assert_eq!(data.read_i32_be().unwrap(), 256);
68    }
69}