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