1use std::{fmt, future::Future, io, pin::Pin, sync::Arc};
4
5use bytes::Bytes;
6use rand::RngExt as _;
7use rustls::{ServerConfig, pki_types::CertificateDer};
8use tokio::io::{AsyncRead, AsyncWrite};
9
10use crate::ServerMiddleware as _;
11use crate::{
12 Conn,
13 auth::{Ready, TlsServerEndPoint},
14 codec::{
15 Backend, BackendMessage, DEFAULT_MAX_FRAME_LEN, Direction as _, Frontend, FrontendMessage,
16 },
17 pre_startup::{DEFAULT_MAX_PRE_STARTUP_PACKET_LEN, PreStartupOffer},
18 server_auth::ServerProtocolOffer,
19 startup::{ProtocolVersion, StartupMessage},
20 tls::ServerTls,
21 transport::Buffered,
22};
23
24fn tls_required_diagnostic() -> BackendMessage {
25 BackendMessage::ErrorResponse(crate::codec::DiagnosticResponse {
26 fields: vec![
27 crate::codec::DiagnosticField {
28 code: b'S',
29 value: Bytes::from_static(b"FATAL"),
30 },
31 crate::codec::DiagnosticField {
32 code: b'V',
33 value: Bytes::from_static(b"FATAL"),
34 },
35 crate::codec::DiagnosticField {
36 code: b'C',
37 value: Bytes::from_static(b"08001"),
38 },
39 crate::codec::DiagnosticField {
40 code: b'M',
41 value: Bytes::from_static(b"Transport Layer Security (TLS) connection is required"),
42 },
43 ],
44 })
45}
46
47#[allow(clippy::type_complexity)]
49pub(crate) trait StartupResolver<State, Peer, Identity> {
50 type Route;
51 type Error;
52
53 fn defer_ready(&self) -> bool {
55 false
56 }
57
58 fn resolve<'a>(
59 &'a mut self,
60 startup: &'a StartupMessage,
61 context: &'a ServerConnectionContext<Peer, Identity>,
62 state: &'a mut State,
63 ) -> Pin<Box<dyn Future<Output = Result<Self::Route, Self::Error>> + 'a>>;
64}
65
66#[derive(Debug)]
68pub(crate) enum RoutedAcceptError<TlsError, AuthenticationError, RouteError> {
69 Accept(AcceptError<TlsError, AuthenticationError>),
70 Route(RouteError),
71}
72
73impl<TlsError, AuthenticationError, RouteError> From<AcceptError<TlsError, AuthenticationError>>
74 for RoutedAcceptError<TlsError, AuthenticationError, RouteError>
75{
76 fn from(error: AcceptError<TlsError, AuthenticationError>) -> Self {
77 Self::Accept(error)
78 }
79}
80
81struct NoStartupRoute;
82
83impl<State, Peer, Identity> StartupResolver<State, Peer, Identity> for NoStartupRoute {
84 type Route = ();
85 type Error = std::convert::Infallible;
86
87 fn resolve<'a>(
88 &'a mut self,
89 _startup: &'a StartupMessage,
90 _context: &'a ServerConnectionContext<Peer, Identity>,
91 _state: &'a mut State,
92 ) -> Pin<Box<dyn Future<Output = Result<(), Self::Error>> + 'a>> {
93 Box::pin(async { Ok(()) })
94 }
95}
96
97#[derive(Clone)]
99pub struct ServerIdentity {
100 config: Arc<ServerConfig>,
101 leaf_certificate: CertificateDer<'static>,
102}
103
104impl ServerIdentity {
105 #[must_use]
111 pub const fn new(config: Arc<ServerConfig>, leaf_certificate: CertificateDer<'static>) -> Self {
112 Self {
113 config,
114 leaf_certificate,
115 }
116 }
117}
118
119impl fmt::Debug for ServerIdentity {
120 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
121 formatter.write_str("ServerIdentity([REDACTED])")
122 }
123}
124
125pub trait ServerIdentityProvider {
127 type Error;
129
130 fn resolve(&self) -> Result<ServerIdentity, Self::Error>;
136}
137
138#[derive(Clone, Debug, Eq, PartialEq)]
140pub enum NegotiatedServerTls {
141 Plaintext,
143 Tls {
145 server_end_point: Bytes,
147 },
148}
149
150pub struct ServerAuthenticationRequest<'a, Peer> {
152 startup: &'a StartupMessage,
153 tls: &'a NegotiatedServerTls,
154 peer: &'a Peer,
155}
156
157impl<Peer> Copy for ServerAuthenticationRequest<'_, Peer> {}
158
159impl<Peer> Clone for ServerAuthenticationRequest<'_, Peer> {
160 fn clone(&self) -> Self {
161 *self
162 }
163}
164
165impl<Peer> ServerAuthenticationRequest<'_, Peer> {
166 #[must_use]
168 pub const fn startup(&self) -> &StartupMessage {
169 self.startup
170 }
171
172 #[must_use]
174 pub const fn tls(&self) -> &NegotiatedServerTls {
175 self.tls
176 }
177
178 #[must_use]
180 pub const fn peer(&self) -> &Peer {
181 self.peer
182 }
183}
184
185#[derive(Clone, Debug, Eq, PartialEq)]
187pub enum ServerAuthenticationAction<Identity> {
188 Accept(Identity),
190 CleartextPassword,
192 Md5Password {
194 salt: [u8; 4],
196 },
197 Sasl {
199 mechanisms: Vec<Bytes>,
201 },
202 SaslContinue(Bytes),
204 SaslFinal {
206 server_final: Bytes,
208 identity: Identity,
210 },
211 KerberosV5,
213 Gss,
215 Sspi,
217 GssContinue(Bytes),
219}
220
221#[derive(Clone, Debug, Eq, PartialEq)]
223pub enum ServerAuthenticationResponse {
224 Password(Bytes),
226 SaslInitial {
228 mechanism: Bytes,
230 response: Option<Bytes>,
232 },
233 Sasl(Bytes),
235 Token(Bytes),
237}
238
239#[allow(async_fn_in_trait)]
243pub trait ServerAuthentication<Peer> {
244 type Identity;
246 type Error;
248
249 async fn start(
251 &mut self,
252 request: ServerAuthenticationRequest<'_, Peer>,
253 ) -> Result<ServerAuthenticationAction<Self::Identity>, Self::Error>;
254
255 async fn respond(
257 &mut self,
258 request: ServerAuthenticationRequest<'_, Peer>,
259 response: ServerAuthenticationResponse,
260 ) -> Result<ServerAuthenticationAction<Self::Identity>, Self::Error>;
261}
262
263pub trait ServerAuthenticationProvider {
265 type Authentication;
267
268 fn create(&self) -> Self::Authentication;
270}
271
272#[derive(Clone, Debug, Eq, PartialEq)]
274pub struct StaticMd5ServerCredentials {
275 username: Bytes,
276 password: Bytes,
277}
278
279impl StaticMd5ServerCredentials {
280 #[must_use]
282 pub fn new(username: impl Into<Bytes>, password: impl Into<Bytes>) -> Self {
283 Self {
284 username: username.into(),
285 password: password.into(),
286 }
287 }
288}
289
290#[derive(Clone, Debug)]
292pub struct StaticMd5ServerCredentialSession {
293 username: Bytes,
294 password: Bytes,
295 salt: [u8; 4],
296}
297
298impl ServerAuthenticationProvider for StaticMd5ServerCredentials {
299 type Authentication = StaticMd5ServerCredentialSession;
300
301 fn create(&self) -> Self::Authentication {
302 let mut salt = [0; 4];
303 rand::rng().fill(&mut salt);
304 StaticMd5ServerCredentialSession {
305 username: self.username.clone(),
306 password: self.password.clone(),
307 salt,
308 }
309 }
310}
311
312impl<Peer> ServerAuthentication<Peer> for StaticMd5ServerCredentialSession {
313 type Identity = ();
314 type Error = crate::StaticCredentialError;
315
316 async fn start(
317 &mut self,
318 _: ServerAuthenticationRequest<'_, Peer>,
319 ) -> Result<ServerAuthenticationAction<()>, Self::Error> {
320 Ok(ServerAuthenticationAction::Md5Password { salt: self.salt })
321 }
322
323 async fn respond(
324 &mut self,
325 _: ServerAuthenticationRequest<'_, Peer>,
326 response: ServerAuthenticationResponse,
327 ) -> Result<ServerAuthenticationAction<()>, Self::Error> {
328 let ServerAuthenticationResponse::Password(received) = response else {
329 return Err(crate::StaticCredentialError::AuthenticationFailed);
330 };
331 if !crate::credentials::verify_md5_response(
332 &received,
333 &self.username,
334 &self.password,
335 self.salt,
336 ) {
337 return Err(crate::StaticCredentialError::AuthenticationFailed);
338 }
339 Ok(ServerAuthenticationAction::Accept(()))
340 }
341}
342
343#[derive(Clone, Copy, Debug, Eq, PartialEq)]
345pub struct TrustIdentity;
346
347#[derive(Clone, Copy, Debug, Eq, PartialEq)]
349pub enum BuildServerError {
350 MissingTlsPolicy,
352 MissingAuthenticationPolicy,
354 InvalidProtocolLimits,
356}
357
358impl fmt::Display for BuildServerError {
359 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
360 formatter.write_str(match self {
361 Self::MissingTlsPolicy => "server TLS policy is required",
362 Self::MissingAuthenticationPolicy => "server authentication policy is required",
363 Self::InvalidProtocolLimits => {
364 "server protocol limits cannot support connection establishment"
365 }
366 })
367 }
368}
369
370impl std::error::Error for BuildServerError {}
371
372#[derive(Debug)]
374pub enum AcceptError<TlsError = NoServerIdentity, AuthenticationError = std::convert::Infallible> {
375 Io(io::Error),
377 UnsupportedProtocolVersion,
379 TlsRequired,
381 TlsIdentity(TlsError),
383 Authentication(AuthenticationError),
385 AuthenticationProtocol,
387}
388
389impl<TlsError, AuthenticationError> fmt::Display for AcceptError<TlsError, AuthenticationError> {
390 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
391 match self {
392 Self::Io(error) => error.fmt(formatter),
393 Self::UnsupportedProtocolVersion => {
394 formatter.write_str("unsupported PostgreSQL protocol version")
395 }
396 Self::TlsRequired => formatter.write_str("TLS is required before startup"),
397 Self::TlsIdentity(_) => formatter.write_str("server TLS identity is unavailable"),
398 Self::Authentication(_) => formatter.write_str("authentication rejected"),
399 Self::AuthenticationProtocol => formatter.write_str("invalid authentication response"),
400 }
401 }
402}
403
404impl<TlsError, AuthenticationError> std::error::Error for AcceptError<TlsError, AuthenticationError>
405where
406 TlsError: std::error::Error + 'static,
407 AuthenticationError: std::error::Error + 'static,
408{
409 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
410 match self {
411 Self::Io(error) => Some(error),
412 Self::TlsIdentity(error) => Some(error),
413 Self::Authentication(error) => Some(error),
414 Self::UnsupportedProtocolVersion | Self::TlsRequired | Self::AuthenticationProtocol => {
415 None
416 }
417 }
418 }
419}
420
421#[derive(Clone, Copy, Debug, Eq, PartialEq)]
423pub struct ServerTlsPolicy;
424
425impl ServerTlsPolicy {
426 #[allow(non_upper_case_globals)]
428 pub const Disabled: DisabledServerTls = DisabledServerTls;
429
430 #[allow(non_snake_case)]
432 pub const fn Optional<Provider>(provider: Provider) -> OptionalServerTls<Provider> {
433 OptionalServerTls(provider)
434 }
435
436 #[allow(non_snake_case)]
438 pub const fn Required<Provider>(provider: Provider) -> RequiredServerTls<Provider> {
439 RequiredServerTls(provider)
440 }
441}
442
443#[derive(Clone, Copy, Debug, Eq, PartialEq)]
445pub struct DisabledServerTls;
446
447#[derive(Clone)]
449pub struct OptionalServerTls<Provider>(Provider);
450
451#[derive(Clone)]
453pub struct RequiredServerTls<Provider>(Provider);
454
455impl<Provider> fmt::Debug for OptionalServerTls<Provider> {
456 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
457 formatter.write_str("OptionalServerTls([REDACTED])")
458 }
459}
460
461impl<Provider> fmt::Debug for RequiredServerTls<Provider> {
462 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
463 formatter.write_str("RequiredServerTls([REDACTED])")
464 }
465}
466
467#[derive(Clone, Copy, Debug, Eq, PartialEq)]
469pub struct NoServerIdentityProvider;
470
471#[derive(Clone, Copy, Debug, Eq, PartialEq)]
473pub struct NoServerIdentity;
474
475impl fmt::Display for NoServerIdentity {
476 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
477 formatter.write_str("disabled TLS has no identity")
478 }
479}
480
481impl std::error::Error for NoServerIdentity {}
482
483impl ServerIdentityProvider for NoServerIdentityProvider {
484 type Error = NoServerIdentity;
485
486 fn resolve(&self) -> Result<ServerIdentity, Self::Error> {
487 Err(NoServerIdentity)
488 }
489}
490
491mod sealed {
492 pub trait Sealed {}
493}
494
495#[doc(hidden)]
497pub trait ServerTlsConfiguration: sealed::Sealed {
498 type Provider: ServerIdentityProvider;
500 fn provider(&self) -> Option<&Self::Provider>;
502 fn required(&self) -> bool;
504 fn category(&self) -> &'static str;
506}
507
508impl sealed::Sealed for DisabledServerTls {}
509impl<Provider> sealed::Sealed for OptionalServerTls<Provider> {}
510impl<Provider> sealed::Sealed for RequiredServerTls<Provider> {}
511
512impl ServerTlsConfiguration for DisabledServerTls {
513 type Provider = NoServerIdentityProvider;
514 fn provider(&self) -> Option<&Self::Provider> {
515 None
516 }
517 fn required(&self) -> bool {
518 false
519 }
520 fn category(&self) -> &'static str {
521 "disabled"
522 }
523}
524
525impl<Provider: ServerIdentityProvider> ServerTlsConfiguration for OptionalServerTls<Provider> {
526 type Provider = Provider;
527 fn provider(&self) -> Option<&Self::Provider> {
528 Some(&self.0)
529 }
530 fn required(&self) -> bool {
531 false
532 }
533 fn category(&self) -> &'static str {
534 "optional"
535 }
536}
537
538impl<Provider: ServerIdentityProvider> ServerTlsConfiguration for RequiredServerTls<Provider> {
539 type Provider = Provider;
540 fn provider(&self) -> Option<&Self::Provider> {
541 Some(&self.0)
542 }
543 fn required(&self) -> bool {
544 true
545 }
546 fn category(&self) -> &'static str {
547 "required"
548 }
549}
550
551#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
553pub struct TrustServerAuthentication;
554
555impl ServerAuthenticationProvider for TrustServerAuthentication {
556 type Authentication = Self;
557
558 fn create(&self) -> Self::Authentication {
559 *self
560 }
561}
562
563impl<Peer> ServerAuthentication<Peer> for TrustServerAuthentication {
564 type Identity = TrustIdentity;
565 type Error = std::convert::Infallible;
566
567 async fn start(
568 &mut self,
569 _request: ServerAuthenticationRequest<'_, Peer>,
570 ) -> Result<ServerAuthenticationAction<Self::Identity>, Self::Error> {
571 Ok(ServerAuthenticationAction::Accept(TrustIdentity))
572 }
573
574 async fn respond(
575 &mut self,
576 _request: ServerAuthenticationRequest<'_, Peer>,
577 _response: ServerAuthenticationResponse,
578 ) -> Result<ServerAuthenticationAction<Self::Identity>, Self::Error> {
579 unreachable!("trust authentication accepts before a response")
580 }
581}
582
583#[derive(Clone, Copy, Debug, Eq, PartialEq)]
585pub struct ServerProtocolLimits {
586 max_frame_len: usize,
587 max_pre_startup_packet_len: usize,
588}
589
590impl ServerProtocolLimits {
591 #[must_use]
593 pub const fn with_max_frame_len(mut self, bytes: usize) -> Self {
594 self.max_frame_len = bytes;
595 self
596 }
597
598 #[must_use]
600 pub const fn with_max_pre_startup_packet_len(mut self, bytes: usize) -> Self {
601 self.max_pre_startup_packet_len = bytes;
602 self
603 }
604
605 const fn is_valid(self) -> bool {
606 self.max_frame_len >= 9
607 && self.max_frame_len <= i32::MAX as usize
608 && self.max_pre_startup_packet_len >= 8
609 && self.max_pre_startup_packet_len <= i32::MAX as usize
610 }
611}
612
613impl Default for ServerProtocolLimits {
614 fn default() -> Self {
615 Self {
616 max_frame_len: DEFAULT_MAX_FRAME_LEN,
617 max_pre_startup_packet_len: DEFAULT_MAX_PRE_STARTUP_PACKET_LEN,
618 }
619 }
620}
621
622#[derive(Clone)]
624pub struct Server<
625 Tls = DisabledServerTls,
626 Authentication = TrustServerAuthentication,
627 Middleware = IdentityServerHandler,
628> {
629 tls: Tls,
630 authentication: Authentication,
631 limits: ServerProtocolLimits,
632 middleware: Middleware,
633}
634
635impl Server {
636 #[must_use]
638 pub fn builder() -> ServerBuilder {
639 ServerBuilder::default()
640 }
641}
642
643impl<Tls: ServerTlsConfiguration, Authentication, Middleware> fmt::Debug
644 for Server<Tls, Authentication, Middleware>
645{
646 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
647 formatter
648 .debug_struct("Server")
649 .field("tls", &self.tls.category())
650 .field("authentication", &"<redacted>")
651 .field("limits", &self.limits)
652 .finish_non_exhaustive()
653 }
654}
655
656#[derive(Clone)]
658pub struct ServerBuilder<Tls = (), Authentication = (), Middleware = IdentityServerHandler> {
659 tls: Option<Tls>,
660 authentication: Option<Authentication>,
661 limits: ServerProtocolLimits,
662 middleware: Middleware,
663}
664
665impl<Tls, Authentication> fmt::Debug for ServerBuilder<Tls, Authentication> {
666 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
667 formatter
668 .debug_struct("ServerBuilder")
669 .field("tls_configured", &self.tls.is_some())
670 .field("authentication_configured", &self.authentication.is_some())
671 .field("limits", &self.limits)
672 .finish()
673 }
674}
675
676impl Default for ServerBuilder {
677 fn default() -> Self {
678 Self {
679 tls: None,
680 authentication: None,
681 limits: ServerProtocolLimits::default(),
682 middleware: IdentityServerHandler,
683 }
684 }
685}
686
687impl<Tls, Authentication, Middleware> ServerBuilder<Tls, Authentication, Middleware> {
688 #[must_use]
690 pub fn tls<Next>(self, policy: Next) -> ServerBuilder<Next, Authentication, Middleware> {
691 ServerBuilder {
692 tls: Some(policy),
693 authentication: self.authentication,
694 limits: self.limits,
695 middleware: self.middleware,
696 }
697 }
698
699 #[must_use]
701 pub fn authentication<Next>(self, policy: Next) -> ServerBuilder<Tls, Next, Middleware> {
702 ServerBuilder {
703 tls: self.tls,
704 authentication: Some(policy),
705 limits: self.limits,
706 middleware: self.middleware,
707 }
708 }
709
710 #[must_use]
712 pub fn limits(mut self, limits: ServerProtocolLimits) -> Self {
713 self.limits = limits;
714 self
715 }
716
717 #[must_use]
719 pub fn middleware<Next>(
720 self,
721 factory: Next,
722 ) -> ServerBuilder<Tls, Authentication, crate::MiddlewareChain<Middleware, Next>> {
723 ServerBuilder {
724 tls: self.tls,
725 authentication: self.authentication,
726 limits: self.limits,
727 middleware: crate::MiddlewareChain(self.middleware, factory),
728 }
729 }
730
731 pub fn build(self) -> Result<Server<Tls, Authentication, Middleware>, BuildServerError> {
738 let tls = self.tls.ok_or(BuildServerError::MissingTlsPolicy)?;
739 let authentication = self
740 .authentication
741 .ok_or(BuildServerError::MissingAuthenticationPolicy)?;
742 if !self.limits.is_valid() {
743 return Err(BuildServerError::InvalidProtocolLimits);
744 }
745 Ok(Server {
746 tls,
747 authentication,
748 limits: self.limits,
749 middleware: self.middleware,
750 })
751 }
752}
753
754#[derive(Clone, Debug, Eq, PartialEq)]
756pub struct ServerConnectionContext<Peer, Identity> {
757 peer: Peer,
758 tls: Option<NegotiatedServerTls>,
759 identity: Option<Identity>,
760}
761
762impl<Peer, Identity> ServerConnectionContext<Peer, Identity> {
763 #[must_use]
765 pub const fn peer(&self) -> &Peer {
766 &self.peer
767 }
768
769 #[must_use]
775 pub const fn tls(&self) -> &NegotiatedServerTls {
776 match &self.tls {
777 Some(tls) => tls,
778 None => panic!("TLS is not known before pre-startup negotiation"),
779 }
780 }
781
782 #[must_use]
784 pub const fn tls_if_known(&self) -> Option<&NegotiatedServerTls> {
785 self.tls.as_ref()
786 }
787
788 #[must_use]
794 pub const fn identity(&self) -> &Identity {
795 match &self.identity {
796 Some(identity) => identity,
797 None => panic!("identity is not known before authentication"),
798 }
799 }
800
801 #[must_use]
803 pub const fn identity_if_known(&self) -> Option<&Identity> {
804 self.identity.as_ref()
805 }
806}
807
808#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
810pub struct IdentityServerHandler;
811impl<C> crate::MiddlewareFactory<C> for IdentityServerHandler {
812 type Handler = Self;
813 fn create(&self, _: &C) -> Self {
814 *self
815 }
816}
817impl<S, C> crate::ServerMiddleware<S, C> for IdentityServerHandler {}
818
819#[derive(Clone, Eq, PartialEq)]
821pub struct CancellationRequest {
822 process_id: u32,
823 secret_key: Bytes,
824}
825
826impl fmt::Debug for CancellationRequest {
827 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
828 formatter
829 .debug_struct("CancellationRequest")
830 .field("process_id", &self.process_id)
831 .field("secret_key", &"[REDACTED]")
832 .finish()
833 }
834}
835
836impl CancellationRequest {
837 #[must_use]
839 pub const fn process_id(&self) -> u32 {
840 self.process_id
841 }
842
843 #[must_use]
845 pub fn secret_key(&self) -> &[u8] {
846 &self.secret_key
847 }
848}
849
850#[derive(Debug)]
852#[allow(clippy::large_enum_variant)]
853pub enum ServerAccept<
854 Transport,
855 State,
856 Peer,
857 Identity = TrustIdentity,
858 Handler = IdentityServerHandler,
859> {
860 Session(ServerConnection<Transport, State, Peer, Identity, Handler>),
862 Cancellation(ServerCancellation<Transport, State, Peer, Handler>),
864}
865
866pub type ServerAcceptFuture<
868 'a,
869 Transport,
870 State,
871 Peer,
872 Identity,
873 Handler,
874 TlsError,
875 AuthenticationError,
876> = Pin<
877 Box<
878 dyn Future<
879 Output = Result<
880 ServerAccept<Transport, State, Peer, Identity, Handler>,
881 AcceptError<TlsError, AuthenticationError>,
882 >,
883 > + 'a,
884 >,
885>;
886
887#[derive(Debug)]
889pub struct ServerConnection<
890 Transport,
891 State,
892 Peer,
893 Identity = TrustIdentity,
894 Handler = IdentityServerHandler,
895> {
896 core: ServerConnectionCore<Transport, Peer, Identity, Handler>,
897 state: State,
898}
899
900#[derive(Debug)]
901pub(crate) struct ServerConnectionCore<Transport, Peer, Identity, Handler> {
902 conn: ServerConnectionInner<Transport>,
903 startup: StartupMessage,
904 handler: Handler,
905 context: ServerConnectionContext<Peer, Identity>,
906}
907
908#[derive(Debug)]
909enum ServerConnectionInner<Transport> {
910 Plaintext(Box<Conn<Buffered<Transport, Frontend>, Ready>>),
911 Tls(Box<Conn<Buffered<ServerTls<Transport>, Frontend>, Ready>>),
912}
913
914#[derive(Debug)]
916pub enum AcceptedServerTransport<Transport> {
917 Plaintext(Transport),
919 Tls(Box<ServerTls<Transport>>),
921}
922
923impl<Transport, State, Peer, Identity, Handler>
924 ServerConnection<Transport, State, Peer, Identity, Handler>
925{
926 #[must_use]
928 pub const fn context(&self) -> &ServerConnectionContext<Peer, Identity> {
929 &self.core.context
930 }
931
932 #[must_use]
935 pub const fn state(&self) -> &State {
936 &self.state
937 }
938
939 pub(crate) fn into_core_and_state(
940 self,
941 ) -> (
942 ServerConnectionCore<Transport, Peer, Identity, Handler>,
943 State,
944 ) {
945 (self.core, self.state)
946 }
947
948 #[must_use]
950 pub const fn startup(&self) -> &StartupMessage {
951 &self.core.startup
952 }
953
954 pub async fn receive_wire(&mut self) -> io::Result<FrontendMessage>
965 where
966 Transport: AsyncRead + AsyncWrite + Unpin,
967 Handler: crate::ServerMiddleware<State, ServerConnectionContext<Peer, Identity>>,
968 {
969 let message = self.core.receive_wire_raw().await?;
970 Ok(self.core.intercept_frontend(&mut self.state, message))
971 }
972
973 pub async fn send_wire(&mut self, message: BackendMessage) -> io::Result<()>
982 where
983 Transport: AsyncRead + AsyncWrite + Unpin,
984 Handler: crate::ServerMiddleware<State, ServerConnectionContext<Peer, Identity>>,
985 {
986 let message = self.core.intercept_backend(&mut self.state, message);
987 self.core.send_wire_raw(message).await
988 }
989
990 #[must_use]
993 pub fn teardown(
994 self,
995 ) -> (
996 AcceptedServerTransport<Transport>,
997 State,
998 Handler,
999 ServerConnectionContext<Peer, Identity>,
1000 ) {
1001 let (transport, handler, context) = self.core.into_parts();
1002 (transport, self.state, handler, context)
1003 }
1004}
1005
1006impl<Transport, State, Peer, Identity, Handler>
1007 ServerConnection<Transport, State, Peer, Identity, Handler>
1008where
1009 Transport: AsyncRead + AsyncWrite + Unpin,
1010 Handler: crate::ServerMiddleware<State, ServerConnectionContext<Peer, Identity>>,
1011{
1012 pub(crate) async fn send_generated_error(&mut self, message: BackendMessage) -> io::Result<()> {
1013 let message = self.core.intercept_backend(&mut self.state, message);
1014 if !matches!(message, BackendMessage::ErrorResponse(_)) {
1015 return Err(io::Error::new(
1016 io::ErrorKind::InvalidData,
1017 "middleware rejected generated diagnostic",
1018 ));
1019 }
1020 self.core.send_wire_raw(message).await
1021 }
1022}
1023
1024impl<Transport, Peer, Identity, Handler> ServerConnectionCore<Transport, Peer, Identity, Handler> {
1025 pub(crate) const fn context(&self) -> &ServerConnectionContext<Peer, Identity> {
1026 &self.context
1027 }
1028
1029 pub(crate) async fn receive_wire_raw(&mut self) -> io::Result<FrontendMessage>
1030 where
1031 Transport: AsyncRead + AsyncWrite + Unpin,
1032 {
1033 match &mut self.conn {
1034 ServerConnectionInner::Plaintext(conn) => conn.receive_frontend_wire().await,
1035 ServerConnectionInner::Tls(conn) => conn.receive_frontend_wire().await,
1036 }
1037 }
1038
1039 pub(crate) fn intercept_frontend<State>(
1040 &mut self,
1041 state: &mut State,
1042 message: FrontendMessage,
1043 ) -> FrontendMessage
1044 where
1045 Handler: crate::ServerMiddleware<State, ServerConnectionContext<Peer, Identity>>,
1046 {
1047 self.handler.frontend(&self.context, state, message)
1048 }
1049
1050 pub(crate) fn intercept_backend<State>(
1051 &mut self,
1052 state: &mut State,
1053 message: BackendMessage,
1054 ) -> BackendMessage
1055 where
1056 Handler: crate::ServerMiddleware<State, ServerConnectionContext<Peer, Identity>>,
1057 {
1058 self.handler.backend(&self.context, state, message)
1059 }
1060
1061 pub(crate) async fn send_wire_raw(&mut self, message: BackendMessage) -> io::Result<()>
1062 where
1063 Transport: AsyncRead + AsyncWrite + Unpin,
1064 {
1065 let frame = message.to_frame()?;
1066 match &mut self.conn {
1067 ServerConnectionInner::Plaintext(conn) => {
1068 conn.push_frame(frame)?;
1069 conn.flush().await
1070 }
1071 ServerConnectionInner::Tls(conn) => {
1072 conn.push_frame(frame)?;
1073 conn.flush().await
1074 }
1075 }
1076 }
1077
1078 pub(crate) fn into_parts(
1079 self,
1080 ) -> (
1081 AcceptedServerTransport<Transport>,
1082 Handler,
1083 ServerConnectionContext<Peer, Identity>,
1084 ) {
1085 let transport = match self.conn {
1086 ServerConnectionInner::Plaintext(conn) => {
1087 AcceptedServerTransport::Plaintext(conn.into_transport().into_inner())
1088 }
1089 ServerConnectionInner::Tls(conn) => {
1090 AcceptedServerTransport::Tls(Box::new(conn.into_transport().into_inner()))
1091 }
1092 };
1093 (transport, self.handler, self.context)
1094 }
1095}
1096
1097#[derive(Debug)]
1099pub struct ServerCancellation<Transport, State, Peer, Handler = IdentityServerHandler> {
1100 transport: AcceptedServerTransport<Transport>,
1101 request: CancellationRequest,
1102 state: State,
1103 handler: Handler,
1104 context: ServerConnectionContext<Peer, ()>,
1105}
1106
1107impl<Transport, State, Peer, Handler> ServerCancellation<Transport, State, Peer, Handler> {
1108 #[must_use]
1110 pub const fn request(&self) -> &CancellationRequest {
1111 &self.request
1112 }
1113
1114 #[must_use]
1116 pub fn teardown(
1117 self,
1118 ) -> (
1119 AcceptedServerTransport<Transport>,
1120 CancellationRequest,
1121 State,
1122 Handler,
1123 ServerConnectionContext<Peer, ()>,
1124 ) {
1125 (
1126 self.transport,
1127 self.request,
1128 self.state,
1129 self.handler,
1130 self.context,
1131 )
1132 }
1133}
1134
1135impl<Tls, Authentication, Middleware> Server<Tls, Authentication, Middleware>
1136where
1137 Tls: ServerTlsConfiguration,
1138 Authentication: ServerAuthenticationProvider,
1139{
1140 #[allow(clippy::type_complexity)]
1156 pub fn accept<'a, Transport, State, Peer>(
1157 &'a self,
1158 transport: Transport,
1159 peer: Peer,
1160 state: State,
1161 ) -> ServerAcceptFuture<
1162 'a,
1163 Transport,
1164 State,
1165 Peer,
1166 <Authentication::Authentication as ServerAuthentication<Peer>>::Identity,
1167 <Middleware as crate::MiddlewareFactory<
1168 ServerConnectionContext<
1169 Peer,
1170 <Authentication::Authentication as ServerAuthentication<Peer>>::Identity,
1171 >,
1172 >>::Handler,
1173 <Tls::Provider as ServerIdentityProvider>::Error,
1174 <Authentication::Authentication as ServerAuthentication<Peer>>::Error,
1175 >
1176 where
1177 Transport: AsyncRead + AsyncWrite + Unpin + 'a,
1178 State: 'a,
1179 Peer: 'a,
1180 Authentication::Authentication: ServerAuthentication<Peer>,
1181 Middleware: crate::MiddlewareFactory<
1182 ServerConnectionContext<
1183 Peer,
1184 <Authentication::Authentication as ServerAuthentication<Peer>>::Identity,
1185 >,
1186 >,
1187 <Middleware as crate::MiddlewareFactory<
1188 ServerConnectionContext<
1189 Peer,
1190 <Authentication::Authentication as ServerAuthentication<Peer>>::Identity,
1191 >,
1192 >>::Handler: crate::ServerMiddleware<
1193 State,
1194 ServerConnectionContext<
1195 Peer,
1196 <Authentication::Authentication as ServerAuthentication<Peer>>::Identity,
1197 >,
1198 >,
1199 {
1200 Box::pin(async move {
1201 let mut resolver = NoStartupRoute;
1202 self.accept_routed(transport, peer, state, &mut resolver)
1203 .await
1204 .map(|(accepted, _)| accepted)
1205 .map_err(|error| match error {
1206 RoutedAcceptError::Accept(error) => error,
1207 RoutedAcceptError::Route(never) => match never {},
1208 })
1209 })
1210 }
1211
1212 #[allow(clippy::too_many_lines)]
1213 pub(crate) async fn accept_routed<Transport, State, Peer, Resolver>(
1214 &self,
1215 transport: Transport,
1216 peer: Peer,
1217 mut state: State,
1218 resolver: &mut Resolver,
1219 ) -> Result<
1220 (
1221 ServerAccept<
1222 Transport,
1223 State,
1224 Peer,
1225 <Authentication::Authentication as ServerAuthentication<Peer>>::Identity,
1226 <Middleware as crate::MiddlewareFactory<
1227 ServerConnectionContext<
1228 Peer,
1229 <Authentication::Authentication as ServerAuthentication<Peer>>::Identity,
1230 >,
1231 >>::Handler,
1232 >,
1233 Option<Resolver::Route>,
1234 ),
1235 RoutedAcceptError<
1236 <Tls::Provider as ServerIdentityProvider>::Error,
1237 <Authentication::Authentication as ServerAuthentication<Peer>>::Error,
1238 Resolver::Error,
1239 >,
1240 >
1241 where
1242 Transport: AsyncRead + AsyncWrite + Unpin,
1243 Authentication::Authentication: ServerAuthentication<Peer>,
1244 Middleware: crate::MiddlewareFactory<
1245 ServerConnectionContext<
1246 Peer,
1247 <Authentication::Authentication as ServerAuthentication<Peer>>::Identity,
1248 >,
1249 >,
1250 <Middleware as crate::MiddlewareFactory<
1251 ServerConnectionContext<
1252 Peer,
1253 <Authentication::Authentication as ServerAuthentication<Peer>>::Identity,
1254 >,
1255 >>::Handler: crate::ServerMiddleware<
1256 State,
1257 ServerConnectionContext<
1258 Peer,
1259 <Authentication::Authentication as ServerAuthentication<Peer>>::Identity,
1260 >,
1261 >,
1262 Resolver: StartupResolver<
1263 State,
1264 Peer,
1265 <Authentication::Authentication as ServerAuthentication<Peer>>::Identity,
1266 >,
1267 {
1268 let mut context = ServerConnectionContext {
1269 peer,
1270 tls: None,
1271 identity: None,
1272 };
1273 let mut handler = self.middleware.create(&context);
1274 let buffered = self
1275 .buffer_transport(transport)
1276 .map_err(AcceptError::Io)
1277 .map_err(RoutedAcceptError::Accept)?;
1278 let mut conn = Conn::new(buffered);
1279
1280 loop {
1281 let message = match conn.receive_pre_startup_wire().await {
1282 Ok(message) => message,
1283 Err(error) => {
1284 let _ = conn.into_transport();
1285 return Err(RoutedAcceptError::Accept(AcceptError::Io(error)));
1286 }
1287 };
1288 let message = handler.pre_startup(&context, &mut state, message);
1289 match conn.offer_pre_startup(message) {
1290 PreStartupOffer::Ssl(decision) => match self.tls.provider() {
1291 None => {
1292 conn = decision.decline_ssl();
1293 conn = flush_or_abort(conn).await?;
1294 }
1295 Some(provider) => {
1296 let identity = match provider.resolve() {
1297 Ok(identity) => identity,
1298 Err(error) => {
1299 let _ = decision.into_transport();
1300 return Err(RoutedAcceptError::Accept(AcceptError::TlsIdentity(
1301 error,
1302 )));
1303 }
1304 };
1305 let handshake = decision.approve_ssl();
1306 let handshake = flush_or_abort(handshake).await?;
1307 let encrypted = handshake
1308 .accept_tls(identity.config, identity.leaf_certificate)
1309 .await
1310 .map_err(AcceptError::Io)?;
1311 return accept_encrypted(
1312 encrypted,
1313 context,
1314 state,
1315 handler,
1316 &self.authentication,
1317 resolver,
1318 )
1319 .await;
1320 }
1321 },
1322 PreStartupOffer::Gss(decision) => {
1323 conn = decision.decline_gss();
1324 conn = flush_or_abort(conn).await?;
1325 }
1326 PreStartupOffer::Cancel {
1327 conn: terminal,
1328 process_id,
1329 secret_key,
1330 } => {
1331 context.tls = Some(NegotiatedServerTls::Plaintext);
1332 let request = handler.cancellation(
1333 &context,
1334 &mut state,
1335 CancellationRequest {
1336 process_id,
1337 secret_key,
1338 },
1339 );
1340 return Ok((
1341 ServerAccept::Cancellation(ServerCancellation {
1342 transport: AcceptedServerTransport::Plaintext(
1343 terminal.into_transport().into_inner(),
1344 ),
1345 request,
1346 state,
1347 handler,
1348 context: ServerConnectionContext {
1349 peer: context.peer,
1350 tls: Some(NegotiatedServerTls::Plaintext),
1351 identity: None,
1352 },
1353 }),
1354 None,
1355 ));
1356 }
1357 PreStartupOffer::Startup {
1358 conn: mut startup_conn,
1359 message,
1360 } => {
1361 if self.tls.required() {
1362 if let Ok(frame) = tls_required_diagnostic().to_frame() {
1363 let _ = startup_conn.push_frame(frame);
1364 let _ = startup_conn.flush().await;
1365 }
1366 let _ = startup_conn.into_transport();
1367 return Err(RoutedAcceptError::Accept(AcceptError::TlsRequired));
1368 }
1369 context.tls = Some(NegotiatedServerTls::Plaintext);
1370 let message = handler.startup(&context, &mut state, message);
1371 let route = resolver
1372 .resolve(&message, &context, &mut state)
1373 .await
1374 .map_err(RoutedAcceptError::Route)?;
1375 let ready = complete_auth(
1376 startup_conn,
1377 &message,
1378 &self.authentication,
1379 &mut context,
1380 &mut state,
1381 &mut handler,
1382 resolver.defer_ready(),
1383 )
1384 .await?;
1385 return Ok((
1386 ServerAccept::Session(ServerConnection {
1387 core: ServerConnectionCore {
1388 conn: ServerConnectionInner::Plaintext(Box::new(ready)),
1389 startup: message,
1390 handler,
1391 context,
1392 },
1393 state,
1394 }),
1395 Some(route),
1396 ));
1397 }
1398 }
1399 }
1400 }
1401
1402 fn buffer_transport<Transport>(
1403 &self,
1404 transport: Transport,
1405 ) -> io::Result<Buffered<Transport, Frontend>> {
1406 Buffered::with_limits_frontend(
1407 transport,
1408 self.limits.max_frame_len,
1409 self.limits.max_pre_startup_packet_len,
1410 )
1411 }
1412}
1413
1414async fn accept_encrypted<Transport, State, Peer, Authentication, TlsError, Handler, Resolver>(
1415 mut conn: Conn<Buffered<ServerTls<Transport>, Frontend>, crate::pre_startup::PreStartup>,
1416 mut context: ServerConnectionContext<
1417 Peer,
1418 <Authentication::Authentication as ServerAuthentication<Peer>>::Identity,
1419 >,
1420 mut state: State,
1421 mut handler: Handler,
1422 authentication: &Authentication,
1423 resolver: &mut Resolver,
1424) -> Result<
1425 (
1426 ServerAccept<
1427 Transport,
1428 State,
1429 Peer,
1430 <Authentication::Authentication as ServerAuthentication<Peer>>::Identity,
1431 Handler,
1432 >,
1433 Option<Resolver::Route>,
1434 ),
1435 RoutedAcceptError<
1436 TlsError,
1437 <Authentication::Authentication as ServerAuthentication<Peer>>::Error,
1438 Resolver::Error,
1439 >,
1440>
1441where
1442 Transport: AsyncRead + AsyncWrite + Unpin,
1443 Authentication: ServerAuthenticationProvider,
1444 Authentication::Authentication: ServerAuthentication<Peer>,
1445 Handler: crate::ServerMiddleware<
1446 State,
1447 ServerConnectionContext<
1448 Peer,
1449 <Authentication::Authentication as ServerAuthentication<Peer>>::Identity,
1450 >,
1451 >,
1452 Resolver: StartupResolver<
1453 State,
1454 Peer,
1455 <Authentication::Authentication as ServerAuthentication<Peer>>::Identity,
1456 >,
1457{
1458 let negotiated_tls = NegotiatedServerTls::Tls {
1459 server_end_point: Bytes::copy_from_slice(conn.transport().get_ref().tls_server_end_point()),
1460 };
1461 context.tls = Some(negotiated_tls.clone());
1462 loop {
1463 let message = match conn.receive_pre_startup_wire().await {
1464 Ok(message) => message,
1465 Err(error) => {
1466 let _ = conn.into_transport();
1467 return Err(RoutedAcceptError::Accept(AcceptError::Io(error)));
1468 }
1469 };
1470 let message = handler.pre_startup(&context, &mut state, message);
1471 match conn.offer_pre_startup(message) {
1472 PreStartupOffer::Ssl(decision) => {
1473 conn = decision.decline_ssl();
1474 conn = flush_or_abort(conn).await?;
1475 }
1476 PreStartupOffer::Gss(decision) => {
1477 conn = decision.decline_gss();
1478 conn = flush_or_abort(conn).await?;
1479 }
1480 PreStartupOffer::Cancel {
1481 conn: terminal,
1482 process_id,
1483 secret_key,
1484 } => {
1485 let request = handler.cancellation(
1486 &context,
1487 &mut state,
1488 CancellationRequest {
1489 process_id,
1490 secret_key,
1491 },
1492 );
1493 return Ok((
1494 ServerAccept::Cancellation(ServerCancellation {
1495 transport: AcceptedServerTransport::Tls(Box::new(
1496 terminal.into_transport().into_inner(),
1497 )),
1498 request,
1499 state,
1500 handler,
1501 context: ServerConnectionContext {
1502 peer: context.peer,
1503 tls: context.tls,
1504 identity: None,
1505 },
1506 }),
1507 None,
1508 ));
1509 }
1510 PreStartupOffer::Startup {
1511 conn: startup_conn,
1512 message,
1513 } => {
1514 let message = handler.startup(&context, &mut state, message);
1515 let route = resolver
1516 .resolve(&message, &context, &mut state)
1517 .await
1518 .map_err(RoutedAcceptError::Route)?;
1519 let ready = complete_auth(
1520 startup_conn,
1521 &message,
1522 authentication,
1523 &mut context,
1524 &mut state,
1525 &mut handler,
1526 resolver.defer_ready(),
1527 )
1528 .await?;
1529 return Ok((
1530 ServerAccept::Session(ServerConnection {
1531 core: ServerConnectionCore {
1532 conn: ServerConnectionInner::Tls(Box::new(ready)),
1533 startup: message,
1534 handler,
1535 context,
1536 },
1537 state,
1538 }),
1539 Some(route),
1540 ));
1541 }
1542 }
1543 }
1544}
1545
1546#[allow(clippy::too_many_arguments, clippy::too_many_lines)]
1547async fn complete_auth<I, Authentication, Peer, TlsError, State, Handler>(
1548 startup_conn: Conn<Buffered<I, Frontend>, crate::pre_startup::Startup>,
1549 message: &StartupMessage,
1550 provider: &Authentication,
1551 context: &mut ServerConnectionContext<
1552 Peer,
1553 <Authentication::Authentication as ServerAuthentication<Peer>>::Identity,
1554 >,
1555 state: &mut State,
1556 handler: &mut Handler,
1557 defer_ready: bool,
1558) -> Result<
1559 Conn<Buffered<I, Frontend>, Ready>,
1560 AcceptError<TlsError, <Authentication::Authentication as ServerAuthentication<Peer>>::Error>,
1561>
1562where
1563 I: AsyncRead + AsyncWrite + Unpin,
1564 Authentication: ServerAuthenticationProvider,
1565 Authentication::Authentication: ServerAuthentication<Peer>,
1566 Handler: crate::ServerMiddleware<
1567 State,
1568 ServerConnectionContext<
1569 Peer,
1570 <Authentication::Authentication as ServerAuthentication<Peer>>::Identity,
1571 >,
1572 >,
1573{
1574 let validated = match startup_conn.validate_protocol(message.clone(), ProtocolVersion::V3_2) {
1575 ServerProtocolOffer::Supported { conn, .. } => conn,
1576 ServerProtocolOffer::Rejected { conn, .. } => {
1577 let _ = conn.into_transport();
1578 return Err(AcceptError::UnsupportedProtocolVersion);
1579 }
1580 };
1581 let mut policy = provider.create();
1582 let auth = validated.begin_server_auth();
1583 let request = ServerAuthenticationRequest {
1584 startup: message,
1585 tls: context.tls(),
1586 peer: context.peer(),
1587 };
1588 let action = match policy.start(request).await {
1589 Ok(action) => action,
1590 Err(error) => {
1591 let _ = auth.into_transport();
1592 return Err(AcceptError::Authentication(error));
1593 }
1594 };
1595 let (auth, identity, final_frame) = match action {
1596 ServerAuthenticationAction::Accept(identity) => (auth, identity, None),
1597 action @ (ServerAuthenticationAction::CleartextPassword
1598 | ServerAuthenticationAction::Md5Password { .. }) => {
1599 let (waiting, frame) = match action {
1600 ServerAuthenticationAction::CleartextPassword => auth.request_cleartext(),
1601 ServerAuthenticationAction::Md5Password { salt } => auth.request_md5(salt),
1602 _ => unreachable!("matched password action"),
1603 }
1604 .map_err(AcceptError::Io)?;
1605 let frame = intercept_server_backend(handler, context, state, frame)
1606 .map_err(AcceptError::Io)?;
1607 let waiting = push_or_abort(waiting, frame)?;
1608 let waiting = flush_or_abort(waiting).await?;
1609 let (waiting, wire) = receive_frontend_or_abort(waiting).await?;
1610 let wire = handler.frontend(context, state, wire);
1611 let (auth, credential) = match waiting.receive_password(wire) {
1612 Ok(response) => response,
1613 Err(rejected) => {
1614 let (waiting, _) = *rejected;
1615 let _ = waiting.into_transport();
1616 return Err(AcceptError::AuthenticationProtocol);
1617 }
1618 };
1619 match policy
1620 .respond(request, ServerAuthenticationResponse::Password(credential))
1621 .await
1622 {
1623 Ok(ServerAuthenticationAction::Accept(identity)) => (auth, identity, None),
1624 Ok(_) => {
1625 let _ = auth.into_transport();
1626 return Err(AcceptError::AuthenticationProtocol);
1627 }
1628 Err(error) => {
1629 let _ = auth.into_transport();
1630 return Err(AcceptError::Authentication(error));
1631 }
1632 }
1633 }
1634 ServerAuthenticationAction::Sasl { mechanisms } => {
1635 authenticate_sasl(
1636 auth,
1637 mechanisms,
1638 &mut policy,
1639 request,
1640 context,
1641 state,
1642 handler,
1643 )
1644 .await?
1645 }
1646 action @ (ServerAuthenticationAction::KerberosV5
1647 | ServerAuthenticationAction::Gss
1648 | ServerAuthenticationAction::Sspi) => {
1649 authenticate_token(auth, action, &mut policy, request, context, state, handler).await?
1650 }
1651 ServerAuthenticationAction::SaslContinue(_)
1652 | ServerAuthenticationAction::SaslFinal { .. }
1653 | ServerAuthenticationAction::GssContinue(_) => {
1654 let _ = auth.into_transport();
1655 return Err(AcceptError::AuthenticationProtocol);
1656 }
1657 };
1658 context.identity = Some(identity);
1659 let (mut startup_ready, _authentication_ok) =
1660 auth.authentication_ok().map_err(AcceptError::Io)?;
1661 if let Some(final_frame) = final_frame {
1662 let final_frame = intercept_server_backend(handler, context, state, final_frame)
1663 .map_err(AcceptError::Io)?;
1664 startup_ready = push_or_abort(startup_ready, final_frame)?;
1665 }
1666 let authentication_ok = handler
1667 .backend(
1668 context,
1669 state,
1670 BackendMessage::Authentication(crate::codec::Authentication::Ok),
1671 )
1672 .to_frame()
1673 .map_err(AcceptError::Io)?;
1674 let startup_ready = push_or_abort(startup_ready, authentication_ok)?;
1675 let (ready, _ready_frame) = startup_ready.ready().map_err(AcceptError::Io)?;
1676 let ready = if defer_ready {
1677 ready
1678 } else {
1679 let ready_frame = handler
1680 .backend(
1681 context,
1682 state,
1683 BackendMessage::ReadyForQuery(crate::codec::TransactionStatus::Idle),
1684 )
1685 .to_frame()
1686 .map_err(AcceptError::Io)?;
1687 push_or_abort(ready, ready_frame)?
1688 };
1689 let ready = flush_or_abort(ready).await?;
1690 Ok(ready)
1691}
1692
1693async fn authenticate_sasl<I, Policy, Peer, TlsError, State, Handler>(
1694 auth: Conn<Buffered<I, Frontend>, crate::server_auth::ServerAuth>,
1695 mechanisms: Vec<Bytes>,
1696 policy: &mut Policy,
1697 request: ServerAuthenticationRequest<'_, Peer>,
1698 context: &ServerConnectionContext<Peer, Policy::Identity>,
1699 state: &mut State,
1700 handler: &mut Handler,
1701) -> Result<
1702 (
1703 Conn<Buffered<I, Frontend>, crate::server_auth::ServerAuth>,
1704 Policy::Identity,
1705 Option<crate::codec::Frame>,
1706 ),
1707 AcceptError<TlsError, Policy::Error>,
1708>
1709where
1710 I: AsyncRead + AsyncWrite + Unpin,
1711 Policy: ServerAuthentication<Peer>,
1712 Handler: crate::ServerMiddleware<State, ServerConnectionContext<Peer, Policy::Identity>>,
1713{
1714 if mechanisms.iter().any(|mechanism| mechanism.contains(&0)) {
1715 let _ = auth.into_transport();
1716 return Err(AcceptError::Io(io::Error::new(
1717 io::ErrorKind::InvalidInput,
1718 "SASL mechanism contains NUL",
1719 )));
1720 }
1721 let (initial, frame) = auth.request_sasl(mechanisms).map_err(AcceptError::Io)?;
1722 let frame =
1723 intercept_server_backend(handler, context, state, frame).map_err(AcceptError::Io)?;
1724 let initial = push_or_abort(initial, frame)?;
1725 let initial = flush_or_abort(initial).await?;
1726 let (initial, wire) = receive_frontend_or_abort(initial).await?;
1727 let wire = handler.frontend(context, state, wire);
1728 let (mut sasl, initial_response) = match initial.receive_initial(wire) {
1729 Ok(response) => response,
1730 Err(rejected) => {
1731 let (initial, _) = *rejected;
1732 let _ = initial.into_transport();
1733 return Err(AcceptError::AuthenticationProtocol);
1734 }
1735 };
1736 let mut action = match policy
1737 .respond(
1738 request,
1739 ServerAuthenticationResponse::SaslInitial {
1740 mechanism: initial_response.mechanism,
1741 response: initial_response.response,
1742 },
1743 )
1744 .await
1745 {
1746 Ok(action) => action,
1747 Err(error) => {
1748 let _ = sasl.into_transport();
1749 return Err(AcceptError::Authentication(error));
1750 }
1751 };
1752 loop {
1753 match action {
1754 ServerAuthenticationAction::SaslContinue(challenge) => {
1755 let (waiting, frame) = sasl.continue_with(challenge).map_err(AcceptError::Io)?;
1756 let frame = intercept_server_backend(handler, context, state, frame)
1757 .map_err(AcceptError::Io)?;
1758 let waiting = push_or_abort(waiting, frame)?;
1759 let waiting = flush_or_abort(waiting).await?;
1760 let (waiting, wire) = receive_frontend_or_abort(waiting).await?;
1761 let wire = handler.frontend(context, state, wire);
1762 let (next, response) = match waiting.receive_response(wire) {
1763 Ok(response) => response,
1764 Err(rejected) => {
1765 let (waiting, _) = *rejected;
1766 let _ = waiting.into_transport();
1767 return Err(AcceptError::AuthenticationProtocol);
1768 }
1769 };
1770 sasl = next;
1771 action = match policy
1772 .respond(request, ServerAuthenticationResponse::Sasl(response))
1773 .await
1774 {
1775 Ok(action) => action,
1776 Err(error) => {
1777 let _ = sasl.into_transport();
1778 return Err(AcceptError::Authentication(error));
1779 }
1780 };
1781 }
1782 ServerAuthenticationAction::SaslFinal {
1783 server_final,
1784 identity,
1785 } => {
1786 let (auth, frame) = sasl.finish(server_final).map_err(AcceptError::Io)?;
1787 return Ok((auth, identity, Some(frame)));
1788 }
1789 _ => {
1790 let _ = sasl.into_transport();
1791 return Err(AcceptError::AuthenticationProtocol);
1792 }
1793 }
1794 }
1795}
1796
1797async fn authenticate_token<I, Policy, Peer, TlsError, State, Handler>(
1798 auth: Conn<Buffered<I, Frontend>, crate::server_auth::ServerAuth>,
1799 initial_action: ServerAuthenticationAction<Policy::Identity>,
1800 policy: &mut Policy,
1801 request: ServerAuthenticationRequest<'_, Peer>,
1802 context: &ServerConnectionContext<Peer, Policy::Identity>,
1803 state: &mut State,
1804 handler: &mut Handler,
1805) -> Result<
1806 (
1807 Conn<Buffered<I, Frontend>, crate::server_auth::ServerAuth>,
1808 Policy::Identity,
1809 Option<crate::codec::Frame>,
1810 ),
1811 AcceptError<TlsError, Policy::Error>,
1812>
1813where
1814 I: AsyncRead + AsyncWrite + Unpin,
1815 Policy: ServerAuthentication<Peer>,
1816 Handler: crate::ServerMiddleware<State, ServerConnectionContext<Peer, Policy::Identity>>,
1817{
1818 let (waiting, frame) = match initial_action {
1819 ServerAuthenticationAction::KerberosV5 => auth.request_kerberos_v5(),
1820 ServerAuthenticationAction::Gss => auth.request_gss(),
1821 ServerAuthenticationAction::Sspi => auth.request_sspi(),
1822 _ => unreachable!("matched initial token action"),
1823 }
1824 .map_err(AcceptError::Io)?;
1825 let frame =
1826 intercept_server_backend(handler, context, state, frame).map_err(AcceptError::Io)?;
1827 let waiting = push_or_abort(waiting, frame)?;
1828 let mut waiting = flush_or_abort(waiting).await?;
1829 loop {
1830 let received = receive_frontend_or_abort(waiting).await?;
1831 waiting = received.0;
1832 let wire = handler.frontend(context, state, received.1);
1833 let (decision, token) = match waiting.receive_response(wire) {
1834 Ok(response) => response,
1835 Err(rejected) => {
1836 let (waiting, _) = *rejected;
1837 let _ = waiting.into_transport();
1838 return Err(AcceptError::AuthenticationProtocol);
1839 }
1840 };
1841 let action = match policy
1842 .respond(request, ServerAuthenticationResponse::Token(token))
1843 .await
1844 {
1845 Ok(action) => action,
1846 Err(error) => {
1847 let _ = decision.into_transport();
1848 return Err(AcceptError::Authentication(error));
1849 }
1850 };
1851 match action {
1852 ServerAuthenticationAction::Accept(identity) => {
1853 return Ok((decision.verified(), identity, None));
1854 }
1855 ServerAuthenticationAction::GssContinue(token) => {
1856 let (next, frame) = decision.continue_gss(token).map_err(AcceptError::Io)?;
1857 let frame = intercept_server_backend(handler, context, state, frame)
1858 .map_err(AcceptError::Io)?;
1859 let next = push_or_abort(next, frame)?;
1860 waiting = flush_or_abort(next).await?;
1861 }
1862 _ => {
1863 let _ = decision.into_transport();
1864 return Err(AcceptError::AuthenticationProtocol);
1865 }
1866 }
1867 }
1868}
1869
1870fn intercept_server_backend<State, Context, Handler>(
1871 handler: &mut Handler,
1872 context: &Context,
1873 state: &mut State,
1874 frame: crate::codec::Frame,
1875) -> io::Result<crate::codec::Frame>
1876where
1877 Handler: crate::ServerMiddleware<State, Context>,
1878{
1879 handler
1880 .backend(context, state, Backend::decode(frame)?)
1881 .to_frame()
1882}
1883
1884async fn flush_or_abort<I, D, Phase, TlsError, AuthenticationError>(
1885 mut conn: Conn<Buffered<I, D>, Phase>,
1886) -> Result<Conn<Buffered<I, D>, Phase>, AcceptError<TlsError, AuthenticationError>>
1887where
1888 I: AsyncWrite + Unpin,
1889{
1890 if let Err(error) = conn.flush().await {
1891 let _ = conn.into_transport();
1892 return Err(AcceptError::Io(error));
1893 }
1894 Ok(conn)
1895}
1896
1897fn push_or_abort<I, D, Phase, TlsError, AuthenticationError>(
1898 mut conn: Conn<Buffered<I, D>, Phase>,
1899 frame: crate::codec::Frame,
1900) -> Result<Conn<Buffered<I, D>, Phase>, AcceptError<TlsError, AuthenticationError>> {
1901 if let Err(error) = conn.push_frame(frame) {
1902 let _ = conn.into_transport();
1903 return Err(AcceptError::Io(error));
1904 }
1905 Ok(conn)
1906}
1907
1908async fn receive_frontend_or_abort<I, Phase, TlsError, AuthenticationError>(
1909 mut conn: Conn<Buffered<I, Frontend>, Phase>,
1910) -> Result<
1911 (Conn<Buffered<I, Frontend>, Phase>, FrontendMessage),
1912 AcceptError<TlsError, AuthenticationError>,
1913>
1914where
1915 I: AsyncRead + Unpin,
1916{
1917 match conn.receive_frontend_wire().await {
1918 Ok(message) => Ok((conn, message)),
1919 Err(error) => {
1920 let _ = conn.into_transport();
1921 Err(AcceptError::Io(error))
1922 }
1923 }
1924}