uqa_pg_wire/
frontend.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7use std::collections::BTreeMap;
8
9use crate::codec::{message_total_len, DecodeLen, Reader, MESSAGE_HEADER_LEN};
10use crate::protocol::{
11    resolve_format_code, resolve_format_codes, CancelKey, DecodeOutcome, FormatCode, PgWireError,
12    ProtocolVersion, CANCEL_REQUEST_CODE, GSSENC_REQUEST_CODE, SSL_REQUEST_CODE,
13};
14
15pub const DEFAULT_MAX_MESSAGE_LEN: usize = 16 * 1024 * 1024;
16
17#[derive(Debug, Clone, PartialEq, Eq)]
18pub enum StartupFrame {
19    Startup(StartupMessage),
20    CancelRequest {
21        process_id: i32,
22        secret_key: CancelKey,
23    },
24    SSLRequest,
25    GSSEncRequest,
26}
27
28#[derive(Debug, Clone, PartialEq, Eq)]
29pub struct StartupMessage {
30    pub version: ProtocolVersion,
31    pub parameters: BTreeMap<String, String>,
32    /// Startup parameter pairs in wire order, including duplicate names.
33    pub parameter_pairs: Vec<(String, String)>,
34}
35
36#[derive(Debug, Clone, PartialEq, Eq)]
37pub struct StartupNegotiation {
38    pub requested_version: ProtocolVersion,
39    pub negotiated_version: ProtocolVersion,
40    pub unrecognized_options: Vec<String>,
41}
42
43impl StartupNegotiation {
44    #[must_use]
45    pub fn requires_response(&self) -> bool {
46        self.requested_version != self.negotiated_version || !self.unrecognized_options.is_empty()
47    }
48
49    #[must_use]
50    pub fn response(&self) -> Option<crate::backend::BackendMessage> {
51        self.requires_response().then(
52            || crate::backend::BackendMessage::NegotiateProtocolVersion {
53                newest_protocol_version: self.negotiated_version,
54                unrecognized_options: self.unrecognized_options.clone(),
55            },
56        )
57    }
58}
59
60impl StartupMessage {
61    pub fn get(&self, key: &str) -> Option<&str> {
62        self.parameters.get(key).map(String::as_str)
63    }
64
65    pub fn user(&self) -> Option<&str> {
66        self.get("user")
67    }
68
69    pub fn database(&self) -> Option<&str> {
70        self.get("database")
71    }
72
73    pub fn application_name(&self) -> Option<&str> {
74        self.get("application_name")
75    }
76
77    /// Negotiate a `PostgreSQL` 3.x minor version and report every `_pq_.`
78    /// startup option the embedding server has not implemented.
79    pub fn negotiate(
80        &self,
81        supported_protocol_options: &[&str],
82    ) -> Result<StartupNegotiation, PgWireError> {
83        self.negotiate_with_max(ProtocolVersion::LATEST, supported_protocol_options)
84    }
85
86    /// Negotiate against the newest protocol version implemented by the
87    /// embedding server.
88    pub fn negotiate_with_max(
89        &self,
90        newest_supported: ProtocolVersion,
91        supported_protocol_options: &[&str],
92    ) -> Result<StartupNegotiation, PgWireError> {
93        let negotiated_version = self.version.negotiate_with_max(newest_supported)?;
94        let unrecognized_options = self
95            .parameter_pairs
96            .iter()
97            .map(|(name, _)| name)
98            .filter(|name| {
99                name.starts_with("_pq_.") && !supported_protocol_options.contains(&name.as_str())
100            })
101            .cloned()
102            .collect();
103        Ok(StartupNegotiation {
104            requested_version: self.version,
105            negotiated_version,
106            unrecognized_options,
107        })
108    }
109}
110
111#[derive(Debug, Clone, PartialEq, Eq)]
112pub enum FrontendMessage {
113    Query(String),
114    Parse(Parse),
115    Bind(Bind),
116    Describe(DescribeTarget),
117    Execute(Execute),
118    Close(CloseTarget),
119    Flush,
120    Sync,
121    Terminate,
122    Password(PasswordMessage),
123    CopyData(Vec<u8>),
124    CopyDone,
125    CopyFail(String),
126    FunctionCall(FunctionCall),
127}
128
129#[derive(Debug, Clone, PartialEq, Eq)]
130pub struct Parse {
131    pub statement: String,
132    pub query: String,
133    pub parameter_type_oids: Vec<u32>,
134}
135
136#[derive(Debug, Clone, PartialEq, Eq)]
137pub struct Bind {
138    pub portal: String,
139    pub statement: String,
140    pub parameter_formats: Vec<FormatCode>,
141    pub parameters: Vec<Option<Vec<u8>>>,
142    pub result_formats: Vec<FormatCode>,
143}
144
145impl Bind {
146    /// Expand `PostgreSQL`'s zero, one, or one-per-parameter format-code forms.
147    pub fn resolved_parameter_formats(&self) -> Result<Vec<FormatCode>, PgWireError> {
148        resolve_format_codes(
149            &self.parameter_formats,
150            self.parameters.len(),
151            |format_count, parameter_count| PgWireError::ParameterFormatCountMismatch {
152                format_count,
153                parameter_count,
154            },
155        )
156    }
157
158    /// Resolve the wire format for one bound parameter.
159    pub fn parameter_format(&self, index: usize) -> Result<FormatCode, PgWireError> {
160        resolve_format_code(
161            &self.parameter_formats,
162            self.parameters.len(),
163            index,
164            "Bind parameter",
165            |format_count, parameter_count| PgWireError::ParameterFormatCountMismatch {
166                format_count,
167                parameter_count,
168            },
169        )
170    }
171
172    /// Expand `PostgreSQL`'s zero, one, or one-per-column result format forms.
173    pub fn resolved_result_formats(
174        &self,
175        column_count: usize,
176    ) -> Result<Vec<FormatCode>, PgWireError> {
177        resolve_format_codes(
178            &self.result_formats,
179            column_count,
180            |format_count, column_count| PgWireError::ResultFormatCountMismatch {
181                format_count,
182                column_count,
183            },
184        )
185    }
186
187    /// Resolve the requested wire format for one result column.
188    pub fn result_format(
189        &self,
190        index: usize,
191        column_count: usize,
192    ) -> Result<FormatCode, PgWireError> {
193        resolve_format_code(
194            &self.result_formats,
195            column_count,
196            index,
197            "Bind result column",
198            |format_count, column_count| PgWireError::ResultFormatCountMismatch {
199                format_count,
200                column_count,
201            },
202        )
203    }
204}
205
206/// The body of the context-dependent frontend message tagged `p`.
207#[derive(Debug, Clone, PartialEq, Eq)]
208pub struct PasswordMessage(Vec<u8>);
209
210impl PasswordMessage {
211    #[must_use]
212    pub fn new(bytes: impl Into<Vec<u8>>) -> Self {
213        Self(bytes.into())
214    }
215
216    #[must_use]
217    pub fn as_bytes(&self) -> &[u8] {
218        &self.0
219    }
220
221    #[must_use]
222    pub fn into_bytes(self) -> Vec<u8> {
223        self.0
224    }
225}
226
227impl AsRef<[u8]> for PasswordMessage {
228    fn as_ref(&self) -> &[u8] {
229        self.as_bytes()
230    }
231}
232
233#[derive(Debug, Clone, PartialEq, Eq)]
234pub enum DescribeTarget {
235    Statement(String),
236    Portal(String),
237}
238
239#[derive(Debug, Clone, PartialEq, Eq)]
240pub struct Execute {
241    pub portal: String,
242    pub max_rows: i32,
243}
244
245#[derive(Debug, Clone, PartialEq, Eq)]
246pub enum CloseTarget {
247    Statement(String),
248    Portal(String),
249}
250
251#[derive(Debug, Clone, PartialEq, Eq)]
252pub struct FunctionCall {
253    pub function_oid: u32,
254    pub argument_formats: Vec<FormatCode>,
255    pub arguments: Vec<Option<Vec<u8>>>,
256    pub result_format: FormatCode,
257}
258
259impl FunctionCall {
260    /// Expand `PostgreSQL`'s zero, one, or one-per-argument format-code forms.
261    pub fn resolved_argument_formats(&self) -> Result<Vec<FormatCode>, PgWireError> {
262        resolve_format_codes(
263            &self.argument_formats,
264            self.arguments.len(),
265            |format_count, argument_count| PgWireError::FunctionArgumentFormatCountMismatch {
266                format_count,
267                argument_count,
268            },
269        )
270    }
271
272    /// Resolve the wire format for one function-call argument.
273    pub fn argument_format(&self, index: usize) -> Result<FormatCode, PgWireError> {
274        resolve_format_code(
275            &self.argument_formats,
276            self.arguments.len(),
277            index,
278            "FunctionCall argument",
279            |format_count, argument_count| PgWireError::FunctionArgumentFormatCountMismatch {
280                format_count,
281                argument_count,
282            },
283        )
284    }
285}
286
287pub fn decode_startup(input: &[u8]) -> DecodeOutcome<StartupFrame> {
288    decode_startup_with_max(input, DEFAULT_MAX_MESSAGE_LEN)
289}
290
291pub fn decode_startup_with_max(input: &[u8], max_len: usize) -> DecodeOutcome<StartupFrame> {
292    let total = match message_total_len(input, false, max_len) {
293        DecodeLen::Complete(total) => total,
294        DecodeLen::Incomplete => return Ok(None),
295        DecodeLen::Error(error) => return Err(error),
296    };
297
298    let body = &input[4..total];
299    let mut reader = Reader::new(body);
300    let version_or_code = reader.read_i32("startup version")?;
301    let frame = match version_or_code {
302        SSL_REQUEST_CODE => {
303            reader.ensure_empty("SSL request")?;
304            StartupFrame::SSLRequest
305        }
306        GSSENC_REQUEST_CODE => {
307            reader.ensure_empty("GSSENC request")?;
308            StartupFrame::GSSEncRequest
309        }
310        CANCEL_REQUEST_CODE => {
311            let process_id = reader.read_i32("cancel request process id")?;
312            let key_length = reader.remaining();
313            let secret_key = CancelKey::new(
314                reader
315                    .read_exact(key_length, "cancel request secret key")?
316                    .to_vec(),
317            )?;
318            reader.ensure_empty("cancel request")?;
319            StartupFrame::CancelRequest {
320                process_id,
321                secret_key,
322            }
323        }
324        other => {
325            let version = ProtocolVersion::from_raw(other);
326            version.negotiate()?;
327            StartupFrame::Startup(parse_startup_message(version, reader)?)
328        }
329    };
330    Ok(Some((frame, total)))
331}
332
333pub fn decode_frontend(input: &[u8]) -> DecodeOutcome<FrontendMessage> {
334    decode_frontend_with_max(input, DEFAULT_MAX_MESSAGE_LEN)
335}
336
337pub fn decode_frontend_with_max(input: &[u8], max_len: usize) -> DecodeOutcome<FrontendMessage> {
338    let total = match message_total_len(input, true, max_len) {
339        DecodeLen::Complete(total) => total,
340        DecodeLen::Incomplete => return Ok(None),
341        DecodeLen::Error(error) => return Err(error),
342    };
343
344    let tag = input[0];
345    let body = &input[MESSAGE_HEADER_LEN..total];
346    let message = parse_frontend_message(tag, body)?;
347    Ok(Some((message, total)))
348}
349
350fn parse_startup_message(
351    version: ProtocolVersion,
352    mut reader: Reader<'_>,
353) -> Result<StartupMessage, PgWireError> {
354    let mut parameters = BTreeMap::new();
355    let mut parameter_pairs = Vec::new();
356    loop {
357        if reader.remaining() == 0 {
358            return Err(PgWireError::MissingNul {
359                context: "startup parameters",
360            });
361        }
362        if reader.remaining() == 1 {
363            let terminator = reader.read_byte("startup terminator")?;
364            if terminator == 0 {
365                break;
366            }
367            return Err(PgWireError::MissingNul {
368                context: "startup parameters",
369            });
370        }
371
372        let key = reader.read_cstring("startup parameter key")?;
373        if key.is_empty() {
374            reader.ensure_empty("startup parameters")?;
375            break;
376        }
377        let value = reader.read_cstring("startup parameter value")?;
378        parameter_pairs.push((key.clone(), value.clone()));
379        parameters.insert(key, value);
380    }
381    Ok(StartupMessage {
382        version,
383        parameters,
384        parameter_pairs,
385    })
386}
387
388fn parse_frontend_message(tag: u8, body: &[u8]) -> Result<FrontendMessage, PgWireError> {
389    let mut reader = Reader::new(body);
390    let message = match tag {
391        b'Q' => FrontendMessage::Query(parse_single_cstring(&mut reader, "Query")?),
392        b'P' => FrontendMessage::Parse(parse_parse(&mut reader)?),
393        b'B' => FrontendMessage::Bind(parse_bind(&mut reader)?),
394        b'D' => FrontendMessage::Describe(parse_describe(&mut reader)?),
395        b'E' => FrontendMessage::Execute(parse_execute(&mut reader)?),
396        b'C' => FrontendMessage::Close(parse_close(&mut reader)?),
397        b'H' => {
398            reader.ensure_empty("Flush")?;
399            FrontendMessage::Flush
400        }
401        b'S' => {
402            reader.ensure_empty("Sync")?;
403            FrontendMessage::Sync
404        }
405        b'X' => {
406            reader.ensure_empty("Terminate")?;
407            FrontendMessage::Terminate
408        }
409        b'p' => FrontendMessage::Password(PasswordMessage::new(body)),
410        b'd' => FrontendMessage::CopyData(body.to_vec()),
411        b'c' => {
412            reader.ensure_empty("CopyDone")?;
413            FrontendMessage::CopyDone
414        }
415        b'f' => FrontendMessage::CopyFail(parse_single_cstring(&mut reader, "CopyFail")?),
416        b'F' => FrontendMessage::FunctionCall(parse_function_call(&mut reader)?),
417        other => return Err(PgWireError::UnknownFrontendTag(other)),
418    };
419    Ok(message)
420}
421
422fn parse_parse(reader: &mut Reader<'_>) -> Result<Parse, PgWireError> {
423    let statement = reader.read_cstring("Parse statement name")?;
424    let query = reader.read_cstring("Parse query")?;
425    let count = read_count(reader, "Parse parameter type count")?;
426    let mut parameter_type_oids = Vec::with_capacity(count);
427    for _ in 0..count {
428        parameter_type_oids.push(reader.read_u32("Parse parameter type oid")?);
429    }
430    reader.ensure_empty("Parse")?;
431    Ok(Parse {
432        statement,
433        query,
434        parameter_type_oids,
435    })
436}
437
438fn parse_bind(reader: &mut Reader<'_>) -> Result<Bind, PgWireError> {
439    let portal = reader.read_cstring("Bind portal name")?;
440    let statement = reader.read_cstring("Bind statement name")?;
441    let parameter_format_count = read_count(reader, "Bind parameter format count")?;
442    let mut parameter_formats = Vec::with_capacity(parameter_format_count);
443    for _ in 0..parameter_format_count {
444        parameter_formats.push(FormatCode::from_i16(
445            reader.read_i16("Bind parameter format code")?,
446        )?);
447    }
448
449    let parameter_count = read_count(reader, "Bind parameter count")?;
450    if parameter_format_count > 1 && parameter_format_count != parameter_count {
451        return Err(PgWireError::ParameterFormatCountMismatch {
452            format_count: parameter_format_count,
453            parameter_count,
454        });
455    }
456    let mut parameters = Vec::with_capacity(parameter_count);
457    for _ in 0..parameter_count {
458        let value = match reader.read_len_i32("Bind parameter value length")? {
459            Some(length) => Some(reader.read_exact(length, "Bind parameter value")?.to_vec()),
460            None => None,
461        };
462        parameters.push(value);
463    }
464
465    let result_format_count = read_count(reader, "Bind result format count")?;
466    let mut result_formats = Vec::with_capacity(result_format_count);
467    for _ in 0..result_format_count {
468        result_formats.push(FormatCode::from_i16(
469            reader.read_i16("Bind result format code")?,
470        )?);
471    }
472    reader.ensure_empty("Bind")?;
473    Ok(Bind {
474        portal,
475        statement,
476        parameter_formats,
477        parameters,
478        result_formats,
479    })
480}
481
482fn parse_describe(reader: &mut Reader<'_>) -> Result<DescribeTarget, PgWireError> {
483    let target = reader.read_byte("Describe target type")?;
484    let name = reader.read_cstring("Describe target name")?;
485    reader.ensure_empty("Describe")?;
486    match target {
487        b'S' => Ok(DescribeTarget::Statement(name)),
488        b'P' => Ok(DescribeTarget::Portal(name)),
489        other => Err(PgWireError::UnknownFrontendTag(other)),
490    }
491}
492
493fn parse_execute(reader: &mut Reader<'_>) -> Result<Execute, PgWireError> {
494    let portal = reader.read_cstring("Execute portal name")?;
495    let max_rows = reader.read_i32("Execute max rows")?;
496    if max_rows < 0 {
497        return Err(PgWireError::NegativeValue {
498            context: "Execute max rows",
499        });
500    }
501    reader.ensure_empty("Execute")?;
502    Ok(Execute { portal, max_rows })
503}
504
505fn parse_close(reader: &mut Reader<'_>) -> Result<CloseTarget, PgWireError> {
506    let target = reader.read_byte("Close target type")?;
507    let name = reader.read_cstring("Close target name")?;
508    reader.ensure_empty("Close")?;
509    match target {
510        b'S' => Ok(CloseTarget::Statement(name)),
511        b'P' => Ok(CloseTarget::Portal(name)),
512        other => Err(PgWireError::UnknownFrontendTag(other)),
513    }
514}
515
516fn parse_function_call(reader: &mut Reader<'_>) -> Result<FunctionCall, PgWireError> {
517    let function_oid = reader.read_u32("FunctionCall function oid")?;
518    let argument_format_count = read_count(reader, "FunctionCall argument format count")?;
519    let mut argument_formats = Vec::with_capacity(argument_format_count);
520    for _ in 0..argument_format_count {
521        argument_formats.push(FormatCode::from_i16(
522            reader.read_i16("FunctionCall argument format code")?,
523        )?);
524    }
525
526    let argument_count = read_count(reader, "FunctionCall argument count")?;
527    if argument_format_count > 1 && argument_format_count != argument_count {
528        return Err(PgWireError::FunctionArgumentFormatCountMismatch {
529            format_count: argument_format_count,
530            argument_count,
531        });
532    }
533    let mut arguments = Vec::with_capacity(argument_count);
534    for _ in 0..argument_count {
535        let value = match reader.read_len_i32("FunctionCall argument value length")? {
536            Some(length) => Some(
537                reader
538                    .read_exact(length, "FunctionCall argument value")?
539                    .to_vec(),
540            ),
541            None => None,
542        };
543        arguments.push(value);
544    }
545
546    let result_format = FormatCode::from_i16(reader.read_i16("FunctionCall result format code")?)?;
547    reader.ensure_empty("FunctionCall")?;
548    Ok(FunctionCall {
549        function_oid,
550        argument_formats,
551        arguments,
552        result_format,
553    })
554}
555
556fn parse_single_cstring(
557    reader: &mut Reader<'_>,
558    context: &'static str,
559) -> Result<String, PgWireError> {
560    let value = reader.read_cstring(context)?;
561    reader.ensure_empty(context)?;
562    Ok(value)
563}
564
565fn read_count(reader: &mut Reader<'_>, context: &'static str) -> Result<usize, PgWireError> {
566    let count = reader.read_i16(context)?;
567    if count < 0 {
568        return Err(PgWireError::NegativeValue { context });
569    }
570    Ok(count as usize)
571}
uqa_pg_wire/frontend.rs

uqa_pg_wire/
frontend.rs
