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