1use 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)]
13pub enum Auth {}
15
16#[derive(Debug)]
17pub enum PasswordResponse {}
19
20#[derive(Debug)]
21pub enum SaslInitial {}
23
24#[derive(Debug)]
25pub enum Sasl {}
27
28#[derive(Debug)]
29pub enum SaslChallenge {}
31
32#[derive(Debug)]
33pub enum SaslFinal {}
35
36#[derive(Debug)]
37pub enum TokenResponse {}
39
40#[derive(Debug)]
41pub enum TokenChallenge {}
43
44#[derive(Debug)]
45pub enum AwaitingAuthOk {}
47
48#[derive(Debug)]
49pub enum Ready {}
51
52#[derive(Debug)]
53pub enum AwaitingStartupReady {}
55
56pub trait TlsServerEndPoint {
58 fn tls_server_end_point(&self) -> &[u8];
60}
61
62impl<S: TlsServerEndPoint, Phase, Cleanliness> Conn<S, Phase, Cleanliness> {
63 #[must_use]
65 pub fn tls_server_end_point(&self) -> &[u8] {
66 self.transport().tls_server_end_point()
67 }
68}
69
70#[derive(Debug)]
72pub enum AuthOffer<S> {
73 Ok(Conn<S, AwaitingStartupReady>),
75 Cleartext(Conn<S, PasswordResponse>),
77 Md5 {
79 conn: Conn<S, PasswordResponse>,
81 salt: [u8; 4],
83 },
84 Sasl {
86 conn: Conn<S, SaslInitial>,
88 mechanisms: Vec<Bytes>,
90 },
91 Gss(Conn<S, TokenResponse>),
93 Sspi(Conn<S, TokenResponse>),
95 KerberosV5(Conn<S, TokenResponse>),
97}
98
99#[derive(Debug)]
101pub enum AuthEvent<S> {
102 Authentication(AuthOffer<S>),
104 Negotiate {
106 conn: Conn<S, Auth>,
108 message: codec::NegotiateProtocolVersion,
110 },
111 Error {
113 conn: Conn<S, Terminated>,
115 error: codec::DiagnosticResponse,
117 },
118}
119
120#[derive(Debug)]
122pub enum SaslEvent<S> {
123 Continue {
125 conn: Conn<S, SaslChallenge>,
127 challenge: Bytes,
129 },
130 Final {
132 conn: Conn<S, SaslFinal>,
134 server_final: Bytes,
136 },
137 Error {
139 conn: Conn<S, Terminated>,
141 error: codec::DiagnosticResponse,
143 },
144}
145
146#[derive(Debug)]
148pub enum AuthCompletion<S> {
149 Ok(Conn<S, AwaitingStartupReady>),
151 Error {
153 conn: Conn<S, Terminated>,
155 error: codec::DiagnosticResponse,
157 },
158}
159
160#[derive(Debug)]
162pub enum TokenAuthEvent<S> {
163 Continue {
165 conn: Conn<S, TokenResponse>,
167 token: Bytes,
169 },
170 Ok(Conn<S, AwaitingStartupReady>),
172 Error {
174 conn: Conn<S, Terminated>,
176 error: codec::DiagnosticResponse,
178 },
179}
180
181impl<S> Conn<S, Startup, Pristine> {
182 pub fn authentication(self) -> Conn<S, Auth> {
184 self.transition()
185 }
186}
187
188impl<S> Conn<S, Auth, Pristine> {
189 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 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 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 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 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 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 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 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 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 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 pub fn verified(self) -> Conn<S, AwaitingAuthOk> {
453 self.transition()
454 }
455}
456
457impl<S> Conn<S, AwaitingAuthOk, Pristine> {
458 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 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}