pg_proto/
codec.rs

1//! Direction-parameterised framing and lossless message decoding.
2
3use std::{io, marker::PhantomData};
4
5use bytes::{Buf, BufMut, Bytes, BytesMut};
6use tokio_util::codec::{Decoder, Encoder};
7
8use crate::startup::ProtocolVersion;
9
10/// Messages sent by a `PostgreSQL` frontend.
11#[derive(Debug)]
12pub enum Frontend {}
13
14/// Messages sent by a `PostgreSQL` backend.
15#[derive(Debug)]
16pub enum Backend {}
17
18/// A validated `PostgreSQL` tagged frame, including its tag but not its length.
19#[derive(Clone, Debug, Eq, PartialEq)]
20pub struct Frame {
21    /// Direction-scoped one-byte message tag.
22    pub tag: u8,
23    /// Message bytes after the tag and four-byte length field.
24    pub body: Bytes,
25}
26
27/// Decoder direction. The same byte tag has different meanings in each direction.
28pub trait Direction {
29    /// Typed message family produced for this direction.
30    type Message;
31
32    /// # Errors
33    ///
34    /// Returns an error when the tag is unknown or its body is malformed.
35    fn decode(frame: Frame) -> io::Result<Self::Message>;
36}
37
38/// A `PostgreSQL` codec which cannot confuse frontend and backend tag alphabets.
39#[derive(Debug)]
40pub struct PgCodec<D> {
41    max_frame_len: usize,
42    _direction: PhantomData<fn() -> D>,
43}
44
45const MAX_PROTOCOL_FRAME_LEN: usize = i32::MAX as usize + 1;
46/// Default total tagged-frame limit, including tag and length.
47pub const DEFAULT_MAX_FRAME_LEN: usize = 16 * 1024 * 1024;
48
49impl<D> Default for PgCodec<D> {
50    fn default() -> Self {
51        Self {
52            max_frame_len: DEFAULT_MAX_FRAME_LEN,
53            _direction: PhantomData,
54        }
55    }
56}
57
58impl<D> PgCodec<D> {
59    /// Creates a codec with a total tagged-frame limit, including tag and length.
60    ///
61    /// # Errors
62    ///
63    /// Rejects limits smaller than an empty frame or larger than `PostgreSQL`'s
64    /// signed int32 length field can represent.
65    pub fn with_max_frame_len(max_frame_len: usize) -> io::Result<Self> {
66        if !(5..=MAX_PROTOCOL_FRAME_LEN).contains(&max_frame_len) {
67            return Err(io::Error::new(
68                io::ErrorKind::InvalidInput,
69                "frame limit is outside PostgreSQL's tagged-frame range",
70            ));
71        }
72        Ok(Self {
73            max_frame_len,
74            _direction: PhantomData,
75        })
76    }
77}
78
79impl<D: Direction> Decoder for PgCodec<D> {
80    type Item = D::Message;
81    type Error = io::Error;
82
83    fn decode(&mut self, source: &mut BytesMut) -> io::Result<Option<Self::Item>> {
84        let Some(frame) = decode_frame(source, self.max_frame_len)? else {
85            return Ok(None);
86        };
87        let tag = frame.tag;
88        D::decode(frame).map(Some).map_err(|error| {
89            io::Error::new(error.kind(), format!("message tag 0x{tag:02x}: {error}"))
90        })
91    }
92}
93
94impl<D> Encoder<Frame> for PgCodec<D> {
95    type Error = io::Error;
96
97    fn encode(&mut self, item: Frame, destination: &mut BytesMut) -> io::Result<()> {
98        let frame_len = item
99            .body
100            .len()
101            .checked_add(5)
102            .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "frame too large"))?;
103        if frame_len > self.max_frame_len {
104            return Err(io::Error::new(
105                io::ErrorKind::InvalidInput,
106                "frame exceeds configured limit",
107            ));
108        }
109        let length = item
110            .body
111            .len()
112            .checked_add(4)
113            .and_then(|length| u32::try_from(length).ok())
114            .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "frame too large"))?;
115        destination.reserve(item.body.len() + 5);
116        destination.put_u8(item.tag);
117        destination.put_u32(length);
118        destination.extend_from_slice(&item.body);
119        Ok(())
120    }
121}
122
123fn decode_frame(source: &mut BytesMut, max_frame_len: usize) -> io::Result<Option<Frame>> {
124    if source.len() < 5 {
125        source.reserve(5 - source.len());
126        return Ok(None);
127    }
128
129    let length = u32::from_be_bytes(source[1..5].try_into().expect("four-byte slice"));
130    if length < 4 {
131        return Err(io::Error::new(
132            io::ErrorKind::InvalidData,
133            "message length is smaller than its length field",
134        ));
135    }
136    let frame_length = usize::try_from(length)
137        .ok()
138        .and_then(|length| length.checked_add(1))
139        .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "message length overflow"))?;
140    if frame_length > MAX_PROTOCOL_FRAME_LEN {
141        return Err(io::Error::new(
142            io::ErrorKind::InvalidData,
143            "message length exceeds PostgreSQL's signed int32 range",
144        ));
145    }
146    if frame_length > max_frame_len {
147        return Err(io::Error::new(
148            io::ErrorKind::InvalidData,
149            "message exceeds configured frame limit",
150        ));
151    }
152    if source.len() < frame_length {
153        source.reserve(frame_length - source.len());
154        return Ok(None);
155    }
156
157    let tag = source[0];
158    let mut bytes = source.split_to(frame_length).freeze();
159    bytes.advance(5);
160    Ok(Some(Frame { tag, body: bytes }))
161}
162
163/// Frontend messages whose contents a rewriting proxy must retain structurally.
164#[derive(Clone, Eq, PartialEq)]
165pub enum FrontendMessage {
166    /// Defines a prepared statement and its parameter OIDs.
167    Parse(Parse),
168    /// Binds parameters and result formats to a portal.
169    Bind(Bind),
170    /// Requests metadata for a statement or portal.
171    Describe(Describe),
172    /// Closes a statement or portal.
173    Close(Close),
174    /// Executes a bound portal.
175    Execute(Execute),
176    /// Invokes the deprecated function-call protocol.
177    FunctionCall(FunctionCall),
178    /// Executes one or more SQL statements through the simple-query protocol.
179    Query(Bytes),
180    /// Requests that buffered backend responses be flushed.
181    Flush,
182    /// Ends an extended-query pipeline and restores error recovery.
183    Sync,
184    /// Gracefully closes the frontend session.
185    Terminate,
186    /// Carries COPY or replication data.
187    CopyData(Bytes),
188    /// Signals successful completion of a COPY input stream.
189    CopyDone,
190    /// Aborts a COPY input stream with a diagnostic message.
191    CopyFail(Bytes),
192    /// Context determines whether this is password, GSS, or a SASL response.
193    PasswordResponse(Bytes),
194}
195
196impl std::fmt::Debug for FrontendMessage {
197    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
198        match self {
199            Self::Parse(value) => formatter.debug_tuple("Parse").field(value).finish(),
200            Self::Bind(value) => formatter.debug_tuple("Bind").field(value).finish(),
201            Self::Describe(value) => formatter.debug_tuple("Describe").field(value).finish(),
202            Self::Close(value) => formatter.debug_tuple("Close").field(value).finish(),
203            Self::Execute(value) => formatter.debug_tuple("Execute").field(value).finish(),
204            Self::FunctionCall(value) => {
205                formatter.debug_tuple("FunctionCall").field(value).finish()
206            }
207            Self::Query(value) => formatter.debug_tuple("Query").field(value).finish(),
208            Self::Flush => formatter.write_str("Flush"),
209            Self::Sync => formatter.write_str("Sync"),
210            Self::Terminate => formatter.write_str("Terminate"),
211            Self::CopyData(value) => formatter.debug_tuple("CopyData").field(value).finish(),
212            Self::CopyDone => formatter.write_str("CopyDone"),
213            Self::CopyFail(value) => formatter.debug_tuple("CopyFail").field(value).finish(),
214            Self::PasswordResponse(value) => formatter
215                .debug_tuple("PasswordResponse")
216                .field(&format_args!("[REDACTED; {} bytes]", value.len()))
217                .finish(),
218        }
219    }
220}
221
222impl FrontendMessage {
223    /// Reconstructs a frontend frame after inspection or modification.
224    ///
225    /// # Errors
226    ///
227    /// Returns an error when a structured message contains invalid values.
228    pub fn to_frame(&self) -> io::Result<Frame> {
229        match self {
230            Self::Parse(message) => message.to_frame(),
231            Self::Bind(message) => message.to_frame(),
232            Self::Describe(message) => message.to_frame(),
233            Self::Close(message) => message.to_frame(),
234            Self::Execute(message) => message.to_frame(),
235            Self::FunctionCall(message) => message.to_frame(),
236            Self::Query(query) => cstr_message(b'Q', query),
237            Self::Flush => Ok(empty_message(b'H')),
238            Self::Sync => Ok(empty_message(b'S')),
239            Self::Terminate => Ok(empty_message(b'X')),
240            Self::CopyData(data) => Ok(Frame {
241                tag: b'd',
242                body: data.clone(),
243            }),
244            Self::CopyDone => Ok(empty_message(b'c')),
245            Self::CopyFail(message) => cstr_message(b'f', message),
246            Self::PasswordResponse(data) => Ok(Frame {
247                tag: b'p',
248                body: data.clone(),
249            }),
250        }
251    }
252}
253
254/// Structured `Parse` message.
255#[derive(Clone, Debug, Eq, PartialEq)]
256pub struct Parse {
257    /// Client-visible prepared-statement name; empty denotes the unnamed statement.
258    pub statement: Bytes,
259    /// SQL text supplied by the frontend.
260    pub query: Bytes,
261    /// Declared parameter type OIDs; zero entries request inference.
262    pub parameter_types: Vec<u32>,
263}
264
265impl Parse {
266    /// Reconstructs a checked Parse frame after inspection or rewriting.
267    ///
268    /// # Errors
269    ///
270    /// Returns an error for NUL-containing strings or too many parameter types.
271    pub fn to_frame(&self) -> io::Result<Frame> {
272        let mut body = BytesMut::new();
273        put_cstr(&self.statement, &mut body)?;
274        put_cstr(&self.query, &mut body)?;
275        put_count(self.parameter_types.len(), &mut body)?;
276        for oid in &self.parameter_types {
277            body.put_u32(*oid);
278        }
279        Ok(Frame {
280            tag: b'P',
281            body: body.freeze(),
282        })
283    }
284}
285
286/// Structured `Bind` message.
287#[derive(Clone, Debug, Eq, PartialEq)]
288pub struct Bind {
289    /// Portal name; empty denotes the unnamed portal.
290    pub portal: Bytes,
291    /// Prepared-statement name referenced by the portal.
292    pub statement: Bytes,
293    /// Parameter format codes using `PostgreSQL`'s zero/one/per-value cardinality rules.
294    pub parameter_formats: Vec<i16>,
295    /// Parameter values, with `None` representing SQL `NULL`.
296    pub parameters: Vec<Option<Bytes>>,
297    /// Requested result-column format codes.
298    pub result_formats: Vec<i16>,
299}
300
301impl Bind {
302    /// Reconstructs a checked Bind frame, retaining every format code and value.
303    ///
304    /// # Errors
305    ///
306    /// Returns an error for invalid names, excessive counts, or oversized values.
307    pub fn to_frame(&self) -> io::Result<Frame> {
308        let mut body = BytesMut::new();
309        put_cstr(&self.portal, &mut body)?;
310        put_cstr(&self.statement, &mut body)?;
311        put_i16_vec(&self.parameter_formats, &mut body)?;
312        put_count(self.parameters.len(), &mut body)?;
313        for parameter in &self.parameters {
314            match parameter {
315                None => body.put_i32(-1),
316                Some(value) => {
317                    let length = i32::try_from(value.len())
318                        .map_err(|_| invalid_input("Bind parameter is too large"))?;
319                    body.put_i32(length);
320                    body.extend_from_slice(value);
321                }
322            }
323        }
324        put_i16_vec(&self.result_formats, &mut body)?;
325        Ok(Frame {
326            tag: b'B',
327            body: body.freeze(),
328        })
329    }
330}
331
332/// Structured `Describe` message.
333#[derive(Clone, Debug, Eq, PartialEq)]
334pub struct Describe {
335    /// Namespace in which `name` is resolved.
336    pub target: DescribeTarget,
337    /// Statement or portal name to describe.
338    pub name: Bytes,
339}
340
341/// Namespace selected by `Describe` and `Close` messages.
342#[derive(Clone, Copy, Debug, Eq, PartialEq)]
343pub enum DescribeTarget {
344    /// Prepared-statement namespace.
345    Statement,
346    /// Portal namespace.
347    Portal,
348}
349
350/// Structured `Close` message.
351#[derive(Clone, Debug, Eq, PartialEq)]
352pub struct Close {
353    /// Namespace in which `name` is resolved.
354    pub target: DescribeTarget,
355    /// Statement or portal name to close.
356    pub name: Bytes,
357}
358
359/// Structured `Execute` message.
360#[derive(Clone, Debug, Eq, PartialEq)]
361pub struct Execute {
362    /// Portal to execute.
363    pub portal: Bytes,
364    /// Maximum rows to return; zero requests all rows.
365    pub max_rows: i32,
366}
367
368/// Structured legacy `FunctionCall` message.
369#[derive(Clone, Debug, Eq, PartialEq)]
370pub struct FunctionCall {
371    /// OID of the function to invoke.
372    pub function_oid: u32,
373    /// Per-argument format codes.
374    pub argument_formats: Vec<i16>,
375    /// Function arguments, with `None` representing SQL `NULL`.
376    pub arguments: Vec<Option<Bytes>>,
377    /// Requested result format code.
378    pub result_format: i16,
379}
380
381impl Describe {
382    /// Reconstructs a checked Describe frame.
383    ///
384    /// # Errors
385    ///
386    /// Returns an error if the name contains a NUL byte.
387    pub fn to_frame(&self) -> io::Result<Frame> {
388        let mut body = BytesMut::new();
389        body.put_u8(match self.target {
390            DescribeTarget::Statement => b'S',
391            DescribeTarget::Portal => b'P',
392        });
393        put_cstr(&self.name, &mut body)?;
394        Ok(Frame {
395            tag: b'D',
396            body: body.freeze(),
397        })
398    }
399}
400
401impl Close {
402    /// # Errors
403    ///
404    /// Returns an error if the name contains a NUL byte.
405    pub fn to_frame(&self) -> io::Result<Frame> {
406        named_target_frame(b'C', self.target, &self.name)
407    }
408}
409
410impl Execute {
411    /// # Errors
412    ///
413    /// Returns an error if the portal name contains a NUL byte.
414    pub fn to_frame(&self) -> io::Result<Frame> {
415        let mut body = BytesMut::new();
416        put_cstr(&self.portal, &mut body)?;
417        body.put_i32(self.max_rows);
418        Ok(Frame {
419            tag: b'E',
420            body: body.freeze(),
421        })
422    }
423}
424
425impl FunctionCall {
426    /// # Errors
427    ///
428    /// Returns an error for excessive counts or oversized argument values.
429    pub fn to_frame(&self) -> io::Result<Frame> {
430        let mut body = BytesMut::new();
431        body.put_u32(self.function_oid);
432        put_i16_vec(&self.argument_formats, &mut body)?;
433        put_count(self.arguments.len(), &mut body)?;
434        for argument in &self.arguments {
435            put_nullable(argument.as_ref(), &mut body)?;
436        }
437        body.put_i16(self.result_format);
438        Ok(Frame {
439            tag: b'F',
440            body: body.freeze(),
441        })
442    }
443}
444
445impl Direction for Frontend {
446    type Message = FrontendMessage;
447
448    fn decode(frame: Frame) -> io::Result<Self::Message> {
449        match frame.tag {
450            b'P' => decode_parse(frame.body).map(FrontendMessage::Parse),
451            b'B' => decode_bind(frame.body).map(FrontendMessage::Bind),
452            b'D' => decode_describe(frame.body).map(FrontendMessage::Describe),
453            b'C' => decode_close(frame.body).map(FrontendMessage::Close),
454            b'E' => decode_execute(frame.body).map(FrontendMessage::Execute),
455            b'F' => decode_function_call(frame.body).map(FrontendMessage::FunctionCall),
456            b'H' => decode_empty(&frame.body).map(|()| FrontendMessage::Flush),
457            b'Q' => decode_cstr_body(frame.body).map(FrontendMessage::Query),
458            b'S' => decode_empty(&frame.body).map(|()| FrontendMessage::Sync),
459            b'X' => decode_empty(&frame.body).map(|()| FrontendMessage::Terminate),
460            b'c' => decode_empty(&frame.body).map(|()| FrontendMessage::CopyDone),
461            b'd' => Ok(FrontendMessage::CopyData(frame.body)),
462            b'f' => decode_cstr_body(frame.body).map(FrontendMessage::CopyFail),
463            b'p' => Ok(FrontendMessage::PasswordResponse(frame.body)),
464            tag => Err(unknown_tag("frontend", tag)),
465        }
466    }
467}
468
469/// Backend row metadata retained in reconstructable form.
470#[derive(Clone, Debug, Eq, PartialEq)]
471pub struct RowDescription {
472    /// Result columns in wire order.
473    pub fields: Vec<FieldDescription>,
474}
475
476/// Metadata for one result column.
477#[derive(Clone, Debug, Eq, PartialEq)]
478pub struct FieldDescription {
479    /// Column label presented to the frontend.
480    pub name: Bytes,
481    /// Source table OID, or zero when not associated with a table.
482    pub table_oid: u32,
483    /// One-based source column number, or zero when not applicable.
484    pub column: i16,
485    /// `PostgreSQL` data-type OID.
486    pub type_oid: u32,
487    /// Fixed type width, or `-1` for variable-width types.
488    pub type_size: i16,
489    /// Type-specific modifier, or `-1` when absent.
490    pub type_modifier: i32,
491    /// Result format: zero for text and one for binary.
492    pub format: i16,
493}
494
495impl RowDescription {
496    /// Reconstructs checked result metadata after proxy rewriting.
497    ///
498    /// # Errors
499    ///
500    /// Returns an error for excessive fields or NUL-containing field names.
501    pub fn to_frame(&self) -> io::Result<Frame> {
502        let mut body = BytesMut::new();
503        put_count(self.fields.len(), &mut body)?;
504        for field in &self.fields {
505            put_cstr(&field.name, &mut body)?;
506            body.put_u32(field.table_oid);
507            body.put_i16(field.column);
508            body.put_u32(field.type_oid);
509            body.put_i16(field.type_size);
510            body.put_i32(field.type_modifier);
511            body.put_i16(field.format);
512        }
513        Ok(Frame {
514            tag: b'T',
515            body: body.freeze(),
516        })
517    }
518}
519
520/// Backend authentication request or continuation message.
521#[derive(Clone, Debug, Eq, PartialEq)]
522pub enum Authentication {
523    /// Authentication completed successfully.
524    Ok,
525    /// Requests Kerberos V5 authentication.
526    KerberosV5,
527    /// Requests a cleartext password response.
528    CleartextPassword,
529    /// Requests a `PostgreSQL` MD5 password response.
530    Md5Password {
531        /// Four-byte server challenge salt.
532        salt: [u8; 4],
533    },
534    /// Begins a GSSAPI token exchange.
535    Gss,
536    /// Continues a GSSAPI token exchange with a server token.
537    GssContinue(Bytes),
538    /// Begins an SSPI token exchange.
539    Sspi,
540    /// Offers SASL authentication mechanisms.
541    Sasl {
542        /// Mechanism names in server preference order.
543        mechanisms: Vec<Bytes>,
544    },
545    /// Carries a SASL server-first challenge.
546    SaslContinue(Bytes),
547    /// Carries the SASL server-final message.
548    SaslFinal(Bytes),
549}
550
551/// Transaction state reported by `ReadyForQuery`.
552#[derive(Clone, Copy, Debug, Eq, PartialEq)]
553pub enum TransactionStatus {
554    /// No transaction is active.
555    Idle,
556    /// A transaction is active and has not failed.
557    InTransaction,
558    /// A transaction is active and failed; only rollback is legal.
559    FailedTransaction,
560}
561
562/// Backend response negotiating a requested protocol minor version and options.
563#[derive(Clone, Debug, Eq, PartialEq)]
564pub struct NegotiateProtocolVersion {
565    /// Newest protocol version supported by the server.
566    pub newest: ProtocolVersion,
567    /// Startup option names the server does not support.
568    pub unsupported_options: Vec<Bytes>,
569}
570
571/// Ordered diagnostic fields from an error or notice response.
572#[derive(Clone, Debug, Eq, PartialEq)]
573pub struct DiagnosticResponse {
574    /// Diagnostic fields in their original wire order.
575    pub fields: Vec<DiagnosticField>,
576}
577
578/// One tagged `PostgreSQL` diagnostic field.
579#[derive(Clone, Debug, Eq, PartialEq)]
580pub struct DiagnosticField {
581    /// `PostgreSQL`'s one-byte diagnostic field code.
582    pub code: u8,
583    /// Field value without its terminating NUL byte.
584    pub value: Bytes,
585}
586
587/// Format metadata which begins a COPY sub-protocol.
588#[derive(Clone, Debug, Eq, PartialEq)]
589pub struct CopyResponse {
590    /// Overall COPY format: zero for text and one for binary.
591    pub overall_format: u8,
592    /// Per-column format codes.
593    pub column_formats: Vec<i16>,
594}
595
596/// One result row retaining raw text or binary column values.
597#[derive(Clone, Debug, Eq, PartialEq)]
598pub struct DataRow {
599    /// Column values in result order; `None` represents SQL `NULL`.
600    pub columns: Vec<Option<Bytes>>,
601}
602
603/// Messages sent by a `PostgreSQL` backend.
604#[derive(Clone, Debug, Eq, PartialEq)]
605pub enum BackendMessage {
606    /// Describes the columns of a result set.
607    RowDescription(RowDescription),
608    /// Requests, continues, or completes authentication.
609    Authentication(Authentication),
610    /// Confirms `Parse` completion.
611    ParseComplete,
612    /// Confirms `Bind` completion.
613    BindComplete,
614    /// Confirms `Close` completion.
615    CloseComplete,
616    /// Reports a completed SQL command tag.
617    CommandComplete(Bytes),
618    /// Carries COPY or replication data.
619    CopyData(Bytes),
620    /// Signals completion of backend COPY output.
621    CopyDone,
622    /// Enters COPY IN mode.
623    CopyInResponse(CopyResponse),
624    /// Enters COPY OUT mode.
625    CopyOutResponse(CopyResponse),
626    /// Enters bidirectional COPY mode.
627    CopyBothResponse(CopyResponse),
628    /// Carries one result row.
629    DataRow(DataRow),
630    /// Reports an empty simple-query string.
631    EmptyQueryResponse,
632    /// Reports an error with structured diagnostics.
633    ErrorResponse(DiagnosticResponse),
634    /// Reports that a described statement or portal has no row metadata.
635    NoData,
636    /// Reports a runtime parameter value.
637    ParameterStatus {
638        /// Parameter name.
639        name: Bytes,
640        /// Current parameter value.
641        value: Bytes,
642    },
643    /// Reports a non-fatal notice.
644    NoticeResponse(DiagnosticResponse),
645    /// Delivers an asynchronous `LISTEN`/`NOTIFY` notification.
646    NotificationResponse {
647        /// Process ID of the notifying backend.
648        process_id: u32,
649        /// Notification channel.
650        channel: Bytes,
651        /// Notification payload.
652        payload: Bytes,
653    },
654    /// Supplies the backend cancellation key.
655    BackendKeyData {
656        /// Backend process ID.
657        process_id: u32,
658        /// Backend secret cancellation key.
659        secret_key: Bytes,
660    },
661    /// Marks an idle command boundary and reports transaction state.
662    ReadyForQuery(TransactionStatus),
663    /// Reports inferred or declared prepared-statement parameter OIDs.
664    ParameterDescription(Vec<u32>),
665    /// Reports that an execution stopped at its row limit and may resume.
666    PortalSuspended,
667    /// Returns the legacy function-call result bytes.
668    FunctionCallResponse(Bytes),
669    /// Negotiates protocol version and unsupported startup options.
670    NegotiateProtocolVersion(NegotiateProtocolVersion),
671}
672
673impl BackendMessage {
674    /// Reconstructs a backend frame after inspection or modification.
675    ///
676    /// # Errors
677    ///
678    /// Returns an error when a structured message contains invalid values.
679    pub fn to_frame(&self) -> io::Result<Frame> {
680        match self {
681            Self::RowDescription(message) => message.to_frame(),
682            Self::Authentication(message) => authentication_frame(message),
683            Self::ParseComplete => Ok(empty_message(b'1')),
684            Self::BindComplete => Ok(empty_message(b'2')),
685            Self::CloseComplete => Ok(empty_message(b'3')),
686            Self::CommandComplete(tag) => cstr_message(b'C', tag),
687            Self::CopyData(data) => Ok(Frame {
688                tag: b'd',
689                body: data.clone(),
690            }),
691            Self::CopyDone => Ok(empty_message(b'c')),
692            Self::CopyInResponse(response) => copy_response_frame(b'G', response),
693            Self::CopyOutResponse(response) => copy_response_frame(b'H', response),
694            Self::CopyBothResponse(response) => copy_response_frame(b'W', response),
695            Self::DataRow(row) => row.to_frame(),
696            Self::EmptyQueryResponse => Ok(empty_message(b'I')),
697            Self::ErrorResponse(response) => diagnostic_frame(b'E', response),
698            Self::NoData => Ok(empty_message(b'n')),
699            Self::ParameterStatus { name, value } => {
700                let mut body = BytesMut::new();
701                put_cstr(name, &mut body)?;
702                put_cstr(value, &mut body)?;
703                Ok(Frame {
704                    tag: b'S',
705                    body: body.freeze(),
706                })
707            }
708            Self::NoticeResponse(response) => diagnostic_frame(b'N', response),
709            Self::NotificationResponse {
710                process_id,
711                channel,
712                payload,
713            } => {
714                let mut body = BytesMut::new();
715                body.put_u32(*process_id);
716                put_cstr(channel, &mut body)?;
717                put_cstr(payload, &mut body)?;
718                Ok(Frame {
719                    tag: b'A',
720                    body: body.freeze(),
721                })
722            }
723            Self::BackendKeyData {
724                process_id,
725                secret_key,
726            } => {
727                if !(4..=256).contains(&secret_key.len()) {
728                    return Err(invalid_input("cancellation key length is outside 4..=256"));
729                }
730                let mut body = BytesMut::with_capacity(4 + secret_key.len());
731                body.put_u32(*process_id);
732                body.extend_from_slice(secret_key);
733                Ok(Frame {
734                    tag: b'K',
735                    body: body.freeze(),
736                })
737            }
738            Self::ReadyForQuery(status) => Ok(Frame {
739                tag: b'Z',
740                body: Bytes::copy_from_slice(&[status.as_byte()]),
741            }),
742            Self::ParameterDescription(types) => {
743                let mut body = BytesMut::new();
744                put_count(types.len(), &mut body)?;
745                for oid in types {
746                    body.put_u32(*oid);
747                }
748                Ok(Frame {
749                    tag: b't',
750                    body: body.freeze(),
751                })
752            }
753            Self::PortalSuspended => Ok(empty_message(b's')),
754            Self::FunctionCallResponse(data) => Ok(Frame {
755                tag: b'V',
756                body: data.clone(),
757            }),
758            Self::NegotiateProtocolVersion(message) => message.to_frame(),
759        }
760    }
761}
762
763impl TransactionStatus {
764    const fn as_byte(self) -> u8 {
765        match self {
766            Self::Idle => b'I',
767            Self::InTransaction => b'T',
768            Self::FailedTransaction => b'E',
769        }
770    }
771}
772
773impl NegotiateProtocolVersion {
774    /// Reconstructs protocol negotiation, including unsupported option names.
775    ///
776    /// # Errors
777    ///
778    /// Returns an error for too many options or NUL-containing names.
779    pub fn to_frame(&self) -> io::Result<Frame> {
780        let mut body = BytesMut::new();
781        body.put_u32((u32::from(self.newest.major) << 16) | u32::from(self.newest.minor));
782        let count = u32::try_from(self.unsupported_options.len())
783            .map_err(|_| invalid_input("unsupported option count exceeds u32"))?;
784        body.put_u32(count);
785        for option in &self.unsupported_options {
786            put_cstr(option, &mut body)?;
787        }
788        Ok(Frame {
789            tag: b'v',
790            body: body.freeze(),
791        })
792    }
793}
794
795impl DataRow {
796    /// # Errors
797    ///
798    /// Returns an error for too many columns or oversized values.
799    pub fn to_frame(&self) -> io::Result<Frame> {
800        let mut body = BytesMut::new();
801        put_count(self.columns.len(), &mut body)?;
802        for column in &self.columns {
803            put_nullable(column.as_ref(), &mut body)?;
804        }
805        Ok(Frame {
806            tag: b'D',
807            body: body.freeze(),
808        })
809    }
810}
811
812fn diagnostic_frame(tag: u8, response: &DiagnosticResponse) -> io::Result<Frame> {
813    let mut body = BytesMut::new();
814    for field in &response.fields {
815        if field.code == 0 {
816            return Err(invalid_input("diagnostic field code cannot be zero"));
817        }
818        body.put_u8(field.code);
819        put_cstr(&field.value, &mut body)?;
820    }
821    body.put_u8(0);
822    Ok(Frame {
823        tag,
824        body: body.freeze(),
825    })
826}
827
828fn copy_response_frame(tag: u8, response: &CopyResponse) -> io::Result<Frame> {
829    let mut body = BytesMut::new();
830    body.put_u8(response.overall_format);
831    put_i16_vec(&response.column_formats, &mut body)?;
832    Ok(Frame {
833        tag,
834        body: body.freeze(),
835    })
836}
837
838fn authentication_frame(authentication: &Authentication) -> io::Result<Frame> {
839    let mut body = BytesMut::new();
840    match authentication {
841        Authentication::Ok => body.put_u32(0),
842        Authentication::KerberosV5 => body.put_u32(2),
843        Authentication::CleartextPassword => body.put_u32(3),
844        Authentication::Md5Password { salt } => {
845            body.put_u32(5);
846            body.extend_from_slice(salt);
847        }
848        Authentication::Gss => body.put_u32(7),
849        Authentication::GssContinue(data) => {
850            body.put_u32(8);
851            body.extend_from_slice(data);
852        }
853        Authentication::Sspi => body.put_u32(9),
854        Authentication::Sasl { mechanisms } => {
855            body.put_u32(10);
856            for mechanism in mechanisms {
857                put_cstr(mechanism, &mut body)?;
858            }
859            body.put_u8(0);
860        }
861        Authentication::SaslContinue(data) => {
862            body.put_u32(11);
863            body.extend_from_slice(data);
864        }
865        Authentication::SaslFinal(data) => {
866            body.put_u32(12);
867            body.extend_from_slice(data);
868        }
869    }
870    Ok(Frame {
871        tag: b'R',
872        body: body.freeze(),
873    })
874}
875
876impl Direction for Backend {
877    type Message = BackendMessage;
878
879    fn decode(frame: Frame) -> io::Result<Self::Message> {
880        match frame.tag {
881            b'1' => decode_empty(&frame.body).map(|()| BackendMessage::ParseComplete),
882            b'2' => decode_empty(&frame.body).map(|()| BackendMessage::BindComplete),
883            b'3' => decode_empty(&frame.body).map(|()| BackendMessage::CloseComplete),
884            b'C' => decode_cstr_body(frame.body).map(BackendMessage::CommandComplete),
885            b'c' => decode_empty(&frame.body).map(|()| BackendMessage::CopyDone),
886            b'd' => Ok(BackendMessage::CopyData(frame.body)),
887            b'D' => decode_data_row(frame.body).map(BackendMessage::DataRow),
888            b'E' => decode_diagnostic(frame.body).map(BackendMessage::ErrorResponse),
889            b'G' => decode_copy_response(frame.body).map(BackendMessage::CopyInResponse),
890            b'H' => decode_copy_response(frame.body).map(BackendMessage::CopyOutResponse),
891            b'I' => decode_empty(&frame.body).map(|()| BackendMessage::EmptyQueryResponse),
892            b'n' => decode_empty(&frame.body).map(|()| BackendMessage::NoData),
893            b's' => decode_empty(&frame.body).map(|()| BackendMessage::PortalSuspended),
894            b't' => {
895                decode_parameter_description(frame.body).map(BackendMessage::ParameterDescription)
896            }
897            b'T' => decode_row_description(frame.body).map(BackendMessage::RowDescription),
898            b'V' => Ok(BackendMessage::FunctionCallResponse(frame.body)),
899            b'W' => decode_copy_response(frame.body).map(BackendMessage::CopyBothResponse),
900            b'R' => decode_authentication(frame.body).map(BackendMessage::Authentication),
901            b'S' => decode_parameter_status(frame.body),
902            b'N' => decode_diagnostic(frame.body).map(BackendMessage::NoticeResponse),
903            b'A' => decode_notification(frame.body),
904            b'K' => decode_backend_key_data(frame.body),
905            b'Z' => decode_ready(frame.body),
906            b'v' => decode_negotiate_protocol_version(frame.body),
907            tag => Err(unknown_tag("backend", tag)),
908        }
909    }
910}
911
912fn decode_authentication(mut body: Bytes) -> io::Result<Authentication> {
913    let kind = take_u32(&mut body)?;
914    let auth = match kind {
915        0 => Authentication::Ok,
916        2 => Authentication::KerberosV5,
917        3 => Authentication::CleartextPassword,
918        5 => {
919            require(&body, 4)?;
920            let salt = body.split_to(4);
921            Authentication::Md5Password {
922                salt: salt[..].try_into().expect("four-byte slice"),
923            }
924        }
925        7 => Authentication::Gss,
926        8 => Authentication::GssContinue(body.split_to(body.len())),
927        9 => Authentication::Sspi,
928        10 => {
929            let mut mechanisms = Vec::new();
930            while !body.is_empty() && body[0] != 0 {
931                mechanisms.push(take_cstr(&mut body)?);
932            }
933            require(&body, 1)?;
934            body.advance(1);
935            Authentication::Sasl { mechanisms }
936        }
937        11 => Authentication::SaslContinue(body.split_to(body.len())),
938        12 => Authentication::SaslFinal(body.split_to(body.len())),
939        _ => return Err(invalid("unknown authentication request")),
940    };
941    require_empty(&body)?;
942    Ok(auth)
943}
944
945fn decode_parameter_status(mut body: Bytes) -> io::Result<BackendMessage> {
946    let name = take_cstr(&mut body)?;
947    let value = take_cstr(&mut body)?;
948    require_empty(&body)?;
949    Ok(BackendMessage::ParameterStatus { name, value })
950}
951
952fn decode_notification(mut body: Bytes) -> io::Result<BackendMessage> {
953    let process_id = take_u32(&mut body)?;
954    let channel = take_cstr(&mut body)?;
955    let payload = take_cstr(&mut body)?;
956    require_empty(&body)?;
957    Ok(BackendMessage::NotificationResponse {
958        process_id,
959        channel,
960        payload,
961    })
962}
963
964fn decode_backend_key_data(mut body: Bytes) -> io::Result<BackendMessage> {
965    let process_id = take_u32(&mut body)?;
966    if !(4..=256).contains(&body.len()) {
967        return Err(invalid("cancellation key length is outside 4..=256"));
968    }
969    let secret_key = body;
970    Ok(BackendMessage::BackendKeyData {
971        process_id,
972        secret_key,
973    })
974}
975
976fn decode_ready(mut body: Bytes) -> io::Result<BackendMessage> {
977    require(&body, 1)?;
978    let status = match body.get_u8() {
979        b'I' => TransactionStatus::Idle,
980        b'T' => TransactionStatus::InTransaction,
981        b'E' => TransactionStatus::FailedTransaction,
982        _ => return Err(invalid("unknown transaction status")),
983    };
984    require_empty(&body)?;
985    Ok(BackendMessage::ReadyForQuery(status))
986}
987
988fn decode_negotiate_protocol_version(mut body: Bytes) -> io::Result<BackendMessage> {
989    let newest = take_u32(&mut body)?;
990    let major = u16::try_from(newest >> 16).map_err(|_| invalid("protocol major overflow"))?;
991    let minor = u16::try_from(newest & 0xffff).map_err(|_| invalid("protocol minor overflow"))?;
992    let count = take_u32(&mut body)?;
993    let capacity = usize::try_from(count).map_err(|_| invalid("option count overflow"))?;
994    let mut unsupported_options = Vec::with_capacity(capacity);
995    for _ in 0..count {
996        unsupported_options.push(take_cstr(&mut body)?);
997    }
998    require_empty(&body)?;
999    Ok(BackendMessage::NegotiateProtocolVersion(
1000        NegotiateProtocolVersion {
1001            newest: ProtocolVersion { major, minor },
1002            unsupported_options,
1003        },
1004    ))
1005}
1006
1007fn decode_parse(mut body: Bytes) -> io::Result<Parse> {
1008    let statement = take_cstr(&mut body)?;
1009    let query = take_cstr(&mut body)?;
1010    let count = take_u16(&mut body)?;
1011    let mut parameter_types = Vec::with_capacity(usize::from(count));
1012    for _ in 0..count {
1013        parameter_types.push(take_u32(&mut body)?);
1014    }
1015    require_empty(&body)?;
1016    Ok(Parse {
1017        statement,
1018        query,
1019        parameter_types,
1020    })
1021}
1022
1023fn decode_bind(mut body: Bytes) -> io::Result<Bind> {
1024    let portal = take_cstr(&mut body)?;
1025    let statement = take_cstr(&mut body)?;
1026    let parameter_formats = take_i16_vec(&mut body)?;
1027    let parameter_count = take_u16(&mut body)?;
1028    let mut parameters = Vec::with_capacity(usize::from(parameter_count));
1029    for _ in 0..parameter_count {
1030        let length = take_i32(&mut body)?;
1031        if length == -1 {
1032            parameters.push(None);
1033        } else {
1034            let length =
1035                usize::try_from(length).map_err(|_| invalid("negative parameter length"))?;
1036            require(&body, length)?;
1037            parameters.push(Some(body.split_to(length)));
1038        }
1039    }
1040    let result_formats = take_i16_vec(&mut body)?;
1041    require_empty(&body)?;
1042    Ok(Bind {
1043        portal,
1044        statement,
1045        parameter_formats,
1046        parameters,
1047        result_formats,
1048    })
1049}
1050
1051fn decode_describe(mut body: Bytes) -> io::Result<Describe> {
1052    let target = take_target(&mut body)?;
1053    let name = take_cstr(&mut body)?;
1054    require_empty(&body)?;
1055    Ok(Describe { target, name })
1056}
1057
1058fn decode_data_row(mut body: Bytes) -> io::Result<DataRow> {
1059    let count = take_u16(&mut body)?;
1060    let mut columns = Vec::with_capacity(usize::from(count));
1061    for _ in 0..count {
1062        columns.push(take_nullable(&mut body)?);
1063    }
1064    require_empty(&body)?;
1065    Ok(DataRow { columns })
1066}
1067
1068fn decode_diagnostic(mut body: Bytes) -> io::Result<DiagnosticResponse> {
1069    let mut fields = Vec::new();
1070    loop {
1071        require(&body, 1)?;
1072        let code = body.get_u8();
1073        if code == 0 {
1074            break;
1075        }
1076        fields.push(DiagnosticField {
1077            code,
1078            value: take_cstr(&mut body)?,
1079        });
1080    }
1081    require_empty(&body)?;
1082    Ok(DiagnosticResponse { fields })
1083}
1084
1085fn decode_copy_response(mut body: Bytes) -> io::Result<CopyResponse> {
1086    require(&body, 1)?;
1087    let overall_format = body.get_u8();
1088    let column_formats = take_i16_vec(&mut body)?;
1089    require_empty(&body)?;
1090    Ok(CopyResponse {
1091        overall_format,
1092        column_formats,
1093    })
1094}
1095
1096fn decode_parameter_description(mut body: Bytes) -> io::Result<Vec<u32>> {
1097    let count = take_u16(&mut body)?;
1098    let mut types = Vec::with_capacity(usize::from(count));
1099    for _ in 0..count {
1100        types.push(take_u32(&mut body)?);
1101    }
1102    require_empty(&body)?;
1103    Ok(types)
1104}
1105
1106fn decode_close(mut body: Bytes) -> io::Result<Close> {
1107    let target = take_target(&mut body)?;
1108    let name = take_cstr(&mut body)?;
1109    require_empty(&body)?;
1110    Ok(Close { target, name })
1111}
1112
1113fn decode_execute(mut body: Bytes) -> io::Result<Execute> {
1114    let portal = take_cstr(&mut body)?;
1115    let max_rows = take_i32(&mut body)?;
1116    require_empty(&body)?;
1117    Ok(Execute { portal, max_rows })
1118}
1119
1120fn decode_function_call(mut body: Bytes) -> io::Result<FunctionCall> {
1121    let function_oid = take_u32(&mut body)?;
1122    let argument_formats = take_i16_vec(&mut body)?;
1123    let count = take_u16(&mut body)?;
1124    let mut arguments = Vec::with_capacity(usize::from(count));
1125    for _ in 0..count {
1126        arguments.push(take_nullable(&mut body)?);
1127    }
1128    let result_format = take_i16(&mut body)?;
1129    require_empty(&body)?;
1130    Ok(FunctionCall {
1131        function_oid,
1132        argument_formats,
1133        arguments,
1134        result_format,
1135    })
1136}
1137
1138fn decode_cstr_body(mut body: Bytes) -> io::Result<Bytes> {
1139    let value = take_cstr(&mut body)?;
1140    require_empty(&body)?;
1141    Ok(value)
1142}
1143
1144fn decode_empty(body: &Bytes) -> io::Result<()> {
1145    require_empty(body)
1146}
1147
1148fn take_target(body: &mut Bytes) -> io::Result<DescribeTarget> {
1149    require(body, 1)?;
1150    match body.get_u8() {
1151        b'S' => Ok(DescribeTarget::Statement),
1152        b'P' => Ok(DescribeTarget::Portal),
1153        _ => Err(invalid("invalid statement or portal target")),
1154    }
1155}
1156
1157fn decode_row_description(mut body: Bytes) -> io::Result<RowDescription> {
1158    let count = take_u16(&mut body)?;
1159    let mut fields = Vec::with_capacity(usize::from(count));
1160    for _ in 0..count {
1161        fields.push(FieldDescription {
1162            name: take_cstr(&mut body)?,
1163            table_oid: take_u32(&mut body)?,
1164            column: take_i16(&mut body)?,
1165            type_oid: take_u32(&mut body)?,
1166            type_size: take_i16(&mut body)?,
1167            type_modifier: take_i32(&mut body)?,
1168            format: take_i16(&mut body)?,
1169        });
1170    }
1171    require_empty(&body)?;
1172    Ok(RowDescription { fields })
1173}
1174
1175fn take_i16_vec(body: &mut Bytes) -> io::Result<Vec<i16>> {
1176    let count = take_u16(body)?;
1177    (0..count).map(|_| take_i16(body)).collect()
1178}
1179
1180fn take_nullable(body: &mut Bytes) -> io::Result<Option<Bytes>> {
1181    let length = take_i32(body)?;
1182    if length == -1 {
1183        return Ok(None);
1184    }
1185    let length = usize::try_from(length).map_err(|_| invalid("negative value length"))?;
1186    require(body, length)?;
1187    Ok(Some(body.split_to(length)))
1188}
1189
1190fn put_i16_vec(values: &[i16], body: &mut BytesMut) -> io::Result<()> {
1191    put_count(values.len(), body)?;
1192    for value in values {
1193        body.put_i16(*value);
1194    }
1195    Ok(())
1196}
1197
1198fn put_nullable(value: Option<&Bytes>, body: &mut BytesMut) -> io::Result<()> {
1199    match value {
1200        None => body.put_i32(-1),
1201        Some(value) => {
1202            let length =
1203                i32::try_from(value.len()).map_err(|_| invalid_input("value is too large"))?;
1204            body.put_i32(length);
1205            body.extend_from_slice(value);
1206        }
1207    }
1208    Ok(())
1209}
1210
1211fn put_count(count: usize, body: &mut BytesMut) -> io::Result<()> {
1212    let count =
1213        u16::try_from(count).map_err(|_| invalid_input("message item count exceeds u16"))?;
1214    body.put_u16(count);
1215    Ok(())
1216}
1217
1218fn put_cstr(value: &[u8], body: &mut BytesMut) -> io::Result<()> {
1219    if value.contains(&0) {
1220        return Err(invalid_input("message string contains a NUL byte"));
1221    }
1222    body.extend_from_slice(value);
1223    body.put_u8(0);
1224    Ok(())
1225}
1226
1227fn named_target_frame(tag: u8, target: DescribeTarget, name: &[u8]) -> io::Result<Frame> {
1228    let mut body = BytesMut::new();
1229    body.put_u8(match target {
1230        DescribeTarget::Statement => b'S',
1231        DescribeTarget::Portal => b'P',
1232    });
1233    put_cstr(name, &mut body)?;
1234    Ok(Frame {
1235        tag,
1236        body: body.freeze(),
1237    })
1238}
1239
1240fn cstr_message(tag: u8, value: &[u8]) -> io::Result<Frame> {
1241    let mut body = BytesMut::new();
1242    put_cstr(value, &mut body)?;
1243    Ok(Frame {
1244        tag,
1245        body: body.freeze(),
1246    })
1247}
1248
1249fn empty_message(tag: u8) -> Frame {
1250    Frame {
1251        tag,
1252        body: Bytes::new(),
1253    }
1254}
1255
1256fn take_cstr(body: &mut Bytes) -> io::Result<Bytes> {
1257    let end = body
1258        .iter()
1259        .position(|byte| *byte == 0)
1260        .ok_or_else(|| invalid("unterminated string"))?;
1261    let value = body.split_to(end);
1262    body.advance(1);
1263    Ok(value)
1264}
1265
1266fn take_u16(body: &mut Bytes) -> io::Result<u16> {
1267    require(body, 2)?;
1268    Ok(body.get_u16())
1269}
1270
1271fn take_i16(body: &mut Bytes) -> io::Result<i16> {
1272    require(body, 2)?;
1273    Ok(body.get_i16())
1274}
1275
1276fn take_u32(body: &mut Bytes) -> io::Result<u32> {
1277    require(body, 4)?;
1278    Ok(body.get_u32())
1279}
1280
1281fn take_i32(body: &mut Bytes) -> io::Result<i32> {
1282    require(body, 4)?;
1283    Ok(body.get_i32())
1284}
1285
1286fn require(body: &Bytes, length: usize) -> io::Result<()> {
1287    if body.len() < length {
1288        Err(invalid("truncated message body"))
1289    } else {
1290        Ok(())
1291    }
1292}
1293
1294fn require_empty(body: &Bytes) -> io::Result<()> {
1295    if body.is_empty() {
1296        Ok(())
1297    } else {
1298        Err(invalid("trailing message bytes"))
1299    }
1300}
1301
1302fn invalid(message: &'static str) -> io::Error {
1303    io::Error::new(io::ErrorKind::InvalidData, message)
1304}
1305
1306fn invalid_input(message: &'static str) -> io::Error {
1307    io::Error::new(io::ErrorKind::InvalidInput, message)
1308}
1309
1310fn unknown_tag(direction: &str, tag: u8) -> io::Error {
1311    io::Error::new(
1312        io::ErrorKind::InvalidData,
1313        format!("unknown {direction} message tag 0x{tag:02x}"),
1314    )
1315}
1316
1317#[cfg(test)]
1318mod tests {
1319    use super::*;
1320
1321    #[test]
1322    fn direction_disambiguates_s_tag() {
1323        let frontend_frame = Frame {
1324            tag: b'S',
1325            body: Bytes::new(),
1326        };
1327        assert!(matches!(
1328            Frontend::decode(frontend_frame),
1329            Ok(FrontendMessage::Sync)
1330        ));
1331        assert_eq!(
1332            Backend::decode(Frame {
1333                tag: b'S',
1334                body: Bytes::from_static(b"client_encoding\0UTF8\0"),
1335            })
1336            .expect("valid ParameterStatus"),
1337            BackendMessage::ParameterStatus {
1338                name: Bytes::from_static(b"client_encoding"),
1339                value: Bytes::from_static(b"UTF8"),
1340            }
1341        );
1342    }
1343
1344    #[test]
1345    fn parse_is_losslessly_structured() {
1346        let mut bytes = BytesMut::from(&b"P\0\0\0\x19stmt\0select $1\0\0\x01\0\0\0\x17"[..]);
1347        let original = bytes.clone();
1348        let message = PgCodec::<Frontend>::default()
1349            .decode(&mut bytes)
1350            .expect("valid frame")
1351            .expect("complete frame");
1352        let expected = FrontendMessage::Parse(Parse {
1353            statement: Bytes::from_static(b"stmt"),
1354            query: Bytes::from_static(b"select $1"),
1355            parameter_types: vec![23],
1356        });
1357        assert_eq!(message, expected);
1358
1359        let FrontendMessage::Parse(parsed) = message else {
1360            unreachable!()
1361        };
1362        let frame = parsed.to_frame().expect("reconstructable Parse");
1363        let mut encoded = BytesMut::new();
1364        PgCodec::<Frontend>::default()
1365            .encode(frame, &mut encoded)
1366            .expect("encodable frame");
1367        assert_eq!(encoded, original);
1368    }
1369
1370    #[test]
1371    fn incomplete_frame_does_not_consume_input() {
1372        let mut bytes = BytesMut::from(&b"S\0\0\0\x04"[..4]);
1373        let original = bytes.clone();
1374        assert!(
1375            PgCodec::<Frontend>::default()
1376                .decode(&mut bytes)
1377                .expect("incomplete input is not an error")
1378                .is_none()
1379        );
1380        assert_eq!(bytes, original);
1381    }
1382
1383    #[test]
1384    fn frame_limits_reject_oversized_input_before_allocation() {
1385        let mut codec = PgCodec::<Frontend>::with_max_frame_len(9).unwrap();
1386        let mut oversized = BytesMut::from(&b"Q\0\0\0\x09"[..]);
1387        let error = codec.decode(&mut oversized).unwrap_err();
1388        assert_eq!(error.kind(), io::ErrorKind::InvalidData);
1389        assert_eq!(oversized.len(), 5);
1390
1391        let mut signed_overflow = BytesMut::from(&[b'Q', 0x80, 0, 0, 0][..]);
1392        let error = PgCodec::<Frontend>::default()
1393            .decode(&mut signed_overflow)
1394            .unwrap_err();
1395        assert_eq!(error.kind(), io::ErrorKind::InvalidData);
1396
1397        let error = codec
1398            .encode(
1399                Frame {
1400                    tag: b'Q',
1401                    body: Bytes::from_static(b"12345"),
1402                },
1403                &mut BytesMut::new(),
1404            )
1405            .unwrap_err();
1406        assert_eq!(error.kind(), io::ErrorKind::InvalidInput);
1407    }
1408
1409    #[test]
1410    fn default_frame_limit_is_network_safe_and_configurable() {
1411        let mut input = BytesMut::from(&[b'Q', 1, 0, 0, 4][..]);
1412        let error = PgCodec::<Frontend>::default()
1413            .decode(&mut input)
1414            .unwrap_err();
1415        assert_eq!(error.kind(), io::ErrorKind::InvalidData);
1416        assert!(error.to_string().contains("configured frame limit"));
1417
1418        assert!(PgCodec::<Frontend>::with_max_frame_len(MAX_PROTOCOL_FRAME_LEN).is_ok());
1419    }
1420
1421    #[test]
1422    fn bind_round_trips_nulls_formats_and_values() {
1423        let bind = Bind {
1424            portal: Bytes::from_static(b"portal"),
1425            statement: Bytes::from_static(b"statement"),
1426            parameter_formats: vec![1, 0],
1427            parameters: vec![None, Some(Bytes::from_static(b"value"))],
1428            result_formats: vec![1],
1429        };
1430        let frame = bind.to_frame().expect("valid Bind");
1431        assert_eq!(
1432            Frontend::decode(frame).expect("decodable Bind"),
1433            FrontendMessage::Bind(bind)
1434        );
1435    }
1436
1437    #[test]
1438    fn row_description_round_trips_all_metadata() {
1439        let description = RowDescription {
1440            fields: vec![FieldDescription {
1441                name: Bytes::from_static(b"answer"),
1442                table_oid: 16_384,
1443                column: 2,
1444                type_oid: 23,
1445                type_size: 4,
1446                type_modifier: -1,
1447                format: 1,
1448            }],
1449        };
1450        let frame = description.to_frame().expect("valid RowDescription");
1451        assert_eq!(
1452            Backend::decode(frame).expect("decodable RowDescription"),
1453            BackendMessage::RowDescription(description)
1454        );
1455    }
1456
1457    #[test]
1458    fn frontend_message_family_round_trips_structurally() {
1459        let messages = vec![
1460            FrontendMessage::Close(Close {
1461                target: DescribeTarget::Portal,
1462                name: Bytes::from_static(b"p"),
1463            }),
1464            FrontendMessage::Execute(Execute {
1465                portal: Bytes::from_static(b"p"),
1466                max_rows: 10,
1467            }),
1468            FrontendMessage::FunctionCall(FunctionCall {
1469                function_oid: 42,
1470                argument_formats: vec![1],
1471                arguments: vec![Some(Bytes::from_static(b"arg")), None],
1472                result_format: 1,
1473            }),
1474            FrontendMessage::Query(Bytes::from_static(b"select 1")),
1475            FrontendMessage::Flush,
1476            FrontendMessage::Sync,
1477            FrontendMessage::Terminate,
1478            FrontendMessage::CopyData(Bytes::from_static(b"row\n")),
1479            FrontendMessage::CopyDone,
1480            FrontendMessage::CopyFail(Bytes::from_static(b"cancelled")),
1481            FrontendMessage::PasswordResponse(Bytes::from_static(b"opaque response")),
1482        ];
1483        for message in messages {
1484            let frame = message
1485                .to_frame()
1486                .expect("reconstructable frontend message");
1487            assert_eq!(
1488                Frontend::decode(frame).expect("decodable frontend message"),
1489                message
1490            );
1491        }
1492    }
1493
1494    #[test]
1495    fn backend_message_family_round_trips_structurally() {
1496        let diagnostic = DiagnosticResponse {
1497            fields: vec![
1498                DiagnosticField {
1499                    code: b'S',
1500                    value: Bytes::from_static(b"ERROR"),
1501                },
1502                DiagnosticField {
1503                    code: b'M',
1504                    value: Bytes::from_static(b"rewritable message"),
1505                },
1506            ],
1507        };
1508        let copy = CopyResponse {
1509            overall_format: 0,
1510            column_formats: vec![0, 1],
1511        };
1512        let messages = vec![
1513            BackendMessage::ParseComplete,
1514            BackendMessage::BindComplete,
1515            BackendMessage::CloseComplete,
1516            BackendMessage::CommandComplete(Bytes::from_static(b"SELECT 1")),
1517            BackendMessage::CopyData(Bytes::from_static(b"row\n")),
1518            BackendMessage::CopyDone,
1519            BackendMessage::CopyInResponse(copy.clone()),
1520            BackendMessage::CopyOutResponse(copy.clone()),
1521            BackendMessage::CopyBothResponse(copy),
1522            BackendMessage::DataRow(DataRow {
1523                columns: vec![Some(Bytes::from_static(b"42")), None],
1524            }),
1525            BackendMessage::EmptyQueryResponse,
1526            BackendMessage::ErrorResponse(diagnostic.clone()),
1527            BackendMessage::NoticeResponse(diagnostic),
1528            BackendMessage::NoData,
1529            BackendMessage::ParameterDescription(vec![23, 25]),
1530            BackendMessage::PortalSuspended,
1531            BackendMessage::FunctionCallResponse(Bytes::from_static(b"result")),
1532        ];
1533        for message in messages {
1534            let frame = message.to_frame().expect("reconstructable backend message");
1535            assert_eq!(
1536                Backend::decode(frame).expect("decodable backend message"),
1537                message
1538            );
1539        }
1540    }
1541}
pg_proto/codec.rs

pg_proto/
codec.rs
