Skip to main content

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("invalid SQLSTATE {code:?}; expected exactly five ASCII letters or digits")]
66    InvalidSqlState { code: String },
67    #[error(
68        "Bind parameter format count {format_count} must be zero, one, or match parameter count {parameter_count}"
69    )]
70    ParameterFormatCountMismatch {
71        format_count: usize,
72        parameter_count: usize,
73    },
74    #[error(
75        "FunctionCall argument format count {format_count} must be zero, one, or match argument count {argument_count}"
76    )]
77    FunctionArgumentFormatCountMismatch {
78        format_count: usize,
79        argument_count: usize,
80    },
81    #[error("text COPY response column {column} uses the binary format")]
82    BinaryColumnInTextCopy { column: usize },
83    #[error("{context} count {count} exceeds representable PostgreSQL i16")]
84    CountTooLarge { context: &'static str, count: usize },
85    #[error("{context} length {length} exceeds representable PostgreSQL i32")]
86    LengthTooLarge {
87        context: &'static str,
88        length: usize,
89    },
90    #[error("{context} cannot be negative")]
91    NegativeValue { context: &'static str },
92}
93
94pub type DecodeError = PgWireError;
95
96#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
97pub struct ProtocolVersion {
98    pub major: u16,
99    pub minor: u16,
100}
101
102impl ProtocolVersion {
103    pub const V3_0: Self = Self { major: 3, minor: 0 };
104    pub const V3_2: Self = Self { major: 3, minor: 2 };
105    pub const LATEST: Self = Self::V3_2;
106
107    pub fn from_raw(raw: i32) -> Self {
108        Self {
109            major: ((raw >> 16) & 0xffff) as u16,
110            minor: (raw & 0xffff) as u16,
111        }
112    }
113
114    pub const fn raw(self) -> i32 {
115        i32::from_be_bytes([
116            (self.major >> 8) as u8,
117            self.major as u8,
118            (self.minor >> 8) as u8,
119            self.minor as u8,
120        ])
121    }
122
123    /// Select the newest protocol version this crate supports without
124    /// negotiating to a version newer than the frontend requested.
125    pub fn negotiate(self) -> Result<Self, PgWireError> {
126        self.negotiate_with_max(Self::LATEST)
127    }
128
129    /// Select a protocol version no newer than either the frontend request or
130    /// the newest version implemented by the embedding server.
131    pub fn negotiate_with_max(self, newest_supported: Self) -> Result<Self, PgWireError> {
132        if self.major != Self::LATEST.major {
133            return Err(PgWireError::UnsupportedProtocolVersion(self.raw()));
134        }
135        if !newest_supported.is_supported_server_max() {
136            return Err(PgWireError::UnsupportedProtocolVersion(
137                newest_supported.raw(),
138            ));
139        }
140        Ok(Self {
141            major: self.major,
142            minor: self.minor.min(newest_supported.minor),
143        })
144    }
145
146    /// Versions an embedding server may configure as its implementation
147    /// maximum. `PostgreSQL` 18 has implementations for 3.0 and 3.2; a 3.1
148    /// frontend request can remain selected, but 3.1 is not a server maximum.
149    #[must_use]
150    pub const fn is_supported_server_max(self) -> bool {
151        matches!(self, Self::V3_0 | Self::V3_2)
152    }
153}
154
155/// Opaque cancellation secret carried by `BackendKeyData` and
156/// `CancelRequest`.
157///
158/// `PostgreSQL` 18 accepts 1 through 256 bytes when decoding a cancel request.
159/// A backend key has the stricter 4 through 256 byte range in protocol 3.2,
160/// and is exactly 4 bytes before protocol 3.2.
161#[derive(Debug, Clone, PartialEq, Eq)]
162pub struct CancelKey(Vec<u8>);
163
164impl CancelKey {
165    pub fn new(bytes: impl Into<Vec<u8>>) -> Result<Self, PgWireError> {
166        let bytes = bytes.into();
167        if !(MIN_CANCEL_REQUEST_KEY_LEN..=MAX_CANCEL_KEY_LEN).contains(&bytes.len()) {
168            return Err(PgWireError::InvalidCancelKeyLength {
169                length: bytes.len(),
170                minimum: MIN_CANCEL_REQUEST_KEY_LEN,
171                maximum: MAX_CANCEL_KEY_LEN,
172            });
173        }
174        Ok(Self(bytes))
175    }
176
177    #[must_use]
178    pub fn from_i32(secret_key: i32) -> Self {
179        Self(secret_key.to_be_bytes().to_vec())
180    }
181
182    #[must_use]
183    pub fn as_bytes(&self) -> &[u8] {
184        &self.0
185    }
186
187    #[must_use]
188    pub fn into_bytes(self) -> Vec<u8> {
189        self.0
190    }
191
192    pub fn validate_for_backend_key_data(
193        &self,
194        version: ProtocolVersion,
195    ) -> Result<(), PgWireError> {
196        let negotiated = version.negotiate()?;
197        if negotiated < ProtocolVersion::V3_2 && self.0.len() != MIN_BACKEND_KEY_DATA_KEY_LEN {
198            return Err(PgWireError::CancelKeyLengthForProtocol {
199                length: self.0.len(),
200                major: negotiated.major,
201                minor: negotiated.minor,
202            });
203        }
204        if self.0.len() < MIN_BACKEND_KEY_DATA_KEY_LEN {
205            return Err(PgWireError::InvalidCancelKeyLength {
206                length: self.0.len(),
207                minimum: MIN_BACKEND_KEY_DATA_KEY_LEN,
208                maximum: MAX_CANCEL_KEY_LEN,
209            });
210        }
211        Ok(())
212    }
213}
214
215impl From<i32> for CancelKey {
216    fn from(secret_key: i32) -> Self {
217        Self::from_i32(secret_key)
218    }
219}
220
221impl AsRef<[u8]> for CancelKey {
222    fn as_ref(&self) -> &[u8] {
223        self.as_bytes()
224    }
225}
226
227#[derive(Debug, Clone, Copy, PartialEq, Eq)]
228pub enum FormatCode {
229    Text,
230    Binary,
231}
232
233impl FormatCode {
234    pub fn from_i16(value: i16) -> Result<Self, PgWireError> {
235        match value {
236            0 => Ok(Self::Text),
237            1 => Ok(Self::Binary),
238            other => Err(PgWireError::InvalidFormatCode(other)),
239        }
240    }
241
242    pub const fn as_i16(self) -> i16 {
243        match self {
244            Self::Text => 0,
245            Self::Binary => 1,
246        }
247    }
248}
249
250#[derive(Debug, Clone, Copy, PartialEq, Eq)]
251pub enum TransactionStatus {
252    Idle,
253    InTransaction,
254    Failed,
255}
256
257impl TransactionStatus {
258    pub fn from_byte(value: u8) -> Result<Self, PgWireError> {
259        match value {
260            b'I' => Ok(Self::Idle),
261            b'T' => Ok(Self::InTransaction),
262            b'E' => Ok(Self::Failed),
263            other => Err(PgWireError::InvalidTransactionStatus(other)),
264        }
265    }
266
267    pub const fn as_byte(self) -> u8 {
268        match self {
269            Self::Idle => b'I',
270            Self::InTransaction => b'T',
271            Self::Failed => b'E',
272        }
273    }
274}