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        tags::EMPTY_QUERY_RESPONSE => BackendMessage::EmptyQueryResponse,
218        tags::NOTIFICATION_RESPONSE => decode_notification_response(msg_data)?,
219        // The COPY family begins a bulk-transfer sub-protocol fraiseql-wire does
220        // not implement. Surfacing an explicit `Unsupported` (rather than letting
221        // the tag fall through to the unknown-tag arm) keeps the distinction
222        // honest at the dispatch table: a recognized-but-unsupported message is
223        // not the same as a malformed one, and neither may be mistaken for
224        // "need more bytes" by the read loop (audit H42).
225        tags::COPY_IN_RESPONSE | tags::COPY_OUT_RESPONSE | tags::COPY_BOTH_RESPONSE => {
226            return Err(io::Error::new(
227                io::ErrorKind::Unsupported,
228                "COPY protocol is not supported by fraiseql-wire",
229            ))
230        }
231        _ => {
232            return Err(io::Error::new(
233                io::ErrorKind::InvalidData,
234                format!("unknown message tag: {}", tag),
235            ))
236        }
237    };
238
239    Ok((msg, len + 1))
240}
241
242fn decode_authentication(data: &[u8]) -> io::Result<BackendMessage> {
243    let mut cur = Cursor::new(data);
244    let auth_type = cur
245        .read_i32_be()
246        .map_err(|_| io::Error::new(io::ErrorKind::UnexpectedEof, "auth type"))?;
247
248    let auth_msg = match auth_type {
249        auth::OK => AuthenticationMessage::Ok,
250        auth::CLEARTEXT_PASSWORD => AuthenticationMessage::CleartextPassword,
251        auth::MD5_PASSWORD => {
252            let salt_slice = cur
253                .read_slice(4)
254                .map_err(|_| io::Error::new(io::ErrorKind::UnexpectedEof, "salt data"))?;
255            let salt: [u8; 4] = salt_slice
256                .try_into()
257                // Reason: provably-safe — `read_slice(4)` returns a 4-byte slice.
258                .expect("slice of length 4 always converts to [u8; 4]");
259            AuthenticationMessage::Md5Password { salt }
260        }
261        auth::SASL => {
262            // SASL: read mechanism list (null-terminated strings)
263            let mut mechanisms = Vec::new();
264            loop {
265                if cur.is_empty() {
266                    break;
267                }
268                let Some(end) = cur.position_of_null() else {
269                    break;
270                };
271                let mech_bytes = cur.read_slice(end).unwrap_or(&[]);
272                let mechanism = String::from_utf8_lossy(mech_bytes).to_string();
273                // Skip the null terminator we just located.
274                let _ = cur.read_u8();
275                if mechanism.is_empty() {
276                    break;
277                }
278                if mechanisms.len() >= MAX_SASL_MECHANISMS {
279                    break;
280                }
281                mechanisms.push(mechanism);
282            }
283            AuthenticationMessage::Sasl { mechanisms }
284        }
285        auth::SASL_CONTINUE => {
286            // SASL continue: read remaining data as bytes
287            let data_vec = cur.remaining().to_vec();
288            AuthenticationMessage::SaslContinue { data: data_vec }
289        }
290        auth::SASL_FINAL => {
291            // SASL final: read remaining data as bytes
292            let data_vec = cur.remaining().to_vec();
293            AuthenticationMessage::SaslFinal { data: data_vec }
294        }
295        _ => {
296            return Err(io::Error::new(
297                io::ErrorKind::Unsupported,
298                format!("unsupported auth type: {}", auth_type),
299            ))
300        }
301    };
302
303    Ok(BackendMessage::Authentication(auth_msg))
304}
305
306fn decode_backend_key_data(data: &[u8]) -> io::Result<BackendMessage> {
307    let mut cur = Cursor::new(data);
308    let process_id = cur
309        .read_i32_be()
310        .map_err(|_| io::Error::new(io::ErrorKind::UnexpectedEof, "backend key data"))?;
311    let secret_key = cur
312        .read_i32_be()
313        .map_err(|_| io::Error::new(io::ErrorKind::UnexpectedEof, "backend key data"))?;
314    Ok(BackendMessage::BackendKeyData {
315        process_id,
316        secret_key,
317    })
318}
319
320fn decode_command_complete(data: &[u8]) -> io::Result<BackendMessage> {
321    let mut cur = Cursor::new(data);
322    let tag_bytes = cur.read_until_null()?;
323    let tag = String::from_utf8_lossy(tag_bytes).to_string();
324    Ok(BackendMessage::CommandComplete(tag))
325}
326
327fn decode_data_row(data: &[u8]) -> io::Result<BackendMessage> {
328    let mut cur = Cursor::new(data);
329    let field_count_i16 = cur
330        .read_i16_be()
331        .map_err(|_| io::Error::new(io::ErrorKind::UnexpectedEof, "field count"))?;
332    if field_count_i16 < 0 {
333        return Err(io::Error::new(
334            io::ErrorKind::InvalidData,
335            "negative field count",
336        ));
337    }
338    let field_count = field_count_i16 as usize;
339    if field_count > MAX_FIELD_COUNT {
340        return Err(io::Error::new(
341            io::ErrorKind::InvalidData,
342            format!("DataRow field count {field_count} exceeds maximum {MAX_FIELD_COUNT}"),
343        ));
344    }
345    let mut fields = Vec::with_capacity(field_count);
346
347    for _ in 0..field_count {
348        let field_len = cur
349            .read_i32_be()
350            .map_err(|_| io::Error::new(io::ErrorKind::UnexpectedEof, "field length"))?;
351
352        let field = if field_len == -1 {
353            None
354        } else if field_len < 0 {
355            return Err(io::Error::new(
356                io::ErrorKind::InvalidData,
357                "negative field length",
358            ));
359        } else {
360            let len = field_len as usize;
361            let field_slice = cur
362                .read_slice(len)
363                .map_err(|_| io::Error::new(io::ErrorKind::UnexpectedEof, "field data"))?;
364            Some(Bytes::copy_from_slice(field_slice))
365        };
366        fields.push(field);
367    }
368
369    Ok(BackendMessage::DataRow(fields))
370}
371
372fn decode_error_response(data: &[u8]) -> io::Result<BackendMessage> {
373    let fields = decode_error_fields(data)?;
374    Ok(BackendMessage::ErrorResponse(fields))
375}
376
377fn decode_notice_response(data: &[u8]) -> io::Result<BackendMessage> {
378    let fields = decode_error_fields(data)?;
379    Ok(BackendMessage::NoticeResponse(fields))
380}
381
382fn decode_error_fields(data: &[u8]) -> io::Result<ErrorFields> {
383    let mut fields = ErrorFields::default();
384    let mut cur = Cursor::new(data);
385
386    loop {
387        if cur.is_empty() {
388            break;
389        }
390        let field_type = cur.read_u8()?;
391        if field_type == 0 {
392            break;
393        }
394
395        let end = cur.position_of_null().ok_or_else(|| {
396            io::Error::new(
397                io::ErrorKind::InvalidData,
398                "missing null terminator in error field",
399            )
400        })?;
401        if end > MAX_ERROR_FIELD_BYTES {
402            return Err(io::Error::new(
403                io::ErrorKind::InvalidData,
404                format!("Error field too large ({end} bytes, max {MAX_ERROR_FIELD_BYTES})"),
405            ));
406        }
407        let value_bytes = cur.read_slice(end).unwrap_or(&[]);
408        let value = String::from_utf8_lossy(value_bytes).to_string();
409        // Skip the null terminator.
410        let _ = cur.read_u8();
411
412        match field_type {
413            b'S' => fields.severity = Some(value),
414            b'C' => fields.code = Some(value),
415            b'M' => fields.message = Some(value),
416            b'D' => fields.detail = Some(value),
417            b'H' => fields.hint = Some(value),
418            b'P' => fields.position = Some(value),
419            _ => {} // Ignore unknown fields
420        }
421    }
422
423    Ok(fields)
424}
425
426fn decode_parameter_status(data: &[u8]) -> io::Result<BackendMessage> {
427    let mut cur = Cursor::new(data);
428
429    let name_end = cur.position_of_null().ok_or_else(|| {
430        io::Error::new(
431            io::ErrorKind::InvalidData,
432            "missing null terminator in parameter name",
433        )
434    })?;
435    if name_end > MAX_PARAMETER_NAME_BYTES {
436        return Err(io::Error::new(
437            io::ErrorKind::InvalidData,
438            format!("Parameter name too long ({name_end} bytes, max {MAX_PARAMETER_NAME_BYTES})"),
439        ));
440    }
441    let name_bytes = cur.read_slice(name_end).unwrap_or(&[]);
442    let name = String::from_utf8_lossy(name_bytes).to_string();
443    // Skip null terminator.
444    let _ = cur.read_u8();
445
446    if cur.is_empty() {
447        return Err(io::Error::new(
448            io::ErrorKind::UnexpectedEof,
449            "parameter value",
450        ));
451    }
452    let value_end = cur.position_of_null().ok_or_else(|| {
453        io::Error::new(
454            io::ErrorKind::InvalidData,
455            "missing null terminator in parameter value",
456        )
457    })?;
458    if value_end > MAX_PARAMETER_VALUE_BYTES {
459        return Err(io::Error::new(
460            io::ErrorKind::InvalidData,
461            format!(
462                "Parameter value too long ({value_end} bytes, max {MAX_PARAMETER_VALUE_BYTES})"
463            ),
464        ));
465    }
466    let value_bytes = cur.read_slice(value_end).unwrap_or(&[]);
467    let value = String::from_utf8_lossy(value_bytes).to_string();
468
469    Ok(BackendMessage::ParameterStatus { name, value })
470}
471
472fn decode_notification_response(data: &[u8]) -> io::Result<BackendMessage> {
473    let mut cur = Cursor::new(data);
474    let process_id = cur
475        .read_i32_be()
476        .map_err(|_| io::Error::new(io::ErrorKind::UnexpectedEof, "notification process id"))?;
477    let channel_bytes = cur.read_until_null()?;
478    let channel = String::from_utf8_lossy(channel_bytes).to_string();
479    let payload_bytes = cur.read_until_null()?;
480    let payload = String::from_utf8_lossy(payload_bytes).to_string();
481    Ok(BackendMessage::NotificationResponse {
482        process_id,
483        channel,
484        payload,
485    })
486}
487
488fn decode_ready_for_query(data: &[u8]) -> io::Result<BackendMessage> {
489    let status = *data
490        .first()
491        .ok_or_else(|| io::Error::new(io::ErrorKind::UnexpectedEof, "status byte"))?;
492    Ok(BackendMessage::ReadyForQuery { status })
493}
494
495fn decode_row_description(data: &[u8]) -> io::Result<BackendMessage> {
496    let mut cur = Cursor::new(data);
497    let field_count_i16 = cur
498        .read_i16_be()
499        .map_err(|_| io::Error::new(io::ErrorKind::UnexpectedEof, "field count"))?;
500    if field_count_i16 < 0 {
501        return Err(io::Error::new(
502            io::ErrorKind::InvalidData,
503            "negative field count",
504        ));
505    }
506    let field_count = field_count_i16 as usize;
507    if field_count > MAX_FIELD_COUNT {
508        return Err(io::Error::new(
509            io::ErrorKind::InvalidData,
510            format!("RowDescription field count {field_count} exceeds maximum {MAX_FIELD_COUNT}"),
511        ));
512    }
513    let mut fields = Vec::with_capacity(field_count);
514
515    for _ in 0..field_count {
516        // Read name (null-terminated string)
517        let name_end = cur.position_of_null().ok_or_else(|| {
518            io::Error::new(
519                io::ErrorKind::InvalidData,
520                "missing null terminator in field name",
521            )
522        })?;
523        let name_bytes = cur.read_slice(name_end).unwrap_or(&[]);
524        let name = String::from_utf8_lossy(name_bytes).to_string();
525        // Skip null terminator.
526        let _ = cur.read_u8();
527
528        // Read field descriptor (18 bytes: 4+2+4+2+4+2)
529        let table_oid = cur
530            .read_i32_be()
531            .map_err(|_| io::Error::new(io::ErrorKind::UnexpectedEof, "field descriptor"))?;
532        let column_attr = cur
533            .read_i16_be()
534            .map_err(|_| io::Error::new(io::ErrorKind::UnexpectedEof, "field descriptor"))?;
535        let type_oid = cur
536            .read_i32_be()
537            .map_err(|_| io::Error::new(io::ErrorKind::UnexpectedEof, "field descriptor"))?
538            as u32;
539        let type_size = cur
540            .read_i16_be()
541            .map_err(|_| io::Error::new(io::ErrorKind::UnexpectedEof, "field descriptor"))?;
542        let type_modifier = cur
543            .read_i32_be()
544            .map_err(|_| io::Error::new(io::ErrorKind::UnexpectedEof, "field descriptor"))?;
545        let format_code = cur
546            .read_i16_be()
547            .map_err(|_| io::Error::new(io::ErrorKind::UnexpectedEof, "field descriptor"))?;
548
549        fields.push(FieldDescription {
550            name,
551            table_oid,
552            column_attr,
553            type_oid,
554            type_size,
555            type_modifier,
556            format_code,
557        });
558    }
559
560    Ok(BackendMessage::RowDescription(fields))
561}
562
563#[cfg(test)]
564mod tests;
fraiseql_wire/protocol/decode/mod.rs

fraiseql_wire/protocol/decode/
mod.rs
