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