pg_proto/
server_auth.rs

1//! Server-role authentication typestates for proxy-side client termination.
2
3use std::io;
4
5use bytes::{Buf, Bytes};
6
7use crate::{
8    Conn,
9    auth::Ready,
10    codec::{
11        Authentication, BackendMessage, DiagnosticResponse, Frame, FrontendMessage,
12        NegotiateProtocolVersion, TransactionStatus,
13    },
14    grammar::server_authentication as auth_grammar,
15    pre_startup::{Startup, Terminated},
16    startup::{ProtocolVersion, StartupMessage},
17};
18
19#[derive(Debug)]
20/// The proxy is selecting how to authenticate its client.
21pub enum ServerAuth {}
22
23#[derive(Debug)]
24/// The client's startup protocol version is supported.
25pub enum ServerStartupValidated {}
26
27#[derive(Debug)]
28/// The client's startup protocol version must be rejected.
29pub enum ServerStartupRejected {}
30
31#[derive(Debug)]
32/// A password response is expected from the client.
33pub enum ServerPassword {}
34
35#[derive(Debug)]
36/// A SASL initial response is expected from the client.
37pub enum ServerSaslInitial {}
38
39#[derive(Debug)]
40/// A recursive SASL response exchange is in progress.
41pub enum ServerSasl {}
42
43#[derive(Debug)]
44/// A GSS, SSPI, or Kerberos response token is expected.
45pub enum ServerAuthResponse {}
46
47#[derive(Debug)]
48/// Authentication succeeded and startup metadata may be sent before readiness.
49pub enum ServerStartupReady {}
50
51/// Result of validating a client's requested protocol version.
52#[derive(Debug)]
53pub enum ServerProtocolOffer<S, C> {
54    /// The major version is supported, with optional minor-version negotiation.
55    Supported {
56        /// Connection authorised to begin authentication.
57        conn: Conn<S, ServerStartupValidated, C>,
58        /// Inspected startup message.
59        message: StartupMessage,
60        /// Highest supported version to advertise when the requested minor is newer.
61        negotiate_to: Option<ProtocolVersion>,
62    },
63    /// The major protocol version is unsupported.
64    Rejected {
65        /// Connection authorised only to send an error and terminate.
66        conn: Conn<S, ServerStartupRejected, C>,
67        /// Rejected startup message.
68        message: StartupMessage,
69    },
70}
71
72/// A decoded SASL initial response selected by the client.
73#[derive(Clone, Eq, PartialEq)]
74pub struct SaslInitialResponse {
75    /// SASL mechanism selected by the client.
76    pub mechanism: Bytes,
77    /// Optional mechanism-specific initial response.
78    pub response: Option<Bytes>,
79}
80
81impl std::fmt::Debug for SaslInitialResponse {
82    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
83        formatter
84            .debug_struct("SaslInitialResponse")
85            .field("mechanism", &self.mechanism)
86            .field("response", &self.response.as_ref().map(Bytes::len))
87            .finish()
88    }
89}
90
91/// A rejected frontend message paired with the unchanged authentication state.
92pub type ServerProjection<T, S, Phase, C> = Result<T, Box<(Conn<S, Phase, C>, FrontendMessage)>>;
93
94/// Projection of a valid password body or the unchanged server-password state.
95pub type PasswordProjection<S, C> =
96    ServerProjection<(Conn<S, ServerAuth, C>, Bytes), S, ServerPassword, C>;
97
98/// Projection of a valid SASL initial response or the unchanged initial state.
99pub type SaslInitialProjection<S, C> =
100    ServerProjection<(Conn<S, ServerSasl, C>, SaslInitialResponse), S, ServerSaslInitial, C>;
101
102impl<S, C> Conn<S, Startup, C> {
103    /// Validates the startup protocol before authentication can begin.
104    pub fn validate_protocol(
105        self,
106        message: StartupMessage,
107        newest: ProtocolVersion,
108    ) -> ServerProtocolOffer<S, C> {
109        if message.version.major == newest.major {
110            let negotiate_to = (message.version.minor > newest.minor).then_some(newest);
111            ServerProtocolOffer::Supported {
112                conn: self.transition(),
113                message,
114                negotiate_to,
115            }
116        } else {
117            ServerProtocolOffer::Rejected {
118                conn: self.transition(),
119                message,
120            }
121        }
122    }
123}
124
125impl<S, C> Conn<S, ServerStartupValidated, C> {
126    /// Begins proxy-side authentication of a protocol-compatible client.
127    pub fn begin_server_auth(self) -> Conn<S, ServerAuth, C> {
128        self.transition()
129    }
130}
131
132impl<S, C> Conn<S, ServerStartupRejected, C> {
133    /// Rejects an unsupported major protocol version and terminates the session.
134    ///
135    /// # Errors
136    ///
137    /// Returns an error if a diagnostic field is invalid.
138    pub fn error(
139        self,
140        response: DiagnosticResponse,
141    ) -> io::Result<(Conn<S, Terminated, C>, Frame)> {
142        Ok((
143            self.transition(),
144            BackendMessage::ErrorResponse(response).to_frame()?,
145        ))
146    }
147}
148
149impl<S, C> Conn<S, ServerAuth, C> {
150    /// Requests a cleartext password from the client.
151    ///
152    /// # Errors
153    ///
154    /// Returns an error only if the fixed authentication message cannot be encoded.
155    pub fn request_cleartext(self) -> io::Result<(Conn<S, ServerPassword, C>, Frame)> {
156        Ok((
157            self.transition(),
158            authentication_frame(Authentication::CleartextPassword)?,
159        ))
160    }
161
162    /// Requests a `PostgreSQL` MD5 password response from the client.
163    ///
164    /// # Errors
165    ///
166    /// Returns an error only if the authentication message cannot be encoded.
167    pub fn request_md5(self, salt: [u8; 4]) -> io::Result<(Conn<S, ServerPassword, C>, Frame)> {
168        Ok((
169            self.transition(),
170            authentication_frame(Authentication::Md5Password { salt })?,
171        ))
172    }
173
174    /// Offers one or more SASL mechanisms to the client.
175    ///
176    /// # Errors
177    ///
178    /// Returns an error if a mechanism contains a NUL byte.
179    pub fn request_sasl(
180        self,
181        mechanisms: Vec<Bytes>,
182    ) -> io::Result<(Conn<S, ServerSaslInitial, C>, Frame)> {
183        Ok((
184            self.transition(),
185            authentication_frame(Authentication::Sasl { mechanisms })?,
186        ))
187    }
188
189    /// Requests Kerberos V5 authentication.
190    ///
191    /// # Errors
192    ///
193    /// Returns an error only if the fixed authentication message cannot be encoded.
194    pub fn request_kerberos_v5(self) -> io::Result<(Conn<S, ServerAuthResponse, C>, Frame)> {
195        Ok((
196            self.transition(),
197            authentication_frame(Authentication::KerberosV5)?,
198        ))
199    }
200
201    /// Requests GSSAPI authentication.
202    ///
203    /// # Errors
204    ///
205    /// Returns an error only if the fixed authentication message cannot be encoded.
206    pub fn request_gss(self) -> io::Result<(Conn<S, ServerAuthResponse, C>, Frame)> {
207        Ok((
208            self.transition(),
209            authentication_frame(Authentication::Gss)?,
210        ))
211    }
212
213    /// Requests SSPI authentication.
214    ///
215    /// # Errors
216    ///
217    /// Returns an error only if the fixed authentication message cannot be encoded.
218    pub fn request_sspi(self) -> io::Result<(Conn<S, ServerAuthResponse, C>, Frame)> {
219        Ok((
220            self.transition(),
221            authentication_frame(Authentication::Sspi)?,
222        ))
223    }
224
225    /// Confirms authentication and enters the startup-completion phase.
226    ///
227    /// # Errors
228    ///
229    /// Returns an error only if the fixed authentication message cannot be encoded.
230    pub fn authentication_ok(self) -> io::Result<(Conn<S, ServerStartupReady, C>, Frame)> {
231        Ok((self.transition(), authentication_frame(Authentication::Ok)?))
232    }
233}
234
235impl<S, C> Conn<S, ServerPassword, C> {
236    /// Projects the inspected password response and returns to policy evaluation.
237    ///
238    /// # Errors
239    ///
240    /// Returns the unchanged state and message if it is not a valid password response.
241    pub fn receive_password(self, message: FrontendMessage) -> PasswordProjection<S, C> {
242        match (
243            auth_grammar::project_external(auth_grammar::RuntimeState::PasswordResponse, &message),
244            message,
245        ) {
246            (Some(auth_grammar::Event::Response), FrontendMessage::PasswordResponse(body)) => {
247                match password_body(body) {
248                    Ok(password) => Ok((self.transition(), password)),
249                    Err(body) => Err(Box::new((self, FrontendMessage::PasswordResponse(body)))),
250                }
251            }
252            (_, other) => Err(Box::new((self, other))),
253        }
254    }
255}
256
257impl<S, C> Conn<S, ServerSaslInitial, C> {
258    /// Projects the client's selected SASL mechanism and initial response.
259    ///
260    /// # Errors
261    ///
262    /// Returns the unchanged state and message if the SASL initial response is malformed.
263    pub fn receive_initial(self, message: FrontendMessage) -> SaslInitialProjection<S, C> {
264        match (
265            auth_grammar::project_external(auth_grammar::RuntimeState::SaslInitial, &message),
266            message,
267        ) {
268            (Some(auth_grammar::Event::Initial), FrontendMessage::PasswordResponse(body)) => {
269                match sasl_initial(body) {
270                    Ok(initial) => Ok((self.transition(), initial)),
271                    Err(body) => Err(Box::new((self, FrontendMessage::PasswordResponse(body)))),
272                }
273            }
274            (_, other) => Err(Box::new((self, other))),
275        }
276    }
277}
278
279impl<S, C> Conn<S, ServerSasl, C> {
280    /// Projects one client SASL response and remains in the recursive exchange.
281    ///
282    /// # Errors
283    ///
284    /// Returns the unchanged state and message if it is not a SASL response.
285    pub fn receive_response(
286        self,
287        message: FrontendMessage,
288    ) -> ServerProjection<(Self, Bytes), S, ServerSasl, C> {
289        match (
290            auth_grammar::project_external(auth_grammar::RuntimeState::SaslResponse, &message),
291            message,
292        ) {
293            (Some(auth_grammar::Event::Response), FrontendMessage::PasswordResponse(response)) => {
294                Ok((self, response))
295            }
296            (_, other) => Err(Box::new((self, other))),
297        }
298    }
299
300    /// Sends a SASL challenge and remains in the recursive exchange.
301    ///
302    /// # Errors
303    ///
304    /// Returns an error only if the authentication message cannot be encoded.
305    pub fn continue_with(self, challenge: Bytes) -> io::Result<(Self, Frame)> {
306        Ok((
307            self,
308            authentication_frame(Authentication::SaslContinue(challenge))?,
309        ))
310    }
311
312    /// Sends verified server-final data and returns to authentication completion.
313    ///
314    /// # Errors
315    ///
316    /// Returns an error only if the authentication message cannot be encoded.
317    pub fn finish(self, server_final: Bytes) -> io::Result<(Conn<S, ServerAuth, C>, Frame)> {
318        Ok((
319            self.transition(),
320            authentication_frame(Authentication::SaslFinal(server_final))?,
321        ))
322    }
323}
324
325impl<S, C> Conn<S, ServerAuthResponse, C> {
326    /// Projects one GSS, SSPI, or Kerberos response token.
327    ///
328    /// # Errors
329    ///
330    /// Returns the unchanged state and message if it is not an authentication token.
331    pub fn receive_response(
332        self,
333        message: FrontendMessage,
334    ) -> ServerProjection<(Self, Bytes), S, ServerAuthResponse, C> {
335        match (
336            auth_grammar::project_external(auth_grammar::RuntimeState::TokenResponse, &message),
337            message,
338        ) {
339            (Some(auth_grammar::Event::Response), FrontendMessage::PasswordResponse(response)) => {
340                Ok((self, response))
341            }
342            (_, other) => Err(Box::new((self, other))),
343        }
344    }
345
346    /// Sends a GSS continuation token and remains in the authentication exchange.
347    ///
348    /// # Errors
349    ///
350    /// Returns an error only if the authentication message cannot be encoded.
351    pub fn continue_gss(self, token: Bytes) -> io::Result<(Self, Frame)> {
352        Ok((
353            self,
354            authentication_frame(Authentication::GssContinue(token))?,
355        ))
356    }
357
358    /// Returns to policy evaluation after the mechanism verifies its response.
359    pub fn verified(self) -> Conn<S, ServerAuth, C> {
360        self.transition()
361    }
362}
363
364impl<S, C> Conn<S, ServerStartupReady, C> {
365    /// Emits a startup parameter while remaining before `ReadyForQuery`.
366    ///
367    /// # Errors
368    ///
369    /// Returns an error if either value contains a NUL byte.
370    pub fn parameter_status(self, name: Bytes, value: Bytes) -> io::Result<(Self, Frame)> {
371        Ok((
372            self,
373            BackendMessage::ParameterStatus { name, value }.to_frame()?,
374        ))
375    }
376
377    /// Emits the proxy-minted cancellation key exposed to this client.
378    ///
379    /// # Errors
380    ///
381    /// Returns an error if the key is outside the protocol's 4–256 byte range.
382    pub fn backend_key_data(self, process_id: u32, secret_key: Bytes) -> io::Result<(Self, Frame)> {
383        if !(4..=256).contains(&secret_key.len()) {
384            return Err(io::Error::new(
385                io::ErrorKind::InvalidInput,
386                "cancellation key length is outside 4..=256",
387            ));
388        }
389        Ok((
390            self,
391            BackendMessage::BackendKeyData {
392                process_id,
393                secret_key,
394            }
395            .to_frame()?,
396        ))
397    }
398
399    /// Responds to unsupported protocol 3.1/3.2 startup options.
400    ///
401    /// # Errors
402    ///
403    /// Returns an error if an option name contains a NUL byte or counts overflow.
404    pub fn negotiate_protocol(
405        self,
406        newest: ProtocolVersion,
407        unsupported_options: Vec<Bytes>,
408    ) -> io::Result<(Self, Frame)> {
409        Ok((
410            self,
411            BackendMessage::NegotiateProtocolVersion(NegotiateProtocolVersion {
412                newest,
413                unsupported_options,
414            })
415            .to_frame()?,
416        ))
417    }
418
419    /// Completes startup with an idle `ReadyForQuery`.
420    ///
421    /// # Errors
422    ///
423    /// Returns an error only if the fixed message cannot be encoded.
424    pub fn ready(self) -> io::Result<(Conn<S, Ready, C>, Frame)> {
425        Ok((
426            self.transition(),
427            BackendMessage::ReadyForQuery(TransactionStatus::Idle).to_frame()?,
428        ))
429    }
430}
431
432fn authentication_frame(authentication: Authentication) -> io::Result<Frame> {
433    BackendMessage::Authentication(authentication).to_frame()
434}
435
436fn password_body(body: Bytes) -> Result<Bytes, Bytes> {
437    if body.last() == Some(&0) && !body[..body.len() - 1].contains(&0) {
438        Ok(body.slice(..body.len() - 1))
439    } else {
440        Err(body)
441    }
442}
443
444fn sasl_initial(mut body: Bytes) -> Result<SaslInitialResponse, Bytes> {
445    let original = body.clone();
446    let Some(nul) = body.iter().position(|byte| *byte == 0) else {
447        return Err(original);
448    };
449    let mechanism = body.split_to(nul);
450    body.advance(1);
451    if body.len() < 4 {
452        return Err(original);
453    }
454    let length = body.get_i32();
455    if length == -1 && body.is_empty() {
456        return Ok(SaslInitialResponse {
457            mechanism,
458            response: None,
459        });
460    }
461    let Ok(length) = usize::try_from(length) else {
462        return Err(original);
463    };
464    if body.len() != length {
465        return Err(original);
466    }
467    Ok(SaslInitialResponse {
468        mechanism,
469        response: Some(body),
470    })
471}
472
473#[cfg(test)]
474mod tests {
475    use super::*;
476    use crate::{
477        Pristine,
478        scram::{SCRAM_SHA_256, ScramServer, ServerChannelBinding},
479    };
480    use bytes::{BufMut as _, BytesMut};
481    use postgres_protocol::authentication::sasl::{ChannelBinding, ScramSha256};
482    use std::collections::BTreeMap;
483
484    fn validated_startup() -> Conn<(), ServerStartupValidated, Pristine> {
485        let startup: Conn<(), Startup, Pristine> = Conn::new(()).transition();
486        let message = StartupMessage {
487            version: ProtocolVersion::V3_2,
488            parameters: BTreeMap::new(),
489        };
490        let ServerProtocolOffer::Supported { conn, .. } =
491            startup.validate_protocol(message, ProtocolVersion::V3_2)
492        else {
493            panic!("supported protocol was rejected")
494        };
495        conn
496    }
497
498    #[test]
499    fn cleartext_response_returns_to_independent_policy_choice() {
500        let (password, request) = validated_startup()
501            .begin_server_auth()
502            .request_cleartext()
503            .unwrap();
504        assert_eq!(
505            request,
506            authentication_frame(Authentication::CleartextPassword).unwrap()
507        );
508
509        let (auth, response) = password
510            .receive_password(FrontendMessage::PasswordResponse(Bytes::from_static(
511                b"secret\0",
512            )))
513            .unwrap();
514        assert_eq!(response, Bytes::from_static(b"secret"));
515        let (ready, ok) = auth.authentication_ok().unwrap();
516        assert_eq!(ok, authentication_frame(Authentication::Ok).unwrap());
517        ready.into_transport();
518    }
519
520    #[test]
521    fn sasl_is_a_recursive_server_sub_session() {
522        let (initial, _) = validated_startup()
523            .begin_server_auth()
524            .request_sasl(vec![Bytes::from_static(b"SCRAM-SHA-256")])
525            .unwrap();
526        let body = Bytes::from_static(b"SCRAM-SHA-256\0\0\0\0\x03one");
527        let (sasl, initial) = initial
528            .receive_initial(FrontendMessage::PasswordResponse(body))
529            .unwrap();
530        assert_eq!(initial.response, Some(Bytes::from_static(b"one")));
531        let (sasl, _) = sasl
532            .continue_with(Bytes::from_static(b"challenge"))
533            .unwrap();
534        let (sasl, response) = sasl
535            .receive_response(FrontendMessage::PasswordResponse(Bytes::from_static(
536                b"two",
537            )))
538            .unwrap();
539        assert_eq!(response, Bytes::from_static(b"two"));
540        let (auth, _) = sasl.finish(Bytes::from_static(b"verified")).unwrap();
541        auth.into_transport();
542    }
543
544    #[test]
545    fn startup_completion_mints_keys_and_requires_ready() {
546        let (startup_ready, _) = validated_startup()
547            .begin_server_auth()
548            .authentication_ok()
549            .unwrap();
550        let (startup_ready, parameter) = startup_ready
551            .parameter_status(
552                Bytes::from_static(b"server_version"),
553                Bytes::from_static(b"18"),
554            )
555            .unwrap();
556        assert_eq!(parameter.tag, b'S');
557        let (startup_ready, key) = startup_ready
558            .backend_key_data(42, Bytes::from_static(b"secret-key"))
559            .unwrap();
560        assert_eq!(key.tag, b'K');
561        let (ready, frame) = startup_ready.ready().unwrap();
562        assert_eq!(frame.body, Bytes::from_static(b"I"));
563        ready.into_transport();
564    }
565
566    #[test]
567    fn protocol_validation_negotiates_minor_and_rejects_major() {
568        let startup: Conn<(), Startup> = Conn::new(()).transition();
569        let message = StartupMessage {
570            version: ProtocolVersion { major: 3, minor: 9 },
571            parameters: BTreeMap::new(),
572        };
573        let ServerProtocolOffer::Supported {
574            conn, negotiate_to, ..
575        } = startup.validate_protocol(message, ProtocolVersion::V3_2)
576        else {
577            panic!("compatible major was rejected")
578        };
579        assert_eq!(negotiate_to, Some(ProtocolVersion::V3_2));
580        conn.into_transport();
581
582        let startup: Conn<(), Startup> = Conn::new(()).transition();
583        let message = StartupMessage {
584            version: ProtocolVersion { major: 4, minor: 0 },
585            parameters: BTreeMap::new(),
586        };
587        let ServerProtocolOffer::Rejected { conn, .. } =
588            startup.validate_protocol(message, ProtocolVersion::V3_2)
589        else {
590            panic!("unsupported major was accepted")
591        };
592        conn.into_transport();
593    }
594
595    #[test]
596    fn scram_server_engine_completes_the_typed_authentication_session() {
597        use crate::grammar::server_authentication::{Event, RuntimeFsm, RuntimeState};
598
599        let mut generated = RuntimeFsm::new();
600        generated.step(Event::Begin).unwrap();
601        let (initial_state, offer_frame) = validated_startup()
602            .begin_server_auth()
603            .request_sasl(vec![Bytes::from_static(SCRAM_SHA_256)])
604            .unwrap();
605        generated.step(Event::Sasl).unwrap();
606        assert_eq!(offer_frame.tag, b'R');
607
608        let mut client = ScramSha256::new(b"secret", ChannelBinding::unsupported());
609        let mut initial_body = BytesMut::new();
610        initial_body.extend_from_slice(SCRAM_SHA_256);
611        initial_body.put_u8(0);
612        initial_body.put_i32(i32::try_from(client.message().len()).unwrap());
613        initial_body.extend_from_slice(client.message());
614        let (sasl, initial) = initial_state
615            .receive_initial(FrontendMessage::PasswordResponse(initial_body.freeze()))
616            .unwrap();
617        generated.step(Event::Initial).unwrap();
618
619        let verifier = ScramServer::with_parameters(
620            b"secret",
621            b"fixed test salt".to_vec(),
622            crate::scram::DEFAULT_ITERATIONS,
623            ServerChannelBinding::None,
624        )
625        .unwrap();
626        let (exchange, challenge) = verifier
627            .start(&initial.mechanism, initial.response.as_deref().unwrap())
628            .unwrap();
629        let (sasl, challenge_frame) = sasl.continue_with(challenge.clone()).unwrap();
630        generated.step(Event::Continue).unwrap();
631        assert_eq!(challenge_frame.tag, b'R');
632        client.update(&challenge).unwrap();
633
634        let (sasl, response) = sasl
635            .receive_response(FrontendMessage::PasswordResponse(Bytes::copy_from_slice(
636                client.message(),
637            )))
638            .unwrap();
639        generated.step(Event::Response).unwrap();
640        let server_final = exchange.finish(&response).unwrap();
641        client.finish(&server_final).unwrap();
642        let (auth, final_frame) = sasl.finish(server_final).unwrap();
643        generated.step(Event::Final).unwrap();
644        assert_eq!(final_frame.tag, b'R');
645        let (startup_ready, ok) = auth.authentication_ok().unwrap();
646        generated.step(Event::Ok).unwrap();
647        assert_eq!(ok.tag, b'R');
648        assert_eq!(generated.state(), RuntimeState::StartupReady);
649        startup_ready.into_transport();
650    }
651}
pg_proto/server_auth.rs

pg_proto/
server_auth.rs
