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