1use std::{collections::VecDeque, fmt, future::Future, io, pin::Pin};
4
5use tokio::io::{AsyncRead, AsyncWrite};
6
7use crate::{
8 ConnectTarget, NoPipeline, StartupParameters,
9 pipeline::{
10 BackendAction, FrontendAction, FrontendHandling, OperationId, Pipeline, PipelinePolicy,
11 },
12};
13
14fn is_backend_batch_barrier(message: &crate::codec::BackendMessage) -> bool {
15 use crate::codec::BackendMessage as B;
16 matches!(
17 message,
18 B::CommandComplete(_)
19 | B::PortalSuspended
20 | B::EmptyQueryResponse
21 | B::ErrorResponse(_)
22 | B::ReadyForQuery(_)
23 | B::CopyInResponse(_)
24 | B::CopyOutResponse(_)
25 | B::CopyBothResponse(_)
26 | B::CopyDone
27 | B::NoticeResponse(_)
28 | B::NotificationResponse { .. }
29 | B::ParameterStatus { .. }
30 | B::BackendKeyData { .. }
31 )
32}
33
34#[derive(Clone, Copy, Debug, Eq, PartialEq)]
39pub enum CancellationPolicy {
40 Reject,
42 Forward,
44}
45
46#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
48pub enum EstablishmentFailurePolicy {
49 #[default]
51 Close,
52 SafeDiagnostic,
54}
55
56fn safe_establishment_diagnostic() -> crate::codec::BackendMessage {
57 crate::codec::BackendMessage::ErrorResponse(crate::codec::DiagnosticResponse {
58 fields: vec![
59 crate::codec::DiagnosticField {
60 code: b'S',
61 value: bytes::Bytes::from_static(b"ERROR"),
62 },
63 crate::codec::DiagnosticField {
64 code: b'M',
65 value: bytes::Bytes::from_static(b"connection establishment failed"),
66 },
67 ],
68 })
69}
70
71#[derive(Clone, Debug, Eq, PartialEq)]
73pub struct CancellationRoute {
74 target: ConnectTarget,
75 upstream: crate::demux::CancelKey,
76}
77
78impl CancellationRoute {
79 #[must_use]
81 pub const fn new(target: ConnectTarget, upstream: crate::demux::CancelKey) -> Self {
82 Self { target, upstream }
83 }
84 #[must_use]
86 pub const fn target(&self) -> &ConnectTarget {
87 &self.target
88 }
89 #[must_use]
91 pub const fn upstream_key(&self) -> &crate::demux::CancelKey {
92 &self.upstream
93 }
94}
95
96pub trait IntermediaryCancellationRegistry {
102 type Error;
104 fn register(&self, route: CancellationRoute) -> Result<crate::demux::CancelKey, Self::Error>;
110 fn resolve(&self, client: &crate::demux::CancelKey) -> Option<CancellationRoute>;
112 fn detach(&self, client: &crate::demux::CancelKey) -> Option<CancellationRoute>;
114}
115
116#[derive(Clone, Copy, Debug, Default)]
118pub struct RejectCancellation;
119impl IntermediaryCancellationRegistry for RejectCancellation {
120 type Error = std::convert::Infallible;
121 fn register(&self, _: CancellationRoute) -> Result<crate::demux::CancelKey, Self::Error> {
122 unreachable!()
123 }
124 fn resolve(&self, _: &crate::demux::CancelKey) -> Option<CancellationRoute> {
125 None
126 }
127 fn detach(&self, _: &crate::demux::CancelKey) -> Option<CancellationRoute> {
128 None
129 }
130}
131
132#[derive(Clone, Copy, Debug, Eq, PartialEq)]
134pub enum IntermediaryBuildError {
135 MissingServer,
137 MissingClient,
139 MissingStartupResolver,
141 MissingCancellationPolicy,
143}
144
145impl fmt::Display for IntermediaryBuildError {
146 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
147 formatter.write_str(match self {
148 Self::MissingServer => "an intermediary server component is required",
149 Self::MissingClient => "an intermediary client component is required",
150 Self::MissingStartupResolver => "an asynchronous startup resolver is required",
151 Self::MissingCancellationPolicy => "an explicit cancellation policy is required",
152 })
153 }
154}
155
156impl std::error::Error for IntermediaryBuildError {}
157
158#[derive(Clone, Copy, Debug)]
160pub struct InitialServerContext<'a, Peer> {
161 peer: &'a Peer,
162 tls: &'a crate::NegotiatedServerTls,
163}
164
165impl<'a, Peer> InitialServerContext<'a, Peer> {
166 pub(crate) const fn new(peer: &'a Peer, tls: &'a crate::NegotiatedServerTls) -> Self {
167 Self { peer, tls }
168 }
169
170 #[must_use]
172 pub const fn peer(&self) -> &Peer {
173 self.peer
174 }
175
176 #[must_use]
178 pub const fn tls(&self) -> &crate::NegotiatedServerTls {
179 self.tls
180 }
181}
182
183#[allow(async_fn_in_trait)]
187pub trait StartupRouteResolver<Peer> {
188 type Error;
190
191 async fn resolve(
193 &self,
194 startup: StartupParameters,
195 context: InitialServerContext<'_, Peer>,
196 ) -> Result<ConnectTarget, Self::Error>;
197}
198
199#[allow(async_fn_in_trait)]
203pub trait AuthenticatedRoutePolicy<Peer, Identity> {
204 type Error;
206 async fn route(
208 &self,
209 target: ConnectTarget,
210 context: AuthenticatedRouteContext<'_, Peer, Identity>,
211 ) -> Result<ConnectTarget, Self::Error>;
212}
213
214#[derive(Clone, Copy, Debug)]
216pub struct AuthenticatedRouteContext<'a, Peer, Identity> {
217 peer: &'a Peer,
218 identity: &'a Identity,
219}
220
221impl<Peer, Identity> AuthenticatedRouteContext<'_, Peer, Identity> {
222 #[must_use]
224 pub const fn peer(&self) -> &Peer {
225 self.peer
226 }
227
228 #[must_use]
230 pub const fn identity(&self) -> &Identity {
231 self.identity
232 }
233}
234
235#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
237pub struct AllowAuthenticatedRoute;
238
239impl<Peer, Identity> AuthenticatedRoutePolicy<Peer, Identity> for AllowAuthenticatedRoute {
240 type Error = std::convert::Infallible;
241 async fn route(
242 &self,
243 target: ConnectTarget,
244 _context: AuthenticatedRouteContext<'_, Peer, Identity>,
245 ) -> Result<ConnectTarget, Self::Error> {
246 Ok(target)
247 }
248}
249
250#[derive(Debug, Eq, PartialEq)]
252pub enum FrontendMiddlewareOutput {
253 Forward(crate::codec::FrontendMessage),
255 Suppress(crate::codec::FrontendMessage),
257 Respond {
259 request: crate::codec::FrontendMessage,
261 responses: Vec<crate::codec::BackendMessage>,
263 },
264}
265
266#[derive(Debug, Eq, PartialEq)]
268pub enum BackendMiddlewareOutput {
269 Forward(crate::codec::BackendMessage),
271 Expand(Vec<crate::codec::BackendMessage>),
273 Suppress(crate::codec::BackendMessage),
275 Hold,
277}
278
279#[derive(Debug, Eq, PartialEq)]
281pub enum BackendBatchOutput {
282 KeepHolding,
284 ReplaceOneToOne(Vec<crate::codec::BackendMessage>),
286}
287
288#[derive(Clone, Copy, Debug, Eq, PartialEq)]
290pub enum BackendFlushReason {
291 Capacity,
293 ProtocolBarrier,
295 Explicit,
297 Teardown,
299}
300
301#[derive(Clone, Copy, Debug, Eq, PartialEq)]
303pub struct BackendHoldLimits {
304 max_messages: usize,
305 max_bytes: usize,
306}
307
308impl BackendHoldLimits {
309 pub const fn new(
314 max_messages: usize,
315 max_bytes: usize,
316 ) -> Result<Self, BackendHoldConfigError> {
317 if max_messages == 0 || max_bytes == 0 {
318 Err(BackendHoldConfigError)
319 } else {
320 Ok(Self {
321 max_messages,
322 max_bytes,
323 })
324 }
325 }
326 #[must_use]
328 pub const fn max_messages(self) -> usize {
329 self.max_messages
330 }
331 #[must_use]
333 pub const fn max_bytes(self) -> usize {
334 self.max_bytes
335 }
336}
337
338#[derive(Clone, Copy, Debug, Eq, PartialEq)]
340pub struct BackendHoldConfigError;
341
342impl fmt::Display for BackendHoldConfigError {
343 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
344 formatter.write_str("backend hold limits must be non-zero")
345 }
346}
347impl std::error::Error for BackendHoldConfigError {}
348
349#[derive(Clone, Copy, Debug)]
351pub struct HeldBackendMessages<'a> {
352 messages: &'a [crate::codec::BackendMessage],
353 bytes: usize,
354}
355
356impl<'a> HeldBackendMessages<'a> {
357 #[must_use]
359 pub const fn len(self) -> usize {
360 self.messages.len()
361 }
362 #[must_use]
364 pub const fn is_empty(self) -> bool {
365 self.messages.is_empty()
366 }
367 #[must_use]
369 pub const fn bytes(self) -> usize {
370 self.bytes
371 }
372 #[must_use]
374 pub fn iter(self) -> impl ExactSizeIterator<Item = &'a crate::codec::BackendMessage> {
375 self.messages.iter()
376 }
377}
378
379#[derive(Clone, Debug)]
381pub struct AttributedBackendMessages<'a> {
382 held: HeldBackendMessages<'a>,
383 operation_ids: Vec<Option<OperationId>>,
384}
385
386impl<'a> AttributedBackendMessages<'a> {
387 #[must_use]
389 pub const fn messages(&self) -> HeldBackendMessages<'a> {
390 self.held
391 }
392
393 #[must_use]
396 pub fn iter(
397 &self,
398 ) -> impl ExactSizeIterator<Item = (Option<OperationId>, &'a crate::codec::BackendMessage)> + '_
399 {
400 self.operation_ids.iter().copied().zip(self.held.messages)
401 }
402}
403
404#[allow(async_fn_in_trait)]
408pub trait IntermediaryMiddleware<State, ServerContext, ClientContext> {
409 type Error;
411
412 async fn frontend(
415 &mut self,
416 _server: &ServerContext,
417 _client: &ClientContext,
418 _state: &mut State,
419 message: crate::codec::FrontendMessage,
420 ) -> Result<FrontendMiddlewareOutput, Self::Error> {
421 Ok(FrontendMiddlewareOutput::Forward(message))
422 }
423
424 async fn frontend_operation(
431 &mut self,
432 server: &ServerContext,
433 client: &ClientContext,
434 state: &mut State,
435 _operation: OperationId,
436 message: crate::codec::FrontendMessage,
437 ) -> Result<FrontendMiddlewareOutput, Self::Error> {
438 self.frontend(server, client, state, message).await
439 }
440
441 async fn backend(
444 &mut self,
445 _server: &ServerContext,
446 _client: &ClientContext,
447 _state: &mut State,
448 message: crate::codec::BackendMessage,
449 ) -> Result<BackendMiddlewareOutput, Self::Error> {
450 Ok(BackendMiddlewareOutput::Forward(message))
451 }
452
453 async fn backend_operation(
459 &mut self,
460 server: &ServerContext,
461 client: &ClientContext,
462 state: &mut State,
463 _operation: Option<OperationId>,
464 message: crate::codec::BackendMessage,
465 ) -> Result<BackendMiddlewareOutput, Self::Error> {
466 self.backend(server, client, state, message).await
467 }
468
469 async fn flush_backend(
471 &mut self,
472 _server: &ServerContext,
473 _client: &ClientContext,
474 _state: &mut State,
475 held: HeldBackendMessages<'_>,
476 _reason: BackendFlushReason,
477 ) -> Result<BackendBatchOutput, Self::Error> {
478 Ok(BackendBatchOutput::ReplaceOneToOne(
479 held.iter().cloned().collect(),
480 ))
481 }
482
483 async fn flush_backend_operations(
487 &mut self,
488 server: &ServerContext,
489 client: &ClientContext,
490 state: &mut State,
491 held: AttributedBackendMessages<'_>,
492 reason: BackendFlushReason,
493 ) -> Result<BackendBatchOutput, Self::Error> {
494 self.flush_backend(server, client, state, held.messages(), reason)
495 .await
496 }
497}
498
499#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
501pub struct IdentityIntermediaryMiddleware;
502
503impl<State, ServerContext, ClientContext>
504 IntermediaryMiddleware<State, ServerContext, ClientContext> for IdentityIntermediaryMiddleware
505{
506 type Error = std::convert::Infallible;
507}
508
509pub trait IntermediaryMiddlewareFactory<ServerContext, ClientContext> {
511 type Handler;
513 fn create(&self, server: &ServerContext, client: &ClientContext) -> Self::Handler;
515}
516
517impl<ServerContext, ClientContext, Handler, Factory>
518 IntermediaryMiddlewareFactory<ServerContext, ClientContext> for Factory
519where
520 Factory: Fn(&ServerContext, &ClientContext) -> Handler,
521{
522 type Handler = Handler;
523 fn create(&self, server: &ServerContext, client: &ClientContext) -> Handler {
524 self(server, client)
525 }
526}
527
528impl<ServerContext, ClientContext> IntermediaryMiddlewareFactory<ServerContext, ClientContext>
529 for IdentityIntermediaryMiddleware
530{
531 type Handler = Self;
532 fn create(&self, _server: &ServerContext, _client: &ClientContext) -> Self {
533 *self
534 }
535}
536
537pub struct Intermediary<
539 Server = (),
540 Client = (),
541 Resolver = (),
542 Route = AllowAuthenticatedRoute,
543 Policy = NoPipeline,
544 Boundary = IdentityIntermediaryMiddleware,
545 Cancellation = RejectCancellation,
546> {
547 pub(crate) server: Server,
548 pub(crate) client: Client,
549 pub(crate) resolver: Resolver,
550 pub(crate) route: Route,
551 pub(crate) pipeline: Policy,
552 pub(crate) boundary: Boundary,
553 pub(crate) cancellation: CancellationPolicy,
554 pub(crate) cancellation_registry: Cancellation,
555 pub(crate) failure_policy: EstablishmentFailurePolicy,
556 pub(crate) backend_hold_limits: Option<BackendHoldLimits>,
557}
558
559impl Intermediary<()> {
560 #[must_use]
562 pub fn builder() -> IntermediaryBuilder {
563 IntermediaryBuilder::default()
564 }
565}
566
567impl<S, C, R, A, P, B, K> fmt::Debug for Intermediary<S, C, R, A, P, B, K> {
568 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
569 formatter
570 .debug_struct("Intermediary")
571 .field("server", &"<configured>")
572 .field("client", &"<configured>")
573 .field("resolver", &"<redacted>")
574 .field("authenticated_route", &"<redacted>")
575 .field("cancellation", &self.cancellation)
576 .finish_non_exhaustive()
577 }
578}
579
580pub struct IntermediaryBuilder<
582 Server = (),
583 Client = (),
584 Resolver = (),
585 Route = AllowAuthenticatedRoute,
586 Policy = NoPipeline,
587 Boundary = IdentityIntermediaryMiddleware,
588 Cancellation = RejectCancellation,
589> {
590 server: Option<Server>,
591 client: Option<Client>,
592 resolver: Option<Resolver>,
593 route: Route,
594 pipeline: Policy,
595 boundary: Boundary,
596 cancellation: Option<CancellationPolicy>,
597 cancellation_registry: Cancellation,
598 failure_policy: EstablishmentFailurePolicy,
599 backend_hold_limits: Option<BackendHoldLimits>,
600}
601
602impl Default for IntermediaryBuilder {
603 fn default() -> Self {
604 Self {
605 server: None,
606 client: None,
607 resolver: None,
608 route: AllowAuthenticatedRoute,
609 pipeline: NoPipeline,
610 boundary: IdentityIntermediaryMiddleware,
611 cancellation: None,
612 cancellation_registry: RejectCancellation,
613 failure_policy: EstablishmentFailurePolicy::Close,
614 backend_hold_limits: None,
615 }
616 }
617}
618
619impl<S, C, R, A, P, B, K> IntermediaryBuilder<S, C, R, A, P, B, K> {
620 #[must_use]
622 pub fn server<Next>(self, server: Next) -> IntermediaryBuilder<Next, C, R, A, P, B, K> {
623 IntermediaryBuilder {
624 server: Some(server),
625 client: self.client,
626 resolver: self.resolver,
627 route: self.route,
628 pipeline: self.pipeline,
629 boundary: self.boundary,
630 cancellation: self.cancellation,
631 cancellation_registry: self.cancellation_registry,
632 failure_policy: self.failure_policy,
633 backend_hold_limits: self.backend_hold_limits,
634 }
635 }
636
637 #[must_use]
639 pub fn client<Next>(self, client: Next) -> IntermediaryBuilder<S, Next, R, A, P, B, K> {
640 IntermediaryBuilder {
641 server: self.server,
642 client: Some(client),
643 resolver: self.resolver,
644 route: self.route,
645 pipeline: self.pipeline,
646 boundary: self.boundary,
647 cancellation: self.cancellation,
648 cancellation_registry: self.cancellation_registry,
649 failure_policy: self.failure_policy,
650 backend_hold_limits: self.backend_hold_limits,
651 }
652 }
653
654 #[must_use]
656 pub fn startup_resolver<Next>(
657 self,
658 resolver: Next,
659 ) -> IntermediaryBuilder<S, C, Next, A, P, B, K> {
660 IntermediaryBuilder {
661 server: self.server,
662 client: self.client,
663 resolver: Some(resolver),
664 route: self.route,
665 pipeline: self.pipeline,
666 boundary: self.boundary,
667 cancellation: self.cancellation,
668 cancellation_registry: self.cancellation_registry,
669 failure_policy: self.failure_policy,
670 backend_hold_limits: self.backend_hold_limits,
671 }
672 }
673
674 #[must_use]
676 pub fn authenticated_route<Next>(
677 self,
678 route: Next,
679 ) -> IntermediaryBuilder<S, C, R, Next, P, B, K> {
680 IntermediaryBuilder {
681 server: self.server,
682 client: self.client,
683 resolver: self.resolver,
684 route,
685 pipeline: self.pipeline,
686 boundary: self.boundary,
687 cancellation: self.cancellation,
688 cancellation_registry: self.cancellation_registry,
689 failure_policy: self.failure_policy,
690 backend_hold_limits: self.backend_hold_limits,
691 }
692 }
693
694 #[must_use]
696 pub fn pipeline<Next: PipelinePolicy>(
697 self,
698 pipeline: Next,
699 ) -> IntermediaryBuilder<S, C, R, A, Next, B, K> {
700 IntermediaryBuilder {
701 server: self.server,
702 client: self.client,
703 resolver: self.resolver,
704 route: self.route,
705 pipeline,
706 boundary: self.boundary,
707 cancellation: self.cancellation,
708 cancellation_registry: self.cancellation_registry,
709 failure_policy: self.failure_policy,
710 backend_hold_limits: self.backend_hold_limits,
711 }
712 }
713
714 #[must_use]
716 pub fn middleware<Next>(self, boundary: Next) -> IntermediaryBuilder<S, C, R, A, P, Next, K> {
717 IntermediaryBuilder {
718 server: self.server,
719 client: self.client,
720 resolver: self.resolver,
721 route: self.route,
722 pipeline: self.pipeline,
723 boundary,
724 cancellation: self.cancellation,
725 cancellation_registry: self.cancellation_registry,
726 failure_policy: self.failure_policy,
727 backend_hold_limits: self.backend_hold_limits,
728 }
729 }
730
731 #[must_use]
733 pub fn cancellation(mut self, cancellation: CancellationPolicy) -> Self {
734 self.cancellation = match cancellation {
735 CancellationPolicy::Reject => Some(CancellationPolicy::Reject),
736 CancellationPolicy::Forward => None,
737 };
738 self
739 }
740
741 #[must_use]
743 pub fn establishment_failure(mut self, policy: EstablishmentFailurePolicy) -> Self {
744 self.failure_policy = policy;
745 self
746 }
747
748 #[must_use]
750 pub fn backend_batching(mut self, limits: BackendHoldLimits) -> Self {
751 self.backend_hold_limits = Some(limits);
752 self
753 }
754
755 #[must_use]
757 pub fn cancellation_registry<Next>(
758 self,
759 registry: Next,
760 ) -> IntermediaryBuilder<S, C, R, A, P, B, Next> {
761 IntermediaryBuilder {
762 server: self.server,
763 client: self.client,
764 resolver: self.resolver,
765 route: self.route,
766 pipeline: self.pipeline,
767 boundary: self.boundary,
768 cancellation: Some(CancellationPolicy::Forward),
769 cancellation_registry: registry,
770 failure_policy: self.failure_policy,
771 backend_hold_limits: self.backend_hold_limits,
772 }
773 }
774
775 #[allow(clippy::type_complexity)]
781 pub fn build(self) -> Result<Intermediary<S, C, R, A, P, B, K>, IntermediaryBuildError> {
782 Ok(Intermediary {
783 server: self.server.ok_or(IntermediaryBuildError::MissingServer)?,
784 client: self.client.ok_or(IntermediaryBuildError::MissingClient)?,
785 resolver: self
786 .resolver
787 .ok_or(IntermediaryBuildError::MissingStartupResolver)?,
788 route: self.route,
789 pipeline: self.pipeline,
790 boundary: self.boundary,
791 cancellation: self
792 .cancellation
793 .ok_or(IntermediaryBuildError::MissingCancellationPolicy)?,
794 cancellation_registry: self.cancellation_registry,
795 failure_policy: self.failure_policy,
796 backend_hold_limits: self.backend_hold_limits,
797 })
798 }
799}
800
801struct StartupResolverAdapter<'a, Resolver> {
802 resolver: &'a Resolver,
803}
804
805#[derive(Debug)]
807pub enum StartupResolutionError<Error> {
808 Parameters(io::Error),
810 Resolver(Error),
812}
813
814impl<Error: fmt::Display> fmt::Display for StartupResolutionError<Error> {
815 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
816 match self {
817 Self::Parameters(error) => error.fmt(formatter),
818 Self::Resolver(error) => error.fmt(formatter),
819 }
820 }
821}
822
823impl<Error: std::error::Error + 'static> std::error::Error for StartupResolutionError<Error> {
824 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
825 match self {
826 Self::Parameters(error) => Some(error),
827 Self::Resolver(error) => Some(error),
828 }
829 }
830}
831
832impl<Resolver, State, Peer, Identity>
833 crate::server_component::StartupResolver<State, Peer, Identity>
834 for StartupResolverAdapter<'_, Resolver>
835where
836 Resolver: StartupRouteResolver<Peer>,
837{
838 type Route = ConnectTarget;
839 type Error = StartupResolutionError<Resolver::Error>;
840
841 fn defer_ready(&self) -> bool {
842 true
843 }
844
845 fn resolve<'a>(
846 &'a mut self,
847 startup: &'a crate::startup::StartupMessage,
848 context: &'a crate::ServerConnectionContext<Peer, Identity>,
849 _state: &'a mut State,
850 ) -> Pin<Box<dyn Future<Output = Result<Self::Route, Self::Error>> + 'a>> {
851 let parameters = StartupParameters::from_wire(startup);
852 let initial = context
853 .tls_if_known()
854 .map(|tls| InitialServerContext::new(context.peer(), tls));
855 let resolver = self.resolver;
856 Box::pin(async move {
857 let parameters = parameters.map_err(StartupResolutionError::Parameters)?;
858 let initial = initial.expect("startup routing runs after TLS negotiation");
859 resolver
860 .resolve(parameters, initial)
861 .await
862 .map_err(StartupResolutionError::Resolver)
863 })
864 }
865}
866
867pub enum IntermediaryAcceptError<
869 ServerError,
870 ResolverError,
871 RouteError,
872 ClientError,
873 RegistryError = std::convert::Infallible,
874 CancellationError = std::convert::Infallible,
875 MiddlewareError = std::convert::Infallible,
876> {
877 Server(ServerError),
879 StartupRoute(StartupResolutionError<ResolverError>),
881 CancellationRejected,
883 AuthenticatedRoute(RouteError),
885 Client(ClientError),
887 CancellationRegistry(RegistryError),
889 ServerOutput(io::Error),
891 Cancellation(CancellationError),
893 Middleware(MiddlewareError),
895}
896
897impl<S, R, A, C, K, X, M> fmt::Debug for IntermediaryAcceptError<S, R, A, C, K, X, M> {
898 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
899 formatter.write_str(match self {
900 Self::Server(_) => "IntermediaryAcceptError::Server([REDACTED])",
901 Self::StartupRoute(_) => "IntermediaryAcceptError::StartupRoute([REDACTED])",
902 Self::CancellationRejected => "IntermediaryAcceptError::CancellationRejected",
903 Self::AuthenticatedRoute(_) => {
904 "IntermediaryAcceptError::AuthenticatedRoute([REDACTED])"
905 }
906 Self::Client(_) => "IntermediaryAcceptError::Client([REDACTED])",
907 Self::CancellationRegistry(_) => {
908 "IntermediaryAcceptError::CancellationRegistry([REDACTED])"
909 }
910 Self::ServerOutput(_) => "IntermediaryAcceptError::ServerOutput([REDACTED])",
911 Self::Cancellation(_) => "IntermediaryAcceptError::Cancellation([REDACTED])",
912 Self::Middleware(_) => "IntermediaryAcceptError::Middleware([REDACTED])",
913 })
914 }
915}
916
917impl<S, R, A, C, K, X, M> fmt::Display for IntermediaryAcceptError<S, R, A, C, K, X, M> {
918 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
919 match self {
920 Self::Server(_) => formatter.write_str("client-facing establishment failed"),
921 Self::StartupRoute(_) => formatter.write_str("startup routing failed"),
922 Self::CancellationRejected => {
923 formatter.write_str("cancellation is explicitly rejected")
924 }
925 Self::AuthenticatedRoute(_) => formatter.write_str("authenticated routing failed"),
926 Self::Client(_) => formatter.write_str("PostgreSQL-facing establishment failed"),
927 Self::CancellationRegistry(_) => {
928 formatter.write_str("cancellation registration failed")
929 }
930 Self::ServerOutput(_) => {
931 formatter.write_str("client-facing establishment output failed")
932 }
933 Self::Cancellation(_) => formatter.write_str("cancellation forwarding failed"),
934 Self::Middleware(_) => {
935 formatter.write_str("forwarding middleware rejected establishment output")
936 }
937 }
938 }
939}
940
941impl<S, R, A, C, K, X, M> std::error::Error for IntermediaryAcceptError<S, R, A, C, K, X, M>
942where
943 S: std::error::Error + 'static,
944 R: std::error::Error + 'static,
945 A: std::error::Error + 'static,
946 C: std::error::Error + 'static,
947 K: std::error::Error + 'static,
948 X: std::error::Error + 'static,
949 M: std::error::Error + 'static,
950{
951 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
952 match self {
953 Self::Server(error) => Some(error),
954 Self::StartupRoute(error) => Some(error),
955 Self::CancellationRejected => None,
956 Self::AuthenticatedRoute(error) => Some(error),
957 Self::Client(error) => Some(error),
958 Self::CancellationRegistry(error) => Some(error),
959 Self::ServerOutput(error) => Some(error),
960 Self::Cancellation(error) => Some(error),
961 Self::Middleware(error) => Some(error),
962 }
963 }
964}
965
966#[derive(Debug)]
968pub struct IntermediaryContexts<ServerContext, ClientContext> {
969 server: ServerContext,
970 client: ClientContext,
971}
972
973impl<ServerContext, ClientContext> IntermediaryContexts<ServerContext, ClientContext> {
974 #[must_use]
976 pub const fn server(&self) -> &ServerContext {
977 &self.server
978 }
979 #[must_use]
981 pub const fn client(&self) -> &ClientContext {
982 &self.client
983 }
984}
985
986pub struct IntermediaryConnection<
988 DT,
989 UT,
990 State,
991 Peer,
992 ServerIdentity,
993 ClientEvidence,
994 ServerHandler,
995 ClientHandler,
996 Boundary,
997 Policy,
998 Cancellation = RejectCancellation,
999> {
1000 downstream:
1001 crate::server_component::ServerConnectionCore<DT, Peer, ServerIdentity, ServerHandler>,
1002 upstream: crate::client_component::ClientConnectionCore<
1003 crate::ClientTransport<UT>,
1004 crate::Pristine,
1005 ClientEvidence,
1006 ClientHandler,
1007 >,
1008 state: State,
1009 boundary: Boundary,
1010 pipeline: Pipeline<Policy>,
1011 target: ConnectTarget,
1012 pending_frontend: Option<crate::codec::FrontendMessage>,
1013 backend_hold: crate::backend_hold::BackendHold,
1014 backend_hold_limits: Option<BackendHoldLimits>,
1015 pending_local: VecDeque<PendingLocalResponses>,
1016 cancellation_registry: Cancellation,
1017 client_cancel_key: Option<crate::demux::CancelKey>,
1018}
1019
1020struct PendingLocalResponses {
1021 operation: crate::pipeline::OperationId,
1022 messages: VecDeque<crate::codec::BackendMessage>,
1023}
1024
1025#[derive(Debug)]
1027pub enum IntermediaryAccept<Connection> {
1028 Session(Connection),
1030 CancellationForwarded,
1032}
1033
1034impl<Connection> IntermediaryAccept<Connection> {
1035 #[must_use]
1040 pub fn into_session(self) -> Connection {
1041 match self {
1042 Self::Session(connection) => connection,
1043 Self::CancellationForwarded => panic!("accepted cancellation has no session"),
1044 }
1045 }
1046}
1047
1048#[derive(Debug)]
1050pub enum ForwardedMessage {
1051 Frontend(crate::codec::FrontendMessage),
1053 Backend(crate::codec::BackendMessage),
1055 BackendExpanded {
1057 source: crate::codec::BackendMessage,
1059 messages: Vec<crate::codec::BackendMessage>,
1061 },
1062 FrontendSuppressed(crate::codec::FrontendMessage),
1064 FrontendLocallyHandled(crate::codec::FrontendMessage),
1066 BackendSuppressed(crate::codec::BackendMessage),
1068 BackendHeld,
1070}
1071
1072#[derive(Debug, Eq, PartialEq)]
1074pub enum FrontendForwarding {
1075 Forwarded(crate::codec::FrontendMessage),
1077 Suppressed(crate::codec::FrontendMessage),
1079 LocallyHandled(crate::codec::FrontendMessage),
1081}
1082
1083impl FrontendForwarding {
1084 #[must_use]
1086 pub fn into_message(self) -> crate::codec::FrontendMessage {
1087 match self {
1088 Self::Forwarded(message)
1089 | Self::Suppressed(message)
1090 | Self::LocallyHandled(message) => message,
1091 }
1092 }
1093}
1094
1095#[derive(Debug, Eq, PartialEq)]
1097pub enum BackendForwarding {
1098 Forwarded(crate::codec::BackendMessage),
1100 Expanded {
1102 source: crate::codec::BackendMessage,
1104 messages: Vec<crate::codec::BackendMessage>,
1106 },
1107 Suppressed(crate::codec::BackendMessage),
1109 Held,
1111}
1112
1113impl BackendForwarding {
1114 #[must_use]
1120 pub fn into_message(self) -> crate::codec::BackendMessage {
1121 match self {
1122 Self::Forwarded(message) | Self::Suppressed(message) => message,
1123 Self::Expanded { source, .. } => source,
1124 Self::Held => panic!("a held response remains owned by the connection"),
1125 }
1126 }
1127}
1128
1129#[derive(Debug, Eq, PartialEq)]
1131pub enum BackendBatchForwarding {
1132 Released(Vec<crate::codec::BackendMessage>),
1134 Kept,
1136 Empty,
1138}
1139
1140#[derive(Clone, Copy, Debug, Eq, PartialEq)]
1142pub enum BackendBatchProjectionError {
1143 Cardinality {
1145 expected: usize,
1147 actual: usize,
1149 },
1150 IllegalSource,
1152 IllegalReplacement,
1154 DifferentSpan,
1156}
1157
1158impl From<crate::pipeline::BackendSequenceError> for BackendBatchProjectionError {
1159 fn from(error: crate::pipeline::BackendSequenceError) -> Self {
1160 match error {
1161 crate::pipeline::BackendSequenceError::Cardinality { expected, actual } => {
1162 Self::Cardinality { expected, actual }
1163 }
1164 crate::pipeline::BackendSequenceError::Source(_) => Self::IllegalSource,
1165 crate::pipeline::BackendSequenceError::Replacement(_) => Self::IllegalReplacement,
1166 crate::pipeline::BackendSequenceError::DifferentSpan => Self::DifferentSpan,
1167 }
1168 }
1169}
1170
1171impl<DT, UT, State, Peer, SI, CE, SH, CH, Boundary, Policy, K>
1172 IntermediaryConnection<DT, UT, State, Peer, SI, CE, SH, CH, Boundary, Policy, K>
1173where
1174 Policy: PipelinePolicy,
1175{
1176 #[must_use]
1178 pub const fn target(&self) -> &ConnectTarget {
1179 &self.target
1180 }
1181 #[must_use]
1183 pub const fn state(&self) -> &State {
1184 &self.state
1185 }
1186 #[must_use]
1188 pub const fn cancellation_key(&self) -> Option<&crate::demux::CancelKey> {
1189 self.client_cancel_key.as_ref()
1190 }
1191
1192 #[must_use]
1194 pub fn held_backend_messages(&self) -> HeldBackendMessages<'_> {
1195 HeldBackendMessages {
1196 messages: self.backend_hold.messages(),
1197 bytes: self.backend_hold.bytes(),
1198 }
1199 }
1200
1201 pub fn detach_cancellation(&mut self) -> Option<CancellationRoute>
1203 where
1204 K: IntermediaryCancellationRegistry,
1205 {
1206 self.client_cancel_key
1207 .take()
1208 .and_then(|key| self.cancellation_registry.detach(&key))
1209 }
1210}
1211
1212impl<
1213 DT,
1214 UT,
1215 State,
1216 Peer,
1217 ServerIdentity,
1218 ClientEvidence,
1219 ServerHandler,
1220 ClientHandler,
1221 Boundary,
1222 Policy,
1223 K,
1224>
1225 IntermediaryConnection<
1226 DT,
1227 UT,
1228 State,
1229 Peer,
1230 ServerIdentity,
1231 ClientEvidence,
1232 ServerHandler,
1233 ClientHandler,
1234 Boundary,
1235 Policy,
1236 K,
1237 >
1238where
1239 DT: AsyncRead + AsyncWrite + Unpin,
1240 UT: AsyncRead + AsyncWrite + Unpin,
1241 ServerHandler:
1242 crate::ServerMiddleware<State, crate::ServerConnectionContext<Peer, ServerIdentity>>,
1243 ClientHandler: crate::ClientMiddleware<State, crate::ClientConnectionContext<ClientEvidence>>,
1244 Boundary: IntermediaryMiddleware<
1245 State,
1246 crate::ServerConnectionContext<Peer, ServerIdentity>,
1247 crate::ClientConnectionContext<ClientEvidence>,
1248 >,
1249 Policy: PipelinePolicy,
1250 K: IntermediaryCancellationRegistry,
1251{
1252 pub async fn forward_frontend(
1259 &mut self,
1260 ) -> Result<FrontendForwarding, ForwardError<Boundary::Error>> {
1261 if let Some(message) = self.pending_frontend.take() {
1262 self.process_frontend(message, false).await
1263 } else {
1264 let message = self.downstream.receive_wire_raw().await?;
1265 self.process_frontend(message, true).await
1266 }
1267 }
1268
1269 async fn process_frontend(
1270 &mut self,
1271 message: crate::codec::FrontendMessage,
1272 intercept_source_and_boundary: bool,
1273 ) -> Result<FrontendForwarding, ForwardError<Boundary::Error>> {
1274 let decision = if intercept_source_and_boundary {
1275 let message = self.downstream.intercept_frontend(&mut self.state, message);
1276 let operation = self.pipeline.next_operation_id();
1277 self.boundary
1278 .frontend_operation(
1279 self.downstream.context(),
1280 self.upstream.context(),
1281 &mut self.state,
1282 operation,
1283 message,
1284 )
1285 .await
1286 .map_err(ForwardError::Middleware)?
1287 } else {
1288 FrontendMiddlewareOutput::Forward(message)
1289 };
1290 let (message, handling) = match decision {
1291 FrontendMiddlewareOutput::Forward(message) => {
1292 let message = if intercept_source_and_boundary {
1293 self.upstream.intercept_frontend(&mut self.state, message)
1294 } else {
1295 message
1296 };
1297 (message, FrontendHandling::Forward)
1298 }
1299 FrontendMiddlewareOutput::Suppress(message) => {
1300 return Ok(FrontendForwarding::Suppressed(message));
1301 }
1302 FrontendMiddlewareOutput::Respond { request, responses } => {
1303 let admission = self
1304 .pipeline
1305 .accept_frontend(request.clone(), FrontendHandling::Local)
1306 .map_err(ForwardError::Frontend)?;
1307 let FrontendAction::Discard { id } = admission.into_action() else {
1308 unreachable!()
1309 };
1310 let messages = responses
1311 .into_iter()
1312 .map(|message| self.downstream.intercept_backend(&mut self.state, message))
1313 .collect();
1314 self.pending_local.push_back(PendingLocalResponses {
1315 operation: id,
1316 messages,
1317 });
1318 self.flush_local_responses().await?;
1319 return Ok(FrontendForwarding::LocallyHandled(request));
1320 }
1321 };
1322 let admission = match self.pipeline.accept_frontend(message.clone(), handling) {
1323 Ok(admission) => admission,
1324 Err(error) => {
1325 self.pending_frontend = Some(message);
1326 return Err(ForwardError::Frontend(error));
1327 }
1328 };
1329 let FrontendAction::Forward { message, .. } = admission.into_action() else {
1330 unreachable!()
1331 };
1332 self.upstream.send_wire_raw(message.clone()).await?;
1333 Ok(FrontendForwarding::Forwarded(message))
1334 }
1335
1336 pub async fn forward_backend(
1349 &mut self,
1350 ) -> Result<BackendForwarding, ForwardError<Boundary::Error>> {
1351 if self.backend_hold.pending().is_some() {
1352 return self.process_pending_backend().await;
1353 }
1354 if self.backend_hold_is_full() {
1355 match self
1356 .flush_backend_hold_for(BackendFlushReason::Capacity)
1357 .await?
1358 {
1359 BackendBatchForwarding::Released(_) | BackendBatchForwarding::Empty => {}
1360 BackendBatchForwarding::Kept => return Err(ForwardError::BackendHoldCapacity),
1361 }
1362 }
1363 let message = self.upstream.receive_wire_raw().await?;
1364 self.process_backend(message).await
1365 }
1366
1367 async fn process_backend(
1368 &mut self,
1369 message: crate::codec::BackendMessage,
1370 ) -> Result<BackendForwarding, ForwardError<Boundary::Error>> {
1371 let message = self.upstream.intercept_backend(&mut self.state, message);
1372 self.backend_hold.set_pending(message);
1373 if self
1374 .backend_hold
1375 .pending()
1376 .is_some_and(is_backend_batch_barrier)
1377 && !self.backend_hold.is_empty()
1378 {
1379 match self
1380 .flush_backend_hold_for(BackendFlushReason::ProtocolBarrier)
1381 .await?
1382 {
1383 BackendBatchForwarding::Released(_) | BackendBatchForwarding::Empty => {}
1384 BackendBatchForwarding::Kept => return Err(ForwardError::BackendHoldRefused),
1385 }
1386 }
1387 self.process_pending_backend().await
1388 }
1389
1390 async fn process_pending_backend(
1391 &mut self,
1392 ) -> Result<BackendForwarding, ForwardError<Boundary::Error>> {
1393 let source = self
1394 .backend_hold
1395 .pending()
1396 .expect("pending backend processing requires a source")
1397 .clone();
1398 let operation = self.pipeline.backend_operation_id(&source);
1399 let decision = self
1400 .boundary
1401 .backend_operation(
1402 self.downstream.context(),
1403 self.upstream.context(),
1404 &mut self.state,
1405 operation,
1406 source.clone(),
1407 )
1408 .await
1409 .map_err(ForwardError::Middleware)?;
1410 let outcome = match decision {
1411 BackendMiddlewareOutput::Forward(message) => {
1412 let _ = self.backend_hold.take_pending();
1413 let message = self.downstream.intercept_backend(&mut self.state, message);
1414 let message = self.emit_backend(message).await?;
1415 BackendForwarding::Forwarded(message)
1416 }
1417 BackendMiddlewareOutput::Suppress(message) => {
1418 let _ = self.backend_hold.take_pending();
1419 let message = self.advance_backend(message)?;
1420 BackendForwarding::Suppressed(message)
1421 }
1422 BackendMiddlewareOutput::Expand(messages) => {
1423 if messages.is_empty() {
1424 return Err(ForwardError::EmptyExpansion(source));
1425 }
1426 let _ = self.backend_hold.take_pending();
1427 let mut emitted = Vec::with_capacity(messages.len());
1428 for message in messages {
1429 let message = self.downstream.intercept_backend(&mut self.state, message);
1430 emitted.push(self.emit_backend(message).await?);
1431 }
1432 BackendForwarding::Expanded {
1433 source,
1434 messages: emitted,
1435 }
1436 }
1437 BackendMiddlewareOutput::Hold => {
1438 if self.backend_hold_limits.is_none() {
1439 return Err(ForwardError::BackendHoldingDisabled(source));
1440 }
1441 self.backend_hold.hold_pending();
1442 BackendForwarding::Held
1443 }
1444 };
1445 self.flush_local_responses().await?;
1446 Ok(outcome)
1447 }
1448
1449 fn backend_hold_is_full(&self) -> bool {
1450 self.backend_hold_limits.is_some_and(|limits| {
1451 self.backend_hold.len() >= limits.max_messages()
1452 || self.backend_hold.bytes() >= limits.max_bytes()
1453 })
1454 }
1455
1456 pub async fn flush_backend_hold(
1462 &mut self,
1463 ) -> Result<BackendBatchForwarding, ForwardError<Boundary::Error>> {
1464 self.flush_backend_hold_for(BackendFlushReason::Explicit)
1465 .await
1466 }
1467
1468 pub async fn prepare_backend_teardown(
1474 &mut self,
1475 ) -> Result<BackendBatchForwarding, ForwardError<Boundary::Error>> {
1476 let outcome = self
1477 .flush_backend_hold_for(BackendFlushReason::Teardown)
1478 .await?;
1479 if matches!(outcome, BackendBatchForwarding::Kept) {
1480 return Err(ForwardError::BackendHoldRefused);
1481 }
1482 Ok(outcome)
1483 }
1484
1485 async fn flush_backend_hold_for(
1486 &mut self,
1487 reason: BackendFlushReason,
1488 ) -> Result<BackendBatchForwarding, ForwardError<Boundary::Error>> {
1489 if self.backend_hold.is_empty() {
1490 return Ok(BackendBatchForwarding::Empty);
1491 }
1492 let held = AttributedBackendMessages {
1493 held: HeldBackendMessages {
1494 messages: self.backend_hold.messages(),
1495 bytes: self.backend_hold.bytes(),
1496 },
1497 operation_ids: self
1498 .pipeline
1499 .backend_operation_ids(self.backend_hold.messages()),
1500 };
1501 let decision = self
1502 .boundary
1503 .flush_backend_operations(
1504 self.downstream.context(),
1505 self.upstream.context(),
1506 &mut self.state,
1507 held,
1508 reason,
1509 )
1510 .await
1511 .map_err(ForwardError::Middleware)?;
1512 let BackendBatchOutput::ReplaceOneToOne(messages) = decision else {
1513 return Ok(BackendBatchForwarding::Kept);
1514 };
1515 let messages: Vec<_> = messages
1516 .into_iter()
1517 .map(|message| self.downstream.intercept_backend(&mut self.state, message))
1518 .collect();
1519 for message in &messages {
1520 message.to_frame().map_err(ForwardError::Io)?;
1521 }
1522 let prepared = self
1523 .pipeline
1524 .prepare_backend_replacements(self.backend_hold.messages(), &messages)
1525 .map_err(|error| ForwardError::BackendBatch {
1526 error: error.into(),
1527 proposed: messages.clone(),
1528 })?;
1529 self.pipeline = prepared;
1530 let _sources = self.backend_hold.clear();
1531 for message in &messages {
1532 self.downstream.send_wire_raw(message.clone()).await?;
1533 }
1534 self.flush_local_responses().await?;
1535 Ok(BackendBatchForwarding::Released(messages))
1536 }
1537
1538 fn advance_backend(
1539 &mut self,
1540 message: crate::codec::BackendMessage,
1541 ) -> Result<crate::codec::BackendMessage, ForwardError<Boundary::Error>> {
1542 match self
1543 .pipeline
1544 .accept_backend(message)
1545 .map_err(ForwardError::Backend)?
1546 {
1547 BackendAction::Emit(message) => Ok(message),
1548 BackendAction::Deferred(message) => Err(ForwardError::Deferred(message)),
1549 }
1550 }
1551
1552 async fn emit_backend(
1553 &mut self,
1554 message: crate::codec::BackendMessage,
1555 ) -> Result<crate::codec::BackendMessage, ForwardError<Boundary::Error>> {
1556 let message = self.advance_backend(message)?;
1557 self.downstream.send_wire_raw(message.clone()).await?;
1558 Ok(message)
1559 }
1560
1561 async fn flush_local_responses(&mut self) -> Result<(), ForwardError<Boundary::Error>> {
1562 loop {
1563 let Some(pending) = self.pending_local.front_mut() else {
1564 return Ok(());
1565 };
1566 let Some(message) = pending.messages.pop_front() else {
1567 self.pending_local.pop_front();
1568 continue;
1569 };
1570 match self.pipeline.try_emit_local(pending.operation, message) {
1571 Ok(BackendAction::Emit(message)) => {
1572 self.downstream.send_wire_raw(message).await?;
1573 }
1574 Ok(BackendAction::Deferred(message)) => {
1575 pending.messages.push_front(message);
1576 return Ok(());
1577 }
1578 Err(error) => return Err(ForwardError::Backend(error)),
1579 }
1580 }
1581 }
1582
1583 pub async fn forward_next(
1594 &mut self,
1595 ) -> Result<ForwardedMessage, ForwardError<Boundary::Error>> {
1596 if self.backend_hold.pending().is_some() {
1597 return self
1598 .process_pending_backend()
1599 .await
1600 .map(|outcome| match outcome {
1601 BackendForwarding::Forwarded(message) => ForwardedMessage::Backend(message),
1602 BackendForwarding::Expanded { source, messages } => {
1603 ForwardedMessage::BackendExpanded { source, messages }
1604 }
1605 BackendForwarding::Suppressed(message) => {
1606 ForwardedMessage::BackendSuppressed(message)
1607 }
1608 BackendForwarding::Held => ForwardedMessage::BackendHeld,
1609 });
1610 }
1611 if self.backend_hold_is_full() {
1612 match self
1613 .flush_backend_hold_for(BackendFlushReason::Capacity)
1614 .await?
1615 {
1616 BackendBatchForwarding::Released(_) | BackendBatchForwarding::Empty => {}
1617 BackendBatchForwarding::Kept => return Err(ForwardError::BackendHoldCapacity),
1618 }
1619 }
1620 if self.pending_frontend.is_some() {
1621 let message = self.upstream.receive_wire_raw().await?;
1622 return self
1623 .process_backend(message)
1624 .await
1625 .map(|outcome| match outcome {
1626 BackendForwarding::Forwarded(message) => ForwardedMessage::Backend(message),
1627 BackendForwarding::Expanded { source, messages } => {
1628 ForwardedMessage::BackendExpanded { source, messages }
1629 }
1630 BackendForwarding::Suppressed(message) => {
1631 ForwardedMessage::BackendSuppressed(message)
1632 }
1633 BackendForwarding::Held => ForwardedMessage::BackendHeld,
1634 });
1635 }
1636 tokio::select! {
1637 result = self.downstream.receive_wire_raw() => {
1638 let message = result?;
1639 self.process_frontend(message, true).await.map(|outcome| match outcome {
1640 FrontendForwarding::Forwarded(message) => ForwardedMessage::Frontend(message),
1641 FrontendForwarding::Suppressed(message) => ForwardedMessage::FrontendSuppressed(message),
1642 FrontendForwarding::LocallyHandled(message) => ForwardedMessage::FrontendLocallyHandled(message),
1643 })
1644 }
1645 result = self.upstream.receive_wire_raw() => {
1646 let message = result?;
1647 self.process_backend(message).await.map(|outcome| match outcome {
1648 BackendForwarding::Forwarded(message) => ForwardedMessage::Backend(message),
1649 BackendForwarding::Expanded { source, messages } => {
1650 ForwardedMessage::BackendExpanded { source, messages }
1651 }
1652 BackendForwarding::Suppressed(message) => ForwardedMessage::BackendSuppressed(message),
1653 BackendForwarding::Held => ForwardedMessage::BackendHeld,
1654 })
1655 }
1656 }
1657 }
1658
1659 #[allow(clippy::type_complexity)]
1666 pub fn teardown(
1667 mut self,
1668 ) -> (
1669 crate::AcceptedServerTransport<DT>,
1670 crate::ClientTransport<UT>,
1671 State,
1672 Boundary,
1673 (ServerHandler, ClientHandler),
1674 IntermediaryContexts<
1675 crate::ServerConnectionContext<Peer, ServerIdentity>,
1676 crate::ClientConnectionContext<ClientEvidence>,
1677 >,
1678 ) {
1679 assert!(
1680 self.backend_hold.is_empty() && self.backend_hold.pending().is_none(),
1681 "backend messages remain held; call prepare_backend_teardown before teardown"
1682 );
1683 let _ = self.detach_cancellation();
1684 let (downstream, server_handler, server_context) = self.downstream.into_parts();
1685 let (upstream, client_handler, client_context) = self.upstream.into_parts();
1686 (
1687 downstream,
1688 upstream,
1689 self.state,
1690 self.boundary,
1691 (server_handler, client_handler),
1692 IntermediaryContexts {
1693 server: server_context,
1694 client: client_context,
1695 },
1696 )
1697 }
1698}
1699
1700#[derive(Debug)]
1702pub enum ForwardError<MiddlewareError = std::convert::Infallible> {
1703 Io(io::Error),
1705 Frontend(crate::pipeline::FrontendProjectionError),
1707 Backend(crate::pipeline::BackendProjectionError),
1709 Deferred(crate::codec::BackendMessage),
1711 EmptyExpansion(crate::codec::BackendMessage),
1713 BackendHoldingDisabled(crate::codec::BackendMessage),
1715 BackendHoldCapacity,
1717 BackendHoldRefused,
1719 BackendBatch {
1721 error: BackendBatchProjectionError,
1723 proposed: Vec<crate::codec::BackendMessage>,
1725 },
1726 Middleware(MiddlewareError),
1728}
1729
1730impl<E> fmt::Display for ForwardError<E> {
1731 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
1732 match self {
1733 Self::Io(error) => error.fmt(formatter),
1734 Self::Frontend(_) => {
1735 formatter.write_str("frontend message violates pipeline legality or capacity")
1736 }
1737 Self::Backend(_) => formatter.write_str("backend message violates pipeline legality"),
1738 Self::Deferred(_) => formatter.write_str("backend response is not yet emittable"),
1739 Self::EmptyExpansion(_) => {
1740 formatter.write_str("backend expansion must contain at least one response")
1741 }
1742 Self::BackendHoldingDisabled(_) => {
1743 formatter.write_str("backend holding is not configured")
1744 }
1745 Self::BackendHoldCapacity => {
1746 formatter.write_str("backend hold is full and batch policy kept holding")
1747 }
1748 Self::BackendHoldRefused => {
1749 formatter.write_str("backend batch policy refused a required release")
1750 }
1751 Self::BackendBatch { .. } => formatter
1752 .write_str("backend batch replacement is not a legal one-to-one protocol span"),
1753 Self::Middleware(_) => formatter.write_str("forwarding middleware rejected a message"),
1754 }
1755 }
1756}
1757
1758impl<E> std::error::Error for ForwardError<E>
1759where
1760 E: std::error::Error + 'static,
1761{
1762 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
1763 match self {
1764 Self::Io(error) => Some(error),
1765 Self::Middleware(error) => Some(error),
1766 Self::Frontend(_)
1767 | Self::Backend(_)
1768 | Self::Deferred(_)
1769 | Self::EmptyExpansion(_)
1770 | Self::BackendHoldingDisabled(_)
1771 | Self::BackendHoldCapacity
1772 | Self::BackendHoldRefused
1773 | Self::BackendBatch { .. } => None,
1774 }
1775 }
1776}
1777
1778impl<E> From<io::Error> for ForwardError<E> {
1779 fn from(error: io::Error) -> Self {
1780 Self::Io(error)
1781 }
1782}
1783
1784impl<ST, SA, SM, Connector, CT, CA, CM, Resolver, Route, Policy, Boundary, K>
1785 Intermediary<
1786 crate::Server<ST, SA, SM>,
1787 crate::Client<Connector, CT, CA, CM>,
1788 Resolver,
1789 Route,
1790 Policy,
1791 Boundary,
1792 K,
1793 >
1794where
1795 ST: crate::ServerTlsConfiguration,
1796 SA: crate::ServerAuthenticationProvider,
1797 CT: crate::client_component::ClientTlsConfiguration,
1798 CA: crate::ClientAuthentication,
1799 CM: crate::MiddlewareFactory<crate::ClientInitialContext>,
1800 Policy: PipelinePolicy,
1801 K: IntermediaryCancellationRegistry + Clone,
1802{
1803 #[allow(clippy::type_complexity, clippy::too_many_lines)]
1810 pub async fn accept<DT, State, Peer, CW, UT, CE>(
1811 &self,
1812 transport: DT,
1813 peer: Peer,
1814 state: State,
1815 ) -> Result<
1816 IntermediaryAccept<
1817 IntermediaryConnection<
1818 DT,
1819 UT,
1820 State,
1821 Peer,
1822 <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1823 CA::Evidence,
1824 <SM as crate::MiddlewareFactory<
1825 crate::ServerConnectionContext<
1826 Peer,
1827 <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1828 >,
1829 >>::Handler,
1830 <CM as crate::MiddlewareFactory<crate::ClientInitialContext>>::Handler,
1831 <Boundary as IntermediaryMiddlewareFactory<
1832 crate::ServerConnectionContext<
1833 Peer,
1834 <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1835 >,
1836 crate::ClientConnectionContext<CA::Evidence>,
1837 >>::Handler,
1838 Policy,
1839 K,
1840 >,
1841 >,
1842 IntermediaryAcceptError<
1843 crate::AcceptError<
1844 <ST::Provider as crate::ServerIdentityProvider>::Error,
1845 <SA::Authentication as crate::ServerAuthentication<Peer>>::Error,
1846 >,
1847 Resolver::Error,
1848 Route::Error,
1849 crate::ConnectError<
1850 CE,
1851 crate::ClientTlsError<<CT::Provider as crate::ClientTlsProvider>::Error>,
1852 crate::ClientAuthenticationError<CA::Error>,
1853 >,
1854 K::Error,
1855 crate::CancelError<CE>,
1856 <<Boundary as IntermediaryMiddlewareFactory<
1857 crate::ServerConnectionContext<
1858 Peer,
1859 <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1860 >,
1861 crate::ClientConnectionContext<CA::Evidence>,
1862 >>::Handler as IntermediaryMiddleware<
1863 State,
1864 crate::ServerConnectionContext<
1865 Peer,
1866 <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1867 >,
1868 crate::ClientConnectionContext<CA::Evidence>,
1869 >>::Error,
1870 >,
1871 >
1872 where
1873 DT: AsyncRead + AsyncWrite + Unpin,
1874 UT: AsyncRead + AsyncWrite + Unpin,
1875 SA::Authentication: crate::ServerAuthentication<Peer>,
1876 SM: crate::MiddlewareFactory<
1877 crate::ServerConnectionContext<
1878 Peer,
1879 <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1880 >,
1881 >,
1882 <SM as crate::MiddlewareFactory<
1883 crate::ServerConnectionContext<
1884 Peer,
1885 <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1886 >,
1887 >>::Handler: crate::ServerMiddleware<
1888 State,
1889 crate::ServerConnectionContext<
1890 Peer,
1891 <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1892 >,
1893 >,
1894 Resolver: StartupRouteResolver<Peer>,
1895 Connector: Fn(&ConnectTarget) -> CW,
1896 CW: Future<Output = Result<UT, CE>>,
1897 <CM as crate::MiddlewareFactory<crate::ClientInitialContext>>::Handler:
1898 crate::ClientMiddleware<State, crate::ClientConnectionContext<CA::Evidence>>,
1899 Route: AuthenticatedRoutePolicy<
1900 Peer,
1901 <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1902 >,
1903 Boundary: IntermediaryMiddlewareFactory<
1904 crate::ServerConnectionContext<
1905 Peer,
1906 <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1907 >,
1908 crate::ClientConnectionContext<CA::Evidence>,
1909 >,
1910 <Boundary as IntermediaryMiddlewareFactory<
1911 crate::ServerConnectionContext<
1912 Peer,
1913 <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1914 >,
1915 crate::ClientConnectionContext<CA::Evidence>,
1916 >>::Handler: IntermediaryMiddleware<
1917 State,
1918 crate::ServerConnectionContext<
1919 Peer,
1920 <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1921 >,
1922 crate::ClientConnectionContext<CA::Evidence>,
1923 >,
1924 {
1925 let mut resolver = StartupResolverAdapter {
1926 resolver: &self.resolver,
1927 };
1928 let (accepted, selected) = self
1929 .server
1930 .accept_routed(transport, peer, state, &mut resolver)
1931 .await
1932 .map_err(|error| match error {
1933 crate::server_component::RoutedAcceptError::Accept(error) => {
1934 IntermediaryAcceptError::Server(error)
1935 }
1936 crate::server_component::RoutedAcceptError::Route(error) => {
1937 IntermediaryAcceptError::StartupRoute(error)
1938 }
1939 })?;
1940 let mut downstream = match accepted {
1941 crate::ServerAccept::Session(downstream) => downstream,
1942 crate::ServerAccept::Cancellation(cancellation) => {
1943 if self.cancellation == CancellationPolicy::Reject {
1944 let _ = cancellation.teardown();
1945 return Err(IntermediaryAcceptError::CancellationRejected);
1946 }
1947 let request = cancellation.request();
1948 let client_key = crate::demux::CancelKey {
1949 process_id: request.process_id(),
1950 secret_key: bytes::Bytes::copy_from_slice(request.secret_key()),
1951 };
1952 let Some(route) = self.cancellation_registry.resolve(&client_key) else {
1953 let _ = cancellation.teardown();
1954 return Err(IntermediaryAcceptError::CancellationRejected);
1955 };
1956 if let Err(error) = self
1957 .client
1958 .cancel(route.target(), route.upstream_key())
1959 .await
1960 {
1961 let _ = cancellation.teardown();
1962 return Err(IntermediaryAcceptError::Cancellation(error));
1963 }
1964 let _ = cancellation.teardown();
1965 return Ok(IntermediaryAccept::CancellationForwarded);
1966 }
1967 };
1968 let startup = match StartupParameters::from_wire(downstream.startup()) {
1969 Ok(startup) => startup,
1970 Err(error) => {
1971 if self.failure_policy == EstablishmentFailurePolicy::SafeDiagnostic {
1972 let _ = downstream
1973 .send_generated_error(safe_establishment_diagnostic())
1974 .await;
1975 }
1976 let _ = downstream.teardown();
1977 return Err(IntermediaryAcceptError::StartupRoute(
1978 StartupResolutionError::Parameters(error),
1979 ));
1980 }
1981 };
1982 let context = AuthenticatedRouteContext {
1983 peer: downstream.context().peer(),
1984 identity: downstream.context().identity(),
1985 };
1986 let Some(selected) = selected else {
1987 let _ = downstream.teardown();
1988 return Err(IntermediaryAcceptError::CancellationRejected);
1989 };
1990 let selected = match self.route.route(selected, context).await {
1991 Ok(target) => target,
1992 Err(error) => {
1993 if self.failure_policy == EstablishmentFailurePolicy::SafeDiagnostic {
1994 let _ = downstream
1995 .send_generated_error(safe_establishment_diagnostic())
1996 .await;
1997 }
1998 let _ = downstream.teardown();
1999 return Err(IntermediaryAcceptError::AuthenticatedRoute(error));
2000 }
2001 };
2002 let (mut downstream, mut state) = downstream.into_core_and_state();
2003 let upstream = match self
2004 .client
2005 .connect_core(selected.clone(), startup, &mut state)
2006 .await
2007 {
2008 Ok(upstream) => upstream,
2009 Err(error) => {
2010 if self.failure_policy == EstablishmentFailurePolicy::SafeDiagnostic {
2011 let diagnostic = safe_establishment_diagnostic();
2012 let diagnostic = downstream.intercept_backend(&mut state, diagnostic);
2013 if matches!(diagnostic, crate::codec::BackendMessage::ErrorResponse(_)) {
2014 let _ = downstream.send_wire_raw(diagnostic).await;
2017 }
2018 }
2019 let _ = downstream.into_parts();
2020 return Err(IntermediaryAcceptError::Client(error));
2021 }
2022 };
2023 let boundary = self
2024 .boundary
2025 .create(downstream.context(), upstream.context());
2026 let (client_cancel_key, backend_key_message) =
2027 match (self.cancellation, upstream.context().backend_key().cloned()) {
2028 (CancellationPolicy::Forward, Some(upstream_key)) => {
2029 let client_key = match self
2030 .cancellation_registry
2031 .register(CancellationRoute::new(selected.clone(), upstream_key))
2032 {
2033 Ok(key) => key,
2034 Err(error) => {
2035 if self.failure_policy == EstablishmentFailurePolicy::SafeDiagnostic {
2036 let diagnostic = downstream
2037 .intercept_backend(&mut state, safe_establishment_diagnostic());
2038 if matches!(
2039 diagnostic,
2040 crate::codec::BackendMessage::ErrorResponse(_)
2041 ) {
2042 let _ = downstream.send_wire_raw(diagnostic).await;
2043 }
2044 }
2045 let _ = downstream.into_parts();
2046 let _ = upstream.into_parts();
2047 return Err(IntermediaryAcceptError::CancellationRegistry(error));
2048 }
2049 };
2050 let message = crate::codec::BackendMessage::BackendKeyData {
2051 process_id: client_key.process_id,
2052 secret_key: client_key.secret_key.clone(),
2053 };
2054 (Some(client_key), Some(message))
2055 }
2056 _ => (None, None),
2057 };
2058 let mut connection = IntermediaryConnection {
2059 downstream,
2060 upstream,
2061 state,
2062 boundary,
2063 pipeline: Pipeline::new(self.pipeline),
2064 target: selected,
2065 pending_frontend: None,
2066 backend_hold: crate::backend_hold::BackendHold::default(),
2067 backend_hold_limits: self.backend_hold_limits,
2068 pending_local: VecDeque::new(),
2069 cancellation_registry: self.cancellation_registry.clone(),
2070 client_cancel_key,
2071 };
2072 if let Some(message) = backend_key_message {
2073 let expected = message.clone();
2074 let message = connection
2075 .boundary
2076 .backend(
2077 connection.downstream.context(),
2078 connection.upstream.context(),
2079 &mut connection.state,
2080 message,
2081 )
2082 .await;
2083 let message = match message {
2084 Ok(BackendMiddlewareOutput::Forward(message)) => message,
2085 Ok(
2086 BackendMiddlewareOutput::Suppress(_)
2087 | BackendMiddlewareOutput::Expand(_)
2088 | BackendMiddlewareOutput::Hold,
2089 ) => {
2090 let _ = connection.detach_cancellation();
2091 let _ = connection.teardown();
2092 return Err(IntermediaryAcceptError::ServerOutput(io::Error::new(
2093 io::ErrorKind::InvalidData,
2094 "middleware suppressed or expanded generated cancellation key",
2095 )));
2096 }
2097 Err(error) => {
2098 let _ = connection.detach_cancellation();
2099 let _ = connection.teardown();
2100 return Err(IntermediaryAcceptError::Middleware(error));
2101 }
2102 };
2103 let message = connection
2104 .downstream
2105 .intercept_backend(&mut connection.state, message);
2106 if message != expected {
2107 let _ = connection.detach_cancellation();
2108 let _ = connection.teardown();
2109 return Err(IntermediaryAcceptError::ServerOutput(io::Error::new(
2110 io::ErrorKind::InvalidData,
2111 "middleware rejected generated cancellation key",
2112 )));
2113 }
2114 if let Err(error) = connection.downstream.send_wire_raw(message).await {
2115 let _ = connection.detach_cancellation();
2116 let _ = connection.teardown();
2117 return Err(IntermediaryAcceptError::ServerOutput(error));
2118 }
2119 }
2120 let ready = connection.downstream.intercept_backend(
2121 &mut connection.state,
2122 crate::codec::BackendMessage::ReadyForQuery(crate::codec::TransactionStatus::Idle),
2123 );
2124 if !matches!(ready, crate::codec::BackendMessage::ReadyForQuery(_)) {
2125 let _ = connection.detach_cancellation();
2126 let _ = connection.teardown();
2127 return Err(IntermediaryAcceptError::ServerOutput(io::Error::new(
2128 io::ErrorKind::InvalidData,
2129 "middleware rejected generated readiness",
2130 )));
2131 }
2132 if let Err(error) = connection.downstream.send_wire_raw(ready).await {
2133 let _ = connection.detach_cancellation();
2134 let _ = connection.teardown();
2135 return Err(IntermediaryAcceptError::ServerOutput(error));
2136 }
2137 Ok(IntermediaryAccept::Session(connection))
2138 }
2139}