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    CancelKey, DecodeOutcome, FormatCode, PgWireError, ProtocolVersion, CANCEL_REQUEST_CODE,
12    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(Vec<u8>),
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
145#[derive(Debug, Clone, PartialEq, Eq)]
146pub enum DescribeTarget {
147    Statement(String),
148    Portal(String),
149}
150
151#[derive(Debug, Clone, PartialEq, Eq)]
152pub struct Execute {
153    pub portal: String,
154    pub max_rows: i32,
155}
156
157#[derive(Debug, Clone, PartialEq, Eq)]
158pub enum CloseTarget {
159    Statement(String),
160    Portal(String),
161}
162
163#[derive(Debug, Clone, PartialEq, Eq)]
164pub struct FunctionCall {
165    pub function_oid: u32,
166    pub argument_formats: Vec<FormatCode>,
167    pub arguments: Vec<Option<Vec<u8>>>,
168    pub result_format: FormatCode,
169}
170
171pub fn decode_startup(input: &[u8]) -> DecodeOutcome<StartupFrame> {
172    decode_startup_with_max(input, DEFAULT_MAX_MESSAGE_LEN)
173}
174
175pub fn decode_startup_with_max(input: &[u8], max_len: usize) -> DecodeOutcome<StartupFrame> {
176    let total = match message_total_len(input, false, max_len) {
177        DecodeLen::Complete(total) => total,
178        DecodeLen::Incomplete => return Ok(None),
179        DecodeLen::Error(error) => return Err(error),
180    };
181
182    let body = &input[4..total];
183    let mut reader = Reader::new(body);
184    let version_or_code = reader.read_i32("startup version")?;
185    let frame = match version_or_code {
186        SSL_REQUEST_CODE => {
187            reader.ensure_empty("SSL request")?;
188            StartupFrame::SSLRequest
189        }
190        GSSENC_REQUEST_CODE => {
191            reader.ensure_empty("GSSENC request")?;
192            StartupFrame::GSSEncRequest
193        }
194        CANCEL_REQUEST_CODE => {
195            let process_id = reader.read_i32("cancel request process id")?;
196            let key_length = reader.remaining();
197            let secret_key = CancelKey::new(
198                reader
199                    .read_exact(key_length, "cancel request secret key")?
200                    .to_vec(),
201            )?;
202            reader.ensure_empty("cancel request")?;
203            StartupFrame::CancelRequest {
204                process_id,
205                secret_key,
206            }
207        }
208        other => {
209            let version = ProtocolVersion::from_raw(other);
210            version.negotiate()?;
211            StartupFrame::Startup(parse_startup_message(version, reader)?)
212        }
213    };
214    Ok(Some((frame, total)))
215}
216
217pub fn decode_frontend(input: &[u8]) -> DecodeOutcome<FrontendMessage> {
218    decode_frontend_with_max(input, DEFAULT_MAX_MESSAGE_LEN)
219}
220
221pub fn decode_frontend_with_max(input: &[u8], max_len: usize) -> DecodeOutcome<FrontendMessage> {
222    let total = match message_total_len(input, true, max_len) {
223        DecodeLen::Complete(total) => total,
224        DecodeLen::Incomplete => return Ok(None),
225        DecodeLen::Error(error) => return Err(error),
226    };
227
228    let tag = input[0];
229    let body = &input[MESSAGE_HEADER_LEN..total];
230    let message = parse_frontend_message(tag, body)?;
231    Ok(Some((message, total)))
232}
233
234fn parse_startup_message(
235    version: ProtocolVersion,
236    mut reader: Reader<'_>,
237) -> Result<StartupMessage, PgWireError> {
238    let mut parameters = BTreeMap::new();
239    let mut parameter_pairs = Vec::new();
240    loop {
241        if reader.remaining() == 0 {
242            return Err(PgWireError::MissingNul {
243                context: "startup parameters",
244            });
245        }
246        if reader.remaining() == 1 {
247            let terminator = reader.read_byte("startup terminator")?;
248            if terminator == 0 {
249                break;
250            }
251            return Err(PgWireError::MissingNul {
252                context: "startup parameters",
253            });
254        }
255
256        let key = reader.read_cstring("startup parameter key")?;
257        if key.is_empty() {
258            reader.ensure_empty("startup parameters")?;
259            break;
260        }
261        let value = reader.read_cstring("startup parameter value")?;
262        parameter_pairs.push((key.clone(), value.clone()));
263        parameters.insert(key, value);
264    }
265    Ok(StartupMessage {
266        version,
267        parameters,
268        parameter_pairs,
269    })
270}
271
272fn parse_frontend_message(tag: u8, body: &[u8]) -> Result<FrontendMessage, PgWireError> {
273    let mut reader = Reader::new(body);
274    let message = match tag {
275        b'Q' => FrontendMessage::Query(parse_single_cstring(&mut reader, "Query")?),
276        b'P' => FrontendMessage::Parse(parse_parse(&mut reader)?),
277        b'B' => FrontendMessage::Bind(parse_bind(&mut reader)?),
278        b'D' => FrontendMessage::Describe(parse_describe(&mut reader)?),
279        b'E' => FrontendMessage::Execute(parse_execute(&mut reader)?),
280        b'C' => FrontendMessage::Close(parse_close(&mut reader)?),
281        b'H' => {
282            reader.ensure_empty("Flush")?;
283            FrontendMessage::Flush
284        }
285        b'S' => {
286            reader.ensure_empty("Sync")?;
287            FrontendMessage::Sync
288        }
289        b'X' => {
290            reader.ensure_empty("Terminate")?;
291            FrontendMessage::Terminate
292        }
293        b'p' => FrontendMessage::Password(body.to_vec()),
294        b'd' => FrontendMessage::CopyData(body.to_vec()),
295        b'c' => {
296            reader.ensure_empty("CopyDone")?;
297            FrontendMessage::CopyDone
298        }
299        b'f' => FrontendMessage::CopyFail(parse_single_cstring(&mut reader, "CopyFail")?),
300        b'F' => FrontendMessage::FunctionCall(parse_function_call(&mut reader)?),
301        other => return Err(PgWireError::UnknownFrontendTag(other)),
302    };
303    Ok(message)
304}
305
306fn parse_parse(reader: &mut Reader<'_>) -> Result<Parse, PgWireError> {
307    let statement = reader.read_cstring("Parse statement name")?;
308    let query = reader.read_cstring("Parse query")?;
309    let count = read_count(reader, "Parse parameter type count")?;
310    let mut parameter_type_oids = Vec::with_capacity(count);
311    for _ in 0..count {
312        parameter_type_oids.push(reader.read_u32("Parse parameter type oid")?);
313    }
314    reader.ensure_empty("Parse")?;
315    Ok(Parse {
316        statement,
317        query,
318        parameter_type_oids,
319    })
320}
321
322fn parse_bind(reader: &mut Reader<'_>) -> Result<Bind, PgWireError> {
323    let portal = reader.read_cstring("Bind portal name")?;
324    let statement = reader.read_cstring("Bind statement name")?;
325    let parameter_format_count = read_count(reader, "Bind parameter format count")?;
326    let mut parameter_formats = Vec::with_capacity(parameter_format_count);
327    for _ in 0..parameter_format_count {
328        parameter_formats.push(FormatCode::from_i16(
329            reader.read_i16("Bind parameter format code")?,
330        )?);
331    }
332
333    let parameter_count = read_count(reader, "Bind parameter count")?;
334    if parameter_format_count > 1 && parameter_format_count != parameter_count {
335        return Err(PgWireError::ParameterFormatCountMismatch {
336            format_count: parameter_format_count,
337            parameter_count,
338        });
339    }
340    let mut parameters = Vec::with_capacity(parameter_count);
341    for _ in 0..parameter_count {
342        let value = match reader.read_len_i32("Bind parameter value length")? {
343            Some(length) => Some(reader.read_exact(length, "Bind parameter value")?.to_vec()),
344            None => None,
345        };
346        parameters.push(value);
347    }
348
349    let result_format_count = read_count(reader, "Bind result format count")?;
350    let mut result_formats = Vec::with_capacity(result_format_count);
351    for _ in 0..result_format_count {
352        result_formats.push(FormatCode::from_i16(
353            reader.read_i16("Bind result format code")?,
354        )?);
355    }
356    reader.ensure_empty("Bind")?;
357    Ok(Bind {
358        portal,
359        statement,
360        parameter_formats,
361        parameters,
362        result_formats,
363    })
364}
365
366fn parse_describe(reader: &mut Reader<'_>) -> Result<DescribeTarget, PgWireError> {
367    let target = reader.read_byte("Describe target type")?;
368    let name = reader.read_cstring("Describe target name")?;
369    reader.ensure_empty("Describe")?;
370    match target {
371        b'S' => Ok(DescribeTarget::Statement(name)),
372        b'P' => Ok(DescribeTarget::Portal(name)),
373        other => Err(PgWireError::UnknownFrontendTag(other)),
374    }
375}
376
377fn parse_execute(reader: &mut Reader<'_>) -> Result<Execute, PgWireError> {
378    let portal = reader.read_cstring("Execute portal name")?;
379    let max_rows = reader.read_i32("Execute max rows")?;
380    if max_rows < 0 {
381        return Err(PgWireError::NegativeValue {
382            context: "Execute max rows",
383        });
384    }
385    reader.ensure_empty("Execute")?;
386    Ok(Execute { portal, max_rows })
387}
388
389fn parse_close(reader: &mut Reader<'_>) -> Result<CloseTarget, PgWireError> {
390    let target = reader.read_byte("Close target type")?;
391    let name = reader.read_cstring("Close target name")?;
392    reader.ensure_empty("Close")?;
393    match target {
394        b'S' => Ok(CloseTarget::Statement(name)),
395        b'P' => Ok(CloseTarget::Portal(name)),
396        other => Err(PgWireError::UnknownFrontendTag(other)),
397    }
398}
399
400fn parse_function_call(reader: &mut Reader<'_>) -> Result<FunctionCall, PgWireError> {
401    let function_oid = reader.read_u32("FunctionCall function oid")?;
402    let argument_format_count = read_count(reader, "FunctionCall argument format count")?;
403    let mut argument_formats = Vec::with_capacity(argument_format_count);
404    for _ in 0..argument_format_count {
405        argument_formats.push(FormatCode::from_i16(
406            reader.read_i16("FunctionCall argument format code")?,
407        )?);
408    }
409
410    let argument_count = read_count(reader, "FunctionCall argument count")?;
411    if argument_format_count > 1 && argument_format_count != argument_count {
412        return Err(PgWireError::FunctionArgumentFormatCountMismatch {
413            format_count: argument_format_count,
414            argument_count,
415        });
416    }
417    let mut arguments = Vec::with_capacity(argument_count);
418    for _ in 0..argument_count {
419        let value = match reader.read_len_i32("FunctionCall argument value length")? {
420            Some(length) => Some(
421                reader
422                    .read_exact(length, "FunctionCall argument value")?
423                    .to_vec(),
424            ),
425            None => None,
426        };
427        arguments.push(value);
428    }
429
430    let result_format = FormatCode::from_i16(reader.read_i16("FunctionCall result format code")?)?;
431    reader.ensure_empty("FunctionCall")?;
432    Ok(FunctionCall {
433        function_oid,
434        argument_formats,
435        arguments,
436        result_format,
437    })
438}
439
440fn parse_single_cstring(
441    reader: &mut Reader<'_>,
442    context: &'static str,
443) -> Result<String, PgWireError> {
444    let value = reader.read_cstring(context)?;
445    reader.ensure_empty(context)?;
446    Ok(value)
447}
448
449fn read_count(reader: &mut Reader<'_>, context: &'static str) -> Result<usize, PgWireError> {
450    let count = reader.read_i16(context)?;
451    if count < 0 {
452        return Err(PgWireError::NegativeValue { context });
453    }
454    Ok(count as usize)
455}
uqa_pg_wire/frontend.rs

uqa_pg_wire/
frontend.rs
