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 let FrontendAction::Forward { message, .. } = admission.into_action() else {
1415 unreachable!()
1416 };
1417 self.upstream.send_wire_raw(message.clone()).await?;
1418 Ok(FrontendForwarding::Forwarded(message))
1419 }
1420
1421 pub async fn forward_backend(
1434 &mut self,
1435 ) -> Result<BackendForwarding, ForwardError<Boundary::Error>> {
1436 if self.backend_hold.pending().is_some() {
1437 return self.process_pending_backend().await;
1438 }
1439 if self.backend_hold_is_full() {
1440 match self
1441 .flush_backend_hold_for(BackendFlushReason::Capacity)
1442 .await?
1443 {
1444 BackendBatchForwarding::Released(_) | BackendBatchForwarding::Empty => {}
1445 BackendBatchForwarding::Kept => return Err(ForwardError::BackendHoldCapacity),
1446 }
1447 }
1448 let message = self.upstream.receive_wire_raw().await?;
1449 self.process_backend(message).await
1450 }
1451
1452 async fn process_backend(
1453 &mut self,
1454 message: crate::codec::BackendMessage,
1455 ) -> Result<BackendForwarding, ForwardError<Boundary::Error>> {
1456 let message = self.upstream.intercept_backend(&mut self.state, message);
1457 self.backend_hold.set_pending(message);
1458 if self
1459 .backend_hold
1460 .pending()
1461 .is_some_and(is_backend_batch_barrier)
1462 && !self.backend_hold.is_empty()
1463 {
1464 match self
1465 .flush_backend_hold_for(BackendFlushReason::ProtocolBarrier)
1466 .await?
1467 {
1468 BackendBatchForwarding::Released(_) | BackendBatchForwarding::Empty => {}
1469 BackendBatchForwarding::Kept => return Err(ForwardError::BackendHoldRefused),
1470 }
1471 }
1472 self.process_pending_backend().await
1473 }
1474
1475 async fn process_pending_backend(
1476 &mut self,
1477 ) -> Result<BackendForwarding, ForwardError<Boundary::Error>> {
1478 let source = self
1479 .backend_hold
1480 .pending()
1481 .expect("pending backend processing requires a source")
1482 .clone();
1483 let operation = self.pipeline.backend_operation_id(&source);
1484 let decision = self
1485 .boundary
1486 .backend_operation(
1487 self.downstream.context(),
1488 self.upstream.context(),
1489 &mut self.state,
1490 operation,
1491 source.clone(),
1492 )
1493 .await
1494 .map_err(ForwardError::Middleware)?;
1495 let outcome = match decision {
1496 BackendMiddlewareOutput::Forward(message) => {
1497 let _ = self.backend_hold.take_pending();
1498 let message = self.downstream.intercept_backend(&mut self.state, message);
1499 let message = self.emit_backend(message).await?;
1500 BackendForwarding::Forwarded(message)
1501 }
1502 BackendMiddlewareOutput::Suppress(message) => {
1503 let _ = self.backend_hold.take_pending();
1504 let message = self.advance_backend(message)?;
1505 BackendForwarding::Suppressed(message)
1506 }
1507 BackendMiddlewareOutput::Expand(messages) => {
1508 if messages.is_empty() {
1509 return Err(ForwardError::EmptyExpansion(source));
1510 }
1511 let _ = self.backend_hold.take_pending();
1512 let mut emitted = Vec::with_capacity(messages.len());
1513 for message in messages {
1514 let message = self.downstream.intercept_backend(&mut self.state, message);
1515 emitted.push(self.emit_backend(message).await?);
1516 }
1517 BackendForwarding::Expanded {
1518 source,
1519 messages: emitted,
1520 }
1521 }
1522 BackendMiddlewareOutput::Hold => {
1523 if self.backend_hold_limits.is_none() {
1524 return Err(ForwardError::BackendHoldingDisabled(source));
1525 }
1526 self.backend_hold.hold_pending();
1527 BackendForwarding::Held
1528 }
1529 };
1530 self.flush_local_responses().await?;
1531 Ok(outcome)
1532 }
1533
1534 fn backend_hold_is_full(&self) -> bool {
1535 self.backend_hold_limits.is_some_and(|limits| {
1536 self.backend_hold.len() >= limits.max_messages()
1537 || self.backend_hold.bytes() >= limits.max_bytes()
1538 })
1539 }
1540
1541 pub async fn flush_backend_hold(
1547 &mut self,
1548 ) -> Result<BackendBatchForwarding, ForwardError<Boundary::Error>> {
1549 self.flush_backend_hold_for(BackendFlushReason::Explicit)
1550 .await
1551 }
1552
1553 pub async fn prepare_backend_teardown(
1559 &mut self,
1560 ) -> Result<BackendBatchForwarding, ForwardError<Boundary::Error>> {
1561 let outcome = self
1562 .flush_backend_hold_for(BackendFlushReason::Teardown)
1563 .await?;
1564 if matches!(outcome, BackendBatchForwarding::Kept) {
1565 return Err(ForwardError::BackendHoldRefused);
1566 }
1567 Ok(outcome)
1568 }
1569
1570 async fn flush_backend_hold_for(
1571 &mut self,
1572 reason: BackendFlushReason,
1573 ) -> Result<BackendBatchForwarding, ForwardError<Boundary::Error>> {
1574 if self.backend_hold.is_empty() {
1575 return Ok(BackendBatchForwarding::Empty);
1576 }
1577 let held = AttributedBackendMessages {
1578 held: HeldBackendMessages {
1579 messages: self.backend_hold.messages(),
1580 bytes: self.backend_hold.bytes(),
1581 },
1582 operation_ids: self
1583 .pipeline
1584 .backend_operation_ids(self.backend_hold.messages()),
1585 };
1586 let decision = self
1587 .boundary
1588 .flush_backend_operations(
1589 self.downstream.context(),
1590 self.upstream.context(),
1591 &mut self.state,
1592 held,
1593 reason,
1594 )
1595 .await
1596 .map_err(ForwardError::Middleware)?;
1597 let BackendBatchOutput::ReplaceOneToOne(messages) = decision else {
1598 return Ok(BackendBatchForwarding::Kept);
1599 };
1600 let messages: Vec<_> = messages
1601 .into_iter()
1602 .map(|message| self.downstream.intercept_backend(&mut self.state, message))
1603 .collect();
1604 for message in &messages {
1605 message.to_frame().map_err(ForwardError::Io)?;
1606 }
1607 let prepared = self
1608 .pipeline
1609 .prepare_backend_replacements(self.backend_hold.messages(), &messages)
1610 .map_err(|error| ForwardError::BackendBatch {
1611 error: error.into(),
1612 proposed: messages.clone(),
1613 })?;
1614 self.pipeline = prepared;
1615 let _sources = self.backend_hold.clear();
1616 for message in &messages {
1617 self.downstream.send_wire_raw(message.clone()).await?;
1618 }
1619 self.flush_local_responses().await?;
1620 Ok(BackendBatchForwarding::Released(messages))
1621 }
1622
1623 fn advance_backend(
1624 &mut self,
1625 message: crate::codec::BackendMessage,
1626 ) -> Result<crate::codec::BackendMessage, ForwardError<Boundary::Error>> {
1627 match self
1628 .pipeline
1629 .accept_backend(message)
1630 .map_err(ForwardError::Backend)?
1631 {
1632 BackendAction::Emit(message) => Ok(message),
1633 BackendAction::Deferred(message) => Err(ForwardError::Deferred(message)),
1634 }
1635 }
1636
1637 async fn emit_backend(
1638 &mut self,
1639 message: crate::codec::BackendMessage,
1640 ) -> Result<crate::codec::BackendMessage, ForwardError<Boundary::Error>> {
1641 let message = self.advance_backend(message)?;
1642 self.downstream.send_wire_raw(message.clone()).await?;
1643 Ok(message)
1644 }
1645
1646 async fn flush_local_responses(&mut self) -> Result<(), ForwardError<Boundary::Error>> {
1647 loop {
1648 let Some(pending) = self.pending_local.front_mut() else {
1649 return Ok(());
1650 };
1651 let Some(message) = pending.messages.pop_front() else {
1652 self.pending_local.pop_front();
1653 continue;
1654 };
1655 match self.pipeline.try_emit_local(pending.operation, message) {
1656 Ok(BackendAction::Emit(message)) => {
1657 self.downstream.send_wire_raw(message).await?;
1658 }
1659 Ok(BackendAction::Deferred(message)) => {
1660 pending.messages.push_front(message);
1661 return Ok(());
1662 }
1663 Err(error) => return Err(ForwardError::Backend(error)),
1664 }
1665 }
1666 }
1667
1668 pub async fn forward_next(
1679 &mut self,
1680 ) -> Result<ForwardedMessage, ForwardError<Boundary::Error>> {
1681 if self.backend_hold.pending().is_some() {
1682 return self
1683 .process_pending_backend()
1684 .await
1685 .map(|outcome| match outcome {
1686 BackendForwarding::Forwarded(message) => ForwardedMessage::Backend(message),
1687 BackendForwarding::Expanded { source, messages } => {
1688 ForwardedMessage::BackendExpanded { source, messages }
1689 }
1690 BackendForwarding::Suppressed(message) => {
1691 ForwardedMessage::BackendSuppressed(message)
1692 }
1693 BackendForwarding::Held => ForwardedMessage::BackendHeld,
1694 });
1695 }
1696 if self.backend_hold_is_full() {
1697 match self
1698 .flush_backend_hold_for(BackendFlushReason::Capacity)
1699 .await?
1700 {
1701 BackendBatchForwarding::Released(_) | BackendBatchForwarding::Empty => {}
1702 BackendBatchForwarding::Kept => return Err(ForwardError::BackendHoldCapacity),
1703 }
1704 }
1705 if self.pending_frontend.is_some() {
1706 let message = self.upstream.receive_wire_raw().await?;
1707 return self
1708 .process_backend(message)
1709 .await
1710 .map(|outcome| match outcome {
1711 BackendForwarding::Forwarded(message) => ForwardedMessage::Backend(message),
1712 BackendForwarding::Expanded { source, messages } => {
1713 ForwardedMessage::BackendExpanded { source, messages }
1714 }
1715 BackendForwarding::Suppressed(message) => {
1716 ForwardedMessage::BackendSuppressed(message)
1717 }
1718 BackendForwarding::Held => ForwardedMessage::BackendHeld,
1719 });
1720 }
1721 tokio::select! {
1722 result = self.downstream.receive_wire_raw() => {
1723 let message = result?;
1724 self.process_frontend(message, true).await.map(|outcome| match outcome {
1725 FrontendForwarding::Forwarded(message) => ForwardedMessage::Frontend(message),
1726 FrontendForwarding::Suppressed(message) => ForwardedMessage::FrontendSuppressed(message),
1727 FrontendForwarding::LocallyHandled(message) => ForwardedMessage::FrontendLocallyHandled(message),
1728 })
1729 }
1730 result = self.upstream.receive_wire_raw() => {
1731 let message = result?;
1732 self.process_backend(message).await.map(|outcome| match outcome {
1733 BackendForwarding::Forwarded(message) => ForwardedMessage::Backend(message),
1734 BackendForwarding::Expanded { source, messages } => {
1735 ForwardedMessage::BackendExpanded { source, messages }
1736 }
1737 BackendForwarding::Suppressed(message) => ForwardedMessage::BackendSuppressed(message),
1738 BackendForwarding::Held => ForwardedMessage::BackendHeld,
1739 })
1740 }
1741 }
1742 }
1743
1744 #[allow(clippy::type_complexity)]
1751 pub fn teardown(
1752 mut self,
1753 ) -> (
1754 crate::AcceptedServerTransport<DT>,
1755 crate::ClientTransport<UT>,
1756 State,
1757 Boundary,
1758 (ServerHandler, ClientHandler),
1759 IntermediaryContexts<
1760 crate::ServerConnectionContext<Peer, ServerIdentity>,
1761 crate::ClientConnectionContext<ClientEvidence>,
1762 >,
1763 ) {
1764 assert!(
1765 self.backend_hold.is_empty() && self.backend_hold.pending().is_none(),
1766 "backend messages remain held; call prepare_backend_teardown before teardown"
1767 );
1768 let _ = self.detach_cancellation();
1769 let (downstream, server_handler, server_context) = self.downstream.into_parts();
1770 let (upstream, client_handler, client_context) = self.upstream.into_parts();
1771 (
1772 downstream,
1773 upstream,
1774 self.state,
1775 self.boundary,
1776 (server_handler, client_handler),
1777 IntermediaryContexts {
1778 server: server_context,
1779 client: client_context,
1780 },
1781 )
1782 }
1783}
1784
1785#[derive(Debug)]
1787pub enum ForwardError<MiddlewareError = std::convert::Infallible> {
1788 Io(io::Error),
1790 Frontend(crate::pipeline::FrontendProjectionError),
1792 Backend(crate::pipeline::BackendProjectionError),
1794 Deferred(crate::codec::BackendMessage),
1796 EmptyExpansion(crate::codec::BackendMessage),
1798 BackendHoldingDisabled(crate::codec::BackendMessage),
1800 BackendHoldCapacity,
1802 BackendHoldRefused,
1804 BackendBatch {
1806 error: BackendBatchProjectionError,
1808 proposed: Vec<crate::codec::BackendMessage>,
1810 },
1811 Middleware(MiddlewareError),
1813}
1814
1815impl<E> fmt::Display for ForwardError<E> {
1816 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
1817 match self {
1818 Self::Io(error) => error.fmt(formatter),
1819 Self::Frontend(_) => {
1820 formatter.write_str("frontend message violates pipeline legality or capacity")
1821 }
1822 Self::Backend(_) => formatter.write_str("backend message violates pipeline legality"),
1823 Self::Deferred(_) => formatter.write_str("backend response is not yet emittable"),
1824 Self::EmptyExpansion(_) => {
1825 formatter.write_str("backend expansion must contain at least one response")
1826 }
1827 Self::BackendHoldingDisabled(_) => {
1828 formatter.write_str("backend holding is not configured")
1829 }
1830 Self::BackendHoldCapacity => {
1831 formatter.write_str("backend hold is full and batch policy kept holding")
1832 }
1833 Self::BackendHoldRefused => {
1834 formatter.write_str("backend batch policy refused a required release")
1835 }
1836 Self::BackendBatch { .. } => formatter
1837 .write_str("backend batch replacement is not a legal one-to-one protocol span"),
1838 Self::Middleware(_) => formatter.write_str("forwarding middleware rejected a message"),
1839 }
1840 }
1841}
1842
1843impl<E> std::error::Error for ForwardError<E>
1844where
1845 E: std::error::Error + 'static,
1846{
1847 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
1848 match self {
1849 Self::Io(error) => Some(error),
1850 Self::Middleware(error) => Some(error),
1851 Self::Frontend(_)
1852 | Self::Backend(_)
1853 | Self::Deferred(_)
1854 | Self::EmptyExpansion(_)
1855 | Self::BackendHoldingDisabled(_)
1856 | Self::BackendHoldCapacity
1857 | Self::BackendHoldRefused
1858 | Self::BackendBatch { .. } => None,
1859 }
1860 }
1861}
1862
1863impl<E> From<io::Error> for ForwardError<E> {
1864 fn from(error: io::Error) -> Self {
1865 Self::Io(error)
1866 }
1867}
1868
1869impl<ST, SA, SM, Connector, CT, CA, CM, Resolver, Route, Policy, Boundary, K>
1870 Intermediary<
1871 crate::Server<ST, SA, SM>,
1872 crate::Client<Connector, CT, CA, CM>,
1873 Resolver,
1874 Route,
1875 Policy,
1876 Boundary,
1877 K,
1878 >
1879where
1880 ST: crate::ServerTlsConfiguration,
1881 SA: crate::ServerAuthenticationProvider,
1882 CT: crate::client_component::ClientTlsConfiguration,
1883 CA: crate::ClientAuthentication,
1884 CM: crate::MiddlewareFactory<crate::ClientInitialContext>,
1885 Policy: PipelinePolicy,
1886 K: IntermediaryCancellationRegistry + Clone,
1887{
1888 #[allow(clippy::type_complexity, clippy::too_many_lines)]
1895 pub async fn accept<DT, State, Peer, CW, UT, CE>(
1896 &self,
1897 transport: DT,
1898 peer: Peer,
1899 state: State,
1900 ) -> Result<
1901 IntermediaryAccept<
1902 IntermediaryConnection<
1903 DT,
1904 UT,
1905 State,
1906 Peer,
1907 <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1908 CA::Evidence,
1909 <SM as crate::MiddlewareFactory<
1910 crate::ServerConnectionContext<
1911 Peer,
1912 <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1913 >,
1914 >>::Handler,
1915 <CM as crate::MiddlewareFactory<crate::ClientInitialContext>>::Handler,
1916 <Boundary as IntermediaryMiddlewareFactory<
1917 crate::ServerConnectionContext<
1918 Peer,
1919 <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1920 >,
1921 crate::ClientConnectionContext<CA::Evidence>,
1922 >>::Handler,
1923 Policy,
1924 K,
1925 >,
1926 >,
1927 IntermediaryAcceptError<
1928 crate::AcceptError<
1929 <ST::Provider as crate::ServerIdentityProvider>::Error,
1930 <SA::Authentication as crate::ServerAuthentication<Peer>>::Error,
1931 >,
1932 Resolver::Error,
1933 Route::Error,
1934 crate::ConnectError<
1935 CE,
1936 crate::ClientTlsError<<CT::Provider as crate::ClientTlsProvider>::Error>,
1937 crate::ClientAuthenticationError<CA::Error>,
1938 >,
1939 K::Error,
1940 crate::CancelError<CE>,
1941 <<Boundary as IntermediaryMiddlewareFactory<
1942 crate::ServerConnectionContext<
1943 Peer,
1944 <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1945 >,
1946 crate::ClientConnectionContext<CA::Evidence>,
1947 >>::Handler as IntermediaryMiddleware<
1948 State,
1949 crate::ServerConnectionContext<
1950 Peer,
1951 <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1952 >,
1953 crate::ClientConnectionContext<CA::Evidence>,
1954 >>::Error,
1955 >,
1956 >
1957 where
1958 DT: AsyncRead + AsyncWrite + Unpin,
1959 UT: AsyncRead + AsyncWrite + Unpin,
1960 SA::Authentication: crate::ServerAuthentication<Peer>,
1961 SM: crate::MiddlewareFactory<
1962 crate::ServerConnectionContext<
1963 Peer,
1964 <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1965 >,
1966 >,
1967 <SM as crate::MiddlewareFactory<
1968 crate::ServerConnectionContext<
1969 Peer,
1970 <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1971 >,
1972 >>::Handler: crate::ServerMiddleware<
1973 State,
1974 crate::ServerConnectionContext<
1975 Peer,
1976 <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1977 >,
1978 >,
1979 Resolver: StartupRouteResolver<Peer>,
1980 Connector: Fn(&ConnectTarget) -> CW,
1981 CW: Future<Output = Result<UT, CE>>,
1982 <CM as crate::MiddlewareFactory<crate::ClientInitialContext>>::Handler:
1983 crate::ClientMiddleware<State, crate::ClientConnectionContext<CA::Evidence>>,
1984 Route: AuthenticatedRoutePolicy<
1985 Peer,
1986 <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1987 >,
1988 Boundary: IntermediaryMiddlewareFactory<
1989 crate::ServerConnectionContext<
1990 Peer,
1991 <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1992 >,
1993 crate::ClientConnectionContext<CA::Evidence>,
1994 >,
1995 <Boundary as IntermediaryMiddlewareFactory<
1996 crate::ServerConnectionContext<
1997 Peer,
1998 <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1999 >,
2000 crate::ClientConnectionContext<CA::Evidence>,
2001 >>::Handler: IntermediaryMiddleware<
2002 State,
2003 crate::ServerConnectionContext<
2004 Peer,
2005 <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
2006 >,
2007 crate::ClientConnectionContext<CA::Evidence>,
2008 >,
2009 {
2010 let mut resolver = StartupResolverAdapter {
2011 resolver: &self.resolver,
2012 };
2013 let (accepted, selected) = self
2014 .server
2015 .accept_routed(transport, peer, state, &mut resolver)
2016 .await
2017 .map_err(|error| match error {
2018 crate::server_component::RoutedAcceptError::Accept(error) => {
2019 IntermediaryAcceptError::Server(error)
2020 }
2021 crate::server_component::RoutedAcceptError::Route(error) => {
2022 IntermediaryAcceptError::StartupRoute(error)
2023 }
2024 })?;
2025 let mut downstream = match accepted {
2026 crate::ServerAccept::Session(downstream) => downstream,
2027 crate::ServerAccept::Cancellation(cancellation) => {
2028 if self.cancellation == CancellationPolicy::Reject {
2029 let _ = cancellation.teardown();
2030 return Err(IntermediaryAcceptError::CancellationRejected);
2031 }
2032 let request = cancellation.request();
2033 let client_key = crate::demux::CancelKey {
2034 process_id: request.process_id(),
2035 secret_key: bytes::Bytes::copy_from_slice(request.secret_key()),
2036 };
2037 let Some(route) = self.cancellation_registry.resolve(&client_key) else {
2038 let _ = cancellation.teardown();
2039 return Err(IntermediaryAcceptError::CancellationRejected);
2040 };
2041 if let Err(error) = self
2042 .client
2043 .cancel(route.target(), route.upstream_key())
2044 .await
2045 {
2046 let _ = cancellation.teardown();
2047 return Err(IntermediaryAcceptError::Cancellation(error));
2048 }
2049 let _ = cancellation.teardown();
2050 return Ok(IntermediaryAccept::CancellationForwarded);
2051 }
2052 };
2053 let startup = match StartupParameters::from_wire(downstream.startup()) {
2054 Ok(startup) => startup,
2055 Err(error) => {
2056 if self.failure_policy == EstablishmentFailurePolicy::SafeDiagnostic {
2057 let _ = downstream
2058 .send_generated_error(safe_establishment_diagnostic())
2059 .await;
2060 }
2061 let _ = downstream.teardown();
2062 return Err(IntermediaryAcceptError::StartupRoute(
2063 StartupResolutionError::Parameters(error),
2064 ));
2065 }
2066 };
2067 let context = AuthenticatedRouteContext {
2068 peer: downstream.context().peer(),
2069 identity: downstream.context().identity(),
2070 };
2071 let Some(selected) = selected else {
2072 let _ = downstream.teardown();
2073 return Err(IntermediaryAcceptError::CancellationRejected);
2074 };
2075 let selected = match self.route.route(selected, context).await {
2076 Ok(target) => target,
2077 Err(error) => {
2078 if self.failure_policy == EstablishmentFailurePolicy::SafeDiagnostic {
2079 let _ = downstream
2080 .send_generated_error(safe_establishment_diagnostic())
2081 .await;
2082 }
2083 let _ = downstream.teardown();
2084 return Err(IntermediaryAcceptError::AuthenticatedRoute(error));
2085 }
2086 };
2087 let (mut downstream, mut state) = downstream.into_core_and_state();
2088 let upstream = match self
2089 .client
2090 .connect_core(selected.clone(), startup, &mut state)
2091 .await
2092 {
2093 Ok(upstream) => upstream,
2094 Err(error) => {
2095 if self.failure_policy == EstablishmentFailurePolicy::SafeDiagnostic {
2096 let diagnostic = safe_establishment_diagnostic();
2097 let diagnostic = downstream.intercept_backend(&mut state, diagnostic);
2098 if matches!(diagnostic, crate::codec::BackendMessage::ErrorResponse(_)) {
2099 let _ = downstream.send_wire_raw(diagnostic).await;
2102 }
2103 }
2104 let _ = downstream.into_parts();
2105 return Err(IntermediaryAcceptError::Client(error));
2106 }
2107 };
2108 let boundary = self
2109 .boundary
2110 .create(downstream.context(), upstream.context());
2111 let (client_cancel_key, backend_key_message) =
2112 match (self.cancellation, upstream.context().backend_key().cloned()) {
2113 (CancellationPolicy::Forward, Some(upstream_key)) => {
2114 let client_key = match self
2115 .cancellation_registry
2116 .register(CancellationRoute::new(selected.clone(), upstream_key))
2117 {
2118 Ok(key) => key,
2119 Err(error) => {
2120 if self.failure_policy == EstablishmentFailurePolicy::SafeDiagnostic {
2121 let diagnostic = downstream
2122 .intercept_backend(&mut state, safe_establishment_diagnostic());
2123 if matches!(
2124 diagnostic,
2125 crate::codec::BackendMessage::ErrorResponse(_)
2126 ) {
2127 let _ = downstream.send_wire_raw(diagnostic).await;
2128 }
2129 }
2130 let _ = downstream.into_parts();
2131 let _ = upstream.into_parts();
2132 return Err(IntermediaryAcceptError::CancellationRegistry(error));
2133 }
2134 };
2135 let message = crate::codec::BackendMessage::BackendKeyData {
2136 process_id: client_key.process_id,
2137 secret_key: client_key.secret_key.clone(),
2138 };
2139 (Some(client_key), Some(message))
2140 }
2141 _ => (None, None),
2142 };
2143 let mut connection = IntermediaryConnection {
2144 downstream,
2145 upstream,
2146 state,
2147 boundary,
2148 pipeline: Pipeline::new(self.pipeline),
2149 target: selected,
2150 pending_frontend: None,
2151 backend_hold: crate::backend_hold::BackendHold::default(),
2152 backend_hold_limits: self.backend_hold_limits,
2153 pending_local: VecDeque::new(),
2154 cancellation_registry: self.cancellation_registry.clone(),
2155 client_cancel_key,
2156 };
2157 if let Some(message) = backend_key_message {
2158 let expected = message.clone();
2159 let message = connection
2160 .boundary
2161 .backend(
2162 connection.downstream.context(),
2163 connection.upstream.context(),
2164 &mut connection.state,
2165 message,
2166 )
2167 .await;
2168 let message = match message {
2169 Ok(BackendMiddlewareOutput::Forward(message)) => message,
2170 Ok(
2171 BackendMiddlewareOutput::Suppress(_)
2172 | BackendMiddlewareOutput::Expand(_)
2173 | BackendMiddlewareOutput::Hold,
2174 ) => {
2175 let _ = connection.detach_cancellation();
2176 let _ = connection.teardown();
2177 return Err(IntermediaryAcceptError::ServerOutput(io::Error::new(
2178 io::ErrorKind::InvalidData,
2179 "middleware suppressed or expanded generated cancellation key",
2180 )));
2181 }
2182 Err(error) => {
2183 let _ = connection.detach_cancellation();
2184 let _ = connection.teardown();
2185 return Err(IntermediaryAcceptError::Middleware(error));
2186 }
2187 };
2188 let message = connection
2189 .downstream
2190 .intercept_backend(&mut connection.state, message);
2191 if message != expected {
2192 let _ = connection.detach_cancellation();
2193 let _ = connection.teardown();
2194 return Err(IntermediaryAcceptError::ServerOutput(io::Error::new(
2195 io::ErrorKind::InvalidData,
2196 "middleware rejected generated cancellation key",
2197 )));
2198 }
2199 if let Err(error) = connection.downstream.send_wire_raw(message).await {
2200 let _ = connection.detach_cancellation();
2201 let _ = connection.teardown();
2202 return Err(IntermediaryAcceptError::ServerOutput(error));
2203 }
2204 }
2205 let ready = connection.downstream.intercept_backend(
2206 &mut connection.state,
2207 crate::codec::BackendMessage::ReadyForQuery(crate::codec::TransactionStatus::Idle),
2208 );
2209 if !matches!(ready, crate::codec::BackendMessage::ReadyForQuery(_)) {
2210 let _ = connection.detach_cancellation();
2211 let _ = connection.teardown();
2212 return Err(IntermediaryAcceptError::ServerOutput(io::Error::new(
2213 io::ErrorKind::InvalidData,
2214 "middleware rejected generated readiness",
2215 )));
2216 }
2217 if let Err(error) = connection.downstream.send_wire_raw(ready).await {
2218 let _ = connection.detach_cancellation();
2219 let _ = connection.teardown();
2220 return Err(IntermediaryAcceptError::ServerOutput(error));
2221 }
2222 Ok(IntermediaryAccept::Session(connection))
2223 }
2224}