use std::{collections::VecDeque, fmt, future::Future, io, pin::Pin};
use tokio::io::{AsyncRead, AsyncWrite};
use crate::{
ConnectTarget, NoPipeline, StartupParameters,
pipeline::{BackendAction, FrontendAction, FrontendHandling, Pipeline, PipelinePolicy},
};
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum CancellationPolicy {
Reject,
Forward,
}
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub enum EstablishmentFailurePolicy {
#[default]
Close,
SafeDiagnostic,
}
fn safe_establishment_diagnostic() -> crate::codec::BackendMessage {
crate::codec::BackendMessage::ErrorResponse(crate::codec::DiagnosticResponse {
fields: vec![
crate::codec::DiagnosticField {
code: b'S',
value: bytes::Bytes::from_static(b"ERROR"),
},
crate::codec::DiagnosticField {
code: b'M',
value: bytes::Bytes::from_static(b"connection establishment failed"),
},
],
})
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct CancellationRoute {
target: ConnectTarget,
upstream: crate::demux::CancelKey,
}
impl CancellationRoute {
#[must_use]
pub const fn new(target: ConnectTarget, upstream: crate::demux::CancelKey) -> Self {
Self { target, upstream }
}
#[must_use]
pub const fn target(&self) -> &ConnectTarget {
&self.target
}
#[must_use]
pub const fn upstream_key(&self) -> &crate::demux::CancelKey {
&self.upstream
}
}
pub trait IntermediaryCancellationRegistry {
type Error;
fn register(&self, route: CancellationRoute) -> Result<crate::demux::CancelKey, Self::Error>;
fn resolve(&self, client: &crate::demux::CancelKey) -> Option<CancellationRoute>;
fn detach(&self, client: &crate::demux::CancelKey) -> Option<CancellationRoute>;
}
#[derive(Clone, Copy, Debug, Default)]
pub struct RejectCancellation;
impl IntermediaryCancellationRegistry for RejectCancellation {
type Error = std::convert::Infallible;
fn register(&self, _: CancellationRoute) -> Result<crate::demux::CancelKey, Self::Error> {
unreachable!()
}
fn resolve(&self, _: &crate::demux::CancelKey) -> Option<CancellationRoute> {
None
}
fn detach(&self, _: &crate::demux::CancelKey) -> Option<CancellationRoute> {
None
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum IntermediaryBuildError {
MissingServer,
MissingClient,
MissingStartupResolver,
MissingCancellationPolicy,
}
impl fmt::Display for IntermediaryBuildError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(match self {
Self::MissingServer => "an intermediary server component is required",
Self::MissingClient => "an intermediary client component is required",
Self::MissingStartupResolver => "an asynchronous startup resolver is required",
Self::MissingCancellationPolicy => "an explicit cancellation policy is required",
})
}
}
impl std::error::Error for IntermediaryBuildError {}
#[derive(Clone, Copy, Debug)]
pub struct InitialServerContext<'a, Peer> {
peer: &'a Peer,
tls: &'a crate::NegotiatedServerTls,
}
impl<'a, Peer> InitialServerContext<'a, Peer> {
pub(crate) const fn new(peer: &'a Peer, tls: &'a crate::NegotiatedServerTls) -> Self {
Self { peer, tls }
}
#[must_use]
pub const fn peer(&self) -> &Peer {
self.peer
}
#[must_use]
pub const fn tls(&self) -> &crate::NegotiatedServerTls {
self.tls
}
}
pub trait StartupRouteResolver<Peer> {
type Error;
fn resolve<'a>(
&'a self,
startup: StartupParameters,
context: InitialServerContext<'a, Peer>,
) -> Pin<Box<dyn Future<Output = Result<ConnectTarget, Self::Error>> + 'a>>;
}
pub trait AuthenticatedRoutePolicy<Peer, Identity> {
type Error;
fn route<'a>(
&'a self,
target: ConnectTarget,
context: AuthenticatedRouteContext<'a, Peer, Identity>,
) -> Pin<Box<dyn Future<Output = Result<ConnectTarget, Self::Error>> + 'a>>;
}
#[derive(Clone, Copy, Debug)]
pub struct AuthenticatedRouteContext<'a, Peer, Identity> {
peer: &'a Peer,
identity: &'a Identity,
}
impl<Peer, Identity> AuthenticatedRouteContext<'_, Peer, Identity> {
#[must_use]
pub const fn peer(&self) -> &Peer {
self.peer
}
#[must_use]
pub const fn identity(&self) -> &Identity {
self.identity
}
}
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub struct AllowAuthenticatedRoute;
impl<Peer, Identity> AuthenticatedRoutePolicy<Peer, Identity> for AllowAuthenticatedRoute {
type Error = std::convert::Infallible;
fn route<'a>(
&'a self,
target: ConnectTarget,
_context: AuthenticatedRouteContext<'a, Peer, Identity>,
) -> Pin<Box<dyn Future<Output = Result<ConnectTarget, Self::Error>> + 'a>> {
Box::pin(async move { Ok(target) })
}
}
#[derive(Debug, Eq, PartialEq)]
pub enum FrontendMiddlewareOutput {
Forward(crate::codec::FrontendMessage),
Suppress(crate::codec::FrontendMessage),
Respond {
request: crate::codec::FrontendMessage,
responses: Vec<crate::codec::BackendMessage>,
},
}
#[derive(Debug, Eq, PartialEq)]
pub enum BackendMiddlewareOutput {
Forward(crate::codec::BackendMessage),
Expand(Vec<crate::codec::BackendMessage>),
Suppress(crate::codec::BackendMessage),
}
pub trait IntermediaryMiddleware<State, ServerContext, ClientContext> {
type Error;
fn frontend<'a>(
&'a mut self,
_server: &'a ServerContext,
_client: &'a ClientContext,
_state: &'a mut State,
message: crate::codec::FrontendMessage,
) -> Pin<Box<dyn Future<Output = Result<FrontendMiddlewareOutput, Self::Error>> + 'a>> {
Box::pin(async move { Ok(FrontendMiddlewareOutput::Forward(message)) })
}
fn backend<'a>(
&'a mut self,
_server: &'a ServerContext,
_client: &'a ClientContext,
_state: &'a mut State,
message: crate::codec::BackendMessage,
) -> Pin<Box<dyn Future<Output = Result<BackendMiddlewareOutput, Self::Error>> + 'a>> {
Box::pin(async move { Ok(BackendMiddlewareOutput::Forward(message)) })
}
}
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub struct IdentityIntermediaryMiddleware;
impl<State, ServerContext, ClientContext>
IntermediaryMiddleware<State, ServerContext, ClientContext> for IdentityIntermediaryMiddleware
{
type Error = std::convert::Infallible;
}
pub trait IntermediaryMiddlewareFactory<ServerContext, ClientContext> {
type Handler;
fn create(&self, server: &ServerContext, client: &ClientContext) -> Self::Handler;
}
impl<ServerContext, ClientContext, Handler, Factory>
IntermediaryMiddlewareFactory<ServerContext, ClientContext> for Factory
where
Factory: Fn(&ServerContext, &ClientContext) -> Handler,
{
type Handler = Handler;
fn create(&self, server: &ServerContext, client: &ClientContext) -> Handler {
self(server, client)
}
}
impl<ServerContext, ClientContext> IntermediaryMiddlewareFactory<ServerContext, ClientContext>
for IdentityIntermediaryMiddleware
{
type Handler = Self;
fn create(&self, _server: &ServerContext, _client: &ClientContext) -> Self {
*self
}
}
pub struct Intermediary<
Server = (),
Client = (),
Resolver = (),
Route = AllowAuthenticatedRoute,
Policy = NoPipeline,
Boundary = IdentityIntermediaryMiddleware,
Cancellation = RejectCancellation,
> {
pub(crate) server: Server,
pub(crate) client: Client,
pub(crate) resolver: Resolver,
pub(crate) route: Route,
pub(crate) pipeline: Policy,
pub(crate) boundary: Boundary,
pub(crate) cancellation: CancellationPolicy,
pub(crate) cancellation_registry: Cancellation,
pub(crate) failure_policy: EstablishmentFailurePolicy,
}
impl Intermediary<()> {
#[must_use]
pub fn builder() -> IntermediaryBuilder {
IntermediaryBuilder::default()
}
}
impl<S, C, R, A, P, B, K> fmt::Debug for Intermediary<S, C, R, A, P, B, K> {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("Intermediary")
.field("server", &"<configured>")
.field("client", &"<configured>")
.field("resolver", &"<redacted>")
.field("authenticated_route", &"<redacted>")
.field("cancellation", &self.cancellation)
.finish_non_exhaustive()
}
}
pub struct IntermediaryBuilder<
Server = (),
Client = (),
Resolver = (),
Route = AllowAuthenticatedRoute,
Policy = NoPipeline,
Boundary = IdentityIntermediaryMiddleware,
Cancellation = RejectCancellation,
> {
server: Option<Server>,
client: Option<Client>,
resolver: Option<Resolver>,
route: Route,
pipeline: Policy,
boundary: Boundary,
cancellation: Option<CancellationPolicy>,
cancellation_registry: Cancellation,
failure_policy: EstablishmentFailurePolicy,
}
impl Default for IntermediaryBuilder {
fn default() -> Self {
Self {
server: None,
client: None,
resolver: None,
route: AllowAuthenticatedRoute,
pipeline: NoPipeline,
boundary: IdentityIntermediaryMiddleware,
cancellation: None,
cancellation_registry: RejectCancellation,
failure_policy: EstablishmentFailurePolicy::Close,
}
}
}
impl<S, C, R, A, P, B, K> IntermediaryBuilder<S, C, R, A, P, B, K> {
#[must_use]
pub fn server<Next>(self, server: Next) -> IntermediaryBuilder<Next, C, R, A, P, B, K> {
IntermediaryBuilder {
server: Some(server),
client: self.client,
resolver: self.resolver,
route: self.route,
pipeline: self.pipeline,
boundary: self.boundary,
cancellation: self.cancellation,
cancellation_registry: self.cancellation_registry,
failure_policy: self.failure_policy,
}
}
#[must_use]
pub fn client<Next>(self, client: Next) -> IntermediaryBuilder<S, Next, R, A, P, B, K> {
IntermediaryBuilder {
server: self.server,
client: Some(client),
resolver: self.resolver,
route: self.route,
pipeline: self.pipeline,
boundary: self.boundary,
cancellation: self.cancellation,
cancellation_registry: self.cancellation_registry,
failure_policy: self.failure_policy,
}
}
#[must_use]
pub fn startup_resolver<Next>(
self,
resolver: Next,
) -> IntermediaryBuilder<S, C, Next, A, P, B, K> {
IntermediaryBuilder {
server: self.server,
client: self.client,
resolver: Some(resolver),
route: self.route,
pipeline: self.pipeline,
boundary: self.boundary,
cancellation: self.cancellation,
cancellation_registry: self.cancellation_registry,
failure_policy: self.failure_policy,
}
}
#[must_use]
pub fn authenticated_route<Next>(
self,
route: Next,
) -> IntermediaryBuilder<S, C, R, Next, P, B, K> {
IntermediaryBuilder {
server: self.server,
client: self.client,
resolver: self.resolver,
route,
pipeline: self.pipeline,
boundary: self.boundary,
cancellation: self.cancellation,
cancellation_registry: self.cancellation_registry,
failure_policy: self.failure_policy,
}
}
#[must_use]
pub fn pipeline<Next: PipelinePolicy>(
self,
pipeline: Next,
) -> IntermediaryBuilder<S, C, R, A, Next, B, K> {
IntermediaryBuilder {
server: self.server,
client: self.client,
resolver: self.resolver,
route: self.route,
pipeline,
boundary: self.boundary,
cancellation: self.cancellation,
cancellation_registry: self.cancellation_registry,
failure_policy: self.failure_policy,
}
}
#[must_use]
pub fn middleware<Next>(self, boundary: Next) -> IntermediaryBuilder<S, C, R, A, P, Next, K> {
IntermediaryBuilder {
server: self.server,
client: self.client,
resolver: self.resolver,
route: self.route,
pipeline: self.pipeline,
boundary,
cancellation: self.cancellation,
cancellation_registry: self.cancellation_registry,
failure_policy: self.failure_policy,
}
}
#[must_use]
pub fn cancellation(mut self, cancellation: CancellationPolicy) -> Self {
self.cancellation = match cancellation {
CancellationPolicy::Reject => Some(CancellationPolicy::Reject),
CancellationPolicy::Forward => None,
};
self
}
#[must_use]
pub fn establishment_failure(mut self, policy: EstablishmentFailurePolicy) -> Self {
self.failure_policy = policy;
self
}
#[must_use]
pub fn cancellation_registry<Next>(
self,
registry: Next,
) -> IntermediaryBuilder<S, C, R, A, P, B, Next> {
IntermediaryBuilder {
server: self.server,
client: self.client,
resolver: self.resolver,
route: self.route,
pipeline: self.pipeline,
boundary: self.boundary,
cancellation: Some(CancellationPolicy::Forward),
cancellation_registry: registry,
failure_policy: self.failure_policy,
}
}
#[allow(clippy::type_complexity)]
pub fn build(self) -> Result<Intermediary<S, C, R, A, P, B, K>, IntermediaryBuildError> {
Ok(Intermediary {
server: self.server.ok_or(IntermediaryBuildError::MissingServer)?,
client: self.client.ok_or(IntermediaryBuildError::MissingClient)?,
resolver: self
.resolver
.ok_or(IntermediaryBuildError::MissingStartupResolver)?,
route: self.route,
pipeline: self.pipeline,
boundary: self.boundary,
cancellation: self
.cancellation
.ok_or(IntermediaryBuildError::MissingCancellationPolicy)?,
cancellation_registry: self.cancellation_registry,
failure_policy: self.failure_policy,
})
}
}
struct StartupResolverAdapter<'a, Resolver> {
resolver: &'a Resolver,
}
#[derive(Debug)]
pub enum StartupResolutionError<Error> {
Parameters(io::Error),
Resolver(Error),
}
impl<Error: fmt::Display> fmt::Display for StartupResolutionError<Error> {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Parameters(error) => error.fmt(formatter),
Self::Resolver(error) => error.fmt(formatter),
}
}
}
impl<Error: std::error::Error + 'static> std::error::Error for StartupResolutionError<Error> {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Self::Parameters(error) => Some(error),
Self::Resolver(error) => Some(error),
}
}
}
impl<Resolver, State, Peer, Identity>
crate::server_component::StartupResolver<State, Peer, Identity>
for StartupResolverAdapter<'_, Resolver>
where
Resolver: StartupRouteResolver<Peer>,
{
type Route = ConnectTarget;
type Error = StartupResolutionError<Resolver::Error>;
fn defer_ready(&self) -> bool {
true
}
fn resolve<'a>(
&'a mut self,
startup: &'a crate::startup::StartupMessage,
context: &'a crate::ServerConnectionContext<Peer, Identity>,
_state: &'a mut State,
) -> Pin<Box<dyn Future<Output = Result<Self::Route, Self::Error>> + 'a>> {
let parameters = StartupParameters::from_wire(startup);
let initial = context
.tls_if_known()
.map(|tls| InitialServerContext::new(context.peer(), tls));
let resolver = self.resolver;
Box::pin(async move {
let parameters = parameters.map_err(StartupResolutionError::Parameters)?;
let initial = initial.expect("startup routing runs after TLS negotiation");
resolver
.resolve(parameters, initial)
.await
.map_err(StartupResolutionError::Resolver)
})
}
}
pub enum IntermediaryAcceptError<
ServerError,
ResolverError,
RouteError,
ClientError,
RegistryError = std::convert::Infallible,
CancellationError = std::convert::Infallible,
MiddlewareError = std::convert::Infallible,
> {
Server(ServerError),
StartupRoute(StartupResolutionError<ResolverError>),
CancellationRejected,
AuthenticatedRoute(RouteError),
Client(ClientError),
CancellationRegistry(RegistryError),
ServerOutput(io::Error),
Cancellation(CancellationError),
Middleware(MiddlewareError),
}
impl<S, R, A, C, K, X, M> fmt::Debug for IntermediaryAcceptError<S, R, A, C, K, X, M> {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(match self {
Self::Server(_) => "IntermediaryAcceptError::Server([REDACTED])",
Self::StartupRoute(_) => "IntermediaryAcceptError::StartupRoute([REDACTED])",
Self::CancellationRejected => "IntermediaryAcceptError::CancellationRejected",
Self::AuthenticatedRoute(_) => {
"IntermediaryAcceptError::AuthenticatedRoute([REDACTED])"
}
Self::Client(_) => "IntermediaryAcceptError::Client([REDACTED])",
Self::CancellationRegistry(_) => {
"IntermediaryAcceptError::CancellationRegistry([REDACTED])"
}
Self::ServerOutput(_) => "IntermediaryAcceptError::ServerOutput([REDACTED])",
Self::Cancellation(_) => "IntermediaryAcceptError::Cancellation([REDACTED])",
Self::Middleware(_) => "IntermediaryAcceptError::Middleware([REDACTED])",
})
}
}
impl<S, R, A, C, K, X, M> fmt::Display for IntermediaryAcceptError<S, R, A, C, K, X, M> {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Server(_) => formatter.write_str("client-facing establishment failed"),
Self::StartupRoute(_) => formatter.write_str("startup routing failed"),
Self::CancellationRejected => {
formatter.write_str("cancellation is explicitly rejected")
}
Self::AuthenticatedRoute(_) => formatter.write_str("authenticated routing failed"),
Self::Client(_) => formatter.write_str("PostgreSQL-facing establishment failed"),
Self::CancellationRegistry(_) => {
formatter.write_str("cancellation registration failed")
}
Self::ServerOutput(_) => {
formatter.write_str("client-facing establishment output failed")
}
Self::Cancellation(_) => formatter.write_str("cancellation forwarding failed"),
Self::Middleware(_) => {
formatter.write_str("forwarding middleware rejected establishment output")
}
}
}
}
impl<S, R, A, C, K, X, M> std::error::Error for IntermediaryAcceptError<S, R, A, C, K, X, M>
where
S: std::error::Error + 'static,
R: std::error::Error + 'static,
A: std::error::Error + 'static,
C: std::error::Error + 'static,
K: std::error::Error + 'static,
X: std::error::Error + 'static,
M: std::error::Error + 'static,
{
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Self::Server(error) => Some(error),
Self::StartupRoute(error) => Some(error),
Self::CancellationRejected => None,
Self::AuthenticatedRoute(error) => Some(error),
Self::Client(error) => Some(error),
Self::CancellationRegistry(error) => Some(error),
Self::ServerOutput(error) => Some(error),
Self::Cancellation(error) => Some(error),
Self::Middleware(error) => Some(error),
}
}
}
#[derive(Debug)]
pub struct IntermediaryContexts<ServerContext, ClientContext> {
server: ServerContext,
client: ClientContext,
}
impl<ServerContext, ClientContext> IntermediaryContexts<ServerContext, ClientContext> {
#[must_use]
pub const fn server(&self) -> &ServerContext {
&self.server
}
#[must_use]
pub const fn client(&self) -> &ClientContext {
&self.client
}
}
pub struct IntermediaryConnection<
DT,
UT,
State,
Peer,
ServerIdentity,
ClientEvidence,
ServerHandler,
ClientHandler,
Boundary,
Policy,
Cancellation = RejectCancellation,
> {
downstream:
crate::server_component::ServerConnectionCore<DT, Peer, ServerIdentity, ServerHandler>,
upstream: crate::client_component::ClientConnectionCore<
crate::ClientTransport<UT>,
crate::Pristine,
ClientEvidence,
ClientHandler,
>,
state: State,
boundary: Boundary,
pipeline: Pipeline<Policy>,
target: ConnectTarget,
pending_frontend: Option<crate::codec::FrontendMessage>,
pending_local: VecDeque<PendingLocalResponses>,
cancellation_registry: Cancellation,
client_cancel_key: Option<crate::demux::CancelKey>,
}
struct PendingLocalResponses {
operation: crate::pipeline::OperationId,
messages: VecDeque<crate::codec::BackendMessage>,
}
#[derive(Debug)]
pub enum IntermediaryAccept<Connection> {
Session(Connection),
CancellationForwarded,
}
impl<Connection> IntermediaryAccept<Connection> {
#[must_use]
pub fn into_session(self) -> Connection {
match self {
Self::Session(connection) => connection,
Self::CancellationForwarded => panic!("accepted cancellation has no session"),
}
}
}
#[derive(Debug)]
pub enum ForwardedMessage {
Frontend(crate::codec::FrontendMessage),
Backend(crate::codec::BackendMessage),
BackendExpanded {
source: crate::codec::BackendMessage,
messages: Vec<crate::codec::BackendMessage>,
},
FrontendSuppressed(crate::codec::FrontendMessage),
FrontendLocallyHandled(crate::codec::FrontendMessage),
BackendSuppressed(crate::codec::BackendMessage),
}
#[derive(Debug, Eq, PartialEq)]
pub enum FrontendForwarding {
Forwarded(crate::codec::FrontendMessage),
Suppressed(crate::codec::FrontendMessage),
LocallyHandled(crate::codec::FrontendMessage),
}
impl FrontendForwarding {
#[must_use]
pub fn into_message(self) -> crate::codec::FrontendMessage {
match self {
Self::Forwarded(message)
| Self::Suppressed(message)
| Self::LocallyHandled(message) => message,
}
}
}
#[derive(Debug, Eq, PartialEq)]
pub enum BackendForwarding {
Forwarded(crate::codec::BackendMessage),
Expanded {
source: crate::codec::BackendMessage,
messages: Vec<crate::codec::BackendMessage>,
},
Suppressed(crate::codec::BackendMessage),
}
impl BackendForwarding {
#[must_use]
pub fn into_message(self) -> crate::codec::BackendMessage {
match self {
Self::Forwarded(message) | Self::Suppressed(message) => message,
Self::Expanded { source, .. } => source,
}
}
}
impl<DT, UT, State, Peer, SI, CE, SH, CH, Boundary, Policy, K>
IntermediaryConnection<DT, UT, State, Peer, SI, CE, SH, CH, Boundary, Policy, K>
where
Policy: PipelinePolicy,
{
#[must_use]
pub const fn target(&self) -> &ConnectTarget {
&self.target
}
#[must_use]
pub const fn state(&self) -> &State {
&self.state
}
#[must_use]
pub const fn cancellation_key(&self) -> Option<&crate::demux::CancelKey> {
self.client_cancel_key.as_ref()
}
pub fn detach_cancellation(&mut self) -> Option<CancellationRoute>
where
K: IntermediaryCancellationRegistry,
{
self.client_cancel_key
.take()
.and_then(|key| self.cancellation_registry.detach(&key))
}
}
impl<
DT,
UT,
State,
Peer,
ServerIdentity,
ClientEvidence,
ServerHandler,
ClientHandler,
Boundary,
Policy,
K,
>
IntermediaryConnection<
DT,
UT,
State,
Peer,
ServerIdentity,
ClientEvidence,
ServerHandler,
ClientHandler,
Boundary,
Policy,
K,
>
where
DT: AsyncRead + AsyncWrite + Unpin,
UT: AsyncRead + AsyncWrite + Unpin,
ServerHandler:
crate::ServerMiddleware<State, crate::ServerConnectionContext<Peer, ServerIdentity>>,
ClientHandler: crate::ClientMiddleware<State, crate::ClientConnectionContext<ClientEvidence>>,
Boundary: IntermediaryMiddleware<
State,
crate::ServerConnectionContext<Peer, ServerIdentity>,
crate::ClientConnectionContext<ClientEvidence>,
>,
Policy: PipelinePolicy,
K: IntermediaryCancellationRegistry,
{
pub async fn forward_frontend(
&mut self,
) -> Result<FrontendForwarding, ForwardError<Boundary::Error>> {
if let Some(message) = self.pending_frontend.take() {
self.process_frontend(message, false).await
} else {
let message = self.downstream.receive_wire_raw().await?;
self.process_frontend(message, true).await
}
}
async fn process_frontend(
&mut self,
message: crate::codec::FrontendMessage,
intercept_source_and_boundary: bool,
) -> Result<FrontendForwarding, ForwardError<Boundary::Error>> {
let decision = if intercept_source_and_boundary {
let message = self.downstream.intercept_frontend(&mut self.state, message);
self.boundary
.frontend(
self.downstream.context(),
self.upstream.context(),
&mut self.state,
message,
)
.await
.map_err(ForwardError::Middleware)?
} else {
FrontendMiddlewareOutput::Forward(message)
};
let (message, handling) = match decision {
FrontendMiddlewareOutput::Forward(message) => {
let message = if intercept_source_and_boundary {
self.upstream.intercept_frontend(&mut self.state, message)
} else {
message
};
(message, FrontendHandling::Forward)
}
FrontendMiddlewareOutput::Suppress(message) => {
return Ok(FrontendForwarding::Suppressed(message));
}
FrontendMiddlewareOutput::Respond { request, responses } => {
let admission = self
.pipeline
.accept_frontend(request.clone(), FrontendHandling::Local)
.map_err(ForwardError::Frontend)?;
let FrontendAction::Discard { id } = admission.into_action() else {
unreachable!()
};
let messages = responses
.into_iter()
.map(|message| self.downstream.intercept_backend(&mut self.state, message))
.collect();
self.pending_local.push_back(PendingLocalResponses {
operation: id,
messages,
});
self.flush_local_responses().await?;
return Ok(FrontendForwarding::LocallyHandled(request));
}
};
let admission = match self.pipeline.accept_frontend(message.clone(), handling) {
Ok(admission) => admission,
Err(error) => {
self.pending_frontend = Some(message);
return Err(ForwardError::Frontend(error));
}
};
let FrontendAction::Forward { message, .. } = admission.into_action() else {
unreachable!()
};
self.upstream.send_wire_raw(message.clone()).await?;
Ok(FrontendForwarding::Forwarded(message))
}
pub async fn forward_backend(
&mut self,
) -> Result<BackendForwarding, ForwardError<Boundary::Error>> {
let message = self.upstream.receive_wire_raw().await?;
self.process_backend(message).await
}
async fn process_backend(
&mut self,
message: crate::codec::BackendMessage,
) -> Result<BackendForwarding, ForwardError<Boundary::Error>> {
let message = self.upstream.intercept_backend(&mut self.state, message);
let source = message.clone();
let decision = self
.boundary
.backend(
self.downstream.context(),
self.upstream.context(),
&mut self.state,
message,
)
.await
.map_err(ForwardError::Middleware)?;
let outcome = match decision {
BackendMiddlewareOutput::Forward(message) => {
let message = self.downstream.intercept_backend(&mut self.state, message);
let message = self.emit_backend(message).await?;
BackendForwarding::Forwarded(message)
}
BackendMiddlewareOutput::Suppress(message) => {
let message = self.advance_backend(message)?;
BackendForwarding::Suppressed(message)
}
BackendMiddlewareOutput::Expand(messages) => {
if messages.is_empty() {
return Err(ForwardError::EmptyExpansion(source));
}
let mut emitted = Vec::with_capacity(messages.len());
for message in messages {
let message = self.downstream.intercept_backend(&mut self.state, message);
emitted.push(self.emit_backend(message).await?);
}
BackendForwarding::Expanded {
source,
messages: emitted,
}
}
};
self.flush_local_responses().await?;
Ok(outcome)
}
fn advance_backend(
&mut self,
message: crate::codec::BackendMessage,
) -> Result<crate::codec::BackendMessage, ForwardError<Boundary::Error>> {
match self
.pipeline
.accept_backend(message)
.map_err(ForwardError::Backend)?
{
BackendAction::Emit(message) => Ok(message),
BackendAction::Deferred(message) => Err(ForwardError::Deferred(message)),
}
}
async fn emit_backend(
&mut self,
message: crate::codec::BackendMessage,
) -> Result<crate::codec::BackendMessage, ForwardError<Boundary::Error>> {
let message = self.advance_backend(message)?;
self.downstream.send_wire_raw(message.clone()).await?;
Ok(message)
}
async fn flush_local_responses(&mut self) -> Result<(), ForwardError<Boundary::Error>> {
loop {
let Some(pending) = self.pending_local.front_mut() else {
return Ok(());
};
let Some(message) = pending.messages.pop_front() else {
self.pending_local.pop_front();
continue;
};
match self.pipeline.try_emit_local(pending.operation, message) {
Ok(BackendAction::Emit(message)) => {
self.downstream.send_wire_raw(message).await?;
}
Ok(BackendAction::Deferred(message)) => {
pending.messages.push_front(message);
return Ok(());
}
Err(error) => return Err(ForwardError::Backend(error)),
}
}
}
pub async fn forward_next(
&mut self,
) -> Result<ForwardedMessage, ForwardError<Boundary::Error>> {
if self.pending_frontend.is_some() {
let message = self.upstream.receive_wire_raw().await?;
return self
.process_backend(message)
.await
.map(|outcome| match outcome {
BackendForwarding::Forwarded(message) => ForwardedMessage::Backend(message),
BackendForwarding::Expanded { source, messages } => {
ForwardedMessage::BackendExpanded { source, messages }
}
BackendForwarding::Suppressed(message) => {
ForwardedMessage::BackendSuppressed(message)
}
});
}
tokio::select! {
result = self.downstream.receive_wire_raw() => {
let message = result?;
self.process_frontend(message, true).await.map(|outcome| match outcome {
FrontendForwarding::Forwarded(message) => ForwardedMessage::Frontend(message),
FrontendForwarding::Suppressed(message) => ForwardedMessage::FrontendSuppressed(message),
FrontendForwarding::LocallyHandled(message) => ForwardedMessage::FrontendLocallyHandled(message),
})
}
result = self.upstream.receive_wire_raw() => {
let message = result?;
self.process_backend(message).await.map(|outcome| match outcome {
BackendForwarding::Forwarded(message) => ForwardedMessage::Backend(message),
BackendForwarding::Expanded { source, messages } => {
ForwardedMessage::BackendExpanded { source, messages }
}
BackendForwarding::Suppressed(message) => ForwardedMessage::BackendSuppressed(message),
})
}
}
}
#[allow(clippy::type_complexity)]
pub fn teardown(
mut self,
) -> (
crate::AcceptedServerTransport<DT>,
crate::ClientTransport<UT>,
State,
Boundary,
(ServerHandler, ClientHandler),
IntermediaryContexts<
crate::ServerConnectionContext<Peer, ServerIdentity>,
crate::ClientConnectionContext<ClientEvidence>,
>,
) {
let _ = self.detach_cancellation();
let (downstream, server_handler, server_context) = self.downstream.into_parts();
let (upstream, client_handler, client_context) = self.upstream.into_parts();
(
downstream,
upstream,
self.state,
self.boundary,
(server_handler, client_handler),
IntermediaryContexts {
server: server_context,
client: client_context,
},
)
}
}
#[derive(Debug)]
pub enum ForwardError<MiddlewareError = std::convert::Infallible> {
Io(io::Error),
Frontend(crate::pipeline::FrontendProjectionError),
Backend(crate::pipeline::BackendProjectionError),
Deferred(crate::codec::BackendMessage),
EmptyExpansion(crate::codec::BackendMessage),
Middleware(MiddlewareError),
}
impl<E> fmt::Display for ForwardError<E> {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Io(error) => error.fmt(formatter),
Self::Frontend(_) => {
formatter.write_str("frontend message violates pipeline legality or capacity")
}
Self::Backend(_) => formatter.write_str("backend message violates pipeline legality"),
Self::Deferred(_) => formatter.write_str("backend response is not yet emittable"),
Self::EmptyExpansion(_) => {
formatter.write_str("backend expansion must contain at least one response")
}
Self::Middleware(_) => formatter.write_str("forwarding middleware rejected a message"),
}
}
}
impl<E> std::error::Error for ForwardError<E>
where
E: std::error::Error + 'static,
{
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Self::Io(error) => Some(error),
Self::Middleware(error) => Some(error),
Self::Frontend(_) | Self::Backend(_) | Self::Deferred(_) | Self::EmptyExpansion(_) => {
None
}
}
}
}
impl<E> From<io::Error> for ForwardError<E> {
fn from(error: io::Error) -> Self {
Self::Io(error)
}
}
impl<ST, SA, SM, Connector, CT, CA, CM, Resolver, Route, Policy, Boundary, K>
Intermediary<
crate::Server<ST, SA, SM>,
crate::Client<Connector, CT, CA, CM>,
Resolver,
Route,
Policy,
Boundary,
K,
>
where
ST: crate::ServerTlsConfiguration,
SA: crate::ServerAuthenticationProvider,
CT: crate::client_component::ClientTlsConfiguration,
CA: crate::ClientAuthentication,
CM: crate::MiddlewareFactory<crate::ClientInitialContext>,
Policy: PipelinePolicy,
K: IntermediaryCancellationRegistry + Clone,
{
#[allow(clippy::type_complexity, clippy::too_many_lines)]
pub async fn accept<DT, State, Peer, CW, UT, CE>(
&self,
transport: DT,
peer: Peer,
state: State,
) -> Result<
IntermediaryAccept<
IntermediaryConnection<
DT,
UT,
State,
Peer,
<SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
CA::Evidence,
<SM as crate::MiddlewareFactory<
crate::ServerConnectionContext<
Peer,
<SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
>,
>>::Handler,
<CM as crate::MiddlewareFactory<crate::ClientInitialContext>>::Handler,
<Boundary as IntermediaryMiddlewareFactory<
crate::ServerConnectionContext<
Peer,
<SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
>,
crate::ClientConnectionContext<CA::Evidence>,
>>::Handler,
Policy,
K,
>,
>,
IntermediaryAcceptError<
crate::AcceptError<
<ST::Provider as crate::ServerIdentityProvider>::Error,
<SA::Authentication as crate::ServerAuthentication<Peer>>::Error,
>,
Resolver::Error,
Route::Error,
crate::ConnectError<
CE,
crate::ClientTlsError<<CT::Provider as crate::ClientTlsProvider>::Error>,
crate::ClientAuthenticationError<CA::Error>,
>,
K::Error,
crate::CancelError<CE>,
<<Boundary as IntermediaryMiddlewareFactory<
crate::ServerConnectionContext<
Peer,
<SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
>,
crate::ClientConnectionContext<CA::Evidence>,
>>::Handler as IntermediaryMiddleware<
State,
crate::ServerConnectionContext<
Peer,
<SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
>,
crate::ClientConnectionContext<CA::Evidence>,
>>::Error,
>,
>
where
DT: AsyncRead + AsyncWrite + Unpin,
UT: AsyncRead + AsyncWrite + Unpin,
SA::Authentication: crate::ServerAuthentication<Peer>,
SM: crate::MiddlewareFactory<
crate::ServerConnectionContext<
Peer,
<SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
>,
>,
<SM as crate::MiddlewareFactory<
crate::ServerConnectionContext<
Peer,
<SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
>,
>>::Handler: crate::ServerMiddleware<
State,
crate::ServerConnectionContext<
Peer,
<SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
>,
>,
Resolver: StartupRouteResolver<Peer>,
Connector: Fn(&ConnectTarget) -> CW,
CW: Future<Output = Result<UT, CE>>,
<CM as crate::MiddlewareFactory<crate::ClientInitialContext>>::Handler:
crate::ClientMiddleware<State, crate::ClientConnectionContext<CA::Evidence>>,
Route: AuthenticatedRoutePolicy<
Peer,
<SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
>,
Boundary: IntermediaryMiddlewareFactory<
crate::ServerConnectionContext<
Peer,
<SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
>,
crate::ClientConnectionContext<CA::Evidence>,
>,
<Boundary as IntermediaryMiddlewareFactory<
crate::ServerConnectionContext<
Peer,
<SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
>,
crate::ClientConnectionContext<CA::Evidence>,
>>::Handler: IntermediaryMiddleware<
State,
crate::ServerConnectionContext<
Peer,
<SA::Authentication as crate::ServerAuthentication<Peer>>::Identity,
>,
crate::ClientConnectionContext<CA::Evidence>,
>,
{
let mut resolver = StartupResolverAdapter {
resolver: &self.resolver,
};
let (accepted, selected) = self
.server
.accept_routed(transport, peer, state, &mut resolver)
.await
.map_err(|error| match error {
crate::server_component::RoutedAcceptError::Accept(error) => {
IntermediaryAcceptError::Server(error)
}
crate::server_component::RoutedAcceptError::Route(error) => {
IntermediaryAcceptError::StartupRoute(error)
}
})?;
let mut downstream = match accepted {
crate::ServerAccept::Session(downstream) => downstream,
crate::ServerAccept::Cancellation(cancellation) => {
if self.cancellation == CancellationPolicy::Reject {
let _ = cancellation.teardown();
return Err(IntermediaryAcceptError::CancellationRejected);
}
let request = cancellation.request();
let client_key = crate::demux::CancelKey {
process_id: request.process_id(),
secret_key: bytes::Bytes::copy_from_slice(request.secret_key()),
};
let Some(route) = self.cancellation_registry.resolve(&client_key) else {
let _ = cancellation.teardown();
return Err(IntermediaryAcceptError::CancellationRejected);
};
if let Err(error) = self
.client
.cancel(route.target(), route.upstream_key())
.await
{
let _ = cancellation.teardown();
return Err(IntermediaryAcceptError::Cancellation(error));
}
let _ = cancellation.teardown();
return Ok(IntermediaryAccept::CancellationForwarded);
}
};
let startup = match StartupParameters::from_wire(downstream.startup()) {
Ok(startup) => startup,
Err(error) => {
if self.failure_policy == EstablishmentFailurePolicy::SafeDiagnostic {
let _ = downstream
.send_generated_error(safe_establishment_diagnostic())
.await;
}
let _ = downstream.teardown();
return Err(IntermediaryAcceptError::StartupRoute(
StartupResolutionError::Parameters(error),
));
}
};
let context = AuthenticatedRouteContext {
peer: downstream.context().peer(),
identity: downstream.context().identity(),
};
let Some(selected) = selected else {
let _ = downstream.teardown();
return Err(IntermediaryAcceptError::CancellationRejected);
};
let selected = match self.route.route(selected, context).await {
Ok(target) => target,
Err(error) => {
if self.failure_policy == EstablishmentFailurePolicy::SafeDiagnostic {
let _ = downstream
.send_generated_error(safe_establishment_diagnostic())
.await;
}
let _ = downstream.teardown();
return Err(IntermediaryAcceptError::AuthenticatedRoute(error));
}
};
let (mut downstream, mut state) = downstream.into_core_and_state();
let upstream = match self
.client
.connect_core(selected.clone(), startup, &mut state)
.await
{
Ok(upstream) => upstream,
Err(error) => {
if self.failure_policy == EstablishmentFailurePolicy::SafeDiagnostic {
let diagnostic = safe_establishment_diagnostic();
let diagnostic = downstream.intercept_backend(&mut state, diagnostic);
if matches!(diagnostic, crate::codec::BackendMessage::ErrorResponse(_)) {
let _ = downstream.send_wire_raw(diagnostic).await;
}
}
let _ = downstream.into_parts();
return Err(IntermediaryAcceptError::Client(error));
}
};
let boundary = self
.boundary
.create(downstream.context(), upstream.context());
let (client_cancel_key, backend_key_message) =
match (self.cancellation, upstream.context().backend_key().cloned()) {
(CancellationPolicy::Forward, Some(upstream_key)) => {
let client_key = match self
.cancellation_registry
.register(CancellationRoute::new(selected.clone(), upstream_key))
{
Ok(key) => key,
Err(error) => {
if self.failure_policy == EstablishmentFailurePolicy::SafeDiagnostic {
let diagnostic = downstream
.intercept_backend(&mut state, safe_establishment_diagnostic());
if matches!(
diagnostic,
crate::codec::BackendMessage::ErrorResponse(_)
) {
let _ = downstream.send_wire_raw(diagnostic).await;
}
}
let _ = downstream.into_parts();
let _ = upstream.into_parts();
return Err(IntermediaryAcceptError::CancellationRegistry(error));
}
};
let message = crate::codec::BackendMessage::BackendKeyData {
process_id: client_key.process_id,
secret_key: client_key.secret_key.clone(),
};
(Some(client_key), Some(message))
}
_ => (None, None),
};
let mut connection = IntermediaryConnection {
downstream,
upstream,
state,
boundary,
pipeline: Pipeline::new(self.pipeline),
target: selected,
pending_frontend: None,
pending_local: VecDeque::new(),
cancellation_registry: self.cancellation_registry.clone(),
client_cancel_key,
};
if let Some(message) = backend_key_message {
let expected = message.clone();
let message = connection
.boundary
.backend(
connection.downstream.context(),
connection.upstream.context(),
&mut connection.state,
message,
)
.await;
let message = match message {
Ok(BackendMiddlewareOutput::Forward(message)) => message,
Ok(BackendMiddlewareOutput::Suppress(_) | BackendMiddlewareOutput::Expand(_)) => {
let _ = connection.detach_cancellation();
let _ = connection.teardown();
return Err(IntermediaryAcceptError::ServerOutput(io::Error::new(
io::ErrorKind::InvalidData,
"middleware suppressed or expanded generated cancellation key",
)));
}
Err(error) => {
let _ = connection.detach_cancellation();
let _ = connection.teardown();
return Err(IntermediaryAcceptError::Middleware(error));
}
};
let message = connection
.downstream
.intercept_backend(&mut connection.state, message);
if message != expected {
let _ = connection.detach_cancellation();
let _ = connection.teardown();
return Err(IntermediaryAcceptError::ServerOutput(io::Error::new(
io::ErrorKind::InvalidData,
"middleware rejected generated cancellation key",
)));
}
if let Err(error) = connection.downstream.send_wire_raw(message).await {
let _ = connection.detach_cancellation();
let _ = connection.teardown();
return Err(IntermediaryAcceptError::ServerOutput(error));
}
}
let ready = connection.downstream.intercept_backend(
&mut connection.state,
crate::codec::BackendMessage::ReadyForQuery(crate::codec::TransactionStatus::Idle),
);
if !matches!(ready, crate::codec::BackendMessage::ReadyForQuery(_)) {
let _ = connection.detach_cancellation();
let _ = connection.teardown();
return Err(IntermediaryAcceptError::ServerOutput(io::Error::new(
io::ErrorKind::InvalidData,
"middleware rejected generated readiness",
)));
}
if let Err(error) = connection.downstream.send_wire_raw(ready).await {
let _ = connection.detach_cancellation();
let _ = connection.teardown();
return Err(IntermediaryAcceptError::ServerOutput(error));
}
Ok(IntermediaryAccept::Session(connection))
}
}