uqa_pg_wire/
backend.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7use crate::codec::{i16_len, i32_len, Writer};
8use crate::protocol::{CancelKey, FormatCode, PgWireError, ProtocolVersion, TransactionStatus};
9
10#[derive(Debug, Clone, PartialEq, Eq)]
11pub enum Authentication {
12    Ok,
13    KerberosV5,
14    CleartextPassword,
15    Md5Password([u8; 4]),
16    Gss,
17    GssContinue(Vec<u8>),
18    Sspi,
19    Sasl { mechanisms: Vec<String> },
20    SaslContinue(Vec<u8>),
21    SaslFinal(Vec<u8>),
22}
23
24#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25pub enum SSLResponse {
26    Accept,
27    Reject,
28}
29
30impl SSLResponse {
31    pub const fn encode(self) -> [u8; 1] {
32        match self {
33            Self::Accept => *b"S",
34            Self::Reject => *b"N",
35        }
36    }
37}
38
39#[derive(Debug, Clone, Copy, PartialEq, Eq)]
40pub enum GSSEncResponse {
41    Accept,
42    Reject,
43}
44
45impl GSSEncResponse {
46    pub const fn encode(self) -> [u8; 1] {
47        match self {
48            Self::Accept => *b"G",
49            Self::Reject => *b"N",
50        }
51    }
52}
53
54#[derive(Debug, Clone, PartialEq, Eq)]
55pub struct BackendKeyData {
56    pub process_id: i32,
57    pub secret_key: CancelKey,
58}
59
60impl BackendKeyData {
61    #[must_use]
62    pub fn legacy(process_id: i32, secret_key: i32) -> Self {
63        Self {
64            process_id,
65            secret_key: CancelKey::from_i32(secret_key),
66        }
67    }
68}
69
70#[derive(Debug, Clone, PartialEq, Eq)]
71pub enum BackendMessage {
72    Authentication(Authentication),
73    BackendKeyData(BackendKeyData),
74    NegotiateProtocolVersion {
75        newest_protocol_version: ProtocolVersion,
76        unrecognized_options: Vec<String>,
77    },
78    ParameterStatus {
79        name: String,
80        value: String,
81    },
82    ReadyForQuery(TransactionStatus),
83    RowDescription(Vec<FieldDescription>),
84    DataRow(Vec<Option<Vec<u8>>>),
85    CommandComplete(String),
86    EmptyQueryResponse,
87    ErrorResponse(ErrorOrNotice),
88    NoticeResponse(ErrorOrNotice),
89    ParseComplete,
90    BindComplete,
91    CloseComplete,
92    NoData,
93    ParameterDescription(Vec<u32>),
94    PortalSuspended,
95    CopyInResponse(CopyResponse),
96    CopyOutResponse(CopyResponse),
97    CopyBothResponse(CopyResponse),
98    CopyData(Vec<u8>),
99    CopyDone,
100    FunctionCallResponse(Option<Vec<u8>>),
101    NotificationResponse(NotificationResponse),
102}
103
104#[derive(Debug, Clone, PartialEq, Eq)]
105pub struct FieldDescription {
106    pub name: String,
107    pub table_oid: u32,
108    pub column_attribute_number: i16,
109    pub type_oid: u32,
110    pub type_size: i16,
111    pub type_modifier: i32,
112    pub format: FormatCode,
113}
114
115impl FieldDescription {
116    pub fn text(name: impl Into<String>, type_oid: u32, type_size: i16) -> Self {
117        Self {
118            name: name.into(),
119            table_oid: 0,
120            column_attribute_number: 0,
121            type_oid,
122            type_size,
123            type_modifier: -1,
124            format: FormatCode::Text,
125        }
126    }
127}
128
129#[derive(Debug, Clone, PartialEq, Eq)]
130pub struct CopyResponse {
131    pub overall_format: FormatCode,
132    pub column_formats: Vec<FormatCode>,
133}
134
135#[derive(Debug, Clone, PartialEq, Eq)]
136pub struct NotificationResponse {
137    pub process_id: i32,
138    pub channel: String,
139    pub payload: String,
140}
141
142#[derive(Debug, Clone, Copy, PartialEq, Eq)]
143pub enum NoticeSeverity {
144    Error,
145    Fatal,
146    Panic,
147    Warning,
148    Notice,
149    Debug,
150    Info,
151    Log,
152}
153
154impl NoticeSeverity {
155    const fn as_str(self) -> &'static str {
156        match self {
157            Self::Error => "ERROR",
158            Self::Fatal => "FATAL",
159            Self::Panic => "PANIC",
160            Self::Warning => "WARNING",
161            Self::Notice => "NOTICE",
162            Self::Debug => "DEBUG",
163            Self::Info => "INFO",
164            Self::Log => "LOG",
165        }
166    }
167}
168
169#[derive(Debug, Clone, PartialEq, Eq)]
170pub struct ErrorOrNotice {
171    pub severity: NoticeSeverity,
172    pub code: String,
173    pub message: String,
174    pub detail: Option<String>,
175    pub hint: Option<String>,
176    pub position: Option<i32>,
177    pub where_: Option<String>,
178    pub schema: Option<String>,
179    pub table: Option<String>,
180    pub column: Option<String>,
181    pub data_type: Option<String>,
182    pub constraint: Option<String>,
183    pub file: Option<String>,
184    pub line: Option<i32>,
185    pub routine: Option<String>,
186}
187
188impl ErrorOrNotice {
189    pub fn error(code: impl Into<String>, message: impl Into<String>) -> Self {
190        Self {
191            severity: NoticeSeverity::Error,
192            code: code.into(),
193            message: message.into(),
194            detail: None,
195            hint: None,
196            position: None,
197            where_: None,
198            schema: None,
199            table: None,
200            column: None,
201            data_type: None,
202            constraint: None,
203            file: None,
204            line: None,
205            routine: None,
206        }
207    }
208}
209
210impl BackendMessage {
211    pub fn encode(&self) -> Result<Vec<u8>, PgWireError> {
212        self.encode_for_protocol(ProtocolVersion::LATEST)
213    }
214
215    pub fn encode_for_protocol(
216        &self,
217        protocol_version: ProtocolVersion,
218    ) -> Result<Vec<u8>, PgWireError> {
219        match self {
220            Self::Authentication(auth) => encode_authentication(auth),
221            Self::BackendKeyData(data) => encode_backend_key_data(data, protocol_version),
222            Self::NegotiateProtocolVersion {
223                newest_protocol_version,
224                unrecognized_options,
225            } => encode_negotiate_protocol_version(*newest_protocol_version, unrecognized_options),
226            Self::ParameterStatus { name, value } => encode_parameter_status(name, value),
227            Self::ReadyForQuery(status) => encode_ready_for_query(*status),
228            Self::RowDescription(fields) => encode_row_description(fields),
229            Self::DataRow(values) => encode_data_row(values),
230            Self::CommandComplete(tag) => encode_command_complete(tag),
231            Self::EmptyQueryResponse => encode_empty_body(b'I'),
232            Self::ErrorResponse(error) => encode_error_or_notice(b'E', error),
233            Self::NoticeResponse(notice) => encode_error_or_notice(b'N', notice),
234            Self::ParseComplete => encode_empty_body(b'1'),
235            Self::BindComplete => encode_empty_body(b'2'),
236            Self::CloseComplete => encode_empty_body(b'3'),
237            Self::NoData => encode_empty_body(b'n'),
238            Self::ParameterDescription(oids) => encode_parameter_description(oids),
239            Self::PortalSuspended => encode_empty_body(b's'),
240            Self::CopyInResponse(response) => encode_copy_response(b'G', response),
241            Self::CopyOutResponse(response) => encode_copy_response(b'H', response),
242            Self::CopyBothResponse(response) => encode_copy_response(b'W', response),
243            Self::CopyData(bytes) => Writer::frame(b'd', bytes),
244            Self::CopyDone => encode_empty_body(b'c'),
245            Self::FunctionCallResponse(value) => encode_function_call_response(value.as_deref()),
246            Self::NotificationResponse(notification) => encode_notification_response(notification),
247        }
248    }
249}
250
251pub fn encode_all(messages: &[BackendMessage]) -> Result<Vec<u8>, PgWireError> {
252    encode_all_for_protocol(messages, ProtocolVersion::LATEST)
253}
254
255pub fn encode_all_for_protocol(
256    messages: &[BackendMessage],
257    protocol_version: ProtocolVersion,
258) -> Result<Vec<u8>, PgWireError> {
259    let mut out = Vec::new();
260    for message in messages {
261        out.extend(message.encode_for_protocol(protocol_version)?);
262    }
263    Ok(out)
264}
265
266pub const fn encode_ssl_response(response: SSLResponse) -> [u8; 1] {
267    response.encode()
268}
269
270pub const fn encode_gssenc_response(response: GSSEncResponse) -> [u8; 1] {
271    response.encode()
272}
273
274pub const TYPE_BOOL: u32 = 16;
275pub const TYPE_BYTEA: u32 = 17;
276pub const TYPE_INT8: u32 = 20;
277pub const TYPE_INT2: u32 = 21;
278pub const TYPE_INT4: u32 = 23;
279pub const TYPE_TEXT: u32 = 25;
280pub const TYPE_FLOAT4: u32 = 700;
281pub const TYPE_FLOAT8: u32 = 701;
282pub const TYPE_VARCHAR: u32 = 1_043;
283pub const TYPE_DATE: u32 = 1_082;
284pub const TYPE_TIMESTAMP: u32 = 1_114;
285pub const TYPE_TIMESTAMPTZ: u32 = 1_184;
286pub const TYPE_JSON: u32 = 114;
287pub const TYPE_JSONB: u32 = 3_802;
288
289pub mod sqlstate {
290    pub const SUCCESSFUL_COMPLETION: &str = "00000";
291    pub const WARNING: &str = "01000";
292    pub const PROTOCOL_VIOLATION: &str = "08P01";
293    pub const FEATURE_NOT_SUPPORTED: &str = "0A000";
294    pub const INVALID_PARAMETER_VALUE: &str = "22023";
295    pub const QUERY_CANCELED: &str = "57014";
296    pub const SYNTAX_ERROR: &str = "42601";
297    pub const UNDEFINED_TABLE: &str = "42P01";
298    pub const INTERNAL_ERROR: &str = "XX000";
299}
300
301fn encode_authentication(auth: &Authentication) -> Result<Vec<u8>, PgWireError> {
302    let mut body = Writer::new();
303    match auth {
304        Authentication::Ok => body.write_i32(0),
305        Authentication::KerberosV5 => body.write_i32(2),
306        Authentication::CleartextPassword => body.write_i32(3),
307        Authentication::Md5Password(salt) => {
308            body.write_i32(5);
309            body.write_bytes(salt);
310        }
311        Authentication::Gss => body.write_i32(7),
312        Authentication::GssContinue(data) => {
313            body.write_i32(8);
314            body.write_bytes(data);
315        }
316        Authentication::Sspi => body.write_i32(9),
317        Authentication::Sasl { mechanisms } => {
318            body.write_i32(10);
319            for mechanism in mechanisms {
320                body.write_cstring(mechanism, "SASL mechanism")?;
321            }
322            body.write_byte(0);
323        }
324        Authentication::SaslContinue(data) => {
325            body.write_i32(11);
326            body.write_bytes(data);
327        }
328        Authentication::SaslFinal(data) => {
329            body.write_i32(12);
330            body.write_bytes(data);
331        }
332    }
333    Writer::frame(b'R', &body.into_inner())
334}
335
336fn encode_backend_key_data(
337    data: &BackendKeyData,
338    protocol_version: ProtocolVersion,
339) -> Result<Vec<u8>, PgWireError> {
340    data.secret_key
341        .validate_for_backend_key_data(protocol_version)?;
342    let mut body = Writer::new();
343    body.write_i32(data.process_id);
344    body.write_bytes(data.secret_key.as_bytes());
345    Writer::frame(b'K', &body.into_inner())
346}
347
348fn encode_negotiate_protocol_version(
349    newest_protocol_version: ProtocolVersion,
350    unrecognized_options: &[String],
351) -> Result<Vec<u8>, PgWireError> {
352    if newest_protocol_version.negotiate()? != newest_protocol_version {
353        return Err(PgWireError::UnsupportedProtocolVersion(
354            newest_protocol_version.raw(),
355        ));
356    }
357    let mut body = Writer::new();
358    body.write_i32(newest_protocol_version.raw());
359    body.write_i32(i32_len(
360        unrecognized_options.len(),
361        "NegotiateProtocolVersion option count",
362    )?);
363    for option in unrecognized_options {
364        body.write_cstring(option, "NegotiateProtocolVersion option")?;
365    }
366    Writer::frame(b'v', &body.into_inner())
367}
368
369fn encode_parameter_status(name: &str, value: &str) -> Result<Vec<u8>, PgWireError> {
370    let mut body = Writer::new();
371    body.write_cstring(name, "ParameterStatus name")?;
372    body.write_cstring(value, "ParameterStatus value")?;
373    Writer::frame(b'S', &body.into_inner())
374}
375
376fn encode_ready_for_query(status: TransactionStatus) -> Result<Vec<u8>, PgWireError> {
377    Writer::frame(b'Z', &[status.as_byte()])
378}
379
380fn encode_row_description(fields: &[FieldDescription]) -> Result<Vec<u8>, PgWireError> {
381    let mut body = Writer::new();
382    body.write_i16(i16_len(fields.len(), "RowDescription field")?);
383    for field in fields {
384        body.write_cstring(&field.name, "RowDescription field name")?;
385        body.write_u32(field.table_oid);
386        body.write_i16(field.column_attribute_number);
387        body.write_u32(field.type_oid);
388        body.write_i16(field.type_size);
389        body.write_i32(field.type_modifier);
390        body.write_format(field.format);
391    }
392    Writer::frame(b'T', &body.into_inner())
393}
394
395fn encode_data_row(values: &[Option<Vec<u8>>]) -> Result<Vec<u8>, PgWireError> {
396    let mut body = Writer::new();
397    body.write_i16(i16_len(values.len(), "DataRow column")?);
398    for value in values {
399        match value {
400            Some(bytes) => {
401                body.write_i32(i32_len(bytes.len(), "DataRow value")?);
402                body.write_bytes(bytes);
403            }
404            None => body.write_i32(-1),
405        }
406    }
407    Writer::frame(b'D', &body.into_inner())
408}
409
410fn encode_command_complete(tag: &str) -> Result<Vec<u8>, PgWireError> {
411    let mut body = Writer::new();
412    body.write_cstring(tag, "CommandComplete tag")?;
413    Writer::frame(b'C', &body.into_inner())
414}
415
416fn encode_error_or_notice(tag: u8, message: &ErrorOrNotice) -> Result<Vec<u8>, PgWireError> {
417    if message.code.len() != 5
418        || !message
419            .code
420            .bytes()
421            .all(|byte| byte.is_ascii_uppercase() || byte.is_ascii_digit())
422    {
423        return Err(PgWireError::InvalidSqlState {
424            code: message.code.clone(),
425        });
426    }
427    let mut body = Writer::new();
428    write_field(&mut body, b'S', message.severity.as_str(), "severity")?;
429    write_field(&mut body, b'V', message.severity.as_str(), "severity")?;
430    write_field(&mut body, b'C', &message.code, "SQLSTATE")?;
431    write_field(&mut body, b'M', &message.message, "error message")?;
432    write_optional_field(&mut body, b'D', message.detail.as_deref(), "error detail")?;
433    write_optional_field(&mut body, b'H', message.hint.as_deref(), "error hint")?;
434    write_optional_i32_field(&mut body, b'P', message.position, "error position")?;
435    write_optional_field(&mut body, b'W', message.where_.as_deref(), "error context")?;
436    write_optional_field(&mut body, b's', message.schema.as_deref(), "error schema")?;
437    write_optional_field(&mut body, b't', message.table.as_deref(), "error table")?;
438    write_optional_field(&mut body, b'c', message.column.as_deref(), "error column")?;
439    write_optional_field(
440        &mut body,
441        b'd',
442        message.data_type.as_deref(),
443        "error data type",
444    )?;
445    write_optional_field(
446        &mut body,
447        b'n',
448        message.constraint.as_deref(),
449        "error constraint",
450    )?;
451    write_optional_field(&mut body, b'F', message.file.as_deref(), "error file")?;
452    write_optional_i32_field(&mut body, b'L', message.line, "error line")?;
453    write_optional_field(&mut body, b'R', message.routine.as_deref(), "error routine")?;
454    body.write_byte(0);
455    Writer::frame(tag, &body.into_inner())
456}
457
458fn encode_parameter_description(oids: &[u32]) -> Result<Vec<u8>, PgWireError> {
459    let mut body = Writer::new();
460    body.write_i16(i16_len(oids.len(), "ParameterDescription parameter")?);
461    for oid in oids {
462        body.write_u32(*oid);
463    }
464    Writer::frame(b't', &body.into_inner())
465}
466
467fn encode_copy_response(tag: u8, response: &CopyResponse) -> Result<Vec<u8>, PgWireError> {
468    if response.overall_format == FormatCode::Text {
469        if let Some((index, _)) = response
470            .column_formats
471            .iter()
472            .enumerate()
473            .find(|(_, format)| **format == FormatCode::Binary)
474        {
475            return Err(PgWireError::BinaryColumnInTextCopy { column: index + 1 });
476        }
477    }
478    let mut body = Writer::new();
479    body.write_byte(match response.overall_format {
480        FormatCode::Text => 0,
481        FormatCode::Binary => 1,
482    });
483    body.write_i16(i16_len(
484        response.column_formats.len(),
485        "CopyResponse column",
486    )?);
487    for format in &response.column_formats {
488        body.write_format(*format);
489    }
490    Writer::frame(tag, &body.into_inner())
491}
492
493fn encode_function_call_response(value: Option<&[u8]>) -> Result<Vec<u8>, PgWireError> {
494    let mut body = Writer::new();
495    match value {
496        Some(bytes) => {
497            body.write_i32(i32_len(bytes.len(), "FunctionCallResponse value")?);
498            body.write_bytes(bytes);
499        }
500        None => body.write_i32(-1),
501    }
502    Writer::frame(b'V', &body.into_inner())
503}
504
505fn encode_notification_response(
506    notification: &NotificationResponse,
507) -> Result<Vec<u8>, PgWireError> {
508    let mut body = Writer::new();
509    body.write_i32(notification.process_id);
510    body.write_cstring(&notification.channel, "NotificationResponse channel")?;
511    body.write_cstring(&notification.payload, "NotificationResponse payload")?;
512    Writer::frame(b'A', &body.into_inner())
513}
514
515fn encode_empty_body(tag: u8) -> Result<Vec<u8>, PgWireError> {
516    Writer::frame(tag, &[])
517}
518
519fn write_field(
520    body: &mut Writer,
521    code: u8,
522    value: &str,
523    context: &'static str,
524) -> Result<(), PgWireError> {
525    body.write_byte(code);
526    body.write_cstring(value, context)
527}
528
529fn write_optional_field(
530    body: &mut Writer,
531    code: u8,
532    value: Option<&str>,
533    context: &'static str,
534) -> Result<(), PgWireError> {
535    if let Some(value) = value {
536        write_field(body, code, value, context)?;
537    }
538    Ok(())
539}
540
541fn write_optional_i32_field(
542    body: &mut Writer,
543    code: u8,
544    value: Option<i32>,
545    context: &'static str,
546) -> Result<(), PgWireError> {
547    if let Some(value) = value {
548        write_field(body, code, &value.to_string(), context)?;
549    }
550    Ok(())
551}
uqa_pg_wire/backend.rs

uqa_pg_wire/
backend.rs
