Skip to main content

pg_proto/
auth.rs

1//! Authentication typestates, including the recursive SASL sub-session.
2
3use bytes::{BufMut, Bytes, BytesMut};
4
5use crate::demux::SessionItem;
6use crate::{
7    Conn, Pristine, codec,
8    grammar::authentication as auth_grammar,
9    pre_startup::{Startup, Terminated},
10};
11
12#[derive(Debug)]
13/// The backend's initial authentication choice.
14pub enum Auth {}
15
16#[derive(Debug)]
17/// A cleartext or MD5 password response is required.
18pub enum PasswordResponse {}
19
20#[derive(Debug)]
21/// A SASL mechanism and optional initial response must be selected.
22pub enum SaslInitial {}
23
24#[derive(Debug)]
25/// The backend must provide the next SASL continuation or final message.
26pub enum Sasl {}
27
28#[derive(Debug)]
29/// A client response to a SASL challenge is required.
30pub enum SaslChallenge {}
31
32#[derive(Debug)]
33/// The received SASL final message must be verified by authentication policy.
34pub enum SaslFinal {}
35
36#[derive(Debug)]
37/// A client token for GSS, SSPI, or Kerberos authentication is required.
38pub enum TokenResponse {}
39
40#[derive(Debug)]
41/// The backend must continue or complete token-based authentication.
42pub enum TokenChallenge {}
43
44#[derive(Debug)]
45/// The selected mechanism has completed and `AuthenticationOk` is required.
46pub enum AwaitingAuthOk {}
47
48#[derive(Debug)]
49/// Authentication and startup have completed and commands may be issued.
50pub enum Ready {}
51
52#[derive(Debug)]
53/// Authentication succeeded and startup messages are being consumed until ready.
54pub enum AwaitingStartupReady {}
55
56/// TLS transports expose the RFC 5929 `tls-server-end-point` binding.
57pub trait TlsServerEndPoint {
58    /// Returns the RFC 5929 channel-binding bytes derived from the peer certificate.
59    fn tls_server_end_point(&self) -> &[u8];
60}
61
62impl<S: TlsServerEndPoint, Phase, Cleanliness> Conn<S, Phase, Cleanliness> {
63    /// Returns the peer-certificate binding for custom authentication policy.
64    #[must_use]
65    pub fn tls_server_end_point(&self) -> &[u8] {
66        self.transport().tls_server_end_point()
67    }
68}
69
70/// External choice offered by the backend during authentication.
71#[derive(Debug)]
72pub enum AuthOffer<S> {
73    /// Authentication completed without a credential exchange.
74    Ok(Conn<S, AwaitingStartupReady>),
75    /// The backend requested a cleartext password.
76    Cleartext(Conn<S, PasswordResponse>),
77    /// The backend requested a `PostgreSQL` MD5 password response.
78    Md5 {
79        /// Connection waiting for the password response.
80        conn: Conn<S, PasswordResponse>,
81        /// Four-byte salt supplied by the backend.
82        salt: [u8; 4],
83    },
84    /// The backend offered a SASL mechanism negotiation.
85    Sasl {
86        /// Connection waiting for mechanism selection.
87        conn: Conn<S, SaslInitial>,
88        /// Mechanism names offered in backend preference order.
89        mechanisms: Vec<Bytes>,
90    },
91    /// The backend requested GSSAPI authentication.
92    Gss(Conn<S, TokenResponse>),
93    /// The backend requested SSPI authentication.
94    Sspi(Conn<S, TokenResponse>),
95    /// The backend requested Kerberos V5 authentication.
96    KerberosV5(Conn<S, TokenResponse>),
97}
98
99/// A startup message that advances or terminates authentication.
100#[derive(Debug)]
101pub enum AuthEvent<S> {
102    /// An authentication request or successful completion.
103    Authentication(AuthOffer<S>),
104    /// A protocol-version negotiation message that leaves authentication active.
105    Negotiate {
106        /// Connection remaining in the authentication phase.
107        conn: Conn<S, Auth>,
108        /// Version and unsupported-option information supplied by the backend.
109        message: codec::NegotiateProtocolVersion,
110    },
111    /// Authentication failed and the connection is terminated.
112    Error {
113        /// Terminated connection.
114        conn: Conn<S, Terminated>,
115        /// Backend diagnostic describing the failure.
116        error: codec::DiagnosticResponse,
117    },
118}
119
120/// An external choice received during a SASL exchange.
121#[derive(Debug)]
122pub enum SaslEvent<S> {
123    /// The backend supplied another challenge.
124    Continue {
125        /// Connection waiting for the corresponding client response.
126        conn: Conn<S, SaslChallenge>,
127        /// Opaque mechanism-specific challenge bytes.
128        challenge: Bytes,
129    },
130    /// The backend supplied its final verifier.
131    Final {
132        /// Connection waiting for verification by authentication policy.
133        conn: Conn<S, SaslFinal>,
134        /// Opaque mechanism-specific server-final bytes.
135        server_final: Bytes,
136    },
137    /// The backend aborted authentication.
138    Error {
139        /// Terminated connection.
140        conn: Conn<S, Terminated>,
141        /// Backend diagnostic describing the failure.
142        error: codec::DiagnosticResponse,
143    },
144}
145
146/// Completion or failure after a credential mechanism finishes.
147#[derive(Debug)]
148pub enum AuthCompletion<S> {
149    /// The backend confirmed authentication.
150    Ok(Conn<S, AwaitingStartupReady>),
151    /// The backend rejected authentication.
152    Error {
153        /// Terminated connection.
154        conn: Conn<S, Terminated>,
155        /// Backend diagnostic describing the failure.
156        error: codec::DiagnosticResponse,
157    },
158}
159
160/// An external choice during GSS, SSPI, or Kerberos token exchange.
161#[derive(Debug)]
162pub enum TokenAuthEvent<S> {
163    /// The backend supplied another token.
164    Continue {
165        /// Connection waiting for the next client token.
166        conn: Conn<S, TokenResponse>,
167        /// Opaque mechanism-specific backend token.
168        token: Bytes,
169    },
170    /// The backend confirmed authentication.
171    Ok(Conn<S, AwaitingStartupReady>),
172    /// The backend rejected authentication.
173    Error {
174        /// Terminated connection.
175        conn: Conn<S, Terminated>,
176        /// Backend diagnostic describing the failure.
177        error: codec::DiagnosticResponse,
178    },
179}
180
181impl<S> Conn<S, Startup, Pristine> {
182    /// Enters backend-driven authentication after sending the startup message.
183    pub fn authentication(self) -> Conn<S, Auth> {
184        self.transition()
185    }
186}
187
188impl<S> Conn<S, Auth, Pristine> {
189    /// Projects either protocol negotiation or an authentication request.
190    ///
191    /// # Errors
192    ///
193    /// Returns an authentication parsing error, or the unchanged connection and
194    /// message when the backend message is unrelated to startup authentication.
195    ///
196    /// # Panics
197    ///
198    /// Panics only if the exhaustive continuation guard above the internal
199    /// projection becomes inconsistent with [`Self::offer`].
200    pub fn offer_backend(
201        self,
202        message: codec::BackendMessage,
203    ) -> Result<AuthEvent<S>, (Self, codec::BackendMessage, Option<std::io::Error>)> {
204        match message {
205            codec::BackendMessage::Authentication(
206                authentication @ (codec::Authentication::GssContinue(_)
207                | codec::Authentication::SaslContinue(_)
208                | codec::Authentication::SaslFinal(_)),
209            ) => Err((
210                self,
211                codec::BackendMessage::Authentication(authentication),
212                Some(std::io::Error::new(
213                    std::io::ErrorKind::InvalidData,
214                    "authentication continuation before mechanism selection",
215                )),
216            )),
217            codec::BackendMessage::Authentication(authentication) => Ok(AuthEvent::Authentication(
218                self.offer(authentication)
219                    .expect("non-continuation authentication is valid in Auth"),
220            )),
221            codec::BackendMessage::NegotiateProtocolVersion(message) => Ok(AuthEvent::Negotiate {
222                conn: self,
223                message,
224            }),
225            codec::BackendMessage::ErrorResponse(error) => Ok(AuthEvent::Error {
226                conn: self.transition(),
227                error,
228            }),
229            message => Err((self, message, None)),
230        }
231    }
232
233    /// Applies one backend authentication request to the session state.
234    ///
235    /// # Errors
236    ///
237    /// SASL continuation/final messages are rejected before SASL is selected.
238    pub fn offer(self, authentication: codec::Authentication) -> std::io::Result<AuthOffer<S>> {
239        match (
240            project_authentication(auth_grammar::RuntimeState::Auth, &authentication),
241            authentication,
242        ) {
243            (Some(auth_grammar::Event::Ok), codec::Authentication::Ok) => {
244                Ok(AuthOffer::Ok(self.transition()))
245            }
246            (Some(auth_grammar::Event::Cleartext), codec::Authentication::CleartextPassword) => {
247                Ok(AuthOffer::Cleartext(self.transition()))
248            }
249            (Some(auth_grammar::Event::Md5), codec::Authentication::Md5Password { salt }) => {
250                Ok(AuthOffer::Md5 {
251                    conn: self.transition(),
252                    salt,
253                })
254            }
255            (Some(auth_grammar::Event::Sasl), codec::Authentication::Sasl { mechanisms }) => {
256                Ok(AuthOffer::Sasl {
257                    conn: self.transition(),
258                    mechanisms,
259                })
260            }
261            (Some(auth_grammar::Event::Gss), codec::Authentication::Gss) => {
262                Ok(AuthOffer::Gss(self.transition()))
263            }
264            (Some(auth_grammar::Event::Sspi), codec::Authentication::Sspi) => {
265                Ok(AuthOffer::Sspi(self.transition()))
266            }
267            (Some(auth_grammar::Event::KerberosV5), codec::Authentication::KerberosV5) => {
268                Ok(AuthOffer::KerberosV5(self.transition()))
269            }
270            _ => Err(std::io::Error::new(
271                std::io::ErrorKind::InvalidData,
272                "authentication continuation before mechanism selection",
273            )),
274        }
275    }
276}
277
278impl<S> Conn<S, TokenResponse, Pristine> {
279    /// Sends a GSS, SSPI, or Kerberos token and waits for continuation or success.
280    pub fn respond(self, token: Bytes) -> (Conn<S, TokenChallenge>, codec::Frame) {
281        (
282            self.transition(),
283            codec::Frame {
284                tag: b'p',
285                body: token,
286            },
287        )
288    }
289}
290
291impl<S> Conn<S, TokenChallenge, Pristine> {
292    /// Projects recursive GSS continuation, successful authentication, or failure.
293    ///
294    /// # Errors
295    ///
296    /// Returns the live connection and message for an illegal response.
297    pub fn offer(
298        self,
299        message: codec::BackendMessage,
300    ) -> Result<TokenAuthEvent<S>, (Self, codec::BackendMessage)> {
301        match (
302            auth_grammar::project_external(auth_grammar::RuntimeState::TokenChallenge, &message),
303            message,
304        ) {
305            (
306                Some(auth_grammar::Event::Continue),
307                codec::BackendMessage::Authentication(codec::Authentication::GssContinue(token)),
308            ) => Ok(TokenAuthEvent::Continue {
309                conn: self.transition(),
310                token,
311            }),
312            (
313                Some(auth_grammar::Event::Ok),
314                codec::BackendMessage::Authentication(codec::Authentication::Ok),
315            ) => Ok(TokenAuthEvent::Ok(self.transition())),
316            (Some(auth_grammar::Event::Error), codec::BackendMessage::ErrorResponse(error)) => {
317                Ok(TokenAuthEvent::Error {
318                    conn: self.transition(),
319                    error,
320                })
321            }
322            (_, message) => Err((self, message)),
323        }
324    }
325}
326
327impl<S> Conn<S, PasswordResponse, Pristine> {
328    /// Sends a cleartext or precomputed MD5 password response.
329    ///
330    /// # Errors
331    ///
332    /// Returns an error if the response contains a NUL byte or is too large.
333    pub fn password(
334        self,
335        password: &[u8],
336    ) -> std::io::Result<(Conn<S, AwaitingAuthOk>, codec::Frame)> {
337        if password.contains(&0) {
338            return Err(std::io::Error::new(
339                std::io::ErrorKind::InvalidInput,
340                "password response contains a NUL byte",
341            ));
342        }
343        let mut body = BytesMut::with_capacity(password.len() + 1);
344        body.extend_from_slice(password);
345        body.put_u8(0);
346        Ok((
347            self.transition(),
348            codec::Frame {
349                tag: b'p',
350                body: body.freeze(),
351            },
352        ))
353    }
354}
355
356impl<S> Conn<S, SaslInitial, Pristine> {
357    /// Selects SCRAM without channel binding.
358    ///
359    /// # Errors
360    ///
361    /// Returns an error if the initial response is too large.
362    pub fn scram_sha_256(self, initial: &[u8]) -> std::io::Result<(Conn<S, Sasl>, codec::Frame)> {
363        sasl_initial(self, b"SCRAM-SHA-256", initial)
364    }
365}
366
367impl<S: TlsServerEndPoint> Conn<S, SaslInitial, Pristine> {
368    /// Selects SCRAM-PLUS. This method is unavailable on transports which cannot
369    /// provide the peer-certificate channel binding.
370    ///
371    /// # Errors
372    ///
373    /// Returns an error if the initial response is too large.
374    pub fn scram_sha_256_plus(
375        self,
376        initial: &[u8],
377    ) -> std::io::Result<(Conn<S, Sasl>, codec::Frame)> {
378        sasl_initial(self, b"SCRAM-SHA-256-PLUS", initial)
379    }
380}
381
382impl<S> Conn<S, Sasl, Pristine> {
383    /// Projects the next server challenge or final verifier.
384    ///
385    /// # Errors
386    ///
387    /// Returns the live connection and authentication message for an illegal branch.
388    pub fn offer(
389        self,
390        authentication: codec::Authentication,
391    ) -> Result<SaslEvent<S>, (Self, codec::Authentication)> {
392        match (
393            project_authentication(auth_grammar::RuntimeState::Sasl, &authentication),
394            authentication,
395        ) {
396            (
397                Some(auth_grammar::Event::Continue),
398                codec::Authentication::SaslContinue(challenge),
399            ) => Ok(SaslEvent::Continue {
400                conn: self.transition(),
401                challenge,
402            }),
403            (Some(auth_grammar::Event::Final), codec::Authentication::SaslFinal(server_final)) => {
404                Ok(SaslEvent::Final {
405                    conn: self.transition(),
406                    server_final,
407                })
408            }
409            (_, authentication) => Err((self, authentication)),
410        }
411    }
412
413    /// Projects an authentication error which terminates an active SASL exchange.
414    ///
415    /// # Errors
416    ///
417    /// Returns the live connection and message for an illegal response.
418    pub fn offer_backend(
419        self,
420        message: codec::BackendMessage,
421    ) -> Result<SaslEvent<S>, (Self, codec::BackendMessage)> {
422        match message {
423            codec::BackendMessage::Authentication(authentication) => self
424                .offer(authentication)
425                .map_err(|(conn, authentication)| {
426                    (conn, codec::BackendMessage::Authentication(authentication))
427                }),
428            codec::BackendMessage::ErrorResponse(error) => Ok(SaslEvent::Error {
429                conn: self.transition(),
430                error,
431            }),
432            message => Err((self, message)),
433        }
434    }
435}
436
437impl<S> Conn<S, SaslChallenge, Pristine> {
438    /// Sends the response to one received challenge and re-enters the SASL loop.
439    pub fn respond(self, response: Bytes) -> (Conn<S, Sasl>, codec::Frame) {
440        (
441            self.transition(),
442            codec::Frame {
443                tag: b'p',
444                body: response,
445            },
446        )
447    }
448}
449
450impl<S> Conn<S, SaslFinal, Pristine> {
451    /// Records that custom SCRAM logic verified the received server-final value.
452    pub fn verified(self) -> Conn<S, AwaitingAuthOk> {
453        self.transition()
454    }
455}
456
457impl<S> Conn<S, AwaitingAuthOk, Pristine> {
458    /// Requires backend evidence that authentication succeeded or failed.
459    ///
460    /// # Errors
461    ///
462    /// Returns the live connection and message for an illegal response.
463    pub fn offer(
464        self,
465        message: codec::BackendMessage,
466    ) -> Result<AuthCompletion<S>, (Self, codec::BackendMessage)> {
467        match (
468            auth_grammar::project_external(auth_grammar::RuntimeState::AwaitingAuthOk, &message),
469            message,
470        ) {
471            (
472                Some(auth_grammar::Event::Ok),
473                codec::BackendMessage::Authentication(codec::Authentication::Ok),
474            ) => Ok(AuthCompletion::Ok(self.transition())),
475            (Some(auth_grammar::Event::Error), codec::BackendMessage::ErrorResponse(error)) => {
476                Ok(AuthCompletion::Error {
477                    conn: self.transition(),
478                    error,
479                })
480            }
481            (_, message) => Err((self, message)),
482        }
483    }
484}
485
486fn project_authentication(
487    state: auth_grammar::RuntimeState,
488    authentication: &codec::Authentication,
489) -> Option<auth_grammar::Event> {
490    auth_grammar::project_external(
491        state,
492        &codec::BackendMessage::Authentication(authentication.clone()),
493    )
494}
495
496impl<S> Conn<S, AwaitingStartupReady, Pristine> {
497    /// Completes startup only when presented with a projected `ReadyForQuery`.
498    ///
499    /// # Errors
500    ///
501    /// Returns the unchanged connection and item when it is not `ReadyForQuery`.
502    pub fn offer_ready(self, item: SessionItem) -> Result<Conn<S, Ready>, (Self, SessionItem)> {
503        if matches!(
504            item,
505            SessionItem::ReadyForQuery {
506                status: codec::TransactionStatus::Idle,
507                parameters_changed: false,
508            }
509        ) {
510            Ok(self.transition())
511        } else {
512            Err((self, item))
513        }
514    }
515}
516
517fn sasl_initial<S>(
518    conn: Conn<S, SaslInitial>,
519    mechanism: &[u8],
520    initial: &[u8],
521) -> std::io::Result<(Conn<S, Sasl>, codec::Frame)> {
522    let length = i32::try_from(initial.len()).map_err(|_| {
523        std::io::Error::new(std::io::ErrorKind::InvalidInput, "SASL response too large")
524    })?;
525    let mut body = BytesMut::with_capacity(mechanism.len() + initial.len() + 5);
526    body.extend_from_slice(mechanism);
527    body.put_u8(0);
528    body.put_i32(length);
529    body.extend_from_slice(initial);
530    Ok((
531        conn.transition(),
532        codec::Frame {
533            tag: b'p',
534            body: body.freeze(),
535        },
536    ))
537}
538
539#[cfg(test)]
540mod tests {
541    use super::*;
542
543    #[derive(Debug)]
544    struct Tls(Vec<u8>);
545
546    impl TlsServerEndPoint for Tls {
547        fn tls_server_end_point(&self) -> &[u8] {
548            &self.0
549        }
550    }
551
552    #[test]
553    fn sasl_continue_alternates_challenge_and_response() {
554        let sasl: Conn<Tls, Sasl> = Conn::new(Tls(vec![1])).transition();
555        let SaslEvent::Continue { conn, challenge } = sasl
556            .offer(codec::Authentication::SaslContinue(Bytes::from_static(
557                b"challenge",
558            )))
559            .unwrap()
560        else {
561            panic!("challenge projected to the wrong branch")
562        };
563        assert_eq!(challenge, Bytes::from_static(b"challenge"));
564        let (sasl, response) = conn.respond(Bytes::from_static(b"response"));
565        assert_eq!(response.body, Bytes::from_static(b"response"));
566
567        let SaslEvent::Final { conn, server_final } = sasl
568            .offer(codec::Authentication::SaslFinal(Bytes::from_static(
569                b"verified",
570            )))
571            .unwrap()
572        else {
573            panic!("server final projected to the wrong branch")
574        };
575        assert_eq!(server_final, Bytes::from_static(b"verified"));
576        conn.verified().into_transport();
577    }
578
579    #[test]
580    fn gss_continuation_is_a_recursive_token_exchange() {
581        let auth: Conn<(), Auth> = Conn::new(()).transition();
582        let AuthOffer::Gss(response) = auth.offer(codec::Authentication::Gss).unwrap() else {
583            panic!("GSS request projected to the wrong branch")
584        };
585        let (waiting, frame) = response.respond(Bytes::from_static(b"client-token-1"));
586        assert_eq!(frame.body, Bytes::from_static(b"client-token-1"));
587
588        let TokenAuthEvent::Continue { conn, token } = waiting
589            .offer(codec::BackendMessage::Authentication(
590                codec::Authentication::GssContinue(Bytes::from_static(b"server-token")),
591            ))
592            .unwrap()
593        else {
594            panic!("GSS continuation projected to the wrong branch")
595        };
596        assert_eq!(token, Bytes::from_static(b"server-token"));
597        let (waiting, _) = conn.respond(Bytes::from_static(b"client-token-2"));
598        let TokenAuthEvent::Ok(awaiting_ready) = waiting
599            .offer(codec::BackendMessage::Authentication(
600                codec::Authentication::Ok,
601            ))
602            .unwrap()
603        else {
604            panic!("authentication success projected to the wrong branch")
605        };
606        awaiting_ready.into_transport();
607    }
608
609    #[test]
610    fn scram_plus_exposes_binding_to_custom_authentication_logic() {
611        let conn: Conn<Tls, SaslInitial> = Conn::new(Tls(vec![1, 2, 3])).transition();
612        assert_eq!(conn.tls_server_end_point(), [1, 2, 3]);
613
614        let (sasl, frame) = conn.scram_sha_256_plus(b"client-first").unwrap();
615        assert_eq!(frame.tag, b'p');
616        assert_eq!(
617            frame.body,
618            Bytes::from_static(b"SCRAM-SHA-256-PLUS\0\0\0\0\x0cclient-first")
619        );
620        let _transport = sasl.into_transport();
621    }
622
623    #[test]
624    fn protocol_negotiation_is_an_auth_self_loop() {
625        let auth: Conn<(), Auth> = Conn::new(()).transition();
626        let negotiation = codec::NegotiateProtocolVersion {
627            newest: crate::startup::ProtocolVersion::V3_2,
628            unsupported_options: vec![Bytes::from_static(b"_pq_.feature")],
629        };
630        let AuthEvent::Negotiate { conn, message } = auth
631            .offer_backend(codec::BackendMessage::NegotiateProtocolVersion(
632                negotiation.clone(),
633            ))
634            .unwrap()
635        else {
636            panic!("negotiation projected to the wrong branch")
637        };
638        assert_eq!(message, negotiation);
639        let AuthEvent::Authentication(AuthOffer::Ok(ready)) = conn
640            .offer_backend(codec::BackendMessage::Authentication(
641                codec::Authentication::Ok,
642            ))
643            .unwrap()
644        else {
645            panic!("authentication projected to the wrong branch")
646        };
647        ready.into_transport();
648    }
649
650    #[test]
651    fn authentication_completion_requires_backend_evidence() {
652        let awaiting: Conn<(), AwaitingAuthOk> = Conn::new(()).transition();
653        let AuthCompletion::Ok(startup) = awaiting
654            .offer(codec::BackendMessage::Authentication(
655                codec::Authentication::Ok,
656            ))
657            .unwrap()
658        else {
659            panic!("AuthenticationOk projected to the wrong branch")
660        };
661        startup.into_transport();
662
663        let awaiting: Conn<(), AwaitingAuthOk> = Conn::new(()).transition();
664        let error = codec::DiagnosticResponse {
665            fields: vec![codec::DiagnosticField {
666                code: b'M',
667                value: Bytes::from_static(b"password authentication failed"),
668            }],
669        };
670        let AuthCompletion::Error {
671            conn,
672            error: projected,
673        } = awaiting
674            .offer(codec::BackendMessage::ErrorResponse(error.clone()))
675            .unwrap()
676        else {
677            panic!("authentication error projected to the wrong branch")
678        };
679        assert_eq!(projected, error);
680        conn.into_transport();
681    }
682}