zero_postgres/protocol/
codec.rs

1//! PostgreSQL wire protocol encoding and decoding primitives.
2//!
3//! PostgreSQL uses big-endian (network byte order) for all integers.
4
5use crate::error::{Error, Result};
6use zerocopy::FromBytes;
7
8use zerocopy::byteorder::big_endian::{I16 as I16BE, I32 as I32BE, U16 as U16BE, U32 as U32BE};
9
10/// Read 1-byte unsigned integer.
11#[inline]
12pub fn read_u8(data: &[u8]) -> Result<(u8, &[u8])> {
13    if data.is_empty() {
14        return Err(Error::Protocol("read_u8: empty buffer".into()));
15    }
16    Ok((data[0], &data[1..]))
17}
18
19/// Read 2-byte big-endian signed integer.
20#[inline]
21pub fn read_i16(data: &[u8]) -> Result<(i16, &[u8])> {
22    if data.len() < 2 {
23        return Err(Error::Protocol(format!(
24            "read_i16: buffer too short: {} < 2",
25            data.len()
26        )));
27    }
28    let value = I16BE::ref_from_bytes(&data[..2])
29        .map_err(|e| Error::Protocol(format!("read_i16: {e:?}")))?
30        .get();
31    Ok((value, &data[2..]))
32}
33
34/// Read 2-byte big-endian unsigned integer.
35#[inline]
36pub fn read_u16(data: &[u8]) -> Result<(u16, &[u8])> {
37    if data.len() < 2 {
38        return Err(Error::Protocol(format!(
39            "read_u16: buffer too short: {} < 2",
40            data.len()
41        )));
42    }
43    let value = U16BE::ref_from_bytes(&data[..2])
44        .map_err(|e| Error::Protocol(format!("read_u16: {e:?}")))?
45        .get();
46    Ok((value, &data[2..]))
47}
48
49/// Read 4-byte big-endian signed integer.
50#[inline]
51pub fn read_i32(data: &[u8]) -> Result<(i32, &[u8])> {
52    if data.len() < 4 {
53        return Err(Error::Protocol(format!(
54            "read_i32: buffer too short: {} < 4",
55            data.len()
56        )));
57    }
58    let value = I32BE::ref_from_bytes(&data[..4])
59        .map_err(|e| Error::Protocol(format!("read_i32: {e:?}")))?
60        .get();
61    Ok((value, &data[4..]))
62}
63
64/// Read 4-byte big-endian unsigned integer.
65#[inline]
66pub fn read_u32(data: &[u8]) -> Result<(u32, &[u8])> {
67    if data.len() < 4 {
68        return Err(Error::Protocol(format!(
69            "read_u32: buffer too short: {} < 4",
70            data.len()
71        )));
72    }
73    let value = U32BE::ref_from_bytes(&data[..4])
74        .map_err(|e| Error::Protocol(format!("read_u32: {e:?}")))?
75        .get();
76    Ok((value, &data[4..]))
77}
78
79/// Read fixed-length bytes.
80#[inline]
81pub fn read_bytes(data: &[u8], len: usize) -> Result<(&[u8], &[u8])> {
82    if data.len() < len {
83        return Err(Error::Protocol(format!(
84            "read_bytes: buffer too short: {} < {}",
85            data.len(),
86            len
87        )));
88    }
89    Ok((&data[..len], &data[len..]))
90}
91
92/// Read null-terminated string (PostgreSQL String type).
93/// Returns the string bytes (without the null terminator) and remaining data.
94#[inline]
95pub fn read_cstring(data: &[u8]) -> Result<(&[u8], &[u8])> {
96    match memchr::memchr(0, data) {
97        Some(pos) => Ok((&data[..pos], &data[pos + 1..])),
98        None => Err(Error::Protocol(
99            "read_cstring: no null terminator found".into(),
100        )),
101    }
102}
103
104/// Read null-terminated string as &str.
105#[inline]
106pub fn read_cstr(data: &[u8]) -> Result<(&str, &[u8])> {
107    let (bytes, rest) = read_cstring(data)?;
108    let s = simdutf8::compat::from_utf8(bytes)
109        .map_err(|e| Error::Protocol(format!("read_cstr: invalid UTF-8: {e}")))?;
110    Ok((s, rest))
111}
112
113/// Write 1-byte unsigned integer.
114#[inline]
115pub fn write_u8(out: &mut Vec<u8>, value: u8) {
116    out.push(value);
117}
118
119/// Write 2-byte big-endian signed integer.
120#[inline]
121pub fn write_i16(out: &mut Vec<u8>, value: i16) {
122    out.extend_from_slice(&value.to_be_bytes());
123}
124
125/// Write 2-byte big-endian unsigned integer.
126#[inline]
127pub fn write_u16(out: &mut Vec<u8>, value: u16) {
128    out.extend_from_slice(&value.to_be_bytes());
129}
130
131/// Write 4-byte big-endian signed integer.
132#[inline]
133pub fn write_i32(out: &mut Vec<u8>, value: i32) {
134    out.extend_from_slice(&value.to_be_bytes());
135}
136
137/// Write 4-byte big-endian unsigned integer.
138#[inline]
139pub fn write_u32(out: &mut Vec<u8>, value: u32) {
140    out.extend_from_slice(&value.to_be_bytes());
141}
142
143/// Write raw bytes.
144#[inline]
145pub fn write_bytes(out: &mut Vec<u8>, data: &[u8]) {
146    out.extend_from_slice(data);
147}
148
149/// Write null-terminated string (PostgreSQL String type).
150#[inline]
151pub fn write_cstring(out: &mut Vec<u8>, s: &[u8]) {
152    out.extend_from_slice(s);
153    out.push(0);
154}
155
156/// Write null-terminated string from &str.
157#[inline]
158pub fn write_cstr(out: &mut Vec<u8>, s: &str) {
159    write_cstring(out, s.as_bytes());
160}
161
162/// Message builder helper that handles the length field.
163///
164/// PostgreSQL message format:
165/// - Type byte (1 byte) - NOT included in length
166/// - Length (4 bytes) - includes itself
167/// - Payload (Length - 4 bytes)
168pub struct MessageBuilder<'a> {
169    buf: &'a mut Vec<u8>,
170    start: usize,
171}
172
173impl<'a> MessageBuilder<'a> {
174    /// Start building a message with a type byte.
175    pub fn new(buf: &'a mut Vec<u8>, type_byte: u8) -> Self {
176        buf.push(type_byte);
177        let start = buf.len();
178        buf.extend_from_slice(&[0, 0, 0, 0]); // Placeholder for length
179        Self { buf, start }
180    }
181
182    /// Start building a startup message (no type byte).
183    pub fn new_startup(buf: &'a mut Vec<u8>) -> Self {
184        let start = buf.len();
185        buf.extend_from_slice(&[0, 0, 0, 0]); // Placeholder for length
186        Self { buf, start }
187    }
188
189    /// Get mutable access to the underlying buffer.
190    pub fn buf(&mut self) -> &mut Vec<u8> {
191        self.buf
192    }
193
194    /// Write a u8.
195    pub fn write_u8(&mut self, value: u8) {
196        write_u8(self.buf, value);
197    }
198
199    /// Write an i16.
200    pub fn write_i16(&mut self, value: i16) {
201        write_i16(self.buf, value);
202    }
203
204    /// Write a u16.
205    pub fn write_u16(&mut self, value: u16) {
206        write_u16(self.buf, value);
207    }
208
209    /// Write an i32.
210    pub fn write_i32(&mut self, value: i32) {
211        write_i32(self.buf, value);
212    }
213
214    /// Write a u32.
215    pub fn write_u32(&mut self, value: u32) {
216        write_u32(self.buf, value);
217    }
218
219    /// Write raw bytes.
220    pub fn write_bytes(&mut self, data: &[u8]) {
221        write_bytes(self.buf, data);
222    }
223
224    /// Write null-terminated string.
225    pub fn write_cstr(&mut self, s: &str) {
226        write_cstr(self.buf, s);
227    }
228
229    /// Finish building the message and fill in the length field.
230    pub fn finish(self) {
231        let len = (self.buf.len() - self.start) as u32;
232        self.buf[self.start..self.start + 4].copy_from_slice(&len.to_be_bytes());
233    }
234}
zero_postgres/protocol/codec.rs

zero_postgres/protocol/
codec.rs
