postgresql_protocol/
frontend.rs

1//! Frontend (client → server) message encoding.
2//!
3//! Messages are appended to a caller-owned [`BytesMut`] so that a full request
4//! (`Parse` + `Describe` + `Bind` + `Execute` + `Sync`) can be assembled in one
5//! buffer and written with a single syscall.
6//!
7//! Counts and lengths taken from the caller are range-checked before being
8//! written: a silently truncated length prefix makes the server reinterpret
9//! the rest of the buffer as new protocol messages.
10
11use bytes::{BufMut, BytesMut};
12
13use crate::{error::Error, oid::Format};
14
15/// Maximum number of values in one `Bind` or `COPY` row.
16///
17/// The protocol encodes these counts as a 16-bit field that PostgreSQL reads
18/// as unsigned, so the limit is 65535 and not 32767.
19pub const MAX_VALUE_COUNT: usize = u16::MAX as usize;
20
21/// Appends a raw C string terminated by a NUL byte.
22pub fn put_cstr(buf: &mut BytesMut, s: &str) {
23    buf.put_slice(s.as_bytes());
24    buf.put_u8(0);
25}
26
27/// Starts a tagged message, writing the tag and a placeholder length.
28///
29/// Returns the offset of the length field, to be passed to [`finish`].
30pub fn begin(buf: &mut BytesMut, tag: u8) -> usize {
31    buf.put_u8(tag);
32    let pos = buf.len();
33    buf.put_i32(0);
34    pos
35}
36
37/// Patches the length field written by [`begin`].
38///
39/// # Errors
40///
41/// Returns [`Error::Encode`] when the message body is larger than 2 GiB and
42/// cannot be expressed in the protocol's 32-bit length prefix. Silently
43/// truncating would make the server reinterpret the rest of the buffer as new
44/// messages.
45pub fn finish(buf: &mut BytesMut, len_pos: usize) -> Result<(), Error> {
46    let len = length32(buf.len() - len_pos)?;
47    buf[len_pos..len_pos + 4].copy_from_slice(&len.to_be_bytes());
48    Ok(())
49}
50
51/// Encodes the startup message (no tag).
52///
53/// # Errors
54///
55/// Returns [`Error::Encode`] if the message is larger than 2 GiB.
56pub fn startup(buf: &mut BytesMut, params: &[(&str, &str)]) -> Result<(), Error> {
57    // Startup parameters are NUL-delimited on the wire: a NUL inside a key or
58    // value would split it and inject extra parameters.
59    for (key, value) in params {
60        if key.contains('\0') || value.contains('\0') {
61            return Err(Error::Encode("startup parameters must not contain a NUL byte".into()));
62        }
63    }
64    let len_pos = buf.len();
65    buf.put_i32(0);
66    buf.put_i32(196608); // protocol 3.0
67    for (k, v) in params {
68        put_cstr(buf, k);
69        put_cstr(buf, v);
70    }
71    buf.put_u8(0);
72    finish(buf, len_pos)
73}
74
75/// Encodes an SSLRequest message.
76pub fn ssl_request(buf: &mut BytesMut) {
77    buf.put_i32(8);
78    buf.put_i32(80877103);
79}
80
81/// Encodes a `CancelRequest` message (no tag; sent on a fresh connection).
82pub fn cancel_request(buf: &mut BytesMut, process_id: i32, secret_key: i32) {
83    buf.put_i32(16);
84    buf.put_i32(80877102);
85    buf.put_i32(process_id);
86    buf.put_i32(secret_key);
87}
88
89/// Encodes a cleartext password message (`PasswordMessage`).
90///
91/// # Errors
92///
93/// Returns [`Error::Encode`] if the message is larger than 2 GiB.
94pub fn password(buf: &mut BytesMut, password: &str) -> Result<(), Error> {
95    let len_pos = begin(buf, b'p');
96    put_cstr(buf, password);
97    finish(buf, len_pos)
98}
99
100/// Encodes a `SASLInitialResponse`.
101///
102/// # Errors
103///
104/// Returns [`Error::Encode`] if the message is larger than 2 GiB.
105pub fn sasl_initial_response(buf: &mut BytesMut, mechanism: &str, data: &[u8]) -> Result<(), Error> {
106    let len_pos = begin(buf, b'p');
107    put_cstr(buf, mechanism);
108    buf.put_i32(length32(data.len())?);
109    buf.put_slice(data);
110    finish(buf, len_pos)
111}
112
113/// Encodes a `SASLResponse`.
114///
115/// # Errors
116///
117/// Returns [`Error::Encode`] if the message is larger than 2 GiB.
118pub fn sasl_response(buf: &mut BytesMut, data: &[u8]) -> Result<(), Error> {
119    let len_pos = begin(buf, b'p');
120    buf.put_slice(data);
121    finish(buf, len_pos)
122}
123
124/// Encodes a simple `Query`.
125///
126/// # Errors
127///
128/// Returns [`Error::Encode`] if the message is larger than 2 GiB.
129pub fn query(buf: &mut BytesMut, sql: &str) -> Result<(), Error> {
130    let len_pos = begin(buf, b'Q');
131    put_cstr(buf, sql);
132    finish(buf, len_pos)
133}
134
135/// Encodes a `Parse`.
136///
137/// # Errors
138///
139/// Returns [`Error::Encode`] when more than [`MAX_VALUE_COUNT`] parameter
140/// types are supplied, or if the message is larger than 2 GiB.
141pub fn parse(buf: &mut BytesMut, name: &str, sql: &str, param_oids: &[u32]) -> Result<(), Error> {
142    let count = count16(param_oids.len())?;
143    let len_pos = begin(buf, b'P');
144    put_cstr(buf, name);
145    put_cstr(buf, sql);
146    buf.put_u16(count);
147    for oid in param_oids {
148        buf.put_u32(*oid);
149    }
150    finish(buf, len_pos)
151}
152
153/// Starts a `Bind` message; call [`write_param`] for each value, then
154/// [`bind_finish`].
155///
156/// All parameters use the binary format, so one format code is written per
157/// parameter. Returns the offset of the length field, to be passed to
158/// [`bind_finish`].
159///
160/// # Errors
161///
162/// Returns [`Error::Encode`] when more than [`MAX_VALUE_COUNT`] parameters are
163/// supplied.
164pub fn bind_start(buf: &mut BytesMut, portal: &str, statement: &str, param_count: usize) -> Result<usize, Error> {
165    let count = count16(param_count)?;
166    let len_pos = begin(buf, b'B');
167    put_cstr(buf, portal);
168    put_cstr(buf, statement);
169    buf.put_u16(count);
170    for _ in 0..param_count {
171        buf.put_i16(Format::Binary.code());
172    }
173    buf.put_u16(count);
174    Ok(len_pos)
175}
176
177/// Writes one `Bind` parameter; `None` is encoded as SQL NULL.
178///
179/// # Errors
180///
181/// Returns [`Error::Encode`] when the value is larger than 2 GiB and cannot be
182/// expressed in the protocol's 32-bit length prefix.
183pub fn write_param(buf: &mut BytesMut, value: Option<&[u8]>) -> Result<(), Error> {
184    match value {
185        None => buf.put_i32(-1),
186        Some(v) => {
187            buf.put_i32(length32(v.len())?);
188            buf.put_slice(v);
189        }
190    }
191    Ok(())
192}
193
194/// Finishes a `Bind` message.
195///
196/// One result-column format code is written, which the server applies to every
197/// result column.
198///
199/// # Errors
200///
201/// Returns [`Error::Encode`] if the message is larger than 2 GiB.
202pub fn bind_finish(buf: &mut BytesMut, result_format: Format, len_pos: usize) -> Result<(), Error> {
203    buf.put_i16(1);
204    buf.put_i16(result_format.code());
205    finish(buf, len_pos)
206}
207
208/// Encodes a `Describe` (`kind` is `b'S'` for a statement or `b'P'` for a portal).
209///
210/// # Errors
211///
212/// Returns [`Error::Encode`] if the message is larger than 2 GiB.
213pub fn describe(buf: &mut BytesMut, kind: u8, name: &str) -> Result<(), Error> {
214    let len_pos = begin(buf, b'D');
215    buf.put_u8(kind);
216    put_cstr(buf, name);
217    finish(buf, len_pos)
218}
219
220/// Encodes an `Execute`.
221///
222/// # Errors
223///
224/// Returns [`Error::Encode`] if the message is larger than 2 GiB.
225pub fn execute(buf: &mut BytesMut, portal: &str, max_rows: i32) -> Result<(), Error> {
226    let len_pos = begin(buf, b'E');
227    put_cstr(buf, portal);
228    buf.put_i32(max_rows);
229    finish(buf, len_pos)
230}
231
232/// Encodes a `Sync`.
233///
234/// # Errors
235///
236/// Returns [`Error::Encode`] if the message is larger than 2 GiB.
237pub fn sync(buf: &mut BytesMut) -> Result<(), Error> {
238    let len_pos = begin(buf, b'S');
239    finish(buf, len_pos)
240}
241
242/// Encodes a `Close` (`kind` is `b'S'` for a statement or `b'P'` for a portal).
243///
244/// # Errors
245///
246/// Returns [`Error::Encode`] if the message is larger than 2 GiB.
247pub fn close(buf: &mut BytesMut, kind: u8, name: &str) -> Result<(), Error> {
248    let len_pos = begin(buf, b'C');
249    buf.put_u8(kind);
250    put_cstr(buf, name);
251    finish(buf, len_pos)
252}
253
254/// Encodes a `Flush`.
255///
256/// # Errors
257///
258/// Returns [`Error::Encode`] if the message is larger than 2 GiB.
259pub fn flush(buf: &mut BytesMut) -> Result<(), Error> {
260    let len_pos = begin(buf, b'H');
261    finish(buf, len_pos)
262}
263
264/// Encodes a `Terminate`.
265///
266/// # Errors
267///
268/// Returns [`Error::Encode`] if the message is larger than 2 GiB.
269pub fn terminate(buf: &mut BytesMut) -> Result<(), Error> {
270    let len_pos = begin(buf, b'X');
271    finish(buf, len_pos)
272}
273
274/// Encodes a `CopyData`.
275///
276/// # Errors
277///
278/// Returns [`Error::Encode`] when the payload is larger than 2 GiB.
279pub fn copy_data(buf: &mut BytesMut, data: &[u8]) -> Result<(), Error> {
280    let len_pos = begin(buf, b'd');
281    buf.put_slice(data);
282    finish(buf, len_pos)
283}
284
285/// Encodes a `CopyDone`.
286///
287/// # Errors
288///
289/// Returns [`Error::Encode`] if the message is larger than 2 GiB.
290pub fn copy_done(buf: &mut BytesMut) -> Result<(), Error> {
291    let len_pos = begin(buf, b'c');
292    finish(buf, len_pos)
293}
294
295/// Encodes a `CopyFail`.
296///
297/// # Errors
298///
299/// Returns [`Error::Encode`] if the message is larger than 2 GiB.
300pub fn copy_fail(buf: &mut BytesMut, message: &str) -> Result<(), Error> {
301    let len_pos = begin(buf, b'f');
302    put_cstr(buf, message);
303    finish(buf, len_pos)
304}
305
306fn count16(count: usize) -> Result<u16, Error> {
307    u16::try_from(count).map_err(|_| Error::Encode(format!("too many values: {count} (max {MAX_VALUE_COUNT})")))
308}
309
310fn length32(len: usize) -> Result<i32, Error> {
311    i32::try_from(len).map_err(|_| Error::Encode(format!("value is too large: {len} bytes")))
312}
313
314#[cfg(test)]
315mod tests {
316    use super::*;
317
318    #[test]
319    fn startup_layout() {
320        let mut buf = BytesMut::new();
321        startup(&mut buf, &[("user", "bob"), ("database", "db")]).unwrap();
322        let len = i32::from_be_bytes([buf[0], buf[1], buf[2], buf[3]]);
323        assert_eq!(usize::try_from(len).unwrap(), buf.len());
324        assert_eq!(&buf[4..8], &196608i32.to_be_bytes());
325        // body ends with a NUL, keys are NUL-terminated
326        assert_eq!(*buf.last().unwrap(), 0);
327    }
328
329    #[test]
330    fn startup_rejects_nul_bytes() {
331        // A NUL would split the parameter and inject extra startup values.
332        let mut buf = BytesMut::new();
333        assert!(startup(&mut buf, &[("application_name", "a\0user\0admin")]).is_err());
334        assert!(startup(&mut buf, &[("us\0er", "x")]).is_err());
335        assert!(buf.is_empty());
336    }
337
338    #[test]
339    fn ssl_request_layout() {
340        let mut buf = BytesMut::new();
341        ssl_request(&mut buf);
342        assert_eq!(buf.len(), 8);
343        assert_eq!(&buf[0..4], &8i32.to_be_bytes());
344        assert_eq!(&buf[4..8], &80877103i32.to_be_bytes());
345    }
346
347    #[test]
348    fn bind_length_is_patched() {
349        let mut buf = BytesMut::new();
350        let pos = bind_start(&mut buf, "", "s1", 1).unwrap();
351        write_param(&mut buf, Some(&[1, 2, 3, 4])).unwrap();
352        bind_finish(&mut buf, Format::Binary, pos).unwrap();
353        let len = i32::from_be_bytes([buf[1], buf[2], buf[3], buf[4]]);
354        assert_eq!(usize::try_from(len).unwrap(), buf.len() - 1);
355    }
356
357    #[test]
358    fn parse_encodes_oids() {
359        let mut buf = BytesMut::new();
360        parse(&mut buf, "", "SELECT $1", &[23]).unwrap();
361        let len = i32::from_be_bytes([buf[1], buf[2], buf[3], buf[4]]);
362        assert_eq!(usize::try_from(len).unwrap(), buf.len() - 1);
363    }
364
365    #[test]
366    fn rejects_impossible_counts_and_lengths() {
367        let mut buf = BytesMut::new();
368        assert!(bind_start(&mut buf, "", "s", MAX_VALUE_COUNT + 1).is_err());
369        assert!(parse(&mut buf, "", "SELECT $1", &vec![23u32; MAX_VALUE_COUNT + 1]).is_err());
370        assert!(write_param(&mut buf, Some(&[0u8; 0])).is_ok());
371        assert!(copy_data(&mut buf, &[1, 2, 3]).is_ok());
372    }
373}
postgresql_protocol/frontend.rs

postgresql_protocol/
frontend.rs
