1use 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)]
20pub enum ServerAuth {}
22
23#[derive(Debug)]
24pub enum ServerStartupValidated {}
26
27#[derive(Debug)]
28pub enum ServerStartupRejected {}
30
31#[derive(Debug)]
32pub enum ServerPassword {}
34
35#[derive(Debug)]
36pub enum ServerSaslInitial {}
38
39#[derive(Debug)]
40pub enum ServerSasl {}
42
43#[derive(Debug)]
44pub enum ServerSaslResponse {}
46
47#[derive(Debug)]
48pub enum ServerAuthResponse {}
50
51#[derive(Debug)]
52pub enum ServerAuthPolicy {}
54
55#[derive(Debug)]
56pub enum ServerStartupReady {}
58
59#[derive(Debug)]
61pub enum ServerProtocolOffer<S, C> {
62 Supported {
64 conn: Conn<S, ServerStartupValidated, C>,
66 message: StartupMessage,
68 negotiate_to: Option<ProtocolVersion>,
70 },
71 Rejected {
73 conn: Conn<S, ServerStartupRejected, C>,
75 message: StartupMessage,
77 },
78}
79
80#[derive(Clone, Eq, PartialEq)]
82pub struct SaslInitialResponse {
83 pub mechanism: Bytes,
85 pub response: Option<Bytes>,
87}
88
89impl std::fmt::Debug for SaslInitialResponse {
90 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
91 formatter
92 .debug_struct("SaslInitialResponse")
93 .field("mechanism", &self.mechanism)
94 .field("response", &self.response.as_ref().map(Bytes::len))
95 .finish()
96 }
97}
98
99pub type ServerProjection<T, S, Phase, C> = Result<T, Box<(Conn<S, Phase, C>, FrontendMessage)>>;
101
102pub type PasswordProjection<S, C> =
104 ServerProjection<(Conn<S, ServerAuth, C>, Bytes), S, ServerPassword, C>;
105
106pub type SaslInitialProjection<S, C> =
108 ServerProjection<(Conn<S, ServerSasl, C>, SaslInitialResponse), S, ServerSaslInitial, C>;
109pub type SaslResponseProjection<S, C> =
111 ServerProjection<(Conn<S, ServerSasl, C>, Bytes), S, ServerSaslResponse, C>;
112pub type TokenResponseProjection<S, C> =
114 ServerProjection<(Conn<S, ServerAuthPolicy, C>, Bytes), S, ServerAuthResponse, C>;
115
116impl<S, C> Conn<S, Startup, C> {
117 pub fn validate_protocol(
119 self,
120 message: StartupMessage,
121 newest: ProtocolVersion,
122 ) -> ServerProtocolOffer<S, C> {
123 if message.version.major == newest.major {
124 let negotiate_to = (message.version.minor > newest.minor).then_some(newest);
125 ServerProtocolOffer::Supported {
126 conn: self.transition(),
127 message,
128 negotiate_to,
129 }
130 } else {
131 ServerProtocolOffer::Rejected {
132 conn: self.transition(),
133 message,
134 }
135 }
136 }
137}
138
139impl<S, C> Conn<S, ServerStartupValidated, C> {
140 pub fn begin_server_auth(self) -> Conn<S, ServerAuth, C> {
142 self.transition()
143 }
144}
145
146impl<S, C> Conn<S, ServerStartupRejected, C> {
147 pub fn error(
153 self,
154 response: DiagnosticResponse,
155 ) -> io::Result<(Conn<S, Terminated, C>, Frame)> {
156 Ok((
157 self.transition(),
158 BackendMessage::ErrorResponse(response).to_frame()?,
159 ))
160 }
161}
162
163impl<S, C> Conn<S, ServerAuth, C> {
164 pub fn request_cleartext(self) -> io::Result<(Conn<S, ServerPassword, C>, Frame)> {
170 Ok((
171 self.transition(),
172 authentication_frame(Authentication::CleartextPassword)?,
173 ))
174 }
175
176 pub fn request_md5(self, salt: [u8; 4]) -> io::Result<(Conn<S, ServerPassword, C>, Frame)> {
182 Ok((
183 self.transition(),
184 authentication_frame(Authentication::Md5Password { salt })?,
185 ))
186 }
187
188 pub fn request_sasl(
194 self,
195 mechanisms: Vec<Bytes>,
196 ) -> io::Result<(Conn<S, ServerSaslInitial, C>, Frame)> {
197 Ok((
198 self.transition(),
199 authentication_frame(Authentication::Sasl { mechanisms })?,
200 ))
201 }
202
203 pub fn request_kerberos_v5(self) -> io::Result<(Conn<S, ServerAuthResponse, C>, Frame)> {
209 Ok((
210 self.transition(),
211 authentication_frame(Authentication::KerberosV5)?,
212 ))
213 }
214
215 pub fn request_gss(self) -> io::Result<(Conn<S, ServerAuthResponse, C>, Frame)> {
221 Ok((
222 self.transition(),
223 authentication_frame(Authentication::Gss)?,
224 ))
225 }
226
227 pub fn request_sspi(self) -> io::Result<(Conn<S, ServerAuthResponse, C>, Frame)> {
233 Ok((
234 self.transition(),
235 authentication_frame(Authentication::Sspi)?,
236 ))
237 }
238
239 pub fn authentication_ok(self) -> io::Result<(Conn<S, ServerStartupReady, C>, Frame)> {
245 Ok((self.transition(), authentication_frame(Authentication::Ok)?))
246 }
247}
248
249impl<S, C> Conn<S, ServerPassword, C> {
250 pub fn receive_password(self, message: FrontendMessage) -> PasswordProjection<S, C> {
256 match (
257 auth_grammar::project_external(auth_grammar::RuntimeState::PasswordResponse, &message),
258 message,
259 ) {
260 (Some(auth_grammar::Event::Response), FrontendMessage::PasswordResponse(body)) => {
261 match password_body(body) {
262 Ok(password) => Ok((self.transition(), password)),
263 Err(body) => Err(Box::new((self, FrontendMessage::PasswordResponse(body)))),
264 }
265 }
266 (_, other) => Err(Box::new((self, other))),
267 }
268 }
269}
270
271impl<S, C> Conn<S, ServerSaslInitial, C> {
272 pub fn receive_initial(self, message: FrontendMessage) -> SaslInitialProjection<S, C> {
278 match (
279 auth_grammar::project_external(auth_grammar::RuntimeState::SaslInitial, &message),
280 message,
281 ) {
282 (Some(auth_grammar::Event::Initial), FrontendMessage::PasswordResponse(body)) => {
283 match sasl_initial(body) {
284 Ok(initial) => Ok((self.transition(), initial)),
285 Err(body) => Err(Box::new((self, FrontendMessage::PasswordResponse(body)))),
286 }
287 }
288 (_, other) => Err(Box::new((self, other))),
289 }
290 }
291}
292
293impl<S, C> Conn<S, ServerSasl, C> {
294 pub fn continue_with(
300 self,
301 challenge: Bytes,
302 ) -> io::Result<(Conn<S, ServerSaslResponse, C>, Frame)> {
303 Ok((
304 self.transition(),
305 authentication_frame(Authentication::SaslContinue(challenge))?,
306 ))
307 }
308
309 pub fn finish(self, server_final: Bytes) -> io::Result<(Conn<S, ServerAuth, C>, Frame)> {
315 Ok((
316 self.transition(),
317 authentication_frame(Authentication::SaslFinal(server_final))?,
318 ))
319 }
320}
321
322impl<S, C> Conn<S, ServerSaslResponse, C> {
323 pub fn receive_response(self, message: FrontendMessage) -> SaslResponseProjection<S, C> {
329 match (
330 auth_grammar::project_external(auth_grammar::RuntimeState::SaslResponse, &message),
331 message,
332 ) {
333 (Some(auth_grammar::Event::Response), FrontendMessage::PasswordResponse(response)) => {
334 Ok((self.transition(), response))
335 }
336 (_, other) => Err(Box::new((self, other))),
337 }
338 }
339}
340
341impl<S, C> Conn<S, ServerAuthResponse, C> {
342 pub fn receive_response(self, message: FrontendMessage) -> TokenResponseProjection<S, C> {
348 match (
349 auth_grammar::project_external(auth_grammar::RuntimeState::TokenResponse, &message),
350 message,
351 ) {
352 (Some(auth_grammar::Event::Response), FrontendMessage::PasswordResponse(response)) => {
353 Ok((self.transition(), response))
354 }
355 (_, other) => Err(Box::new((self, other))),
356 }
357 }
358}
359
360impl<S, C> Conn<S, ServerAuthPolicy, C> {
361 pub fn continue_gss(self, token: Bytes) -> io::Result<(Conn<S, ServerAuthResponse, C>, Frame)> {
367 Ok((
368 self.transition(),
369 authentication_frame(Authentication::GssContinue(token))?,
370 ))
371 }
372
373 pub fn verified(self) -> Conn<S, ServerAuth, C> {
375 self.transition()
376 }
377}
378
379impl<S, C> Conn<S, ServerStartupReady, C> {
380 pub fn parameter_status(self, name: Bytes, value: Bytes) -> io::Result<(Self, Frame)> {
386 Ok((
387 self,
388 BackendMessage::ParameterStatus { name, value }.to_frame()?,
389 ))
390 }
391
392 pub fn backend_key_data(self, process_id: u32, secret_key: Bytes) -> io::Result<(Self, Frame)> {
398 if !(4..=256).contains(&secret_key.len()) {
399 return Err(io::Error::new(
400 io::ErrorKind::InvalidInput,
401 "cancellation key length is outside 4..=256",
402 ));
403 }
404 Ok((
405 self,
406 BackendMessage::BackendKeyData {
407 process_id,
408 secret_key,
409 }
410 .to_frame()?,
411 ))
412 }
413
414 pub fn negotiate_protocol(
420 self,
421 newest: ProtocolVersion,
422 unsupported_options: Vec<Bytes>,
423 ) -> io::Result<(Self, Frame)> {
424 Ok((
425 self,
426 BackendMessage::NegotiateProtocolVersion(NegotiateProtocolVersion {
427 newest,
428 unsupported_options,
429 })
430 .to_frame()?,
431 ))
432 }
433
434 pub fn ready(self) -> io::Result<(Conn<S, Ready, C>, Frame)> {
440 Ok((
441 self.transition(),
442 BackendMessage::ReadyForQuery(TransactionStatus::Idle).to_frame()?,
443 ))
444 }
445}
446
447fn authentication_frame(authentication: Authentication) -> io::Result<Frame> {
448 BackendMessage::Authentication(authentication).to_frame()
449}
450
451fn password_body(body: Bytes) -> Result<Bytes, Bytes> {
452 if body.last() == Some(&0) && !body[..body.len() - 1].contains(&0) {
453 Ok(body.slice(..body.len() - 1))
454 } else {
455 Err(body)
456 }
457}
458
459fn sasl_initial(mut body: Bytes) -> Result<SaslInitialResponse, Bytes> {
460 let original = body.clone();
461 let Some(nul) = body.iter().position(|byte| *byte == 0) else {
462 return Err(original);
463 };
464 let mechanism = body.split_to(nul);
465 body.advance(1);
466 if body.len() < 4 {
467 return Err(original);
468 }
469 let length = body.get_i32();
470 if length == -1 && body.is_empty() {
471 return Ok(SaslInitialResponse {
472 mechanism,
473 response: None,
474 });
475 }
476 let Ok(length) = usize::try_from(length) else {
477 return Err(original);
478 };
479 if body.len() != length {
480 return Err(original);
481 }
482 Ok(SaslInitialResponse {
483 mechanism,
484 response: Some(body),
485 })
486}
487
488#[cfg(test)]
489mod tests {
490 use super::*;
491 use crate::{
492 Pristine,
493 grammar::server_authentication,
494 middleware::{Identity, Middleware, ServerRole, TypedBackendMessage},
495 scram::{SCRAM_SHA_256, ScramServer, ServerChannelBinding},
496 };
497 use bytes::{BufMut as _, BytesMut};
498 use postgres_protocol::authentication::sasl::{ChannelBinding, ScramSha256};
499 use std::collections::BTreeMap;
500
501 fn validated_startup() -> Conn<(), ServerStartupValidated, Pristine> {
502 let startup: Conn<(), Startup, Pristine> = Conn::new(()).transition();
503 let message = StartupMessage {
504 version: ProtocolVersion::V3_2,
505 parameters: BTreeMap::new(),
506 };
507 let ServerProtocolOffer::Supported { conn, .. } =
508 startup.validate_protocol(message, ProtocolVersion::V3_2)
509 else {
510 panic!("supported protocol was rejected")
511 };
512 conn
513 }
514
515 #[test]
516 fn cleartext_response_returns_to_independent_policy_choice() {
517 let (password, request) = validated_startup()
518 .begin_server_auth()
519 .request_cleartext()
520 .unwrap();
521 assert_eq!(
522 request,
523 authentication_frame(Authentication::CleartextPassword).unwrap()
524 );
525
526 let (auth, response) = password
527 .receive_password(FrontendMessage::PasswordResponse(Bytes::from_static(
528 b"secret\0",
529 )))
530 .unwrap();
531 assert_eq!(response, Bytes::from_static(b"secret"));
532 let (ready, ok) = auth.authentication_ok().unwrap();
533 assert_eq!(ok, authentication_frame(Authentication::Ok).unwrap());
534 ready.into_transport();
535 }
536
537 #[tokio::test]
538 async fn server_outbound_middleware_includes_asynchronous_messages() {
539 let conn = validated_startup().begin_server_auth();
540 let message = TypedBackendMessage::<server_authentication::AuthInternalMessage>::try_from(
541 BackendMessage::NoticeResponse(DiagnosticResponse { fields: vec![] }),
542 )
543 .expect("NoticeResponse is legal without advancing authentication");
544 let mut middleware = Middleware::new((), Identity);
545
546 let output = conn
547 .intercept_outbound_typed::<ServerRole, BackendMessage, _, _>(&mut middleware, message)
548 .await
549 .expect("asynchronous server traffic remains phase legal");
550
551 assert!(matches!(output, TypedBackendMessage::Asynchronous(_)));
552 conn.into_transport();
553 }
554
555 #[test]
556 fn sasl_is_a_recursive_server_sub_session() {
557 let (initial, _) = validated_startup()
558 .begin_server_auth()
559 .request_sasl(vec![Bytes::from_static(b"SCRAM-SHA-256")])
560 .unwrap();
561 let body = Bytes::from_static(b"SCRAM-SHA-256\0\0\0\0\x03one");
562 let (sasl, initial) = initial
563 .receive_initial(FrontendMessage::PasswordResponse(body))
564 .unwrap();
565 assert_eq!(initial.response, Some(Bytes::from_static(b"one")));
566 let (sasl, _) = sasl
567 .continue_with(Bytes::from_static(b"challenge"))
568 .unwrap();
569 let (sasl, response) = sasl
570 .receive_response(FrontendMessage::PasswordResponse(Bytes::from_static(
571 b"two",
572 )))
573 .unwrap();
574 assert_eq!(response, Bytes::from_static(b"two"));
575 let (auth, _) = sasl.finish(Bytes::from_static(b"verified")).unwrap();
576 auth.into_transport();
577 }
578
579 #[test]
580 fn startup_completion_mints_keys_and_requires_ready() {
581 let (startup_ready, _) = validated_startup()
582 .begin_server_auth()
583 .authentication_ok()
584 .unwrap();
585 let (startup_ready, parameter) = startup_ready
586 .parameter_status(
587 Bytes::from_static(b"server_version"),
588 Bytes::from_static(b"18"),
589 )
590 .unwrap();
591 assert_eq!(parameter.tag, b'S');
592 let (startup_ready, key) = startup_ready
593 .backend_key_data(42, Bytes::from_static(b"secret-key"))
594 .unwrap();
595 assert_eq!(key.tag, b'K');
596 let (ready, frame) = startup_ready.ready().unwrap();
597 assert_eq!(frame.body, Bytes::from_static(b"I"));
598 ready.into_transport();
599 }
600
601 #[test]
602 fn protocol_validation_negotiates_minor_and_rejects_major() {
603 let startup: Conn<(), Startup> = Conn::new(()).transition();
604 let message = StartupMessage {
605 version: ProtocolVersion { major: 3, minor: 9 },
606 parameters: BTreeMap::new(),
607 };
608 let ServerProtocolOffer::Supported {
609 conn, negotiate_to, ..
610 } = startup.validate_protocol(message, ProtocolVersion::V3_2)
611 else {
612 panic!("compatible major was rejected")
613 };
614 assert_eq!(negotiate_to, Some(ProtocolVersion::V3_2));
615 conn.into_transport();
616
617 let startup: Conn<(), Startup> = Conn::new(()).transition();
618 let message = StartupMessage {
619 version: ProtocolVersion { major: 4, minor: 0 },
620 parameters: BTreeMap::new(),
621 };
622 let ServerProtocolOffer::Rejected { conn, .. } =
623 startup.validate_protocol(message, ProtocolVersion::V3_2)
624 else {
625 panic!("unsupported major was accepted")
626 };
627 conn.into_transport();
628 }
629
630 #[test]
631 fn scram_server_engine_completes_the_typed_authentication_session() {
632 use crate::grammar::server_authentication::{Event, RuntimeFsm, RuntimeState};
633
634 let mut generated = RuntimeFsm::new();
635 generated.step(Event::Begin).unwrap();
636 let (initial_state, offer_frame) = validated_startup()
637 .begin_server_auth()
638 .request_sasl(vec![Bytes::from_static(SCRAM_SHA_256)])
639 .unwrap();
640 generated.step(Event::Sasl).unwrap();
641 assert_eq!(offer_frame.tag, b'R');
642
643 let mut client = ScramSha256::new(b"secret", ChannelBinding::unsupported());
644 let mut initial_body = BytesMut::new();
645 initial_body.extend_from_slice(SCRAM_SHA_256);
646 initial_body.put_u8(0);
647 initial_body.put_i32(i32::try_from(client.message().len()).unwrap());
648 initial_body.extend_from_slice(client.message());
649 let (sasl, initial) = initial_state
650 .receive_initial(FrontendMessage::PasswordResponse(initial_body.freeze()))
651 .unwrap();
652 generated.step(Event::Initial).unwrap();
653
654 let verifier = ScramServer::with_parameters(
655 b"secret",
656 b"fixed test salt".to_vec(),
657 crate::scram::DEFAULT_ITERATIONS,
658 ServerChannelBinding::None,
659 )
660 .unwrap();
661 let (exchange, challenge) = verifier
662 .start(&initial.mechanism, initial.response.as_deref().unwrap())
663 .unwrap();
664 let (sasl, challenge_frame) = sasl.continue_with(challenge.clone()).unwrap();
665 generated.step(Event::Continue).unwrap();
666 assert_eq!(challenge_frame.tag, b'R');
667 client.update(&challenge).unwrap();
668
669 let (sasl, response) = sasl
670 .receive_response(FrontendMessage::PasswordResponse(Bytes::copy_from_slice(
671 client.message(),
672 )))
673 .unwrap();
674 generated.step(Event::Response).unwrap();
675 let server_final = exchange.finish(&response).unwrap();
676 client.finish(&server_final).unwrap();
677 let (auth, final_frame) = sasl.finish(server_final).unwrap();
678 generated.step(Event::Final).unwrap();
679 assert_eq!(final_frame.tag, b'R');
680 let (startup_ready, ok) = auth.authentication_ok().unwrap();
681 generated.step(Event::Ok).unwrap();
682 assert_eq!(ok.tag, b'R');
683 assert_eq!(generated.state(), RuntimeState::StartupReady);
684 startup_ready.into_transport();
685 }
686}