uqa_pg_wire/
protocol.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7use thiserror::Error;
8
9pub const PROTOCOL_VERSION_3_0: i32 = 196_608;
10pub const PROTOCOL_VERSION_3_2: i32 = 196_610;
11pub const CANCEL_REQUEST_CODE: i32 = 80_877_102;
12pub const SSL_REQUEST_CODE: i32 = 80_877_103;
13pub const GSSENC_REQUEST_CODE: i32 = 80_877_104;
14/// Minimum cancellation key accepted in a `CancelRequest` packet.
15pub const MIN_CANCEL_REQUEST_KEY_LEN: usize = 1;
16/// Minimum cancellation key emitted in `BackendKeyData` under protocol 3.2.
17pub const MIN_BACKEND_KEY_DATA_KEY_LEN: usize = 4;
18pub const MAX_CANCEL_KEY_LEN: usize = 256;
19
20pub type DecodeOutcome<T> = Result<Option<(T, usize)>, PgWireError>;
21
22#[derive(Debug, Error, PartialEq, Eq)]
23pub enum PgWireError {
24    #[error("invalid PostgreSQL wire message length {length}; minimum is {minimum}")]
25    InvalidLength { length: i32, minimum: i32 },
26    #[error("PostgreSQL wire message length {length} exceeds configured maximum {maximum}")]
27    MessageTooLarge { length: i32, maximum: usize },
28    #[error("invalid UTF-8 in {context}")]
29    InvalidUtf8 { context: &'static str },
30    #[error("missing nul terminator in {context}")]
31    MissingNul { context: &'static str },
32    #[error("trailing bytes in {context}: {remaining}")]
33    TrailingBytes {
34        context: &'static str,
35        remaining: usize,
36    },
37    #[error("unexpected end of {context}")]
38    UnexpectedEof { context: &'static str },
39    #[error("unsupported PostgreSQL protocol version {0}")]
40    UnsupportedProtocolVersion(i32),
41    #[error(
42        "invalid PostgreSQL cancellation key length {length}; expected {minimum} through {maximum} bytes"
43    )]
44    InvalidCancelKeyLength {
45        length: usize,
46        minimum: usize,
47        maximum: usize,
48    },
49    #[error(
50        "PostgreSQL protocol {major}.{minor} requires a 4-byte cancellation key, got {length} bytes"
51    )]
52    CancelKeyLengthForProtocol {
53        length: usize,
54        major: u16,
55        minor: u16,
56    },
57    #[error("unknown frontend message tag {0:?}")]
58    UnknownFrontendTag(u8),
59    #[error("invalid format code {0}")]
60    InvalidFormatCode(i16),
61    #[error("invalid transaction status byte {0:?}")]
62    InvalidTransactionStatus(u8),
63    #[error("embedded nul byte in {context}")]
64    EmbeddedNul { context: &'static str },
65    #[error("SASL mechanism names cannot be empty")]
66    EmptySaslMechanism,
67    #[error("AuthenticationSASL must advertise at least one mechanism")]
68    EmptySaslMechanismList,
69    #[error("invalid authentication sequence: cannot process {message} while {state}")]
70    InvalidAuthenticationSequence {
71        state: &'static str,
72        message: &'static str,
73    },
74    #[error("invalid SQLSTATE {code:?}; expected exactly five ASCII letters or digits")]
75    InvalidSqlState { code: String },
76    #[error(
77        "Bind parameter format count {format_count} must be zero, one, or match parameter count {parameter_count}"
78    )]
79    ParameterFormatCountMismatch {
80        format_count: usize,
81        parameter_count: usize,
82    },
83    #[error(
84        "FunctionCall argument format count {format_count} must be zero, one, or match argument count {argument_count}"
85    )]
86    FunctionArgumentFormatCountMismatch {
87        format_count: usize,
88        argument_count: usize,
89    },
90    #[error(
91        "Bind result format count {format_count} must be zero, one, or match result column count {column_count}"
92    )]
93    ResultFormatCountMismatch {
94        format_count: usize,
95        column_count: usize,
96    },
97    #[error("{context} index {index} is out of range for {count} value(s)")]
98    FormatIndexOutOfRange {
99        context: &'static str,
100        index: usize,
101        count: usize,
102    },
103    #[error(
104        "cannot remove a {layer_length}-byte middleware cancellation-key prefix from a {key_length}-byte key"
105    )]
106    InvalidCancelKeyLayerLength {
107        layer_length: usize,
108        key_length: usize,
109    },
110    #[error("text COPY response column {column} uses the binary format")]
111    BinaryColumnInTextCopy { column: usize },
112    #[error("{context} count {count} exceeds representable PostgreSQL i16")]
113    CountTooLarge { context: &'static str, count: usize },
114    #[error("{context} length {length} exceeds representable PostgreSQL i32")]
115    LengthTooLarge {
116        context: &'static str,
117        length: usize,
118    },
119    #[error("{context} cannot be negative")]
120    NegativeValue { context: &'static str },
121}
122
123pub type DecodeError = PgWireError;
124
125#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
126pub struct ProtocolVersion {
127    pub major: u16,
128    pub minor: u16,
129}
130
131impl ProtocolVersion {
132    pub const V3_0: Self = Self { major: 3, minor: 0 };
133    pub const V3_2: Self = Self { major: 3, minor: 2 };
134    pub const LATEST: Self = Self::V3_2;
135
136    pub fn from_raw(raw: i32) -> Self {
137        Self {
138            major: ((raw >> 16) & 0xffff) as u16,
139            minor: (raw & 0xffff) as u16,
140        }
141    }
142
143    pub const fn raw(self) -> i32 {
144        i32::from_be_bytes([
145            (self.major >> 8) as u8,
146            self.major as u8,
147            (self.minor >> 8) as u8,
148            self.minor as u8,
149        ])
150    }
151
152    /// Select the newest protocol version this crate supports without
153    /// negotiating to a version newer than the frontend requested.
154    pub fn negotiate(self) -> Result<Self, PgWireError> {
155        self.negotiate_with_max(Self::LATEST)
156    }
157
158    /// Select a protocol version no newer than either the frontend request or
159    /// the newest version implemented by the embedding server.
160    pub fn negotiate_with_max(self, newest_supported: Self) -> Result<Self, PgWireError> {
161        if self.major != Self::LATEST.major {
162            return Err(PgWireError::UnsupportedProtocolVersion(self.raw()));
163        }
164        if !newest_supported.is_supported_server_max() {
165            return Err(PgWireError::UnsupportedProtocolVersion(
166                newest_supported.raw(),
167            ));
168        }
169        Ok(Self {
170            major: self.major,
171            minor: self.minor.min(newest_supported.minor),
172        })
173    }
174
175    /// Versions an embedding server may configure as its implementation
176    /// maximum. `PostgreSQL` 18 has implementations for 3.0 and 3.2; a 3.1
177    /// frontend request can remain selected, but 3.1 is not a server maximum.
178    #[must_use]
179    pub const fn is_supported_server_max(self) -> bool {
180        matches!(self, Self::V3_0 | Self::V3_2)
181    }
182}
183
184/// Opaque cancellation secret carried by `BackendKeyData` and
185/// `CancelRequest`.
186///
187/// `PostgreSQL` 18 accepts 1 through 256 bytes when decoding a cancel request.
188/// A backend key has the stricter 4 through 256 byte range in protocol 3.2,
189/// and is exactly 4 bytes before protocol 3.2.
190#[derive(Debug, Clone, PartialEq, Eq)]
191pub struct CancelKey(Vec<u8>);
192
193impl CancelKey {
194    pub fn new(bytes: impl Into<Vec<u8>>) -> Result<Self, PgWireError> {
195        let bytes = bytes.into();
196        if !(MIN_CANCEL_REQUEST_KEY_LEN..=MAX_CANCEL_KEY_LEN).contains(&bytes.len()) {
197            return Err(PgWireError::InvalidCancelKeyLength {
198                length: bytes.len(),
199                minimum: MIN_CANCEL_REQUEST_KEY_LEN,
200                maximum: MAX_CANCEL_KEY_LEN,
201            });
202        }
203        Ok(Self(bytes))
204    }
205
206    #[must_use]
207    pub fn from_i32(secret_key: i32) -> Self {
208        Self(secret_key.to_be_bytes().to_vec())
209    }
210
211    #[must_use]
212    pub fn as_bytes(&self) -> &[u8] {
213        &self.0
214    }
215
216    #[must_use]
217    pub fn into_bytes(self) -> Vec<u8> {
218        self.0
219    }
220
221    /// Prefix middleware-owned routing data while retaining the complete
222    /// downstream cancellation secret.
223    ///
224    /// `PostgreSQL` does not prescribe a framing format for middleware data.
225    /// Each layer therefore owns the prefix length it adds and passes that
226    /// same length to [`Self::remove_middleware_prefix`] on the return path.
227    pub fn with_middleware_prefix(&self, prefix: &[u8]) -> Result<Self, PgWireError> {
228        let length = prefix
229            .len()
230            .checked_add(self.0.len())
231            .ok_or(PgWireError::LengthTooLarge {
232                context: "middleware cancellation key",
233                length: prefix.len(),
234            })?;
235        if length > MAX_CANCEL_KEY_LEN {
236            return Err(PgWireError::InvalidCancelKeyLength {
237                length,
238                minimum: MIN_CANCEL_REQUEST_KEY_LEN,
239                maximum: MAX_CANCEL_KEY_LEN,
240            });
241        }
242        let mut bytes = Vec::with_capacity(length);
243        bytes.extend_from_slice(prefix);
244        bytes.extend_from_slice(&self.0);
245        Self::new(bytes)
246    }
247
248    /// Remove one middleware prefix and return both the layer data and the
249    /// downstream cancellation secret.
250    pub fn remove_middleware_prefix(
251        &self,
252        prefix_length: usize,
253    ) -> Result<(Vec<u8>, Self), PgWireError> {
254        if prefix_length >= self.0.len() {
255            return Err(PgWireError::InvalidCancelKeyLayerLength {
256                layer_length: prefix_length,
257                key_length: self.0.len(),
258            });
259        }
260        let (prefix, downstream) = self.0.split_at(prefix_length);
261        Ok((prefix.to_vec(), Self::new(downstream.to_vec())?))
262    }
263
264    pub fn validate_for_backend_key_data(
265        &self,
266        version: ProtocolVersion,
267    ) -> Result<(), PgWireError> {
268        let negotiated = version.negotiate()?;
269        if negotiated < ProtocolVersion::V3_2 && self.0.len() != MIN_BACKEND_KEY_DATA_KEY_LEN {
270            return Err(PgWireError::CancelKeyLengthForProtocol {
271                length: self.0.len(),
272                major: negotiated.major,
273                minor: negotiated.minor,
274            });
275        }
276        if self.0.len() < MIN_BACKEND_KEY_DATA_KEY_LEN {
277            return Err(PgWireError::InvalidCancelKeyLength {
278                length: self.0.len(),
279                minimum: MIN_BACKEND_KEY_DATA_KEY_LEN,
280                maximum: MAX_CANCEL_KEY_LEN,
281            });
282        }
283        Ok(())
284    }
285}
286
287impl From<i32> for CancelKey {
288    fn from(secret_key: i32) -> Self {
289        Self::from_i32(secret_key)
290    }
291}
292
293impl AsRef<[u8]> for CancelKey {
294    fn as_ref(&self) -> &[u8] {
295        self.as_bytes()
296    }
297}
298
299#[derive(Debug, Clone, Copy, PartialEq, Eq)]
300pub enum FormatCode {
301    Text,
302    Binary,
303}
304
305impl FormatCode {
306    pub fn from_i16(value: i16) -> Result<Self, PgWireError> {
307        match value {
308            0 => Ok(Self::Text),
309            1 => Ok(Self::Binary),
310            other => Err(PgWireError::InvalidFormatCode(other)),
311        }
312    }
313
314    pub const fn as_i16(self) -> i16 {
315        match self {
316            Self::Text => 0,
317            Self::Binary => 1,
318        }
319    }
320}
321
322pub(crate) fn resolve_format_codes(
323    formats: &[FormatCode],
324    value_count: usize,
325    mismatch: impl FnOnce(usize, usize) -> PgWireError,
326) -> Result<Vec<FormatCode>, PgWireError> {
327    match formats {
328        [] => Ok(vec![FormatCode::Text; value_count]),
329        [format] => Ok(vec![*format; value_count]),
330        formats if formats.len() == value_count => Ok(formats.to_vec()),
331        formats => Err(mismatch(formats.len(), value_count)),
332    }
333}
334
335pub(crate) fn resolve_format_code(
336    formats: &[FormatCode],
337    value_count: usize,
338    index: usize,
339    context: &'static str,
340    mismatch: impl FnOnce(usize, usize) -> PgWireError,
341) -> Result<FormatCode, PgWireError> {
342    if index >= value_count {
343        return Err(PgWireError::FormatIndexOutOfRange {
344            context,
345            index,
346            count: value_count,
347        });
348    }
349    match formats {
350        [] => Ok(FormatCode::Text),
351        [format] => Ok(*format),
352        formats if formats.len() == value_count => Ok(formats[index]),
353        formats => Err(mismatch(formats.len(), value_count)),
354    }
355}
356
357#[derive(Debug, Clone, Copy, PartialEq, Eq)]
358pub enum TransactionStatus {
359    Idle,
360    InTransaction,
361    Failed,
362}
363
364impl TransactionStatus {
365    pub fn from_byte(value: u8) -> Result<Self, PgWireError> {
366        match value {
367            b'I' => Ok(Self::Idle),
368            b'T' => Ok(Self::InTransaction),
369            b'E' => Ok(Self::Failed),
370            other => Err(PgWireError::InvalidTransactionStatus(other)),
371        }
372    }
373
374    pub const fn as_byte(self) -> u8 {
375        match self {
376            Self::Idle => b'I',
377            Self::InTransaction => b'T',
378            Self::Failed => b'E',
379        }
380    }
381}
uqa_pg_wire/protocol.rs

uqa_pg_wire/
protocol.rs
