1use std::{
4 collections::BTreeMap,
5 convert::Infallible,
6 fmt,
7 future::Future,
8 io,
9 pin::Pin,
10 sync::Arc,
11 task::{Context, Poll},
12};
13
14use bytes::Bytes;
15use tokio::io::ReadBuf;
16use tokio::io::{AsyncRead, AsyncReadExt as _, AsyncWrite, AsyncWriteExt as _};
17
18use crate::ClientMiddleware as _;
19use crate::{
20 Conn, Pristine,
21 auth::{AuthEvent, AuthOffer, Ready},
22 codec::{Backend, FrontendMessage},
23 demux::SessionItem,
24 session::{ReadyState, SimpleTransition},
25 startup::{ProtocolVersion, StartupMessage},
26 transport::Buffered,
27};
28
29#[derive(Clone, Copy, Debug, Eq, PartialEq)]
31pub enum BuildError {
32 MissingConnector,
34 MissingTls,
36 MissingAuthentication,
38}
39
40impl fmt::Display for BuildError {
41 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
42 formatter.write_str(match self {
43 Self::MissingConnector => "client connector is required",
44 Self::MissingTls => "an explicit client TLS policy is required",
45 Self::MissingAuthentication => "an explicit client authentication policy is required",
46 })
47 }
48}
49
50impl std::error::Error for BuildError {}
51
52#[derive(Clone, Copy, Debug, Eq, PartialEq)]
54pub enum ClientTlsPolicy {
55 Disabled,
57}
58
59impl ClientTlsPolicy {
60 #[must_use]
62 pub fn libpq<Provider>(
63 mode: crate::pre_startup::SslMode,
64 provider: Provider,
65 ) -> ReloadableClientTls<Provider> {
66 ReloadableClientTls { mode, provider }
67 }
68}
69
70#[derive(Clone)]
72pub struct ReloadableClientTls<Provider> {
73 mode: crate::pre_startup::SslMode,
74 provider: Provider,
75}
76
77impl<Provider> fmt::Debug for ReloadableClientTls<Provider> {
78 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
79 formatter
80 .debug_struct("Libpq")
81 .field("mode", &self.mode)
82 .field("provider", &"<redacted>")
83 .finish()
84 }
85}
86
87pub trait ClientTlsConfiguration {
89 type Provider: ClientTlsProvider;
91
92 fn configured(&self) -> Option<(crate::pre_startup::SslMode, &Self::Provider)>;
94}
95
96impl ClientTlsConfiguration for ClientTlsPolicy {
97 type Provider = ();
98
99 fn configured(&self) -> Option<(crate::pre_startup::SslMode, &Self::Provider)> {
100 None
101 }
102}
103
104impl<Provider: ClientTlsProvider> ClientTlsConfiguration for ReloadableClientTls<Provider> {
105 type Provider = Provider;
106
107 fn configured(&self) -> Option<(crate::pre_startup::SslMode, &Self::Provider)> {
108 Some((self.mode, &self.provider))
109 }
110}
111
112#[derive(Clone)]
119pub struct ClientTlsConfig {
120 server_name: rustls::pki_types::ServerName<'static>,
121 roots: rustls::RootCertStore,
122}
123
124impl ClientTlsConfig {
125 #[must_use]
127 pub fn new(
128 server_name: rustls::pki_types::ServerName<'static>,
129 roots: rustls::RootCertStore,
130 ) -> Self {
131 Self { server_name, roots }
132 }
133}
134
135impl fmt::Debug for ClientTlsConfig {
136 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
137 formatter.write_str("ClientTlsConfig(<redacted>)")
138 }
139}
140
141#[allow(async_fn_in_trait)]
145pub trait ClientTlsProvider {
146 type Error;
148
149 async fn resolve(&self, target: &ConnectTarget) -> Result<ClientTlsConfig, Self::Error>;
151}
152
153#[derive(Debug)]
155pub enum ClientTlsError<ProviderError> {
156 Provider(ProviderError),
158 Handshake(io::Error),
160}
161
162impl<ProviderError: fmt::Display> fmt::Display for ClientTlsError<ProviderError> {
163 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
164 match self {
165 Self::Provider(error) => error.fmt(formatter),
166 Self::Handshake(error) => error.fmt(formatter),
167 }
168 }
169}
170
171impl<ProviderError: std::error::Error + 'static> std::error::Error
172 for ClientTlsError<ProviderError>
173{
174 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
175 match self {
176 Self::Provider(error) => Some(error),
177 Self::Handshake(error) => Some(error),
178 }
179 }
180}
181
182impl ClientTlsProvider for () {
183 type Error = Infallible;
184
185 async fn resolve(&self, _target: &ConnectTarget) -> Result<ClientTlsConfig, Self::Error> {
186 unreachable!("disabled TLS never resolves a provider")
187 }
188}
189
190#[derive(Debug)]
192pub enum ClientTransport<Transport> {
193 Plain(Transport),
195 Tls(Box<crate::tls::ClientTls<Transport>>),
197}
198
199impl<Transport: AsyncRead + AsyncWrite + Unpin> AsyncRead for ClientTransport<Transport> {
200 fn poll_read(
201 mut self: Pin<&mut Self>,
202 cx: &mut Context<'_>,
203 buffer: &mut ReadBuf<'_>,
204 ) -> Poll<io::Result<()>> {
205 match &mut *self {
206 Self::Plain(stream) => Pin::new(stream).poll_read(cx, buffer),
207 Self::Tls(stream) => Pin::new(stream).poll_read(cx, buffer),
208 }
209 }
210}
211
212impl<Transport: AsyncRead + AsyncWrite + Unpin> AsyncWrite for ClientTransport<Transport> {
213 fn poll_write(
214 mut self: Pin<&mut Self>,
215 cx: &mut Context<'_>,
216 buffer: &[u8],
217 ) -> Poll<io::Result<usize>> {
218 match &mut *self {
219 Self::Plain(stream) => Pin::new(stream).poll_write(cx, buffer),
220 Self::Tls(stream) => Pin::new(stream).poll_write(cx, buffer),
221 }
222 }
223
224 fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
225 match &mut *self {
226 Self::Plain(stream) => Pin::new(stream).poll_flush(cx),
227 Self::Tls(stream) => Pin::new(stream).poll_flush(cx),
228 }
229 }
230
231 fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
232 match &mut *self {
233 Self::Plain(stream) => Pin::new(stream).poll_shutdown(cx),
234 Self::Tls(stream) => Pin::new(stream).poll_shutdown(cx),
235 }
236 }
237}
238
239async fn negotiate_client_tls<Transport, Provider, State, Handler, Evidence>(
240 mut transport: Transport,
241 mode: crate::pre_startup::SslMode,
242 provider: &Provider,
243 target: &ConnectTarget,
244 context: &mut ClientConnectionContext<Evidence>,
245 state: &mut State,
246 handler: &mut Handler,
247) -> Result<ClientTransport<Transport>, ClientTlsError<Provider::Error>>
248where
249 Transport: AsyncRead + AsyncWrite + Unpin,
250 Provider: ClientTlsProvider,
251 Handler: crate::ClientMiddleware<State, ClientConnectionContext<Evidence>>,
252{
253 if !mode.strategy().request_on_first_connection {
254 context.tls = Some(ClientTlsStatus::Plaintext);
255 return Ok(ClientTransport::Plain(transport));
256 }
257 let request = handler.pre_startup(
258 context,
259 state,
260 crate::pre_startup::PreStartupMessage::SslRequest,
261 );
262 let crate::pre_startup::PreStartupMessage::SslRequest = request else {
263 return Err(ClientTlsError::Handshake(io::Error::new(
264 io::ErrorKind::InvalidData,
265 "middleware replaced SSLRequest with an incompatible packet",
266 )));
267 };
268 transport
269 .write_all(&[0, 0, 0, 8, 4, 210, 22, 47])
270 .await
271 .map_err(ClientTlsError::Handshake)?;
272 transport.flush().await.map_err(ClientTlsError::Handshake)?;
273 match transport
274 .read_u8()
275 .await
276 .map_err(ClientTlsError::Handshake)?
277 {
278 b'S' => {
279 let resolved = provider
280 .resolve(target)
281 .await
282 .map_err(ClientTlsError::Provider)?;
283 let config = Arc::new(crate::tls::client_config(mode, resolved.roots));
284 let stream = crate::tls::connect(transport, resolved.server_name, config)
285 .await
286 .map_err(ClientTlsError::Handshake)?;
287 context.tls = Some(ClientTlsStatus::Encrypted);
288 Ok(ClientTransport::Tls(Box::new(stream)))
289 }
290 b'N' if mode.strategy().allow_server_rejection => {
291 context.tls = Some(ClientTlsStatus::Plaintext);
292 Ok(ClientTransport::Plain(transport))
293 }
294 b'N' => Err(ClientTlsError::Handshake(io::Error::new(
295 io::ErrorKind::PermissionDenied,
296 "server rejected required TLS",
297 ))),
298 b'E' => Err(ClientTlsError::Handshake(io::Error::new(
299 io::ErrorKind::ConnectionAborted,
300 "server terminated TLS negotiation",
301 ))),
302 _ => Err(ClientTlsError::Handshake(io::Error::new(
303 io::ErrorKind::InvalidData,
304 "invalid TLS negotiation response",
305 ))),
306 }
307}
308
309#[derive(Clone, Copy, Default, Eq, PartialEq)]
311pub struct TrustClientAuthentication;
312
313#[derive(Clone, Debug, Eq, PartialEq)]
315#[non_exhaustive]
316pub enum ClientAuthenticationChallenge {
317 CleartextPassword,
319 Md5Password([u8; 4]),
321 Sasl(Vec<Bytes>),
323 SaslContinue(Bytes),
325 SaslFinal(Bytes),
327 Gss,
329 Sspi,
331 KerberosV5,
333 TokenContinue(Bytes),
335}
336
337#[derive(Clone, Debug, Eq, PartialEq)]
339#[non_exhaustive]
340pub enum ClientAuthenticationResponse {
341 Password(Bytes),
343 SaslInitial {
345 mechanism: Bytes,
347 response: Bytes,
349 },
350 Sasl(Bytes),
352 Token(Bytes),
354 Verified,
356}
357
358#[allow(async_fn_in_trait)]
362pub trait ClientAuthentication {
363 type Evidence;
365 type Session: ClientAuthenticationSession<Evidence = Self::Evidence, Error = Self::Error>;
367 type Error;
369
370 async fn begin(&self, target: &ConnectTarget) -> Result<Self::Session, Self::Error>;
372}
373
374#[allow(async_fn_in_trait)]
378pub trait ClientAuthenticationSession {
379 type Evidence;
381 type Error;
383
384 async fn respond(
386 &mut self,
387 challenge: ClientAuthenticationChallenge,
388 ) -> Result<ClientAuthenticationResponse, Self::Error>;
389
390 async fn authenticated(self) -> Result<Self::Evidence, Self::Error>;
392}
393
394#[derive(Debug)]
396pub enum ClientAuthenticationError<PolicyError> {
397 Policy(PolicyError),
399 Rejected,
401 InvalidResponse,
403}
404
405impl<PolicyError: fmt::Display> fmt::Display for ClientAuthenticationError<PolicyError> {
406 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
407 match self {
408 Self::Policy(error) => error.fmt(formatter),
409 Self::Rejected => formatter.write_str("server rejected authentication"),
410 Self::InvalidResponse => {
411 formatter.write_str("authentication policy rejected the credential challenge")
412 }
413 }
414 }
415}
416
417impl<PolicyError: std::error::Error + 'static> std::error::Error
418 for ClientAuthenticationError<PolicyError>
419{
420 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
421 match self {
422 Self::Policy(error) => Some(error),
423 Self::Rejected | Self::InvalidResponse => None,
424 }
425 }
426}
427
428enum AuthenticationDriveError<PolicyError> {
429 Policy(PolicyError),
430 Rejected,
431 InvalidResponse,
432 Protocol(io::Error),
433}
434
435impl ClientAuthentication for TrustClientAuthentication {
436 type Evidence = ();
437 type Session = Self;
438 type Error = Infallible;
439
440 async fn begin(&self, _target: &ConnectTarget) -> Result<Self::Session, Self::Error> {
441 Ok(Self)
442 }
443}
444
445impl ClientAuthenticationSession for TrustClientAuthentication {
446 type Evidence = ();
447 type Error = Infallible;
448
449 async fn respond(
450 &mut self,
451 _challenge: ClientAuthenticationChallenge,
452 ) -> Result<ClientAuthenticationResponse, Self::Error> {
453 Ok(ClientAuthenticationResponse::Verified)
454 }
455
456 async fn authenticated(self) -> Result<Self::Evidence, Self::Error> {
457 Ok(())
458 }
459}
460
461impl fmt::Debug for TrustClientAuthentication {
462 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
463 formatter.write_str("Trust")
464 }
465}
466
467#[derive(Clone, Debug, Eq, PartialEq)]
469pub struct StaticClientCredentials {
470 username: Bytes,
471 password: Bytes,
472}
473
474impl StaticClientCredentials {
475 #[must_use]
477 pub fn new(username: impl Into<Bytes>, password: impl Into<Bytes>) -> Self {
478 Self {
479 username: username.into(),
480 password: password.into(),
481 }
482 }
483}
484
485pub struct StaticClientCredentialSession {
487 username: Bytes,
488 password: Bytes,
489 scram: Option<postgres_protocol::authentication::sasl::ScramSha256>,
490}
491
492#[derive(Debug)]
494pub enum StaticCredentialError {
495 UnsupportedAuthentication,
497 InvalidChallengeSequence,
499 Scram(io::Error),
501 AuthenticationFailed,
503}
504
505impl fmt::Display for StaticCredentialError {
506 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
507 match self {
508 Self::UnsupportedAuthentication => {
509 formatter.write_str("unsupported authentication mechanism")
510 }
511 Self::InvalidChallengeSequence => {
512 formatter.write_str("authentication challenge is out of sequence")
513 }
514 Self::Scram(error) => error.fmt(formatter),
515 Self::AuthenticationFailed => formatter.write_str("authentication failed"),
516 }
517 }
518}
519
520impl std::error::Error for StaticCredentialError {
521 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
522 match self {
523 Self::Scram(error) => Some(error),
524 _ => None,
525 }
526 }
527}
528
529impl ClientAuthentication for StaticClientCredentials {
530 type Evidence = ();
531 type Session = StaticClientCredentialSession;
532 type Error = StaticCredentialError;
533
534 async fn begin(&self, _: &ConnectTarget) -> Result<Self::Session, Self::Error> {
535 Ok(StaticClientCredentialSession {
536 username: self.username.clone(),
537 password: self.password.clone(),
538 scram: None,
539 })
540 }
541}
542
543impl ClientAuthenticationSession for StaticClientCredentialSession {
544 type Evidence = ();
545 type Error = StaticCredentialError;
546
547 async fn respond(
548 &mut self,
549 challenge: ClientAuthenticationChallenge,
550 ) -> Result<ClientAuthenticationResponse, Self::Error> {
551 use postgres_protocol::authentication::sasl::{ChannelBinding, ScramSha256};
552 match challenge {
553 ClientAuthenticationChallenge::CleartextPassword => Ok(
554 ClientAuthenticationResponse::Password(self.password.clone()),
555 ),
556 ClientAuthenticationChallenge::Md5Password(salt) => {
557 Ok(ClientAuthenticationResponse::Password(Bytes::from(
558 crate::credentials::md5_response(&self.username, &self.password, salt),
559 )))
560 }
561 ClientAuthenticationChallenge::Sasl(mechanisms)
562 if mechanisms
563 .iter()
564 .any(|mechanism| mechanism.as_ref() == b"SCRAM-SHA-256") =>
565 {
566 let scram = self.scram.insert(ScramSha256::new(
567 &self.password,
568 ChannelBinding::unsupported(),
569 ));
570 Ok(ClientAuthenticationResponse::SaslInitial {
571 mechanism: Bytes::from_static(b"SCRAM-SHA-256"),
572 response: Bytes::copy_from_slice(scram.message()),
573 })
574 }
575 ClientAuthenticationChallenge::Sasl(_) => {
576 Err(StaticCredentialError::UnsupportedAuthentication)
577 }
578 ClientAuthenticationChallenge::SaslContinue(message) => {
579 let scram = self
580 .scram
581 .as_mut()
582 .ok_or(StaticCredentialError::InvalidChallengeSequence)?;
583 scram
584 .update(&message)
585 .map_err(StaticCredentialError::Scram)?;
586 Ok(ClientAuthenticationResponse::Sasl(Bytes::copy_from_slice(
587 scram.message(),
588 )))
589 }
590 ClientAuthenticationChallenge::SaslFinal(message) => {
591 let scram = self
592 .scram
593 .as_mut()
594 .ok_or(StaticCredentialError::InvalidChallengeSequence)?;
595 scram
596 .finish(&message)
597 .map_err(StaticCredentialError::Scram)?;
598 Ok(ClientAuthenticationResponse::Verified)
599 }
600 _ => Err(StaticCredentialError::UnsupportedAuthentication),
601 }
602 }
603
604 async fn authenticated(self) -> Result<(), Self::Error> {
605 Ok(())
606 }
607}
608
609#[cfg(test)]
610mod static_credential_tests {
611 use super::*;
612
613 #[tokio::test]
614 async fn answers_cleartext_and_md5_challenges() {
615 let credentials = StaticClientCredentials::new("alice", "secret");
616 let mut session = credentials
617 .begin(&ConnectTarget::new("database"))
618 .await
619 .unwrap();
620 assert_eq!(
621 session
622 .respond(ClientAuthenticationChallenge::CleartextPassword)
623 .await
624 .unwrap(),
625 ClientAuthenticationResponse::Password(Bytes::from_static(b"secret"))
626 );
627
628 let salt = [1, 2, 3, 4];
629 assert_eq!(
630 session
631 .respond(ClientAuthenticationChallenge::Md5Password(salt))
632 .await
633 .unwrap(),
634 ClientAuthenticationResponse::Password(Bytes::from(crate::credentials::md5_response(
635 b"alice", b"secret", salt
636 )))
637 );
638 }
639}
640
641#[derive(Clone, Eq, PartialEq)]
643pub struct ConnectTarget {
644 name: String,
645 metadata: BTreeMap<String, String>,
646}
647
648impl fmt::Debug for ConnectTarget {
649 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
650 formatter.write_str("ConnectTarget(<redacted>)")
651 }
652}
653
654impl ConnectTarget {
655 #[must_use]
657 pub fn new(name: impl Into<String>) -> Self {
658 Self {
659 name: name.into(),
660 metadata: BTreeMap::new(),
661 }
662 }
663
664 #[must_use]
666 pub fn name(&self) -> &str {
667 &self.name
668 }
669
670 #[must_use]
672 pub fn with_metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
673 self.metadata.insert(key.into(), value.into());
674 self
675 }
676
677 #[must_use]
679 pub const fn metadata(&self) -> &BTreeMap<String, String> {
680 &self.metadata
681 }
682}
683
684#[derive(Clone, Default, Eq, PartialEq)]
686pub struct StartupParameters {
687 user: Option<String>,
688 database: Option<String>,
689 extensions: BTreeMap<String, String>,
690}
691
692impl fmt::Debug for StartupParameters {
693 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
694 formatter
695 .debug_struct("StartupParameters")
696 .field("user", &self.user.as_ref().map(|_| "<redacted>"))
697 .field("database", &self.database.as_ref().map(|_| "<redacted>"))
698 .field("extensions", &"<redacted>")
699 .finish()
700 }
701}
702
703impl StartupParameters {
704 #[must_use]
706 pub fn new(user: impl Into<String>) -> Self {
707 Self {
708 user: Some(user.into()),
709 ..Self::default()
710 }
711 }
712
713 #[must_use]
715 pub fn user(&self) -> Option<&str> {
716 self.user.as_deref()
717 }
718
719 #[must_use]
721 pub fn database_name(&self) -> Option<&str> {
722 self.database.as_deref()
723 }
724
725 pub(crate) fn from_wire(message: &StartupMessage) -> io::Result<Self> {
726 let mut parameters = Self::default();
727 for (name, value) in &message.parameters {
728 let name = std::str::from_utf8(name).map_err(|_| {
729 io::Error::new(
730 io::ErrorKind::InvalidData,
731 "startup parameter name is not UTF-8",
732 )
733 })?;
734 let value = std::str::from_utf8(value).map_err(|_| {
735 io::Error::new(
736 io::ErrorKind::InvalidData,
737 "startup parameter value is not UTF-8",
738 )
739 })?;
740 match name {
741 "user" => parameters.user = Some(value.to_owned()),
742 "database" => parameters.database = Some(value.to_owned()),
743 _ => {
744 parameters
745 .extensions
746 .insert(name.to_owned(), value.to_owned());
747 }
748 }
749 }
750 Ok(parameters)
751 }
752
753 #[must_use]
755 pub fn database(mut self, database: impl Into<String>) -> Self {
756 self.database = Some(database.into());
757 self
758 }
759
760 pub fn extension(
766 mut self,
767 name: impl Into<String>,
768 value: impl Into<String>,
769 ) -> Result<Self, StartupParameterError> {
770 let name = name.into();
771 if matches!(name.as_str(), "user" | "database") {
772 return Err(StartupParameterError::ReservedExtension(name));
773 }
774 self.extensions.insert(name, value.into());
775 Ok(self)
776 }
777
778 fn merged_with(mut self, overrides: Self) -> Self {
779 if overrides.user.is_some() {
780 self.user = overrides.user;
781 }
782 if overrides.database.is_some() {
783 self.database = overrides.database;
784 }
785 self.extensions.extend(overrides.extensions);
786 self
787 }
788
789 fn into_message(self) -> Result<StartupMessage, StartupParameterError> {
790 let user = self.user.ok_or(StartupParameterError::MissingUser)?;
791 let mut parameters = self
792 .extensions
793 .into_iter()
794 .map(|(key, value)| (Bytes::from(key), Bytes::from(value)))
795 .collect::<BTreeMap<_, _>>();
796 parameters.insert(Bytes::from_static(b"user"), Bytes::from(user));
797 if let Some(database) = self.database {
798 parameters.insert(Bytes::from_static(b"database"), Bytes::from(database));
799 }
800 Ok(StartupMessage {
801 version: ProtocolVersion::V3_2,
802 parameters,
803 })
804 }
805}
806
807#[derive(Clone, Eq, PartialEq)]
809pub enum StartupParameterError {
810 ReservedExtension(String),
812 MissingUser,
814}
815
816impl fmt::Debug for StartupParameterError {
817 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
818 formatter.write_str(match self {
819 Self::ReservedExtension(_) => "ReservedExtension(<redacted>)",
820 Self::MissingUser => "MissingUser",
821 })
822 }
823}
824
825impl fmt::Display for StartupParameterError {
826 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
827 match self {
828 Self::ReservedExtension(name) => {
829 write!(formatter, "startup extension name `{name}` is reserved")
830 }
831 Self::MissingUser => formatter.write_str("startup user is required"),
832 }
833 }
834}
835
836impl std::error::Error for StartupParameterError {}
837
838#[derive(Clone, Copy, Debug, Eq, PartialEq)]
840pub struct ProtocolLimits {
841 max_frame_len: usize,
842}
843
844impl Default for ProtocolLimits {
845 fn default() -> Self {
846 Self {
847 max_frame_len: 1024 * 1024,
848 }
849 }
850}
851
852impl ProtocolLimits {
853 pub fn max_frame_len(mut self, limit: usize) -> Result<Self, ProtocolLimitError> {
859 if !(5..=i32::MAX as usize + 1).contains(&limit) {
860 return Err(ProtocolLimitError);
861 }
862 self.max_frame_len = limit;
863 Ok(self)
864 }
865
866 #[must_use]
868 pub fn without_frame_limit(mut self) -> Self {
869 self.max_frame_len = i32::MAX as usize + 1;
870 self
871 }
872}
873
874#[derive(Clone, Copy, Debug, Eq, PartialEq)]
876pub struct ProtocolLimitError;
877
878impl fmt::Display for ProtocolLimitError {
879 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
880 formatter.write_str("frame limit is outside PostgreSQL's tagged-frame range")
881 }
882}
883
884impl std::error::Error for ProtocolLimitError {}
885
886#[derive(Clone, Debug, Eq, PartialEq)]
888pub struct ClientConnectionContext<Evidence = ()> {
889 target: ConnectTarget,
890 tls: Option<ClientTlsStatus>,
891 identity: Option<Evidence>,
892 backend_key: Option<crate::demux::CancelKey>,
893}
894
895#[derive(Clone, Copy, Debug, Eq, PartialEq)]
897pub enum ClientTlsStatus {
898 Plaintext,
900 Encrypted,
902}
903
904impl<Evidence> ClientConnectionContext<Evidence> {
905 #[must_use]
907 pub const fn target(&self) -> &ConnectTarget {
908 &self.target
909 }
910
911 #[must_use]
913 pub const fn tls(&self) -> Option<ClientTlsStatus> {
914 self.tls
915 }
916
917 #[must_use]
923 pub const fn identity(&self) -> &Evidence {
924 match &self.identity {
925 Some(identity) => identity,
926 None => panic!("identity is not known before authentication"),
927 }
928 }
929
930 #[must_use]
932 pub const fn identity_if_known(&self) -> Option<&Evidence> {
933 self.identity.as_ref()
934 }
935
936 #[must_use]
938 pub const fn backend_key(&self) -> Option<&crate::demux::CancelKey> {
939 self.backend_key.as_ref()
940 }
941}
942
943#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
945pub struct IdentityHandler;
946impl<C> crate::MiddlewareFactory<C> for IdentityHandler {
947 type Handler = Self;
948 fn create(&self, _: &C) -> Self {
949 *self
950 }
951}
952impl<S, C> crate::ClientMiddleware<S, C> for IdentityHandler {}
953
954pub struct ClientInitialContext {
956 target: ConnectTarget,
957}
958impl ClientInitialContext {
959 #[must_use]
961 pub const fn target(&self) -> &ConnectTarget {
962 &self.target
963 }
964}
965
966pub struct Client<
968 Connector = (),
969 Tls = ClientTlsPolicy,
970 Authentication = TrustClientAuthentication,
971 Middleware = IdentityHandler,
972> {
973 connector: Connector,
974 tls: Tls,
975 authentication: Authentication,
976 defaults: StartupParameters,
977 limits: ProtocolLimits,
978 middleware: Middleware,
979}
980
981impl<Connector: Clone, Tls: Clone, Authentication: Clone, Middleware: Clone> Clone
982 for Client<Connector, Tls, Authentication, Middleware>
983{
984 fn clone(&self) -> Self {
985 Self {
986 connector: self.connector.clone(),
987 tls: self.tls.clone(),
988 authentication: self.authentication.clone(),
989 defaults: self.defaults.clone(),
990 limits: self.limits,
991 middleware: self.middleware.clone(),
992 }
993 }
994}
995
996impl Client<()> {
997 #[must_use]
999 pub fn builder() -> ClientBuilder {
1000 ClientBuilder::default()
1001 }
1002}
1003
1004impl<Connector, Tls: fmt::Debug, Authentication: fmt::Debug, Middleware> fmt::Debug
1005 for Client<Connector, Tls, Authentication, Middleware>
1006{
1007 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
1008 formatter
1009 .debug_struct("Client")
1010 .field("tls", &self.tls)
1011 .field("authentication", &self.authentication)
1012 .finish_non_exhaustive()
1013 }
1014}
1015
1016pub struct ClientBuilder<
1018 Connector = (),
1019 Tls = (),
1020 Authentication = (),
1021 Middleware = IdentityHandler,
1022> {
1023 connector: Option<Connector>,
1024 tls: Option<Tls>,
1025 authentication: Option<Authentication>,
1026 defaults: StartupParameters,
1027 limits: ProtocolLimits,
1028 middleware: Middleware,
1029}
1030
1031impl Default for ClientBuilder<()> {
1032 fn default() -> Self {
1033 Self {
1034 connector: None,
1035 tls: None,
1036 authentication: None,
1037 defaults: StartupParameters::default(),
1038 limits: ProtocolLimits::default(),
1039 middleware: IdentityHandler,
1040 }
1041 }
1042}
1043
1044impl ClientBuilder<()> {
1045 #[must_use]
1047 pub fn connector<Next, Work, Transport, Error>(
1048 self,
1049 connector: Next,
1050 ) -> ClientBuilder<Next, (), ()>
1051 where
1052 Next: Fn(&ConnectTarget) -> Work,
1053 Work: Future<Output = Result<Transport, Error>>,
1054 {
1055 ClientBuilder {
1056 connector: Some(connector),
1057 tls: self.tls,
1058 authentication: self.authentication,
1059 defaults: self.defaults,
1060 limits: self.limits,
1061 middleware: self.middleware,
1062 }
1063 }
1064}
1065
1066impl<Connector, Tls, Authentication, Middleware>
1067 ClientBuilder<Connector, Tls, Authentication, Middleware>
1068{
1069 #[must_use]
1071 pub fn tls<Next>(
1072 self,
1073 tls: Next,
1074 ) -> ClientBuilder<Connector, Next, Authentication, Middleware> {
1075 ClientBuilder {
1076 connector: self.connector,
1077 tls: Some(tls),
1078 authentication: self.authentication,
1079 defaults: self.defaults,
1080 limits: self.limits,
1081 middleware: self.middleware,
1082 }
1083 }
1084
1085 #[must_use]
1087 pub fn authentication<Next>(
1088 self,
1089 authentication: Next,
1090 ) -> ClientBuilder<Connector, Tls, Next, Middleware> {
1091 ClientBuilder {
1092 connector: self.connector,
1093 tls: self.tls,
1094 authentication: Some(authentication),
1095 defaults: self.defaults,
1096 limits: self.limits,
1097 middleware: self.middleware,
1098 }
1099 }
1100
1101 #[must_use]
1103 pub fn startup_parameters(mut self, defaults: StartupParameters) -> Self {
1104 self.defaults = defaults;
1105 self
1106 }
1107
1108 #[must_use]
1110 pub fn protocol_limits(mut self, limits: ProtocolLimits) -> Self {
1111 self.limits = limits;
1112 self
1113 }
1114
1115 #[must_use]
1117 pub fn middleware<Next>(
1118 self,
1119 factory: Next,
1120 ) -> ClientBuilder<Connector, Tls, Authentication, crate::MiddlewareChain<Middleware, Next>>
1121 {
1122 ClientBuilder {
1123 connector: self.connector,
1124 tls: self.tls,
1125 authentication: self.authentication,
1126 defaults: self.defaults,
1127 limits: self.limits,
1128 middleware: crate::MiddlewareChain(self.middleware, factory),
1129 }
1130 }
1131
1132 pub fn build(self) -> Result<Client<Connector, Tls, Authentication, Middleware>, BuildError> {
1138 Ok(Client {
1139 connector: self.connector.ok_or(BuildError::MissingConnector)?,
1140 tls: self.tls.ok_or(BuildError::MissingTls)?,
1141 authentication: self
1142 .authentication
1143 .ok_or(BuildError::MissingAuthentication)?,
1144 defaults: self.defaults,
1145 limits: self.limits,
1146 middleware: self.middleware,
1147 })
1148 }
1149}
1150
1151#[derive(Debug)]
1153pub enum ConnectError<ConnectorError, TlsError = Infallible, AuthenticationError = Infallible> {
1154 Connector(ConnectorError),
1156 Tls(TlsError),
1158 Authentication(AuthenticationError),
1160 Startup(StartupParameterError),
1162 Protocol(io::Error),
1164}
1165
1166#[derive(Debug)]
1168pub enum CancelError<ConnectorError> {
1169 Connector(ConnectorError),
1171 Protocol(io::Error),
1173}
1174
1175impl<E: fmt::Display> fmt::Display for CancelError<E> {
1176 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
1177 match self {
1178 Self::Connector(error) => error.fmt(formatter),
1179 Self::Protocol(error) => error.fmt(formatter),
1180 }
1181 }
1182}
1183impl<E: std::error::Error + 'static> std::error::Error for CancelError<E> {
1184 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
1185 match self {
1186 Self::Connector(error) => Some(error),
1187 Self::Protocol(error) => Some(error),
1188 }
1189 }
1190}
1191
1192impl<ConnectorError: fmt::Display, TlsError: fmt::Display, AuthenticationError: fmt::Display>
1193 fmt::Display for ConnectError<ConnectorError, TlsError, AuthenticationError>
1194{
1195 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1196 match self {
1197 Self::Connector(error) => error.fmt(f),
1198 Self::Tls(error) => error.fmt(f),
1199 Self::Authentication(error) => error.fmt(f),
1200 Self::Startup(error) => error.fmt(f),
1201 Self::Protocol(error) => error.fmt(f),
1202 }
1203 }
1204}
1205impl<ConnectorError, TlsError, AuthenticationError> std::error::Error
1206 for ConnectError<ConnectorError, TlsError, AuthenticationError>
1207where
1208 ConnectorError: std::error::Error + 'static,
1209 TlsError: std::error::Error + 'static,
1210 AuthenticationError: std::error::Error + 'static,
1211{
1212 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
1213 match self {
1214 Self::Connector(error) => Some(error),
1215 Self::Tls(error) => Some(error),
1216 Self::Authentication(error) => Some(error),
1217 Self::Startup(error) => Some(error),
1218 Self::Protocol(error) => Some(error),
1219 }
1220 }
1221}
1222
1223#[derive(Debug)]
1225pub enum ConnectionClean {}
1226
1227#[derive(Debug)]
1229pub enum ConnectionChanged {}
1230
1231pub struct ClientConnection<
1233 Transport,
1234 State,
1235 Cleanliness = ConnectionClean,
1236 Evidence = (),
1237 Handler = IdentityHandler,
1238> {
1239 core: ClientConnectionCore<Transport, Cleanliness, Evidence, Handler>,
1240 state: State,
1241}
1242
1243pub(crate) struct ClientConnectionCore<Transport, Cleanliness, Evidence, Handler> {
1244 connection: Conn<Buffered<Transport, Backend>, Ready, Cleanliness>,
1245 handler: Handler,
1246 context: ClientConnectionContext<Evidence>,
1247}
1248
1249impl<Transport, State, Cleanliness, Evidence, Handler>
1250 ClientConnection<Transport, State, Cleanliness, Evidence, Handler>
1251{
1252 #[must_use]
1254 pub const fn context(&self) -> &ClientConnectionContext<Evidence> {
1255 &self.core.context
1256 }
1257
1258 #[must_use]
1261 pub const fn state(&self) -> &State {
1262 &self.state
1263 }
1264
1265 pub async fn receive_wire(&mut self) -> io::Result<crate::codec::BackendMessage>
1272 where
1273 Transport: AsyncRead + Unpin,
1274 Handler: crate::ClientMiddleware<State, ClientConnectionContext<Evidence>>,
1275 {
1276 let message = self.core.receive_wire_raw().await?;
1277 Ok(self.core.intercept_backend(&mut self.state, message))
1278 }
1279
1280 pub fn into_parts(self) -> (Transport, State, Handler, ClientConnectionContext<Evidence>) {
1283 (
1284 self.core.connection.into_transport().into_inner(),
1285 self.state,
1286 self.core.handler,
1287 self.core.context,
1288 )
1289 }
1290}
1291
1292impl<Transport, Cleanliness, Evidence, Handler>
1293 ClientConnectionCore<Transport, Cleanliness, Evidence, Handler>
1294{
1295 pub(crate) const fn context(&self) -> &ClientConnectionContext<Evidence> {
1296 &self.context
1297 }
1298 pub(crate) async fn receive_wire_raw(&mut self) -> io::Result<crate::codec::BackendMessage>
1299 where
1300 Transport: AsyncRead + Unpin,
1301 {
1302 self.connection.receive_backend_wire().await
1303 }
1304
1305 pub(crate) fn intercept_backend<State>(
1306 &mut self,
1307 state: &mut State,
1308 message: crate::codec::BackendMessage,
1309 ) -> crate::codec::BackendMessage
1310 where
1311 Handler: crate::ClientMiddleware<State, ClientConnectionContext<Evidence>>,
1312 {
1313 self.handler.backend(&self.context, state, message)
1314 }
1315
1316 pub(crate) fn intercept_frontend<State>(
1317 &mut self,
1318 state: &mut State,
1319 message: FrontendMessage,
1320 ) -> FrontendMessage
1321 where
1322 Handler: crate::ClientMiddleware<State, ClientConnectionContext<Evidence>>,
1323 {
1324 self.handler.frontend(&self.context, state, message)
1325 }
1326
1327 pub(crate) async fn send_wire_raw(&mut self, message: FrontendMessage) -> io::Result<()>
1328 where
1329 Transport: AsyncWrite + Unpin,
1330 {
1331 self.connection.push_frame(message.to_frame()?)?;
1332 self.connection.flush().await
1333 }
1334
1335 pub(crate) fn into_parts(self) -> (Transport, Handler, ClientConnectionContext<Evidence>) {
1336 (
1337 self.connection.into_transport().into_inner(),
1338 self.handler,
1339 self.context,
1340 )
1341 }
1342}
1343
1344fn intercept_auth_response<State, Evidence, Handler>(
1345 handler: &mut Handler,
1346 context: &ClientConnectionContext<Evidence>,
1347 state: &mut State,
1348 response: Bytes,
1349) -> Result<Bytes, AuthenticationDriveError<Infallible>>
1350where
1351 Handler: crate::ClientMiddleware<State, ClientConnectionContext<Evidence>>,
1352{
1353 match handler.frontend(
1354 context,
1355 state,
1356 crate::codec::FrontendMessage::PasswordResponse(response),
1357 ) {
1358 crate::codec::FrontendMessage::PasswordResponse(response) => Ok(response),
1359 _ => Err(AuthenticationDriveError::InvalidResponse),
1360 }
1361}
1362
1363async fn complete_password<Transport, Policy, State, Handler>(
1364 connection: Conn<Buffered<Transport, Backend>, crate::auth::PasswordResponse>,
1365 challenge: ClientAuthenticationChallenge,
1366 policy: &mut Policy,
1367 context: &ClientConnectionContext<Policy::Evidence>,
1368 state: &mut State,
1369 handler: &mut Handler,
1370) -> Result<
1371 Conn<Buffered<Transport, Backend>, crate::auth::AwaitingStartupReady>,
1372 AuthenticationDriveError<Policy::Error>,
1373>
1374where
1375 Transport: AsyncRead + AsyncWrite + Unpin,
1376 Policy: ClientAuthenticationSession,
1377 Handler: crate::ClientMiddleware<State, ClientConnectionContext<Policy::Evidence>>,
1378{
1379 let response = policy
1380 .respond(challenge)
1381 .await
1382 .map_err(AuthenticationDriveError::Policy)?;
1383 let ClientAuthenticationResponse::Password(password) = response else {
1384 let _ = connection.into_transport();
1385 return Err(AuthenticationDriveError::InvalidResponse);
1386 };
1387 let password = intercept_auth_response(handler, context, state, password)
1388 .map_err(|_| AuthenticationDriveError::InvalidResponse)?;
1389 let (mut awaiting, frame) = connection
1390 .password(&password)
1391 .map_err(AuthenticationDriveError::Protocol)?;
1392 awaiting
1393 .push_frame(frame)
1394 .map_err(AuthenticationDriveError::Protocol)?;
1395 awaiting
1396 .flush()
1397 .await
1398 .map_err(AuthenticationDriveError::Protocol)?;
1399 let message = awaiting
1400 .receive_backend_wire()
1401 .await
1402 .map_err(AuthenticationDriveError::Protocol)?;
1403 let message = handler.backend(context, state, message);
1404 match awaiting.offer(message) {
1405 Ok(crate::auth::AuthCompletion::Ok(connection)) => Ok(connection),
1406 Ok(crate::auth::AuthCompletion::Error { conn, .. }) => {
1407 let _ = conn.into_transport();
1408 Err(AuthenticationDriveError::Rejected)
1409 }
1410 Err((conn, _)) => {
1411 let _ = conn.into_transport();
1412 Err(AuthenticationDriveError::Protocol(io::Error::new(
1413 io::ErrorKind::InvalidData,
1414 "illegal authentication completion",
1415 )))
1416 }
1417 }
1418}
1419
1420async fn complete_sasl<Transport, Policy, State, Handler>(
1421 connection: Conn<Buffered<Transport, Backend>, crate::auth::SaslInitial>,
1422 mechanisms: Vec<Bytes>,
1423 policy: &mut Policy,
1424 context: &ClientConnectionContext<Policy::Evidence>,
1425 state: &mut State,
1426 handler: &mut Handler,
1427) -> Result<
1428 Conn<Buffered<Transport, Backend>, crate::auth::AwaitingStartupReady>,
1429 AuthenticationDriveError<Policy::Error>,
1430>
1431where
1432 Transport: AsyncRead + AsyncWrite + Unpin,
1433 Policy: ClientAuthenticationSession,
1434 Handler: crate::ClientMiddleware<State, ClientConnectionContext<Policy::Evidence>>,
1435{
1436 let response = policy
1437 .respond(ClientAuthenticationChallenge::Sasl(mechanisms))
1438 .await
1439 .map_err(AuthenticationDriveError::Policy)?;
1440 let ClientAuthenticationResponse::SaslInitial {
1441 mechanism,
1442 response,
1443 } = response
1444 else {
1445 let _ = connection.into_transport();
1446 return Err(AuthenticationDriveError::InvalidResponse);
1447 };
1448 let response = intercept_auth_response(handler, context, state, response)
1449 .map_err(|_| AuthenticationDriveError::InvalidResponse)?;
1450 let (mut sasl, frame) = connection
1451 .sasl(&mechanism, &response)
1452 .map_err(AuthenticationDriveError::Protocol)?;
1453 sasl.push_frame(frame)
1454 .map_err(AuthenticationDriveError::Protocol)?;
1455 sasl.flush()
1456 .await
1457 .map_err(AuthenticationDriveError::Protocol)?;
1458 loop {
1459 let message = sasl
1460 .receive_backend_wire()
1461 .await
1462 .map_err(AuthenticationDriveError::Protocol)?;
1463 let message = handler.backend(context, state, message);
1464 match sasl.offer_backend(message) {
1465 Ok(crate::auth::SaslEvent::Continue { conn, challenge }) => {
1466 let response = policy
1467 .respond(ClientAuthenticationChallenge::SaslContinue(challenge))
1468 .await
1469 .map_err(AuthenticationDriveError::Policy)?;
1470 let ClientAuthenticationResponse::Sasl(response) = response else {
1471 let _ = conn.into_transport();
1472 return Err(AuthenticationDriveError::InvalidResponse);
1473 };
1474 let response = intercept_auth_response(handler, context, state, response)
1475 .map_err(|_| AuthenticationDriveError::InvalidResponse)?;
1476 let (mut next, frame) = conn.respond(response);
1477 next.push_frame(frame)
1478 .map_err(AuthenticationDriveError::Protocol)?;
1479 next.flush()
1480 .await
1481 .map_err(AuthenticationDriveError::Protocol)?;
1482 sasl = next;
1483 }
1484 Ok(crate::auth::SaslEvent::Final { conn, server_final }) => {
1485 let response = policy
1486 .respond(ClientAuthenticationChallenge::SaslFinal(server_final))
1487 .await
1488 .map_err(AuthenticationDriveError::Policy)?;
1489 if response != ClientAuthenticationResponse::Verified {
1490 let _ = conn.into_transport();
1491 return Err(AuthenticationDriveError::InvalidResponse);
1492 }
1493 let mut awaiting = conn.verified();
1494 let message = awaiting
1495 .receive_backend_wire()
1496 .await
1497 .map_err(AuthenticationDriveError::Protocol)?;
1498 let message = handler.backend(context, state, message);
1499 return match awaiting.offer(message) {
1500 Ok(crate::auth::AuthCompletion::Ok(connection)) => Ok(connection),
1501 Ok(crate::auth::AuthCompletion::Error { conn, .. }) => {
1502 let _ = conn.into_transport();
1503 Err(AuthenticationDriveError::Rejected)
1504 }
1505 Err((conn, _)) => {
1506 let _ = conn.into_transport();
1507 Err(AuthenticationDriveError::Protocol(io::Error::new(
1508 io::ErrorKind::InvalidData,
1509 "illegal SASL authentication completion",
1510 )))
1511 }
1512 };
1513 }
1514 Ok(crate::auth::SaslEvent::Error { conn, .. }) => {
1515 let _ = conn.into_transport();
1516 return Err(AuthenticationDriveError::Rejected);
1517 }
1518 Err((conn, _)) => {
1519 let _ = conn.into_transport();
1520 return Err(AuthenticationDriveError::Protocol(io::Error::new(
1521 io::ErrorKind::InvalidData,
1522 "illegal SASL authentication message",
1523 )));
1524 }
1525 }
1526 }
1527}
1528
1529async fn complete_token<Transport, Policy, State, Handler>(
1530 connection: Conn<Buffered<Transport, Backend>, crate::auth::TokenResponse>,
1531 challenge: ClientAuthenticationChallenge,
1532 policy: &mut Policy,
1533 context: &ClientConnectionContext<Policy::Evidence>,
1534 state: &mut State,
1535 handler: &mut Handler,
1536) -> Result<
1537 Conn<Buffered<Transport, Backend>, crate::auth::AwaitingStartupReady>,
1538 AuthenticationDriveError<Policy::Error>,
1539>
1540where
1541 Transport: AsyncRead + AsyncWrite + Unpin,
1542 Policy: ClientAuthenticationSession,
1543 Handler: crate::ClientMiddleware<State, ClientConnectionContext<Policy::Evidence>>,
1544{
1545 let response = policy
1546 .respond(challenge)
1547 .await
1548 .map_err(AuthenticationDriveError::Policy)?;
1549 let ClientAuthenticationResponse::Token(token) = response else {
1550 let _ = connection.into_transport();
1551 return Err(AuthenticationDriveError::InvalidResponse);
1552 };
1553 let token = intercept_auth_response(handler, context, state, token)
1554 .map_err(|_| AuthenticationDriveError::InvalidResponse)?;
1555 let (mut waiting, frame) = connection.respond(token);
1556 waiting
1557 .push_frame(frame)
1558 .map_err(AuthenticationDriveError::Protocol)?;
1559 waiting
1560 .flush()
1561 .await
1562 .map_err(AuthenticationDriveError::Protocol)?;
1563 loop {
1564 let message = waiting
1565 .receive_backend_wire()
1566 .await
1567 .map_err(AuthenticationDriveError::Protocol)?;
1568 let message = handler.backend(context, state, message);
1569 match waiting.offer(message) {
1570 Ok(crate::auth::TokenAuthEvent::Continue { conn, token }) => {
1571 let response = policy
1572 .respond(ClientAuthenticationChallenge::TokenContinue(token))
1573 .await
1574 .map_err(AuthenticationDriveError::Policy)?;
1575 let ClientAuthenticationResponse::Token(token) = response else {
1576 let _ = conn.into_transport();
1577 return Err(AuthenticationDriveError::InvalidResponse);
1578 };
1579 let token = intercept_auth_response(handler, context, state, token)
1580 .map_err(|_| AuthenticationDriveError::InvalidResponse)?;
1581 let (mut next, frame) = conn.respond(token);
1582 next.push_frame(frame)
1583 .map_err(AuthenticationDriveError::Protocol)?;
1584 next.flush()
1585 .await
1586 .map_err(AuthenticationDriveError::Protocol)?;
1587 waiting = next;
1588 }
1589 Ok(crate::auth::TokenAuthEvent::Ok(connection)) => return Ok(connection),
1590 Ok(crate::auth::TokenAuthEvent::Error { conn, .. }) => {
1591 let _ = conn.into_transport();
1592 return Err(AuthenticationDriveError::Rejected);
1593 }
1594 Err((conn, _)) => {
1595 let _ = conn.into_transport();
1596 return Err(AuthenticationDriveError::Protocol(io::Error::new(
1597 io::ErrorKind::InvalidData,
1598 "illegal token authentication message",
1599 )));
1600 }
1601 }
1602 }
1603}
1604
1605enum SessionEstablishError<AuthenticationError> {
1606 Authentication(ClientAuthenticationError<AuthenticationError>),
1607 Protocol(io::Error),
1608}
1609
1610fn replace_session_item(
1611 item: SessionItem,
1612 replacement: crate::codec::BackendMessage,
1613) -> Option<SessionItem> {
1614 match (item, replacement) {
1615 (SessionItem::Message(_), message) => Some(SessionItem::Message(message)),
1616 (
1617 SessionItem::ReadyForQuery {
1618 parameters_changed, ..
1619 },
1620 crate::codec::BackendMessage::ReadyForQuery(status),
1621 ) => Some(SessionItem::ReadyForQuery {
1622 status,
1623 parameters_changed,
1624 }),
1625 (
1626 SessionItem::CommandComplete {
1627 command, notices, ..
1628 },
1629 crate::codec::BackendMessage::CommandComplete(tag),
1630 ) => Some(SessionItem::CommandComplete {
1631 tag,
1632 command,
1633 notices,
1634 }),
1635 _ => None,
1636 }
1637}
1638
1639#[allow(clippy::too_many_arguments, clippy::too_many_lines)] async fn establish_client_session<Transport, Authentication, State, Handler>(
1641 transport: ClientTransport<Transport>,
1642 startup: &StartupMessage,
1643 target: &ConnectTarget,
1644 authentication_policy: &Authentication,
1645 max_frame_len: usize,
1646 context: &ClientConnectionContext<Authentication::Evidence>,
1647 state: &mut State,
1648 handler: &mut Handler,
1649) -> Result<
1650 (
1651 Conn<Buffered<ClientTransport<Transport>, Backend>, Ready>,
1652 Authentication::Evidence,
1653 Option<crate::demux::CancelKey>,
1654 ),
1655 SessionEstablishError<Authentication::Error>,
1656>
1657where
1658 Transport: AsyncRead + AsyncWrite + Unpin,
1659 Authentication: ClientAuthentication,
1660 Handler: crate::ClientMiddleware<State, ClientConnectionContext<Authentication::Evidence>>,
1661{
1662 let mut policy = authentication_policy.begin(target).await.map_err(|error| {
1663 SessionEstablishError::Authentication(ClientAuthenticationError::Policy(error))
1664 })?;
1665 let buffered = Buffered::with_max_frame_len(transport, max_frame_len)
1666 .map_err(SessionEstablishError::Protocol)?;
1667 let (mut startup_connection, packet) = Conn::new(buffered)
1668 .startup(startup)
1669 .map_err(SessionEstablishError::Protocol)?;
1670 startup_connection.push_startup_packet(&packet);
1671 let mut authentication = startup_connection.authentication();
1672 if let Err(error) = authentication.flush().await {
1673 let _ = authentication.into_transport();
1674 return Err(SessionEstablishError::Protocol(error));
1675 }
1676 let awaiting_ready = loop {
1677 let message = match authentication.receive_backend_wire().await {
1678 Ok(message) => message,
1679 Err(error) => {
1680 let _ = authentication.into_transport();
1681 return Err(SessionEstablishError::Protocol(error));
1682 }
1683 };
1684 let message = handler.backend(context, state, message);
1685 match authentication.offer_backend(message) {
1686 Ok(AuthEvent::Authentication(AuthOffer::Ok(connection))) => break connection,
1687 Ok(AuthEvent::Negotiate { conn, .. }) => authentication = conn,
1688 Ok(AuthEvent::Authentication(AuthOffer::Cleartext(connection))) => {
1689 break complete_password(
1690 connection,
1691 ClientAuthenticationChallenge::CleartextPassword,
1692 &mut policy,
1693 context,
1694 state,
1695 handler,
1696 )
1697 .await
1698 .map_err(session_authentication_error)?;
1699 }
1700 Ok(AuthEvent::Authentication(AuthOffer::Md5 { conn, salt })) => {
1701 break complete_password(
1702 conn,
1703 ClientAuthenticationChallenge::Md5Password(salt),
1704 &mut policy,
1705 context,
1706 state,
1707 handler,
1708 )
1709 .await
1710 .map_err(session_authentication_error)?;
1711 }
1712 Ok(AuthEvent::Authentication(AuthOffer::Sasl { conn, mechanisms })) => {
1713 break complete_sasl(conn, mechanisms, &mut policy, context, state, handler)
1714 .await
1715 .map_err(session_authentication_error)?;
1716 }
1717 Ok(AuthEvent::Authentication(AuthOffer::Gss(conn))) => {
1718 break complete_token(
1719 conn,
1720 ClientAuthenticationChallenge::Gss,
1721 &mut policy,
1722 context,
1723 state,
1724 handler,
1725 )
1726 .await
1727 .map_err(session_authentication_error)?;
1728 }
1729 Ok(AuthEvent::Authentication(AuthOffer::Sspi(conn))) => {
1730 break complete_token(
1731 conn,
1732 ClientAuthenticationChallenge::Sspi,
1733 &mut policy,
1734 context,
1735 state,
1736 handler,
1737 )
1738 .await
1739 .map_err(session_authentication_error)?;
1740 }
1741 Ok(AuthEvent::Authentication(AuthOffer::KerberosV5(conn))) => {
1742 break complete_token(
1743 conn,
1744 ClientAuthenticationChallenge::KerberosV5,
1745 &mut policy,
1746 context,
1747 state,
1748 handler,
1749 )
1750 .await
1751 .map_err(session_authentication_error)?;
1752 }
1753 Ok(AuthEvent::Error { conn, .. }) => {
1754 let _ = conn.into_transport();
1755 return Err(SessionEstablishError::Authentication(
1756 ClientAuthenticationError::Rejected,
1757 ));
1758 }
1759 Err((conn, _, source)) => {
1760 authentication = conn;
1761 if let Some(source) = source {
1762 let _ = authentication.into_transport();
1763 return Err(SessionEstablishError::Protocol(source));
1764 }
1765 }
1766 }
1767 };
1768 let mut awaiting_ready = awaiting_ready;
1769 let mut backend_key = None;
1770 let ready = loop {
1771 let item = match awaiting_ready.receive().await {
1772 Ok(item) => item,
1773 Err(error) => {
1774 let _ = awaiting_ready.into_transport();
1775 return Err(SessionEstablishError::Protocol(error));
1776 }
1777 };
1778 if let crate::codec::BackendMessage::BackendKeyData {
1779 process_id,
1780 secret_key,
1781 } = item.clone().into_backend_message()
1782 {
1783 backend_key = Some(crate::demux::CancelKey {
1784 process_id,
1785 secret_key,
1786 });
1787 }
1788 let replacement = handler.backend(context, state, item.clone().into_backend_message());
1789 let Some(item) = replace_session_item(item, replacement) else {
1790 let _ = awaiting_ready.into_transport();
1791 return Err(SessionEstablishError::Protocol(io::Error::new(
1792 io::ErrorKind::InvalidData,
1793 "middleware replacement during startup readiness is not phase-compatible",
1794 )));
1795 };
1796 match awaiting_ready.offer_ready(item) {
1797 Ok(ready) => break ready,
1798 Err((connection, SessionItem::Message(_))) => awaiting_ready = connection,
1799 Err((connection, _)) => {
1800 let _ = connection.into_transport();
1801 return Err(SessionEstablishError::Protocol(io::Error::new(
1802 io::ErrorKind::InvalidData,
1803 "startup did not reach an idle operational phase",
1804 )));
1805 }
1806 }
1807 };
1808 let identity = policy.authenticated().await.map_err(|error| {
1809 SessionEstablishError::Authentication(ClientAuthenticationError::Policy(error))
1810 })?;
1811 Ok((ready, identity, backend_key))
1812}
1813
1814fn session_authentication_error<Error>(
1815 error: AuthenticationDriveError<Error>,
1816) -> SessionEstablishError<Error> {
1817 match error {
1818 AuthenticationDriveError::Policy(error) => {
1819 SessionEstablishError::Authentication(ClientAuthenticationError::Policy(error))
1820 }
1821 AuthenticationDriveError::Rejected => {
1822 SessionEstablishError::Authentication(ClientAuthenticationError::Rejected)
1823 }
1824 AuthenticationDriveError::InvalidResponse => {
1825 SessionEstablishError::Authentication(ClientAuthenticationError::InvalidResponse)
1826 }
1827 AuthenticationDriveError::Protocol(error) => SessionEstablishError::Protocol(error),
1828 }
1829}
1830
1831fn map_session_error<ConnectorError, TlsError, AuthenticationError>(
1832 error: SessionEstablishError<AuthenticationError>,
1833) -> ConnectError<ConnectorError, TlsError, ClientAuthenticationError<AuthenticationError>> {
1834 match error {
1835 SessionEstablishError::Authentication(error) => ConnectError::Authentication(error),
1836 SessionEstablishError::Protocol(error) => ConnectError::Protocol(error),
1837 }
1838}
1839
1840#[derive(Debug)]
1842pub struct QueryError(io::Error);
1843
1844impl fmt::Display for QueryError {
1845 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
1846 self.0.fmt(formatter)
1847 }
1848}
1849
1850impl std::error::Error for QueryError {
1851 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
1852 Some(&self.0)
1853 }
1854}
1855
1856impl<Transport, State, Cleanliness, Evidence, Handler>
1857 ClientConnection<Transport, State, Cleanliness, Evidence, Handler>
1858where
1859 Transport: AsyncRead + AsyncWrite + Unpin,
1860 Handler: crate::ClientMiddleware<State, ClientConnectionContext<Evidence>>,
1861{
1862 #[allow(clippy::too_many_lines)]
1873 pub async fn simple_query(
1874 self,
1875 query: &[u8],
1876 ) -> Result<
1877 (
1878 ClientConnection<Transport, State, ConnectionChanged, Evidence, Handler>,
1879 Vec<crate::codec::BackendMessage>,
1880 ),
1881 QueryError,
1882 > {
1883 let Self {
1884 core:
1885 ClientConnectionCore {
1886 connection,
1887 mut handler,
1888 context,
1889 },
1890 mut state,
1891 } = self;
1892 let outbound = handler.frontend(
1893 &context,
1894 &mut state,
1895 crate::codec::FrontendMessage::Query(Bytes::copy_from_slice(query)),
1896 );
1897 let crate::codec::FrontendMessage::Query(query) = outbound else {
1898 let _ = connection.into_transport();
1899 return Err(QueryError(io::Error::new(
1900 io::ErrorKind::InvalidData,
1901 "middleware replaced Query with an incompatible message",
1902 )));
1903 };
1904 let (mut query_connection, frame) = connection.push_query(&query).map_err(QueryError)?;
1905 if let Err(error) = query_connection.push_frame(frame) {
1906 let _ = query_connection.into_transport();
1907 return Err(QueryError(error));
1908 }
1909 if let Err(error) = query_connection.flush().await {
1910 let _ = query_connection.into_transport();
1911 return Err(QueryError(error));
1912 }
1913 let mut messages = Vec::new();
1914 loop {
1915 let item = match query_connection.receive().await {
1916 Ok(item) => item,
1917 Err(error) => {
1918 let _ = query_connection.into_transport();
1919 return Err(QueryError(error));
1920 }
1921 };
1922 let observed =
1923 handler.backend(&context, &mut state, item.clone().into_backend_message());
1924 let Some(item) = replace_session_item(item, observed.clone()) else {
1925 let _ = query_connection.into_transport();
1926 return Err(QueryError(io::Error::new(
1927 io::ErrorKind::InvalidData,
1928 "middleware replacement during simple query is not phase-compatible",
1929 )));
1930 };
1931 messages.push(observed);
1932 match query_connection.offer(item) {
1933 Ok(SimpleTransition::Continue(connection, _)) => query_connection = connection,
1934 Ok(SimpleTransition::Ready(
1935 ReadyState::Clean(connection)
1936 | ReadyState::Dirty {
1937 conn: connection, ..
1938 },
1939 )) => {
1940 return Ok((
1941 ClientConnection {
1942 core: ClientConnectionCore {
1943 connection: connection.transition(),
1944 handler,
1945 context,
1946 },
1947 state,
1948 },
1949 messages,
1950 ));
1951 }
1952 Ok(SimpleTransition::Error(connection, _)) => {
1953 let _ = connection.into_transport();
1954 return Err(QueryError(io::Error::other(
1955 "backend rejected simple query",
1956 )));
1957 }
1958 Ok(SimpleTransition::CopyIn(connection, _)) => {
1959 let _ = connection.into_transport();
1960 return Err(QueryError(io::Error::other("simple query entered COPY IN")));
1961 }
1962 Ok(SimpleTransition::CopyOut(connection, _)) => {
1963 let _ = connection.into_transport();
1964 return Err(QueryError(io::Error::other(
1965 "simple query entered COPY OUT",
1966 )));
1967 }
1968 Ok(SimpleTransition::CopyBoth(connection, _)) => {
1969 let _ = connection.into_transport();
1970 return Err(QueryError(io::Error::other(
1971 "simple query entered COPY BOTH",
1972 )));
1973 }
1974 Err((connection, _)) => {
1975 let _ = connection.into_transport();
1976 return Err(QueryError(io::Error::new(
1977 io::ErrorKind::InvalidData,
1978 "illegal simple-query response",
1979 )));
1980 }
1981 }
1982 }
1983 }
1984}
1985
1986impl<Connector, Tls, Authentication, Middleware, Work, Transport, Error>
1987 Client<Connector, Tls, Authentication, Middleware>
1988where
1989 Connector: Fn(&ConnectTarget) -> Work,
1990 Work: Future<Output = Result<Transport, Error>>,
1991 Transport: AsyncRead + AsyncWrite + Unpin,
1992 Authentication: ClientAuthentication,
1993 Tls: ClientTlsConfiguration,
1994 Middleware: crate::MiddlewareFactory<ClientInitialContext>,
1995{
1996 pub async fn connect<State>(
2005 &self,
2006 target: ConnectTarget,
2007 overrides: StartupParameters,
2008 mut state: State,
2009 ) -> Result<
2010 ClientConnection<
2011 ClientTransport<Transport>,
2012 State,
2013 ConnectionClean,
2014 Authentication::Evidence,
2015 <Middleware as crate::MiddlewareFactory<ClientInitialContext>>::Handler,
2016 >,
2017 ConnectError<
2018 Error,
2019 ClientTlsError<<Tls::Provider as ClientTlsProvider>::Error>,
2020 ClientAuthenticationError<Authentication::Error>,
2021 >,
2022 >
2023 where
2024 <Middleware as crate::MiddlewareFactory<ClientInitialContext>>::Handler:
2025 crate::ClientMiddleware<State, ClientConnectionContext<Authentication::Evidence>>,
2026 {
2027 let core = self.connect_core(target, overrides, &mut state).await?;
2028 Ok(ClientConnection {
2029 core: ClientConnectionCore {
2030 connection: core.connection.transition(),
2031 handler: core.handler,
2032 context: core.context,
2033 },
2034 state,
2035 })
2036 }
2037
2038 pub(crate) async fn connect_core<State>(
2041 &self,
2042 target: ConnectTarget,
2043 overrides: StartupParameters,
2044 state: &mut State,
2045 ) -> Result<
2046 ClientConnectionCore<
2047 ClientTransport<Transport>,
2048 Pristine,
2049 Authentication::Evidence,
2050 <Middleware as crate::MiddlewareFactory<ClientInitialContext>>::Handler,
2051 >,
2052 ConnectError<
2053 Error,
2054 ClientTlsError<<Tls::Provider as ClientTlsProvider>::Error>,
2055 ClientAuthenticationError<Authentication::Error>,
2056 >,
2057 >
2058 where
2059 <Middleware as crate::MiddlewareFactory<ClientInitialContext>>::Handler:
2060 crate::ClientMiddleware<State, ClientConnectionContext<Authentication::Evidence>>,
2061 {
2062 let mut handler = self.middleware.create(&ClientInitialContext {
2063 target: target.clone(),
2064 });
2065 let startup = self
2066 .defaults
2067 .clone()
2068 .merged_with(overrides)
2069 .into_message()
2070 .map_err(ConnectError::Startup)?;
2071 let mut context = ClientConnectionContext {
2072 target: target.clone(),
2073 tls: None,
2074 identity: None,
2075 backend_key: None,
2076 };
2077 let transport = (self.connector)(&target)
2078 .await
2079 .map_err(ConnectError::Connector)?;
2080 let configured_tls = self.tls.configured();
2081 let transport = match configured_tls {
2082 None => {
2083 context.tls = Some(ClientTlsStatus::Plaintext);
2084 ClientTransport::Plain(transport)
2085 }
2086 Some((mode, provider)) => negotiate_client_tls(
2087 transport,
2088 mode,
2089 provider,
2090 &target,
2091 &mut context,
2092 state,
2093 &mut handler,
2094 )
2095 .await
2096 .map_err(ConnectError::Tls)?,
2097 };
2098 let first_startup = handler.startup(&context, state, startup.clone());
2099 let first = establish_client_session(
2100 transport,
2101 &first_startup,
2102 &target,
2103 &self.authentication,
2104 self.limits.max_frame_len,
2105 &context,
2106 state,
2107 &mut handler,
2108 )
2109 .await;
2110 let retry_provider = match configured_tls {
2111 Some((crate::pre_startup::SslMode::Allow, provider)) if first.is_err() => {
2112 Some(provider)
2113 }
2114 _ => None,
2115 };
2116 let (ready, identity, backend_key) = if let Some(provider) = retry_provider {
2117 let transport = (self.connector)(&target)
2118 .await
2119 .map_err(ConnectError::Connector)?;
2120 context.tls = None;
2124 let transport = negotiate_client_tls(
2125 transport,
2126 crate::pre_startup::SslMode::Require,
2127 provider,
2128 &target,
2129 &mut context,
2130 state,
2131 &mut handler,
2132 )
2133 .await
2134 .map_err(ConnectError::Tls)?;
2135 let retry_startup = handler.startup(&context, state, startup);
2136 establish_client_session(
2137 transport,
2138 &retry_startup,
2139 &target,
2140 &self.authentication,
2141 self.limits.max_frame_len,
2142 &context,
2143 state,
2144 &mut handler,
2145 )
2146 .await
2147 .map_err(map_session_error)?
2148 } else {
2149 first.map_err(map_session_error)?
2150 };
2151 Ok(ClientConnectionCore {
2152 connection: ready,
2153 handler,
2154 context: {
2155 context.identity = Some(identity);
2156 context.backend_key = backend_key;
2157 context
2158 },
2159 })
2160 }
2161}
2162
2163impl<Connector, Tls, Authentication, Middleware, Work, Transport, Error>
2164 Client<Connector, Tls, Authentication, Middleware>
2165where
2166 Connector: Fn(&ConnectTarget) -> Work,
2167 Work: Future<Output = Result<Transport, Error>>,
2168 Transport: AsyncWrite + Unpin,
2169{
2170 pub async fn cancel(
2179 &self,
2180 target: &ConnectTarget,
2181 key: &crate::demux::CancelKey,
2182 ) -> Result<(), CancelError<Error>> {
2183 let mut transport = (self.connector)(target)
2184 .await
2185 .map_err(CancelError::Connector)?;
2186 let packet = crate::pre_startup::PreStartupMessage::CancelRequest {
2187 process_id: key.process_id,
2188 secret_key: key.secret_key.clone(),
2189 }
2190 .to_packet()
2191 .map_err(CancelError::Protocol)?;
2192 tokio::io::AsyncWriteExt::write_all(&mut transport, &packet)
2193 .await
2194 .map_err(CancelError::Protocol)?;
2195 tokio::io::AsyncWriteExt::shutdown(&mut transport)
2196 .await
2197 .map_err(CancelError::Protocol)
2198 }
2199}