fraiseql_wire/protocol/encode/
mod.rs

1//! Protocol message encoding
2
3use super::message::FrontendMessage;
4use bytes::{BufMut, BytesMut};
5use std::io;
6
7/// Write a 4-byte big-endian length back into a previously-reserved slot.
8///
9/// Every encoder writes a placeholder `put_i32(0)` early, captures the offset
10/// in `len_pos`, then continues writing the body. This helper backfills the
11/// real length using a bounds-checked `.get_mut()` so the file stays panic-free
12/// under `#![deny(clippy::indexing_slicing)]`.
13fn fill_length(buf: &mut BytesMut, len_pos: usize, len: usize) {
14    let bytes = (len as i32).to_be_bytes();
15    if let Some(slot) = buf.get_mut(len_pos..len_pos + 4) {
16        slot.copy_from_slice(&bytes);
17    }
18    // If the slot is missing the buffer was corrupted upstream; the encoder
19    // contract guarantees the caller reserved exactly 4 bytes at `len_pos`
20    // before appending the body, so the `if let` always succeeds in practice.
21}
22
23/// Encode a frontend message into bytes
24///
25/// # Errors
26///
27/// Returns `io::Error` if the message contains invalid UTF-8 or cannot be serialized.
28pub fn encode_message(msg: &FrontendMessage) -> io::Result<BytesMut> {
29    let mut buf = BytesMut::new();
30
31    match msg {
32        FrontendMessage::Startup { version, params } => {
33            encode_startup(&mut buf, *version, params)?;
34        }
35        FrontendMessage::Password(password) => {
36            encode_password(&mut buf, password)?;
37        }
38        FrontendMessage::Query(query) => {
39            encode_query(&mut buf, query)?;
40        }
41        FrontendMessage::Terminate => {
42            encode_terminate(&mut buf)?;
43        }
44        FrontendMessage::SaslInitialResponse { mechanism, data } => {
45            encode_sasl_initial_response(&mut buf, mechanism, data)?;
46        }
47        FrontendMessage::SaslResponse { data } => {
48            encode_sasl_response(&mut buf, data)?;
49        }
50    }
51
52    Ok(buf)
53}
54
55fn encode_startup(buf: &mut BytesMut, version: i32, params: &[(String, String)]) -> io::Result<()> {
56    // Startup messages don't have a type byte
57    // Reserve space for length (will be filled at end)
58    let len_pos = buf.len();
59    buf.put_i32(0);
60
61    // Protocol version
62    buf.put_i32(version);
63
64    // Parameters (key-value pairs, null-terminated)
65    for (key, value) in params {
66        buf.put(key.as_bytes());
67        buf.put_u8(0);
68        buf.put(value.as_bytes());
69        buf.put_u8(0);
70    }
71
72    // Final null terminator
73    buf.put_u8(0);
74
75    // Fill in length
76    let len = buf.len() - len_pos;
77    fill_length(buf, len_pos, len);
78
79    Ok(())
80}
81
82fn encode_password(buf: &mut BytesMut, password: &str) -> io::Result<()> {
83    buf.put_u8(b'p');
84    let len_pos = buf.len();
85    buf.put_i32(0);
86
87    buf.put(password.as_bytes());
88    buf.put_u8(0);
89
90    let len = buf.len() - len_pos;
91    fill_length(buf, len_pos, len);
92
93    Ok(())
94}
95
96fn encode_query(buf: &mut BytesMut, query: &str) -> io::Result<()> {
97    buf.put_u8(b'Q');
98    let len_pos = buf.len();
99    buf.put_i32(0);
100
101    buf.put(query.as_bytes());
102    buf.put_u8(0);
103
104    let len = buf.len() - len_pos;
105    fill_length(buf, len_pos, len);
106
107    Ok(())
108}
109
110fn encode_terminate(buf: &mut BytesMut) -> io::Result<()> {
111    buf.put_u8(b'X');
112    buf.put_i32(4); // Length includes itself
113    Ok(())
114}
115
116fn encode_sasl_initial_response(
117    buf: &mut BytesMut,
118    mechanism: &str,
119    data: &[u8],
120) -> io::Result<()> {
121    buf.put_u8(b'p');
122    let len_pos = buf.len();
123    buf.put_i32(0);
124
125    // Mechanism name (null-terminated)
126    buf.put(mechanism.as_bytes());
127    buf.put_u8(0);
128
129    // SASL data (as length-prefixed bytes)
130    buf.put_i32(data.len() as i32);
131    buf.put_slice(data);
132
133    let len = buf.len() - len_pos;
134    fill_length(buf, len_pos, len);
135
136    Ok(())
137}
138
139fn encode_sasl_response(buf: &mut BytesMut, data: &[u8]) -> io::Result<()> {
140    buf.put_u8(b'p');
141    let len_pos = buf.len();
142    buf.put_i32(0);
143
144    // SASL data (length-prefixed)
145    buf.put_slice(data);
146
147    let len = buf.len() - len_pos;
148    fill_length(buf, len_pos, len);
149
150    Ok(())
151}
152
153#[cfg(test)]
154mod tests;
fraiseql_wire/protocol/encode/mod.rs

fraiseql_wire/protocol/encode/
mod.rs
