uqa_pg_wire/
auth.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7use crate::backend::Authentication;
8use crate::codec::Reader;
9use crate::frontend::PasswordMessage;
10use crate::protocol::PgWireError;
11
12#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13pub enum AuthenticationResponseKind {
14    Password,
15    KerberosV5,
16    Gss,
17    Sspi,
18    SaslInitial,
19    Sasl,
20}
21
22#[derive(Debug, Clone, PartialEq, Eq)]
23pub enum AuthenticationResponse {
24    Password(String),
25    KerberosV5(Vec<u8>),
26    Gss(Vec<u8>),
27    Sspi(Vec<u8>),
28    SaslInitial {
29        mechanism: String,
30        initial_response: Option<Vec<u8>>,
31    },
32    Sasl(Vec<u8>),
33}
34
35#[derive(Debug, Clone, Copy, PartialEq, Eq)]
36enum AuthenticationFamily {
37    Password,
38    KerberosV5,
39    Gss,
40    Sspi,
41    Sasl,
42    SaslFinal,
43}
44
45#[derive(Debug, Clone, Copy, PartialEq, Eq)]
46enum ExchangeState {
47    Ready,
48    AwaitingFrontend(AuthenticationResponseKind),
49    AwaitingBackend(AuthenticationFamily),
50    Complete,
51    Failed,
52}
53
54impl ExchangeState {
55    const fn description(self) -> &'static str {
56        match self {
57            Self::Ready => "ready for an authentication request",
58            Self::AwaitingFrontend(_) => "awaiting a frontend authentication response",
59            Self::AwaitingBackend(_) => "awaiting the next backend authentication message",
60            Self::Complete => "authentication is complete",
61            Self::Failed => "authentication has failed",
62        }
63    }
64}
65
66/// Validates one `PostgreSQL` authentication exchange while leaving credential
67/// verification and secret storage to the embedding server.
68#[derive(Debug, Clone, PartialEq, Eq)]
69pub struct AuthenticationExchange {
70    state: ExchangeState,
71}
72
73impl Default for AuthenticationExchange {
74    fn default() -> Self {
75        Self::new()
76    }
77}
78
79impl AuthenticationExchange {
80    #[must_use]
81    pub const fn new() -> Self {
82        Self {
83            state: ExchangeState::Ready,
84        }
85    }
86
87    /// Record an authentication message before it is sent to the frontend.
88    pub fn send(&mut self, authentication: &Authentication) -> Result<(), PgWireError> {
89        let next = match (self.state, authentication) {
90            (ExchangeState::Ready | ExchangeState::AwaitingBackend(_), Authentication::Ok) => {
91                ExchangeState::Complete
92            }
93            (ExchangeState::Ready, Authentication::KerberosV5) => {
94                ExchangeState::AwaitingFrontend(AuthenticationResponseKind::KerberosV5)
95            }
96            (
97                ExchangeState::Ready,
98                Authentication::CleartextPassword | Authentication::Md5Password(_),
99            ) => ExchangeState::AwaitingFrontend(AuthenticationResponseKind::Password),
100            (ExchangeState::Ready, Authentication::Gss) => {
101                ExchangeState::AwaitingFrontend(AuthenticationResponseKind::Gss)
102            }
103            (
104                ExchangeState::AwaitingBackend(AuthenticationFamily::Gss),
105                Authentication::GssContinue(_),
106            ) => ExchangeState::AwaitingFrontend(AuthenticationResponseKind::Gss),
107            (
108                ExchangeState::AwaitingBackend(AuthenticationFamily::Sspi),
109                Authentication::GssContinue(_),
110            ) => ExchangeState::AwaitingFrontend(AuthenticationResponseKind::Sspi),
111            (ExchangeState::Ready, Authentication::Sspi) => {
112                ExchangeState::AwaitingFrontend(AuthenticationResponseKind::Sspi)
113            }
114            (ExchangeState::Ready, Authentication::Sasl { .. }) => {
115                ExchangeState::AwaitingFrontend(AuthenticationResponseKind::SaslInitial)
116            }
117            (
118                ExchangeState::AwaitingBackend(AuthenticationFamily::Sasl),
119                Authentication::SaslContinue(_),
120            ) => ExchangeState::AwaitingFrontend(AuthenticationResponseKind::Sasl),
121            (
122                ExchangeState::AwaitingBackend(AuthenticationFamily::Sasl),
123                Authentication::SaslFinal(_),
124            ) => ExchangeState::AwaitingBackend(AuthenticationFamily::SaslFinal),
125            (state, authentication) => {
126                return Err(PgWireError::InvalidAuthenticationSequence {
127                    state: state.description(),
128                    message: authentication.description(),
129                });
130            }
131        };
132        self.state = next;
133        Ok(())
134    }
135
136    /// Decode the context-dependent frontend message tagged `p` and advance
137    /// the exchange to its next backend decision point.
138    pub fn receive(
139        &mut self,
140        message: &PasswordMessage,
141    ) -> Result<AuthenticationResponse, PgWireError> {
142        let ExchangeState::AwaitingFrontend(kind) = self.state else {
143            return Err(PgWireError::InvalidAuthenticationSequence {
144                state: self.state.description(),
145                message: "a frontend authentication response",
146            });
147        };
148        let response = decode_response(kind, message)?;
149        self.state = ExchangeState::AwaitingBackend(match kind {
150            AuthenticationResponseKind::Password => AuthenticationFamily::Password,
151            AuthenticationResponseKind::KerberosV5 => AuthenticationFamily::KerberosV5,
152            AuthenticationResponseKind::Gss => AuthenticationFamily::Gss,
153            AuthenticationResponseKind::Sspi => AuthenticationFamily::Sspi,
154            AuthenticationResponseKind::SaslInitial | AuthenticationResponseKind::Sasl => {
155                AuthenticationFamily::Sasl
156            }
157        });
158        Ok(response)
159    }
160
161    pub fn fail(&mut self) -> Result<(), PgWireError> {
162        if matches!(self.state, ExchangeState::Complete | ExchangeState::Failed) {
163            return Err(PgWireError::InvalidAuthenticationSequence {
164                state: self.state.description(),
165                message: "authentication failure",
166            });
167        }
168        self.state = ExchangeState::Failed;
169        Ok(())
170    }
171
172    #[must_use]
173    pub const fn is_complete(&self) -> bool {
174        matches!(self.state, ExchangeState::Complete)
175    }
176
177    #[must_use]
178    pub const fn is_failed(&self) -> bool {
179        matches!(self.state, ExchangeState::Failed)
180    }
181
182    #[must_use]
183    pub const fn awaiting_response(&self) -> Option<AuthenticationResponseKind> {
184        match self.state {
185            ExchangeState::AwaitingFrontend(kind) => Some(kind),
186            _ => None,
187        }
188    }
189}
190
191fn decode_response(
192    kind: AuthenticationResponseKind,
193    message: &PasswordMessage,
194) -> Result<AuthenticationResponse, PgWireError> {
195    match kind {
196        AuthenticationResponseKind::Password => {
197            let mut reader = Reader::new(message.as_bytes());
198            let password = reader.read_cstring("PasswordMessage password")?;
199            reader.ensure_empty("PasswordMessage")?;
200            Ok(AuthenticationResponse::Password(password))
201        }
202        AuthenticationResponseKind::KerberosV5 => Ok(AuthenticationResponse::KerberosV5(
203            message.as_bytes().to_vec(),
204        )),
205        AuthenticationResponseKind::Gss => {
206            Ok(AuthenticationResponse::Gss(message.as_bytes().to_vec()))
207        }
208        AuthenticationResponseKind::Sspi => {
209            Ok(AuthenticationResponse::Sspi(message.as_bytes().to_vec()))
210        }
211        AuthenticationResponseKind::SaslInitial => decode_sasl_initial(message),
212        AuthenticationResponseKind::Sasl => {
213            Ok(AuthenticationResponse::Sasl(message.as_bytes().to_vec()))
214        }
215    }
216}
217
218fn decode_sasl_initial(message: &PasswordMessage) -> Result<AuthenticationResponse, PgWireError> {
219    let mut reader = Reader::new(message.as_bytes());
220    let mechanism = reader.read_cstring("SASLInitialResponse mechanism")?;
221    if mechanism.is_empty() {
222        return Err(PgWireError::EmptySaslMechanism);
223    }
224    let length = reader.read_i32("SASLInitialResponse data length")?;
225    let initial_response = match length {
226        -1 => None,
227        length if length < -1 => {
228            return Err(PgWireError::NegativeValue {
229                context: "SASLInitialResponse data length",
230            });
231        }
232        length => Some(
233            reader
234                .read_exact(length as usize, "SASLInitialResponse data")?
235                .to_vec(),
236        ),
237    };
238    reader.ensure_empty("SASLInitialResponse")?;
239    Ok(AuthenticationResponse::SaslInitial {
240        mechanism,
241        initial_response,
242    })
243}
uqa_pg_wire/auth.rs

uqa_pg_wire/
auth.rs
