1use std::{fmt, future::Future, io, pin::Pin};
4
5use tokio::io::{AsyncRead, AsyncWrite};
6
7use crate::{
8 ConnectTarget, NoPipeline, StartupParameters,
9 pipeline::{BackendAction, FrontendAction, FrontendHandling, Pipeline, PipelinePolicy},
10};
11
12#[derive(Clone, Copy, Debug, Eq, PartialEq)]
17pub enum CancellationPolicy {
18 Reject,
20 Forward,
22}
23
24#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
26pub enum EstablishmentFailurePolicy {
27 #[default]
29 Close,
30 SafeDiagnostic,
32}
33
34fn safe_establishment_diagnostic() -> crate::codec::BackendMessage {
35 crate::codec::BackendMessage::ErrorResponse(crate::codec::DiagnosticResponse {
36 fields: vec![
37 crate::codec::DiagnosticField {
38 code: b'S',
39 value: bytes::Bytes::from_static(b"ERROR"),
40 },
41 crate::codec::DiagnosticField {
42 code: b'M',
43 value: bytes::Bytes::from_static(b"connection establishment failed"),
44 },
45 ],
46 })
47}
48
49#[derive(Clone, Debug, Eq, PartialEq)]
51pub struct CancellationRoute {
52 target: ConnectTarget,
53 upstream: crate::demux::CancelKey,
54}
55
56impl CancellationRoute {
57 #[must_use]
59 pub const fn new(target: ConnectTarget, upstream: crate::demux::CancelKey) -> Self {
60 Self { target, upstream }
61 }
62 #[must_use]
64 pub const fn target(&self) -> &ConnectTarget {
65 &self.target
66 }
67 #[must_use]
69 pub const fn upstream_key(&self) -> &crate::demux::CancelKey {
70 &self.upstream
71 }
72}
73
74pub trait IntermediaryCancellationRegistry {
80 type Error;
82 fn register(&self, route: CancellationRoute) -> Result<crate::demux::CancelKey, Self::Error>;
88 fn resolve(&self, client: &crate::demux::CancelKey) -> Option<CancellationRoute>;
90 fn detach(&self, client: &crate::demux::CancelKey) -> Option<CancellationRoute>;
92}
93
94#[derive(Clone, Copy, Debug, Default)]
96pub struct RejectCancellation;
97impl IntermediaryCancellationRegistry for RejectCancellation {
98 type Error = std::convert::Infallible;
99 fn register(&self, _: CancellationRoute) -> Result<crate::demux::CancelKey, Self::Error> {
100 unreachable!()
101 }
102 fn resolve(&self, _: &crate::demux::CancelKey) -> Option<CancellationRoute> {
103 None
104 }
105 fn detach(&self, _: &crate::demux::CancelKey) -> Option<CancellationRoute> {
106 None
107 }
108}
109
110#[derive(Clone, Copy, Debug, Eq, PartialEq)]
112pub enum IntermediaryBuildError {
113 MissingServer,
115 MissingClient,
117 MissingStartupResolver,
119 MissingCancellationPolicy,
121}
122
123impl fmt::Display for IntermediaryBuildError {
124 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
125 formatter.write_str(match self {
126 Self::MissingServer => "an intermediary server component is required",
127 Self::MissingClient => "an intermediary client component is required",
128 Self::MissingStartupResolver => "an asynchronous startup resolver is required",
129 Self::MissingCancellationPolicy => "an explicit cancellation policy is required",
130 })
131 }
132}
133
134impl std::error::Error for IntermediaryBuildError {}
135
136#[derive(Clone, Copy, Debug)]
138pub struct InitialServerContext<'a, Peer> {
139 peer: &'a Peer,
140 tls: &'a crate::NegotiatedServerTls,
141}
142
143impl<'a, Peer> InitialServerContext<'a, Peer> {
144 pub(crate) const fn new(peer: &'a Peer, tls: &'a crate::NegotiatedServerTls) -> Self {
145 Self { peer, tls }
146 }
147
148 #[must_use]
150 pub const fn peer(&self) -> &Peer {
151 self.peer
152 }
153
154 #[must_use]
156 pub const fn tls(&self) -> &crate::NegotiatedServerTls {
157 self.tls
158 }
159}
160
161pub trait StartupRouteResolver<Peer> {
163 type Error;
165
166 fn resolve<'a>(
168 &'a self,
169 startup: StartupParameters,
170 context: InitialServerContext<'a, Peer>,
171 ) -> Pin<Box<dyn Future<Output = Result<ConnectTarget, Self::Error>> + 'a>>;
172}
173
174pub trait AuthenticatedRoutePolicy<Peer, Identity> {
176 type Error;
178 fn route<'a>(
180 &'a self,
181 target: ConnectTarget,
182 context: AuthenticatedRouteContext<'a, Peer, Identity>,
183 ) -> Pin<Box<dyn Future<Output = Result<ConnectTarget, Self::Error>> + 'a>>;
184}
185
186#[derive(Clone, Copy, Debug)]
188pub struct AuthenticatedRouteContext<'a, Peer, Identity> {
189 peer: &'a Peer,
190 identity: &'a Identity,
191}
192
193impl<Peer, Identity> AuthenticatedRouteContext<'_, Peer, Identity> {
194 #[must_use]
196 pub const fn peer(&self) -> &Peer {
197 self.peer
198 }
199
200 #[must_use]
202 pub const fn identity(&self) -> &Identity {
203 self.identity
204 }
205}
206
207#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
209pub struct AllowAuthenticatedRoute;
210
211impl<Peer, Identity> AuthenticatedRoutePolicy<Peer, Identity> for AllowAuthenticatedRoute {
212 type Error = std::convert::Infallible;
213 fn route<'a>(
214 &'a self,
215 target: ConnectTarget,
216 _context: AuthenticatedRouteContext<'a, Peer, Identity>,
217 ) -> Pin<Box<dyn Future<Output = Result<ConnectTarget, Self::Error>> + 'a>> {
218 Box::pin(async move { Ok(target) })
219 }
220}
221
222pub trait IntermediaryMiddleware<State, ServerContext, ClientContext> {
224 fn frontend(
227 &mut self,
228 _server: &ServerContext,
229 _client: &ClientContext,
230 _state: &mut State,
231 message: crate::codec::FrontendMessage,
232 ) -> crate::codec::FrontendMessage {
233 message
234 }
235
236 fn backend(
239 &mut self,
240 _server: &ServerContext,
241 _client: &ClientContext,
242 _state: &mut State,
243 message: crate::codec::BackendMessage,
244 ) -> crate::codec::BackendMessage {
245 message
246 }
247}
248
249#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
251pub struct IdentityIntermediaryMiddleware;
252
253impl<State, ServerContext, ClientContext>
254 IntermediaryMiddleware<State, ServerContext, ClientContext> for IdentityIntermediaryMiddleware
255{
256}
257
258pub trait IntermediaryMiddlewareFactory<ServerContext, ClientContext> {
260 type Handler;
262 fn create(&self, server: &ServerContext, client: &ClientContext) -> Self::Handler;
264}
265
266impl<ServerContext, ClientContext, Handler, Factory>
267 IntermediaryMiddlewareFactory<ServerContext, ClientContext> for Factory
268where
269 Factory: Fn(&ServerContext, &ClientContext) -> Handler,
270{
271 type Handler = Handler;
272 fn create(&self, server: &ServerContext, client: &ClientContext) -> Handler {
273 self(server, client)
274 }
275}
276
277impl<ServerContext, ClientContext> IntermediaryMiddlewareFactory<ServerContext, ClientContext>
278 for IdentityIntermediaryMiddleware
279{
280 type Handler = Self;
281 fn create(&self, _server: &ServerContext, _client: &ClientContext) -> Self {
282 *self
283 }
284}
285
286pub struct Intermediary<
288 Server = (),
289 Client = (),
290 Resolver = (),
291 Route = AllowAuthenticatedRoute,
292 Policy = NoPipeline,
293 Boundary = IdentityIntermediaryMiddleware,
294 Cancellation = RejectCancellation,
295> {
296 pub(crate) server: Server,
297 pub(crate) client: Client,
298 pub(crate) resolver: Resolver,
299 pub(crate) route: Route,
300 pub(crate) pipeline: Policy,
301 pub(crate) boundary: Boundary,
302 pub(crate) cancellation: CancellationPolicy,
303 pub(crate) cancellation_registry: Cancellation,
304 pub(crate) failure_policy: EstablishmentFailurePolicy,
305}
306
307impl Intermediary<()> {
308 #[must_use]
310 pub fn builder() -> IntermediaryBuilder {
311 IntermediaryBuilder::default()
312 }
313}
314
315impl<S, C, R, A, P, B, K> fmt::Debug for Intermediary<S, C, R, A, P, B, K> {
316 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
317 formatter
318 .debug_struct("Intermediary")
319 .field("server", &"<configured>")
320 .field("client", &"<configured>")
321 .field("resolver", &"<redacted>")
322 .field("authenticated_route", &"<redacted>")
323 .field("cancellation", &self.cancellation)
324 .finish_non_exhaustive()
325 }
326}
327
328pub struct IntermediaryBuilder<
330 Server = (),
331 Client = (),
332 Resolver = (),
333 Route = AllowAuthenticatedRoute,
334 Policy = NoPipeline,
335 Boundary = IdentityIntermediaryMiddleware,
336 Cancellation = RejectCancellation,
337> {
338 server: Option<Server>,
339 client: Option<Client>,
340 resolver: Option<Resolver>,
341 route: Route,
342 pipeline: Policy,
343 boundary: Boundary,
344 cancellation: Option<CancellationPolicy>,
345 cancellation_registry: Cancellation,
346 failure_policy: EstablishmentFailurePolicy,
347}
348
349impl Default for IntermediaryBuilder {
350 fn default() -> Self {
351 Self {
352 server: None,
353 client: None,
354 resolver: None,
355 route: AllowAuthenticatedRoute,
356 pipeline: NoPipeline,
357 boundary: IdentityIntermediaryMiddleware,
358 cancellation: None,
359 cancellation_registry: RejectCancellation,
360 failure_policy: EstablishmentFailurePolicy::Close,
361 }
362 }
363}
364
365impl<S, C, R, A, P, B, K> IntermediaryBuilder<S, C, R, A, P, B, K> {
366 #[must_use]
368 pub fn server<Next>(self, server: Next) -> IntermediaryBuilder<Next, C, R, A, P, B, K> {
369 IntermediaryBuilder {
370 server: Some(server),
371 client: self.client,
372 resolver: self.resolver,
373 route: self.route,
374 pipeline: self.pipeline,
375 boundary: self.boundary,
376 cancellation: self.cancellation,
377 cancellation_registry: self.cancellation_registry,
378 failure_policy: self.failure_policy,
379 }
380 }
381
382 #[must_use]
384 pub fn client<Next>(self, client: Next) -> IntermediaryBuilder<S, Next, R, A, P, B, K> {
385 IntermediaryBuilder {
386 server: self.server,
387 client: Some(client),
388 resolver: self.resolver,
389 route: self.route,
390 pipeline: self.pipeline,
391 boundary: self.boundary,
392 cancellation: self.cancellation,
393 cancellation_registry: self.cancellation_registry,
394 failure_policy: self.failure_policy,
395 }
396 }
397
398 #[must_use]
400 pub fn startup_resolver<Next>(
401 self,
402 resolver: Next,
403 ) -> IntermediaryBuilder<S, C, Next, A, P, B, K> {
404 IntermediaryBuilder {
405 server: self.server,
406 client: self.client,
407 resolver: Some(resolver),
408 route: self.route,
409 pipeline: self.pipeline,
410 boundary: self.boundary,
411 cancellation: self.cancellation,
412 cancellation_registry: self.cancellation_registry,
413 failure_policy: self.failure_policy,
414 }
415 }
416
417 #[must_use]
419 pub fn authenticated_route<Next>(
420 self,
421 route: Next,
422 ) -> IntermediaryBuilder<S, C, R, Next, P, B, K> {
423 IntermediaryBuilder {
424 server: self.server,
425 client: self.client,
426 resolver: self.resolver,
427 route,
428 pipeline: self.pipeline,
429 boundary: self.boundary,
430 cancellation: self.cancellation,
431 cancellation_registry: self.cancellation_registry,
432 failure_policy: self.failure_policy,
433 }
434 }
435
436 #[must_use]
438 pub fn pipeline<Next: PipelinePolicy>(
439 self,
440 pipeline: Next,
441 ) -> IntermediaryBuilder<S, C, R, A, Next, B, K> {
442 IntermediaryBuilder {
443 server: self.server,
444 client: self.client,
445 resolver: self.resolver,
446 route: self.route,
447 pipeline,
448 boundary: self.boundary,
449 cancellation: self.cancellation,
450 cancellation_registry: self.cancellation_registry,
451 failure_policy: self.failure_policy,
452 }
453 }
454
455 #[must_use]
457 pub fn middleware<Next>(self, boundary: Next) -> IntermediaryBuilder<S, C, R, A, P, Next, K> {
458 IntermediaryBuilder {
459 server: self.server,
460 client: self.client,
461 resolver: self.resolver,
462 route: self.route,
463 pipeline: self.pipeline,
464 boundary,
465 cancellation: self.cancellation,
466 cancellation_registry: self.cancellation_registry,
467 failure_policy: self.failure_policy,
468 }
469 }
470
471 #[must_use]
473 pub fn cancellation(mut self, cancellation: CancellationPolicy) -> Self {
474 self.cancellation = match cancellation {
475 CancellationPolicy::Reject => Some(CancellationPolicy::Reject),
476 CancellationPolicy::Forward => None,
477 };
478 self
479 }
480
481 #[must_use]
483 pub fn establishment_failure(mut self, policy: EstablishmentFailurePolicy) -> Self {
484 self.failure_policy = policy;
485 self
486 }
487
488 #[must_use]
490 pub fn cancellation_registry<Next>(
491 self,
492 registry: Next,
493 ) -> IntermediaryBuilder<S, C, R, A, P, B, Next> {
494 IntermediaryBuilder {
495 server: self.server,
496 client: self.client,
497 resolver: self.resolver,
498 route: self.route,
499 pipeline: self.pipeline,
500 boundary: self.boundary,
501 cancellation: Some(CancellationPolicy::Forward),
502 cancellation_registry: registry,
503 failure_policy: self.failure_policy,
504 }
505 }
506
507 #[allow(clippy::type_complexity)]
513 pub fn build(self) -> Result<Intermediary<S, C, R, A, P, B, K>, IntermediaryBuildError> {
514 Ok(Intermediary {
515 server: self.server.ok_or(IntermediaryBuildError::MissingServer)?,
516 client: self.client.ok_or(IntermediaryBuildError::MissingClient)?,
517 resolver: self
518 .resolver
519 .ok_or(IntermediaryBuildError::MissingStartupResolver)?,
520 route: self.route,
521 pipeline: self.pipeline,
522 boundary: self.boundary,
523 cancellation: self
524 .cancellation
525 .ok_or(IntermediaryBuildError::MissingCancellationPolicy)?,
526 cancellation_registry: self.cancellation_registry,
527 failure_policy: self.failure_policy,
528 })
529 }
530}
531
532struct StartupResolverAdapter<'a, Resolver> {
533 resolver: &'a Resolver,
534}
535
536#[derive(Debug)]
538pub enum StartupResolutionError<Error> {
539 Parameters(io::Error),
541 Resolver(Error),
543}
544
545impl<Error: fmt::Display> fmt::Display for StartupResolutionError<Error> {
546 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
547 match self {
548 Self::Parameters(error) => error.fmt(formatter),
549 Self::Resolver(error) => error.fmt(formatter),
550 }
551 }
552}
553
554impl<Error: std::error::Error + 'static> std::error::Error for StartupResolutionError<Error> {
555 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
556 match self {
557 Self::Parameters(error) => Some(error),
558 Self::Resolver(error) => Some(error),
559 }
560 }
561}
562
563impl<Resolver, State, Peer, Identity>
564 crate::server_component::StartupResolver<State, Peer, Identity>
565 for StartupResolverAdapter<'_, Resolver>
566where
567 Resolver: StartupRouteResolver<Peer>,
568{
569 type Route = ConnectTarget;
570 type Error = StartupResolutionError<Resolver::Error>;
571
572 fn defer_ready(&self) -> bool {
573 true
574 }
575
576 fn resolve<'a>(
577 &'a mut self,
578 startup: &'a crate::startup::StartupMessage,
579 context: &'a crate::ServerConnectionContext<Peer, Identity>,
580 _state: &'a mut State,
581 ) -> Pin<Box<dyn Future<Output = Result<Self::Route, Self::Error>> + 'a>> {
582 let parameters = StartupParameters::from_wire(startup);
583 let initial = context
584 .tls_if_known()
585 .map(|tls| InitialServerContext::new(context.peer(), tls));
586 let resolver = self.resolver;
587 Box::pin(async move {
588 let parameters = parameters.map_err(StartupResolutionError::Parameters)?;
589 let initial = initial.expect("startup routing runs after TLS negotiation");
590 resolver
591 .resolve(parameters, initial)
592 .await
593 .map_err(StartupResolutionError::Resolver)
594 })
595 }
596}
597
598pub enum IntermediaryAcceptError<
600 ServerError,
601 ResolverError,
602 RouteError,
603 ClientError,
604 RegistryError = std::convert::Infallible,
605 CancellationError = std::convert::Infallible,
606> {
607 Server(ServerError),
609 StartupRoute(StartupResolutionError<ResolverError>),
611 CancellationRejected,
613 AuthenticatedRoute(RouteError),
615 Client(ClientError),
617 CancellationRegistry(RegistryError),
619 ServerOutput(io::Error),
621 Cancellation(CancellationError),
623}
624
625impl<S, R, A, C, K, X> fmt::Debug for IntermediaryAcceptError<S, R, A, C, K, X> {
626 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
627 formatter.write_str(match self {
628 Self::Server(_) => "IntermediaryAcceptError::Server([REDACTED])",
629 Self::StartupRoute(_) => "IntermediaryAcceptError::StartupRoute([REDACTED])",
630 Self::CancellationRejected => "IntermediaryAcceptError::CancellationRejected",
631 Self::AuthenticatedRoute(_) => {
632 "IntermediaryAcceptError::AuthenticatedRoute([REDACTED])"
633 }
634 Self::Client(_) => "IntermediaryAcceptError::Client([REDACTED])",
635 Self::CancellationRegistry(_) => {
636 "IntermediaryAcceptError::CancellationRegistry([REDACTED])"
637 }
638 Self::ServerOutput(_) => "IntermediaryAcceptError::ServerOutput([REDACTED])",
639 Self::Cancellation(_) => "IntermediaryAcceptError::Cancellation([REDACTED])",
640 })
641 }
642}
643
644impl<S, R, A, C, K, X> fmt::Display for IntermediaryAcceptError<S, R, A, C, K, X> {
645 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
646 match self {
647 Self::Server(_) => formatter.write_str("client-facing establishment failed"),
648 Self::StartupRoute(_) => formatter.write_str("startup routing failed"),
649 Self::CancellationRejected => {
650 formatter.write_str("cancellation is explicitly rejected")
651 }
652 Self::AuthenticatedRoute(_) => formatter.write_str("authenticated routing failed"),
653 Self::Client(_) => formatter.write_str("PostgreSQL-facing establishment failed"),
654 Self::CancellationRegistry(_) => {
655 formatter.write_str("cancellation registration failed")
656 }
657 Self::ServerOutput(_) => {
658 formatter.write_str("client-facing establishment output failed")
659 }
660 Self::Cancellation(_) => formatter.write_str("cancellation forwarding failed"),
661 }
662 }
663}
664
665impl<S, R, A, C, K, X> std::error::Error for IntermediaryAcceptError<S, R, A, C, K, X>
666where
667 S: std::error::Error + 'static,
668 R: std::error::Error + 'static,
669 A: std::error::Error + 'static,
670 C: std::error::Error + 'static,
671 K: std::error::Error + 'static,
672 X: std::error::Error + 'static,
673{
674 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
675 match self {
676 Self::Server(error) => Some(error),
677 Self::StartupRoute(error) => Some(error),
678 Self::CancellationRejected => None,
679 Self::AuthenticatedRoute(error) => Some(error),
680 Self::Client(error) => Some(error),
681 Self::CancellationRegistry(error) => Some(error),
682 Self::ServerOutput(error) => Some(error),
683 Self::Cancellation(error) => Some(error),
684 }
685 }
686}
687
688#[derive(Debug)]
690pub struct IntermediaryContexts<ServerContext, ClientContext> {
691 server: ServerContext,
692 client: ClientContext,
693}
694
695impl<ServerContext, ClientContext> IntermediaryContexts<ServerContext, ClientContext> {
696 #[must_use]
698 pub const fn server(&self) -> &ServerContext {
699 &self.server
700 }
701 #[must_use]
703 pub const fn client(&self) -> &ClientContext {
704 &self.client
705 }
706}
707
708pub struct IntermediaryConnection<
710 DT,
711 UT,
712 State,
713 Peer,
714 ServerIdentity,
715 ClientEvidence,
716 ServerHandler,
717 ClientHandler,
718 Boundary,
719 Policy,
720 Cancellation = RejectCancellation,
721> {
722 downstream:
723 crate::server_component::ServerConnectionCore<DT, Peer, ServerIdentity, ServerHandler>,
724 upstream: crate::client_component::ClientConnectionCore<
725 crate::ClientTransport<UT>,
726 crate::Pristine,
727 ClientEvidence,
728 ClientHandler,
729 >,
730 state: State,
731 boundary: Boundary,
732 pipeline: Pipeline<Policy>,
733 target: ConnectTarget,
734 pending_frontend: Option<crate::codec::FrontendMessage>,
735 cancellation_registry: Cancellation,
736 client_cancel_key: Option<crate::demux::CancelKey>,
737}
738
739#[derive(Debug)]
741pub enum IntermediaryAccept<Connection> {
742 Session(Connection),
744 CancellationForwarded,
746}
747
748impl<Connection> IntermediaryAccept<Connection> {
749 #[must_use]
754 pub fn into_session(self) -> Connection {
755 match self {
756 Self::Session(connection) => connection,
757 Self::CancellationForwarded => panic!("accepted cancellation has no session"),
758 }
759 }
760}
761
762#[derive(Debug)]
764pub enum ForwardedMessage {
765 Frontend(crate::codec::FrontendMessage),
767 Backend(crate::codec::BackendMessage),
769}
770
771impl<DT, UT, State, Peer, SI, CE, SH, CH, Boundary, Policy, K>
772 IntermediaryConnection<DT, UT, State, Peer, SI, CE, SH, CH, Boundary, Policy, K>
773where
774 Policy: PipelinePolicy,
775{
776 #[must_use]
778 pub const fn target(&self) -> &ConnectTarget {
779 &self.target
780 }
781 #[must_use]
783 pub const fn state(&self) -> &State {
784 &self.state
785 }
786 #[must_use]
788 pub const fn cancellation_key(&self) -> Option<&crate::demux::CancelKey> {
789 self.client_cancel_key.as_ref()
790 }
791
792 pub fn detach_cancellation(&mut self) -> Option<CancellationRoute>
794 where
795 K: IntermediaryCancellationRegistry,
796 {
797 self.client_cancel_key
798 .take()
799 .and_then(|key| self.cancellation_registry.detach(&key))
800 }
801}
802
803impl<
804 DT,
805 UT,
806 State,
807 Peer,
808 ServerIdentity,
809 ClientEvidence,
810 ServerHandler,
811 ClientHandler,
812 Boundary,
813 Policy,
814 K,
815>
816 IntermediaryConnection<
817 DT,
818 UT,
819 State,
820 Peer,
821 ServerIdentity,
822 ClientEvidence,
823 ServerHandler,
824 ClientHandler,
825 Boundary,
826 Policy,
827 K,
828 >
829where
830 DT: AsyncRead + AsyncWrite + Unpin,
831 UT: AsyncRead + AsyncWrite + Unpin,
832 ServerHandler:
833 crate::ServerMiddleware<State, crate::ServerConnectionContext<Peer, ServerIdentity>>,
834 ClientHandler: crate::ClientMiddleware<State, crate::ClientConnectionContext<ClientEvidence>>,
835 Boundary: IntermediaryMiddleware<
836 State,
837 crate::ServerConnectionContext<Peer, ServerIdentity>,
838 crate::ClientConnectionContext<ClientEvidence>,
839 >,
840 Policy: PipelinePolicy,
841 K: IntermediaryCancellationRegistry,
842{
843 pub async fn forward_frontend(
850 &mut self,
851 ) -> Result<crate::codec::FrontendMessage, ForwardError> {
852 if let Some(message) = self.pending_frontend.take() {
853 self.process_frontend(message, false).await
854 } else {
855 let message = self.downstream.receive_wire_raw().await?;
856 self.process_frontend(message, true).await
857 }
858 }
859
860 async fn process_frontend(
861 &mut self,
862 message: crate::codec::FrontendMessage,
863 intercept_source_and_boundary: bool,
864 ) -> Result<crate::codec::FrontendMessage, ForwardError> {
865 let message = if intercept_source_and_boundary {
866 let message = self.downstream.intercept_frontend(&mut self.state, message);
867 let message = self.boundary.frontend(
868 self.downstream.context(),
869 self.upstream.context(),
870 &mut self.state,
871 message,
872 );
873 self.upstream.intercept_frontend(&mut self.state, message)
874 } else {
875 message
876 };
877 let admission = match self
878 .pipeline
879 .accept_frontend(message.clone(), FrontendHandling::Forward)
880 {
881 Ok(admission) => admission,
882 Err(error) => {
883 self.pending_frontend = Some(message);
884 return Err(ForwardError::Frontend(error));
885 }
886 };
887 let FrontendAction::Forward { message, .. } = admission.into_action() else {
888 unreachable!()
889 };
890 self.upstream.send_wire_raw(message.clone()).await?;
891 Ok(message)
892 }
893
894 pub async fn forward_backend(&mut self) -> Result<crate::codec::BackendMessage, ForwardError> {
901 let message = self.upstream.receive_wire_raw().await?;
902 self.process_backend(message).await
903 }
904
905 async fn process_backend(
906 &mut self,
907 message: crate::codec::BackendMessage,
908 ) -> Result<crate::codec::BackendMessage, ForwardError> {
909 let message = self.upstream.intercept_backend(&mut self.state, message);
910 let message = self.boundary.backend(
911 self.downstream.context(),
912 self.upstream.context(),
913 &mut self.state,
914 message,
915 );
916 let message = self.downstream.intercept_backend(&mut self.state, message);
917 let message = match self
918 .pipeline
919 .accept_backend(message)
920 .map_err(ForwardError::Backend)?
921 {
922 BackendAction::Emit(message) => message,
923 BackendAction::Deferred(message) => return Err(ForwardError::Deferred(message)),
924 };
925 self.downstream.send_wire_raw(message.clone()).await?;
926 Ok(message)
927 }
928
929 pub async fn forward_next(&mut self) -> Result<ForwardedMessage, ForwardError> {
940 if self.pending_frontend.is_some() {
941 let message = self.upstream.receive_wire_raw().await?;
942 return self
943 .process_backend(message)
944 .await
945 .map(ForwardedMessage::Backend);
946 }
947 tokio::select! {
948 result = self.downstream.receive_wire_raw() => {
949 let message = result?;
950 self.process_frontend(message, true).await.map(ForwardedMessage::Frontend)
951 }
952 result = self.upstream.receive_wire_raw() => {
953 let message = result?;
954 self.process_backend(message).await.map(ForwardedMessage::Backend)
955 }
956 }
957 }
958
959 #[allow(clippy::type_complexity)]
962 pub fn teardown(
963 mut self,
964 ) -> (
965 crate::AcceptedServerTransport<DT>,
966 crate::ClientTransport<UT>,
967 State,
968 Boundary,
969 (ServerHandler, ClientHandler),
970 IntermediaryContexts<
971 crate::ServerConnectionContext<Peer, ServerIdentity>,
972 crate::ClientConnectionContext<ClientEvidence>,
973 >,
974 ) {
975 let _ = self.detach_cancellation();
976 let (downstream, server_handler, server_context) = self.downstream.into_parts();
977 let (upstream, client_handler, client_context) = self.upstream.into_parts();
978 (
979 downstream,
980 upstream,
981 self.state,
982 self.boundary,
983 (server_handler, client_handler),
984 IntermediaryContexts {
985 server: server_context,
986 client: client_context,
987 },
988 )
989 }
990}
991
992#[derive(Debug)]
994pub enum ForwardError {
995 Io(io::Error),
997 Frontend(crate::pipeline::FrontendProjectionError),
999 Backend(crate::pipeline::BackendProjectionError),
1001 Deferred(crate::codec::BackendMessage),
1003}
1004
1005impl fmt::Display for ForwardError {
1006 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
1007 match self {
1008 Self::Io(error) => error.fmt(formatter),
1009 Self::Frontend(_) => {
1010 formatter.write_str("frontend message violates pipeline legality or capacity")
1011 }
1012 Self::Backend(_) => formatter.write_str("backend message violates pipeline legality"),
1013 Self::Deferred(_) => formatter.write_str("backend response is not yet emittable"),
1014 }
1015 }
1016}
1017
1018impl std::error::Error for ForwardError {
1019 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
1020 match self {
1021 Self::Io(error) => Some(error),
1022 _ => None,
1023 }
1024 }
1025}
1026
1027impl From<io::Error> for ForwardError {
1028 fn from(error: io::Error) -> Self {
1029 Self::Io(error)
1030 }
1031}
1032
1033impl<ST, SA, SM, Connector, CT, CA, CM, Resolver, Route, Policy, Boundary, K>
1034 Intermediary<
1035 crate::Server<ST, SA, SM>,
1036 crate::Client<Connector, CT, CA, CM>,
1037 Resolver,
1038 Route,
1039 Policy,
1040 Boundary,
1041 K,
1042 >
1043where
1044 ST: crate::ServerTlsConfiguration,
1045 SA: crate::ServerAuthenticationProvider,
1046 CT: crate::client_component::ClientTlsConfiguration,
1047 CA: crate::ClientAuthentication,
1048 CM: crate::MiddlewareFactory<crate::ClientInitialContext>,
1049 Policy: PipelinePolicy,
1050 K: IntermediaryCancellationRegistry + Clone,
1051{
1052 #[allow(clippy::type_complexity, clippy::too_many_lines)]
1059 pub async fn accept<DT, State, Peer, CW, UT, CE>(
1060 &self,
1061 transport: DT,
1062 peer: Peer,
1063 state: State,
1064 ) -> Result<
1065 IntermediaryAccept<
1066 IntermediaryConnection<
1067 DT,
1068 UT,
1069 State,
1070 Peer,
1071 <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1072 CA::Evidence,
1073 <SM as crate::MiddlewareFactory<
1074 crate::ServerConnectionContext<
1075 Peer,
1076 <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1077 >,
1078 >>::Handler,
1079 <CM as crate::MiddlewareFactory<crate::ClientInitialContext>>::Handler,
1080 <Boundary as IntermediaryMiddlewareFactory<
1081 crate::ServerConnectionContext<
1082 Peer,
1083 <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1084 >,
1085 crate::ClientConnectionContext<CA::Evidence>,
1086 >>::Handler,
1087 Policy,
1088 K,
1089 >,
1090 >,
1091 IntermediaryAcceptError<
1092 crate::AcceptError<
1093 <ST::Provider as crate::ServerIdentityProvider>::Error,
1094 <SA::Authentication as crate::ServerAuthentication<Peer>>::Error,
1095 >,
1096 Resolver::Error,
1097 Route::Error,
1098 crate::ConnectError<
1099 CE,
1100 crate::ClientTlsError<<CT::Provider as crate::ClientTlsProvider>::Error>,
1101 crate::ClientAuthenticationError<CA::Error>,
1102 >,
1103 K::Error,
1104 crate::CancelError<CE>,
1105 >,
1106 >
1107 where
1108 DT: AsyncRead + AsyncWrite + Unpin,
1109 UT: AsyncRead + AsyncWrite + Unpin,
1110 SA::Authentication: crate::ServerAuthentication<Peer>,
1111 SM: crate::MiddlewareFactory<
1112 crate::ServerConnectionContext<
1113 Peer,
1114 <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1115 >,
1116 >,
1117 <SM as crate::MiddlewareFactory<
1118 crate::ServerConnectionContext<
1119 Peer,
1120 <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1121 >,
1122 >>::Handler: crate::ServerMiddleware<
1123 State,
1124 crate::ServerConnectionContext<
1125 Peer,
1126 <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1127 >,
1128 >,
1129 Resolver: StartupRouteResolver<Peer>,
1130 Connector: Fn(&ConnectTarget) -> CW,
1131 CW: Future<Output = Result<UT, CE>>,
1132 <CM as crate::MiddlewareFactory<crate::ClientInitialContext>>::Handler:
1133 crate::ClientMiddleware<State, crate::ClientConnectionContext<CA::Evidence>>,
1134 Route: AuthenticatedRoutePolicy<
1135 Peer,
1136 <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1137 >,
1138 Boundary: IntermediaryMiddlewareFactory<
1139 crate::ServerConnectionContext<
1140 Peer,
1141 <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1142 >,
1143 crate::ClientConnectionContext<CA::Evidence>,
1144 >,
1145 <Boundary as IntermediaryMiddlewareFactory<
1146 crate::ServerConnectionContext<
1147 Peer,
1148 <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1149 >,
1150 crate::ClientConnectionContext<CA::Evidence>,
1151 >>::Handler: IntermediaryMiddleware<
1152 State,
1153 crate::ServerConnectionContext<
1154 Peer,
1155 <SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
1156 >,
1157 crate::ClientConnectionContext<CA::Evidence>,
1158 >,
1159 {
1160 let mut resolver = StartupResolverAdapter {
1161 resolver: &self.resolver,
1162 };
1163 let (accepted, selected) = self
1164 .server
1165 .accept_routed(transport, peer, state, &mut resolver)
1166 .await
1167 .map_err(|error| match error {
1168 crate::server_component::RoutedAcceptError::Accept(error) => {
1169 IntermediaryAcceptError::Server(error)
1170 }
1171 crate::server_component::RoutedAcceptError::Route(error) => {
1172 IntermediaryAcceptError::StartupRoute(error)
1173 }
1174 })?;
1175 let mut downstream = match accepted {
1176 crate::ServerAccept::Session(downstream) => downstream,
1177 crate::ServerAccept::Cancellation(cancellation) => {
1178 if self.cancellation == CancellationPolicy::Reject {
1179 let _ = cancellation.teardown();
1180 return Err(IntermediaryAcceptError::CancellationRejected);
1181 }
1182 let request = cancellation.request();
1183 let client_key = crate::demux::CancelKey {
1184 process_id: request.process_id(),
1185 secret_key: bytes::Bytes::copy_from_slice(request.secret_key()),
1186 };
1187 let Some(route) = self.cancellation_registry.resolve(&client_key) else {
1188 let _ = cancellation.teardown();
1189 return Err(IntermediaryAcceptError::CancellationRejected);
1190 };
1191 if let Err(error) = self
1192 .client
1193 .cancel(route.target(), route.upstream_key())
1194 .await
1195 {
1196 let _ = cancellation.teardown();
1197 return Err(IntermediaryAcceptError::Cancellation(error));
1198 }
1199 let _ = cancellation.teardown();
1200 return Ok(IntermediaryAccept::CancellationForwarded);
1201 }
1202 };
1203 let startup = match StartupParameters::from_wire(downstream.startup()) {
1204 Ok(startup) => startup,
1205 Err(error) => {
1206 if self.failure_policy == EstablishmentFailurePolicy::SafeDiagnostic {
1207 let _ = downstream
1208 .send_generated_error(safe_establishment_diagnostic())
1209 .await;
1210 }
1211 let _ = downstream.teardown();
1212 return Err(IntermediaryAcceptError::StartupRoute(
1213 StartupResolutionError::Parameters(error),
1214 ));
1215 }
1216 };
1217 let context = AuthenticatedRouteContext {
1218 peer: downstream.context().peer(),
1219 identity: downstream.context().identity(),
1220 };
1221 let Some(selected) = selected else {
1222 let _ = downstream.teardown();
1223 return Err(IntermediaryAcceptError::CancellationRejected);
1224 };
1225 let selected = match self.route.route(selected, context).await {
1226 Ok(target) => target,
1227 Err(error) => {
1228 if self.failure_policy == EstablishmentFailurePolicy::SafeDiagnostic {
1229 let _ = downstream
1230 .send_generated_error(safe_establishment_diagnostic())
1231 .await;
1232 }
1233 let _ = downstream.teardown();
1234 return Err(IntermediaryAcceptError::AuthenticatedRoute(error));
1235 }
1236 };
1237 let (mut downstream, mut state) = downstream.into_core_and_state();
1238 let upstream = match self
1239 .client
1240 .connect_core(selected.clone(), startup, &mut state)
1241 .await
1242 {
1243 Ok(upstream) => upstream,
1244 Err(error) => {
1245 if self.failure_policy == EstablishmentFailurePolicy::SafeDiagnostic {
1246 let diagnostic = safe_establishment_diagnostic();
1247 let diagnostic = downstream.intercept_backend(&mut state, diagnostic);
1248 if matches!(diagnostic, crate::codec::BackendMessage::ErrorResponse(_)) {
1249 let _ = downstream.send_wire_raw(diagnostic).await;
1252 }
1253 }
1254 let _ = downstream.into_parts();
1255 return Err(IntermediaryAcceptError::Client(error));
1256 }
1257 };
1258 let boundary = self
1259 .boundary
1260 .create(downstream.context(), upstream.context());
1261 let (client_cancel_key, backend_key_message) =
1262 match (self.cancellation, upstream.context().backend_key().cloned()) {
1263 (CancellationPolicy::Forward, Some(upstream_key)) => {
1264 let client_key = match self
1265 .cancellation_registry
1266 .register(CancellationRoute::new(selected.clone(), upstream_key))
1267 {
1268 Ok(key) => key,
1269 Err(error) => {
1270 if self.failure_policy == EstablishmentFailurePolicy::SafeDiagnostic {
1271 let diagnostic = downstream
1272 .intercept_backend(&mut state, safe_establishment_diagnostic());
1273 if matches!(
1274 diagnostic,
1275 crate::codec::BackendMessage::ErrorResponse(_)
1276 ) {
1277 let _ = downstream.send_wire_raw(diagnostic).await;
1278 }
1279 }
1280 let _ = downstream.into_parts();
1281 let _ = upstream.into_parts();
1282 return Err(IntermediaryAcceptError::CancellationRegistry(error));
1283 }
1284 };
1285 let message = crate::codec::BackendMessage::BackendKeyData {
1286 process_id: client_key.process_id,
1287 secret_key: client_key.secret_key.clone(),
1288 };
1289 (Some(client_key), Some(message))
1290 }
1291 _ => (None, None),
1292 };
1293 let mut connection = IntermediaryConnection {
1294 downstream,
1295 upstream,
1296 state,
1297 boundary,
1298 pipeline: Pipeline::new(self.pipeline),
1299 target: selected,
1300 pending_frontend: None,
1301 cancellation_registry: self.cancellation_registry.clone(),
1302 client_cancel_key,
1303 };
1304 if let Some(message) = backend_key_message {
1305 let expected = message.clone();
1306 let message = connection.boundary.backend(
1307 connection.downstream.context(),
1308 connection.upstream.context(),
1309 &mut connection.state,
1310 message,
1311 );
1312 let message = connection
1313 .downstream
1314 .intercept_backend(&mut connection.state, message);
1315 if message != expected {
1316 let _ = connection.detach_cancellation();
1317 let _ = connection.teardown();
1318 return Err(IntermediaryAcceptError::ServerOutput(io::Error::new(
1319 io::ErrorKind::InvalidData,
1320 "middleware rejected generated cancellation key",
1321 )));
1322 }
1323 if let Err(error) = connection.downstream.send_wire_raw(message).await {
1324 let _ = connection.detach_cancellation();
1325 let _ = connection.teardown();
1326 return Err(IntermediaryAcceptError::ServerOutput(error));
1327 }
1328 }
1329 let ready = connection.downstream.intercept_backend(
1330 &mut connection.state,
1331 crate::codec::BackendMessage::ReadyForQuery(crate::codec::TransactionStatus::Idle),
1332 );
1333 if !matches!(ready, crate::codec::BackendMessage::ReadyForQuery(_)) {
1334 let _ = connection.detach_cancellation();
1335 let _ = connection.teardown();
1336 return Err(IntermediaryAcceptError::ServerOutput(io::Error::new(
1337 io::ErrorKind::InvalidData,
1338 "middleware rejected generated readiness",
1339 )));
1340 }
1341 if let Err(error) = connection.downstream.send_wire_raw(ready).await {
1342 let _ = connection.detach_cancellation();
1343 let _ = connection.teardown();
1344 return Err(IntermediaryAcceptError::ServerOutput(error));
1345 }
1346 Ok(IntermediaryAccept::Session(connection))
1347 }
1348}