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

uqa_pg_wire/
backend.rs
