fraiseql_wire/protocol/decode/
mod.rs

1//! Protocol message decoding
2
3use super::constants::{auth, tags};
4use super::message::{AuthenticationMessage, BackendMessage, ErrorFields, FieldDescription};
5use bytes::{Bytes, BytesMut};
6use std::io;
7
8/// Bounds-checked read cursor over a byte slice.
9///
10/// All accessors return `io::Result` so this whole file can stay panic-free
11/// under `#![deny(clippy::indexing_slicing)]`. Each method advances `offset`
12/// only on success.
13struct Cursor<'a> {
14    data: &'a [u8],
15    offset: usize,
16}
17
18impl<'a> Cursor<'a> {
19    const fn new(data: &'a [u8]) -> Self {
20        Self { data, offset: 0 }
21    }
22
23    fn remaining(&self) -> &'a [u8] {
24        // `self.offset` is monotonically advanced only by successful reads,
25        // each of which ensures the offset stays `<= self.data.len()`.
26        self.data.get(self.offset..).unwrap_or(&[])
27    }
28
29    const fn is_empty(&self) -> bool {
30        self.offset >= self.data.len()
31    }
32
33    fn read_u8(&mut self) -> io::Result<u8> {
34        let byte = *self
35            .data
36            .get(self.offset)
37            .ok_or_else(|| io::Error::new(io::ErrorKind::UnexpectedEof, "byte"))?;
38        self.offset += 1;
39        Ok(byte)
40    }
41
42    fn read_i16_be(&mut self) -> io::Result<i16> {
43        let bytes: [u8; 2] = self
44            .data
45            .get(self.offset..self.offset + 2)
46            .ok_or_else(|| io::Error::new(io::ErrorKind::UnexpectedEof, "i16"))?
47            .try_into()
48            // Reason: provably-safe — `.get(offset..offset+2)` returned a
49            // 2-byte slice, and `<[u8; 2]>::try_from(&[u8])` cannot fail on
50            // a slice of the exact length.
51            .expect("slice of length 2 always converts to [u8; 2]");
52        self.offset += 2;
53        Ok(i16::from_be_bytes(bytes))
54    }
55
56    fn read_i32_be(&mut self) -> io::Result<i32> {
57        let bytes: [u8; 4] = self
58            .data
59            .get(self.offset..self.offset + 4)
60            .ok_or_else(|| io::Error::new(io::ErrorKind::UnexpectedEof, "i32"))?
61            .try_into()
62            // Reason: provably-safe — slice length 4 always converts to [u8; 4].
63            .expect("slice of length 4 always converts to [u8; 4]");
64        self.offset += 4;
65        Ok(i32::from_be_bytes(bytes))
66    }
67
68    fn read_slice(&mut self, n: usize) -> io::Result<&'a [u8]> {
69        let slice = self
70            .data
71            .get(self.offset..self.offset + n)
72            .ok_or_else(|| io::Error::new(io::ErrorKind::UnexpectedEof, "slice"))?;
73        self.offset += n;
74        Ok(slice)
75    }
76
77    /// Read until the next `0x00` byte (exclusive), advancing past the null.
78    /// Returns the bytes before the null terminator.
79    fn read_until_null(&mut self) -> io::Result<&'a [u8]> {
80        let tail = self.remaining();
81        let end = tail.iter().position(|&b| b == 0).ok_or_else(|| {
82            io::Error::new(
83                io::ErrorKind::InvalidData,
84                "missing null terminator in string",
85            )
86        })?;
87        let bytes = tail.get(..end).unwrap_or(&[]);
88        // Advance past `end` bytes plus the null terminator.
89        self.offset += end + 1;
90        Ok(bytes)
91    }
92
93    /// Find the next `0x00` byte in the remaining slice without advancing.
94    fn position_of_null(&self) -> Option<usize> {
95        self.remaining().iter().position(|&b| b == 0)
96    }
97}
98
99/// Maximum number of fields accepted in a single DataRow or RowDescription message.
100///
101/// PostgreSQL's protocol allows up to 1600 columns per table (hard limit enforced by
102/// the server), so 2048 is a generous cap that prevents an attacker-supplied message
103/// from triggering a huge `Vec::with_capacity` before any bounds are checked.
104pub(crate) const MAX_FIELD_COUNT: usize = 2048;
105
106/// Maximum byte length of a single error/notice field string (severity, message, etc.).
107///
108/// A 64 KiB cap is generous for any human-readable error message. Without this limit a
109/// malicious server can send a single oversized field and drive unbounded allocation
110/// in `String::from_utf8_lossy` before the string is ever stored.
111pub(crate) const MAX_ERROR_FIELD_BYTES: usize = 64 * 1024; // 64 KiB
112
113/// Maximum number of SASL mechanism names accepted in an Authentication message.
114///
115/// Real providers offer one or two mechanisms (e.g. SCRAM-SHA-256).  Capping at 32
116/// prevents a rogue server from flooding the `Vec<String>` until memory is exhausted.
117pub(crate) const MAX_SASL_MECHANISMS: usize = 32;
118
119/// Maximum byte length of a ParameterStatus name (e.g. `"server_version"`).
120///
121/// PostgreSQL parameter names are short identifiers; 256 bytes is more than enough.
122pub(crate) const MAX_PARAMETER_NAME_BYTES: usize = 256;
123
124/// Maximum byte length of a ParameterStatus value.
125///
126/// 64 KiB covers realistic values (long `TimeZone` strings, etc.) while preventing
127/// a malicious server from inflating memory with an oversized value string.
128pub(crate) const MAX_PARAMETER_VALUE_BYTES: usize = 64 * 1024; // 64 KiB
129
130/// Maximum total length (bytes) of a single protocol message body the decoder
131/// accepts.
132///
133/// PostgreSQL encodes message length as an `i32`, so without a cap a malicious or
134/// compromised peer (or a non-TLS MITM) can declare a length up to ~2 GiB and
135/// force the connection read buffer to grow that large before any per-field cap
136/// runs — a memory-exhaustion `DoS` (audit M-wire-msg-cap). DataRow column values
137/// in particular carry no per-column bound, only the field-count cap. 256 MiB is
138/// far above any legitimate FraiseQL message (other per-field caps are 64 KiB)
139/// while making the unbounded-allocation attack impossible.
140pub(crate) const MAX_MESSAGE_LEN: usize = 256 * 1024 * 1024; // 256 MiB
141
142/// Decode a backend message from `BytesMut` without cloning
143///
144/// This version decodes in-place from a mutable `BytesMut` buffer and returns
145/// the number of bytes consumed. The caller must advance the buffer after calling this.
146///
147/// # Errors
148///
149/// Returns `io::Error` with `UnexpectedEof` if the buffer does not contain a complete
150/// message. Returns `io::Error` with `InvalidData` if the message tag or length is invalid.
151///
152/// # Returns
153/// `Ok((msg, consumed))` - Message and number of bytes consumed
154/// `Err(e)` - IO error if message is incomplete or invalid
155///
156/// # Performance
157/// This version avoids the expensive `buf.clone().freeze()` call by working directly
158/// with references, reducing allocations and copies in the hot path.
159pub fn decode_message(data: &mut BytesMut) -> io::Result<(BackendMessage, usize)> {
160    if data.len() < 5 {
161        return Err(io::Error::new(
162            io::ErrorKind::UnexpectedEof,
163            "incomplete message header",
164        ));
165    }
166
167    let mut header = Cursor::new(data);
168    let tag = header.read_u8()?;
169    let len_i32 = header.read_i32_be()?;
170
171    // PostgreSQL message length includes the 4 length bytes but not the tag byte.
172    // Minimum valid length is 4 (just the length field itself).
173    if len_i32 < 4 {
174        return Err(io::Error::new(
175            io::ErrorKind::InvalidData,
176            "message length too small",
177        ));
178    }
179
180    let len = len_i32 as usize;
181
182    // Upper-bound the declared length BEFORE the "incomplete body" check below,
183    // so an oversized declaration is a fatal `InvalidData` error rather than an
184    // `UnexpectedEof` that the read loop would treat as "need more bytes" and
185    // keep buffering toward (audit M-wire-msg-cap).
186    if len > MAX_MESSAGE_LEN {
187        return Err(io::Error::new(
188            io::ErrorKind::InvalidData,
189            format!("message length {len} exceeds maximum {MAX_MESSAGE_LEN}"),
190        ));
191    }
192
193    if data.len() < len + 1 {
194        return Err(io::Error::new(
195            io::ErrorKind::UnexpectedEof,
196            "incomplete message body",
197        ));
198    }
199
200    // Create a temporary slice starting after the tag and length
201    let msg_start = 5;
202    let msg_end = len + 1;
203    let msg_data = data
204        .get(msg_start..msg_end)
205        .ok_or_else(|| io::Error::new(io::ErrorKind::UnexpectedEof, "message body slice"))?;
206
207    let msg = match tag {
208        tags::AUTHENTICATION => decode_authentication(msg_data)?,
209        tags::BACKEND_KEY_DATA => decode_backend_key_data(msg_data)?,
210        tags::COMMAND_COMPLETE => decode_command_complete(msg_data)?,
211        tags::DATA_ROW => decode_data_row(msg_data)?,
212        tags::ERROR_RESPONSE => decode_error_response(msg_data)?,
213        tags::NOTICE_RESPONSE => decode_notice_response(msg_data)?,
214        tags::PARAMETER_STATUS => decode_parameter_status(msg_data)?,
215        tags::READY_FOR_QUERY => decode_ready_for_query(msg_data)?,
216        tags::ROW_DESCRIPTION => decode_row_description(msg_data)?,
217        _ => {
218            return Err(io::Error::new(
219                io::ErrorKind::InvalidData,
220                format!("unknown message tag: {}", tag),
221            ))
222        }
223    };
224
225    Ok((msg, len + 1))
226}
227
228fn decode_authentication(data: &[u8]) -> io::Result<BackendMessage> {
229    let mut cur = Cursor::new(data);
230    let auth_type = cur
231        .read_i32_be()
232        .map_err(|_| io::Error::new(io::ErrorKind::UnexpectedEof, "auth type"))?;
233
234    let auth_msg = match auth_type {
235        auth::OK => AuthenticationMessage::Ok,
236        auth::CLEARTEXT_PASSWORD => AuthenticationMessage::CleartextPassword,
237        auth::MD5_PASSWORD => {
238            let salt_slice = cur
239                .read_slice(4)
240                .map_err(|_| io::Error::new(io::ErrorKind::UnexpectedEof, "salt data"))?;
241            let salt: [u8; 4] = salt_slice
242                .try_into()
243                // Reason: provably-safe — `read_slice(4)` returns a 4-byte slice.
244                .expect("slice of length 4 always converts to [u8; 4]");
245            AuthenticationMessage::Md5Password { salt }
246        }
247        auth::SASL => {
248            // SASL: read mechanism list (null-terminated strings)
249            let mut mechanisms = Vec::new();
250            loop {
251                if cur.is_empty() {
252                    break;
253                }
254                let Some(end) = cur.position_of_null() else {
255                    break;
256                };
257                let mech_bytes = cur.read_slice(end).unwrap_or(&[]);
258                let mechanism = String::from_utf8_lossy(mech_bytes).to_string();
259                // Skip the null terminator we just located.
260                let _ = cur.read_u8();
261                if mechanism.is_empty() {
262                    break;
263                }
264                if mechanisms.len() >= MAX_SASL_MECHANISMS {
265                    break;
266                }
267                mechanisms.push(mechanism);
268            }
269            AuthenticationMessage::Sasl { mechanisms }
270        }
271        auth::SASL_CONTINUE => {
272            // SASL continue: read remaining data as bytes
273            let data_vec = cur.remaining().to_vec();
274            AuthenticationMessage::SaslContinue { data: data_vec }
275        }
276        auth::SASL_FINAL => {
277            // SASL final: read remaining data as bytes
278            let data_vec = cur.remaining().to_vec();
279            AuthenticationMessage::SaslFinal { data: data_vec }
280        }
281        _ => {
282            return Err(io::Error::new(
283                io::ErrorKind::Unsupported,
284                format!("unsupported auth type: {}", auth_type),
285            ))
286        }
287    };
288
289    Ok(BackendMessage::Authentication(auth_msg))
290}
291
292fn decode_backend_key_data(data: &[u8]) -> io::Result<BackendMessage> {
293    let mut cur = Cursor::new(data);
294    let process_id = cur
295        .read_i32_be()
296        .map_err(|_| io::Error::new(io::ErrorKind::UnexpectedEof, "backend key data"))?;
297    let secret_key = cur
298        .read_i32_be()
299        .map_err(|_| io::Error::new(io::ErrorKind::UnexpectedEof, "backend key data"))?;
300    Ok(BackendMessage::BackendKeyData {
301        process_id,
302        secret_key,
303    })
304}
305
306fn decode_command_complete(data: &[u8]) -> io::Result<BackendMessage> {
307    let mut cur = Cursor::new(data);
308    let tag_bytes = cur.read_until_null()?;
309    let tag = String::from_utf8_lossy(tag_bytes).to_string();
310    Ok(BackendMessage::CommandComplete(tag))
311}
312
313fn decode_data_row(data: &[u8]) -> io::Result<BackendMessage> {
314    let mut cur = Cursor::new(data);
315    let field_count_i16 = cur
316        .read_i16_be()
317        .map_err(|_| io::Error::new(io::ErrorKind::UnexpectedEof, "field count"))?;
318    if field_count_i16 < 0 {
319        return Err(io::Error::new(
320            io::ErrorKind::InvalidData,
321            "negative field count",
322        ));
323    }
324    let field_count = field_count_i16 as usize;
325    if field_count > MAX_FIELD_COUNT {
326        return Err(io::Error::new(
327            io::ErrorKind::InvalidData,
328            format!("DataRow field count {field_count} exceeds maximum {MAX_FIELD_COUNT}"),
329        ));
330    }
331    let mut fields = Vec::with_capacity(field_count);
332
333    for _ in 0..field_count {
334        let field_len = cur
335            .read_i32_be()
336            .map_err(|_| io::Error::new(io::ErrorKind::UnexpectedEof, "field length"))?;
337
338        let field = if field_len == -1 {
339            None
340        } else if field_len < 0 {
341            return Err(io::Error::new(
342                io::ErrorKind::InvalidData,
343                "negative field length",
344            ));
345        } else {
346            let len = field_len as usize;
347            let field_slice = cur
348                .read_slice(len)
349                .map_err(|_| io::Error::new(io::ErrorKind::UnexpectedEof, "field data"))?;
350            Some(Bytes::copy_from_slice(field_slice))
351        };
352        fields.push(field);
353    }
354
355    Ok(BackendMessage::DataRow(fields))
356}
357
358fn decode_error_response(data: &[u8]) -> io::Result<BackendMessage> {
359    let fields = decode_error_fields(data)?;
360    Ok(BackendMessage::ErrorResponse(fields))
361}
362
363fn decode_notice_response(data: &[u8]) -> io::Result<BackendMessage> {
364    let fields = decode_error_fields(data)?;
365    Ok(BackendMessage::NoticeResponse(fields))
366}
367
368fn decode_error_fields(data: &[u8]) -> io::Result<ErrorFields> {
369    let mut fields = ErrorFields::default();
370    let mut cur = Cursor::new(data);
371
372    loop {
373        if cur.is_empty() {
374            break;
375        }
376        let field_type = cur.read_u8()?;
377        if field_type == 0 {
378            break;
379        }
380
381        let end = cur.position_of_null().ok_or_else(|| {
382            io::Error::new(
383                io::ErrorKind::InvalidData,
384                "missing null terminator in error field",
385            )
386        })?;
387        if end > MAX_ERROR_FIELD_BYTES {
388            return Err(io::Error::new(
389                io::ErrorKind::InvalidData,
390                format!("Error field too large ({end} bytes, max {MAX_ERROR_FIELD_BYTES})"),
391            ));
392        }
393        let value_bytes = cur.read_slice(end).unwrap_or(&[]);
394        let value = String::from_utf8_lossy(value_bytes).to_string();
395        // Skip the null terminator.
396        let _ = cur.read_u8();
397
398        match field_type {
399            b'S' => fields.severity = Some(value),
400            b'C' => fields.code = Some(value),
401            b'M' => fields.message = Some(value),
402            b'D' => fields.detail = Some(value),
403            b'H' => fields.hint = Some(value),
404            b'P' => fields.position = Some(value),
405            _ => {} // Ignore unknown fields
406        }
407    }
408
409    Ok(fields)
410}
411
412fn decode_parameter_status(data: &[u8]) -> io::Result<BackendMessage> {
413    let mut cur = Cursor::new(data);
414
415    let name_end = cur.position_of_null().ok_or_else(|| {
416        io::Error::new(
417            io::ErrorKind::InvalidData,
418            "missing null terminator in parameter name",
419        )
420    })?;
421    if name_end > MAX_PARAMETER_NAME_BYTES {
422        return Err(io::Error::new(
423            io::ErrorKind::InvalidData,
424            format!("Parameter name too long ({name_end} bytes, max {MAX_PARAMETER_NAME_BYTES})"),
425        ));
426    }
427    let name_bytes = cur.read_slice(name_end).unwrap_or(&[]);
428    let name = String::from_utf8_lossy(name_bytes).to_string();
429    // Skip null terminator.
430    let _ = cur.read_u8();
431
432    if cur.is_empty() {
433        return Err(io::Error::new(
434            io::ErrorKind::UnexpectedEof,
435            "parameter value",
436        ));
437    }
438    let value_end = cur.position_of_null().ok_or_else(|| {
439        io::Error::new(
440            io::ErrorKind::InvalidData,
441            "missing null terminator in parameter value",
442        )
443    })?;
444    if value_end > MAX_PARAMETER_VALUE_BYTES {
445        return Err(io::Error::new(
446            io::ErrorKind::InvalidData,
447            format!(
448                "Parameter value too long ({value_end} bytes, max {MAX_PARAMETER_VALUE_BYTES})"
449            ),
450        ));
451    }
452    let value_bytes = cur.read_slice(value_end).unwrap_or(&[]);
453    let value = String::from_utf8_lossy(value_bytes).to_string();
454
455    Ok(BackendMessage::ParameterStatus { name, value })
456}
457
458fn decode_ready_for_query(data: &[u8]) -> io::Result<BackendMessage> {
459    let status = *data
460        .first()
461        .ok_or_else(|| io::Error::new(io::ErrorKind::UnexpectedEof, "status byte"))?;
462    Ok(BackendMessage::ReadyForQuery { status })
463}
464
465fn decode_row_description(data: &[u8]) -> io::Result<BackendMessage> {
466    let mut cur = Cursor::new(data);
467    let field_count_i16 = cur
468        .read_i16_be()
469        .map_err(|_| io::Error::new(io::ErrorKind::UnexpectedEof, "field count"))?;
470    if field_count_i16 < 0 {
471        return Err(io::Error::new(
472            io::ErrorKind::InvalidData,
473            "negative field count",
474        ));
475    }
476    let field_count = field_count_i16 as usize;
477    if field_count > MAX_FIELD_COUNT {
478        return Err(io::Error::new(
479            io::ErrorKind::InvalidData,
480            format!("RowDescription field count {field_count} exceeds maximum {MAX_FIELD_COUNT}"),
481        ));
482    }
483    let mut fields = Vec::with_capacity(field_count);
484
485    for _ in 0..field_count {
486        // Read name (null-terminated string)
487        let name_end = cur.position_of_null().ok_or_else(|| {
488            io::Error::new(
489                io::ErrorKind::InvalidData,
490                "missing null terminator in field name",
491            )
492        })?;
493        let name_bytes = cur.read_slice(name_end).unwrap_or(&[]);
494        let name = String::from_utf8_lossy(name_bytes).to_string();
495        // Skip null terminator.
496        let _ = cur.read_u8();
497
498        // Read field descriptor (18 bytes: 4+2+4+2+4+2)
499        let table_oid = cur
500            .read_i32_be()
501            .map_err(|_| io::Error::new(io::ErrorKind::UnexpectedEof, "field descriptor"))?;
502        let column_attr = cur
503            .read_i16_be()
504            .map_err(|_| io::Error::new(io::ErrorKind::UnexpectedEof, "field descriptor"))?;
505        let type_oid = cur
506            .read_i32_be()
507            .map_err(|_| io::Error::new(io::ErrorKind::UnexpectedEof, "field descriptor"))?
508            as u32;
509        let type_size = cur
510            .read_i16_be()
511            .map_err(|_| io::Error::new(io::ErrorKind::UnexpectedEof, "field descriptor"))?;
512        let type_modifier = cur
513            .read_i32_be()
514            .map_err(|_| io::Error::new(io::ErrorKind::UnexpectedEof, "field descriptor"))?;
515        let format_code = cur
516            .read_i16_be()
517            .map_err(|_| io::Error::new(io::ErrorKind::UnexpectedEof, "field descriptor"))?;
518
519        fields.push(FieldDescription {
520            name,
521            table_oid,
522            column_attr,
523            type_oid,
524            type_size,
525            type_modifier,
526            format_code,
527        });
528    }
529
530    Ok(BackendMessage::RowDescription(fields))
531}
532
533#[cfg(test)]
534mod tests;
fraiseql_wire/protocol/decode/mod.rs

fraiseql_wire/protocol/decode/
mod.rs
