pg_proto/
middleware.rs

1//! Stateful, composable interception of owned protocol messages.
2//!
3//! Middleware receives ownership of a decoded message and a mutable reference to
4//! caller-defined state. Returning the input unchanged is a no-op; implementations
5//! may instead mutate it or return another message of the same type. Protocol
6//! session APIs remain responsible for checking that the result is legal in their
7//! current state before advancing.
8
9use std::marker::PhantomData;
10use std::{convert::Infallible, io};
11
12use crate::{
13    codec::{BackendMessage, FrontendMessage},
14    demux::Demux,
15    grammar::{
16        authentication, backend, frontend, pre_startup, server_authentication, server_pre_startup,
17    },
18    pre_startup::{EncryptionReply, PreStartupMessage},
19};
20
21/// State-aware validation of one directional protocol message type.
22pub trait AcceptsMessage<Message> {
23    /// Reports whether `message` is legal without advancing this state.
24    fn accepts(&self, message: &Message) -> bool;
25}
26
27/// A protocol message which can verify that it has a valid wire representation.
28pub trait ReconstructableMessage {
29    /// Reports whether this typed value can be encoded on the wire.
30    fn is_reconstructable(&self) -> bool;
31}
32
33impl ReconstructableMessage for FrontendMessage {
34    fn is_reconstructable(&self) -> bool {
35        self.to_frame().is_ok()
36    }
37}
38
39impl ReconstructableMessage for BackendMessage {
40    fn is_reconstructable(&self) -> bool {
41        self.to_frame().is_ok()
42    }
43}
44
45impl ReconstructableMessage for PreStartupMessage {
46    fn is_reconstructable(&self) -> bool {
47        self.to_packet().is_ok()
48    }
49}
50
51impl ReconstructableMessage for EncryptionReply {
52    fn is_reconstructable(&self) -> bool {
53        true
54    }
55}
56
57/// Validated backend traffic which does not advance the current protocol phase.
58pub struct AsynchronousBackendMessage(BackendMessage);
59
60impl AsynchronousBackendMessage {
61    /// Borrows the decoded asynchronous backend message.
62    #[must_use]
63    pub const fn as_wire(&self) -> &BackendMessage {
64        &self.0
65    }
66
67    /// Returns the decoded asynchronous backend message.
68    #[must_use]
69    pub fn into_wire(self) -> BackendMessage {
70        self.0
71    }
72}
73
74impl TryFrom<BackendMessage> for AsynchronousBackendMessage {
75    type Error = BackendMessage;
76
77    fn try_from(message: BackendMessage) -> Result<Self, Self::Error> {
78        if Demux::is_asynchronous(&message) {
79            Ok(Self(message))
80        } else {
81            Err(message)
82        }
83    }
84}
85
86/// Any server message legal in a phase, including non-advancing asynchronous traffic.
87pub enum TypedBackendMessage<ProtocolMessage> {
88    /// A message represented by a transition in the current grammar phase.
89    Protocol(ProtocolMessage),
90    /// An asynchronous message which leaves the current grammar phase unchanged.
91    Asynchronous(AsynchronousBackendMessage),
92}
93
94impl<ProtocolMessage> AsRef<BackendMessage> for TypedBackendMessage<ProtocolMessage>
95where
96    ProtocolMessage: AsRef<BackendMessage>,
97{
98    fn as_ref(&self) -> &BackendMessage {
99        match self {
100            Self::Protocol(message) => message.as_ref(),
101            Self::Asynchronous(message) => message.as_wire(),
102        }
103    }
104}
105
106impl<ProtocolMessage> TryFrom<BackendMessage> for TypedBackendMessage<ProtocolMessage>
107where
108    ProtocolMessage: TryFrom<BackendMessage, Error = BackendMessage>,
109{
110    type Error = BackendMessage;
111
112    fn try_from(message: BackendMessage) -> Result<Self, Self::Error> {
113        match AsynchronousBackendMessage::try_from(message) {
114            Ok(message) => Ok(Self::Asynchronous(message)),
115            Err(message) => ProtocolMessage::try_from(message).map(Self::Protocol),
116        }
117    }
118}
119
120impl<ProtocolMessage> From<TypedBackendMessage<ProtocolMessage>> for BackendMessage
121where
122    ProtocolMessage: Into<Self>,
123{
124    fn from(message: TypedBackendMessage<ProtocolMessage>) -> Self {
125        match message {
126            TypedBackendMessage::Protocol(message) => message.into(),
127            TypedBackendMessage::Asynchronous(message) => message.into_wire(),
128        }
129    }
130}
131
132macro_rules! projected_messages {
133    ($state:path, $internal:ty, $external:ty, $project_internal:path, $project_external:path) => {
134        impl AcceptsMessage<$internal> for $state {
135            fn accepts(&self, message: &$internal) -> bool {
136                $project_internal(*self, message).is_some()
137            }
138        }
139
140        impl AcceptsMessage<$external> for $state {
141            fn accepts(&self, message: &$external) -> bool {
142                $project_external(*self, message).is_some()
143            }
144        }
145    };
146}
147
148projected_messages!(
149    pre_startup::RuntimeState,
150    PreStartupMessage,
151    EncryptionReply,
152    pre_startup::project_internal,
153    pre_startup::project_external
154);
155projected_messages!(
156    server_pre_startup::RuntimeState,
157    EncryptionReply,
158    PreStartupMessage,
159    server_pre_startup::project_internal,
160    server_pre_startup::project_external
161);
162projected_messages!(
163    authentication::RuntimeState,
164    FrontendMessage,
165    BackendMessage,
166    authentication::project_internal,
167    authentication::project_external
168);
169projected_messages!(
170    server_authentication::RuntimeState,
171    BackendMessage,
172    FrontendMessage,
173    server_authentication::project_internal,
174    server_authentication::project_external
175);
176
177impl AcceptsMessage<FrontendMessage> for frontend::RuntimeState {
178    fn accepts(&self, message: &FrontendMessage) -> bool {
179        frontend::project_internal(*self, message).is_some()
180    }
181}
182
183impl AcceptsMessage<BackendMessage> for frontend::RuntimeState {
184    fn accepts(&self, message: &BackendMessage) -> bool {
185        Demux::is_asynchronous(message) || frontend::project_external(*self, message).is_some()
186    }
187}
188
189impl AcceptsMessage<BackendMessage> for backend::RuntimeState {
190    fn accepts(&self, message: &BackendMessage) -> bool {
191        Demux::is_asynchronous(message) || backend::project_internal(*self, message).is_some()
192    }
193}
194
195impl AcceptsMessage<FrontendMessage> for backend::RuntimeState {
196    fn accepts(&self, message: &FrontendMessage) -> bool {
197        backend::project_external(*self, message).is_some()
198    }
199}
200
201/// Asynchronously intercepts an owned message with access to caller-defined state.
202///
203/// The message type determines the direction at compile time: middleware over
204/// `FrontendMessage` cannot accidentally return a `BackendMessage`, and vice
205/// versa.
206#[allow(async_fn_in_trait)]
207pub trait MessageMiddleware<Message, State> {
208    /// An error which prevents the message from continuing through the chain.
209    type Error;
210
211    /// Observes, mutates, or replaces one message and may await external policy,
212    /// storage, or telemetry work before returning it.
213    ///
214    /// # Errors
215    ///
216    /// Returns a policy-defined error to stop message processing.
217    async fn intercept(
218        &mut self,
219        state: &mut State,
220        message: Message,
221    ) -> Result<Message, Self::Error>;
222}
223
224/// Marker for middleware handling messages sent by a PostgreSQL client.
225#[derive(Clone, Copy, Debug, Eq, PartialEq)]
226pub enum ClientRole {}
227
228/// Marker for middleware handling messages sent by a PostgreSQL server.
229#[derive(Clone, Copy, Debug, Eq, PartialEq)]
230pub enum ServerRole {}
231
232/// Associates a connection typestate with its generated legal message type.
233///
234/// Implementations are provided only for matching sender roles and decoded wire
235/// directions. This is the bridge which lets [`crate::Conn`] infer middleware's
236/// `Role`, `ProtocolPhase`, and `Message` indices from its own phase parameter.
237pub trait TypedPhase<Role, Wire> {
238    /// Generated grammar phase corresponding to the connection typestate.
239    type ProtocolPhase;
240    /// Opaque set of decoded messages legal for this role and phase.
241    type Message: AsRef<Wire> + TryFrom<Wire, Error = Wire> + Into<Wire>;
242}
243
244/// Associates a connection typestate with messages its local role may send.
245///
246/// Unlike [`TypedPhase`], which describes peer-selected input, this trait indexes
247/// the generated internal message set used before a locally generated value is
248/// encoded and sent.
249pub trait TypedOutboundPhase<Role, Wire> {
250    /// Generated grammar phase corresponding to the connection typestate.
251    type ProtocolPhase;
252    /// Opaque set of locally generated messages legal in this phase.
253    type Message: AsRef<Wire> + TryFrom<Wire, Error = Wire> + Into<Wire>;
254}
255
256macro_rules! typed_outbound_phase {
257    ($role:ty, $wire:ty; $($connection:ty => $protocol:path, $message:path);+ $(;)?) => {
258        $(
259            impl TypedOutboundPhase<$role, $wire> for $connection {
260                type ProtocolPhase = $protocol;
261                type Message = $message;
262            }
263        )+
264    };
265}
266
267macro_rules! typed_outbound_backend_phase {
268    ($($connection:ty => $protocol:path, $message:path);+ $(;)?) => {
269        $(
270            impl TypedOutboundPhase<ServerRole, BackendMessage> for $connection {
271                type ProtocolPhase = $protocol;
272                type Message = TypedBackendMessage<$message>;
273            }
274        )+
275    };
276}
277
278/// Authoritative connection-typestate to generated-grammar association catalogue.
279mod grammar_associations {
280    use super::{
281        BackendMessage, ClientRole, EncryptionReply, FrontendMessage, PreStartupMessage,
282        ServerRole, TypedBackendMessage, TypedOutboundPhase, TypedPhase, authentication, backend,
283        frontend, pre_startup, server_authentication, server_pre_startup,
284    };
285
286    typed_outbound_phase!(ClientRole, PreStartupMessage;
287        crate::pre_startup::PreStartup => pre_startup::PreStartup, pre_startup::PreStartupInternalMessage;
288    );
289
290    typed_outbound_phase!(ClientRole, FrontendMessage;
291        crate::auth::PasswordResponse => authentication::PasswordResponse, authentication::PasswordResponseInternalMessage;
292        crate::auth::TokenResponse => authentication::TokenResponse, authentication::TokenResponseInternalMessage;
293        crate::auth::SaslInitial => authentication::SaslInitial, authentication::SaslInitialInternalMessage;
294        crate::auth::SaslChallenge => authentication::SaslChallenge, authentication::SaslChallengeInternalMessage;
295        crate::auth::Ready => frontend::Ready, frontend::ReadyInternalMessage;
296        crate::session::Building => frontend::Building, frontend::BuildingInternalMessage;
297        crate::session::BoundBuilding => frontend::BoundBuilding, frontend::BoundBuildingInternalMessage;
298        crate::session::CopyIn => frontend::CopyIn, frontend::CopyInInternalMessage;
299        crate::session::CopyBoth => frontend::CopyBoth, frontend::CopyBothInternalMessage;
300        crate::session::CopyBothServerDone => frontend::CopyBothServerDone, frontend::CopyBothServerDoneInternalMessage;
301    );
302
303    typed_outbound_phase!(ServerRole, EncryptionReply;
304        crate::pre_startup::ServerSslDecision => server_pre_startup::SslDecision, server_pre_startup::SslDecisionInternalMessage;
305        crate::pre_startup::ServerGssDecision => server_pre_startup::GssDecision, server_pre_startup::GssDecisionInternalMessage;
306    );
307
308    typed_outbound_backend_phase!(
309        crate::server_auth::ServerStartupRejected => server_authentication::Startup, server_authentication::StartupInternalMessage;
310        crate::server_auth::ServerAuth => server_authentication::Auth, server_authentication::AuthInternalMessage;
311        crate::server_auth::ServerPassword => server_authentication::PasswordResponse, server_authentication::PasswordResponseInternalMessage;
312        crate::server_auth::ServerSaslInitial => server_authentication::SaslInitial, server_authentication::SaslInitialInternalMessage;
313        crate::server_auth::ServerSasl => server_authentication::Sasl, server_authentication::SaslInternalMessage;
314        crate::server_auth::ServerSaslResponse => server_authentication::SaslResponse, server_authentication::SaslResponseInternalMessage;
315        crate::server_auth::ServerAuthResponse => server_authentication::TokenResponse, server_authentication::TokenResponseInternalMessage;
316        crate::server_auth::ServerAuthPolicy => server_authentication::TokenPolicy, server_authentication::TokenPolicyInternalMessage;
317        crate::server_auth::ServerStartupReady => server_authentication::StartupReady, server_authentication::StartupReadyInternalMessage;
318        crate::server_session::ServerSimpleQuery => backend::Simple, backend::SimpleInternalMessage;
319        crate::server_session::ServerSimpleError => backend::SimpleError, backend::SimpleErrorInternalMessage;
320        crate::server_session::ServerFunctionCall => backend::FunctionResponse, backend::FunctionResponseInternalMessage;
321        crate::server_session::ServerFunctionCallDone => backend::FunctionReady, backend::FunctionReadyInternalMessage;
322        crate::server_session::ServerFunctionCallError => backend::FunctionReady, backend::FunctionReadyInternalMessage;
323        crate::server_session::ServerParse => backend::ParseResponse, backend::ParseResponseInternalMessage;
324        crate::server_session::ServerBind => backend::BindResponse, backend::BindResponseInternalMessage;
325        crate::server_session::ServerDescribe => backend::DescribeResponse, backend::DescribeResponseInternalMessage;
326        crate::server_session::ServerExecute => backend::ExecuteResponse, backend::ExecuteResponseInternalMessage;
327        crate::server_session::ServerClose => backend::CloseResponse, backend::CloseResponseInternalMessage;
328        crate::server_session::ServerSync => backend::SyncResponse, backend::SyncResponseInternalMessage;
329        crate::server_session::ServerBuilding => backend::Building, backend::BuildingInternalMessage;
330        crate::server_session::ServerExtendedError => backend::ExtendedError, backend::ExtendedErrorInternalMessage;
331        crate::server_session::ServerCopyIn<crate::server_session::CopySimple> => backend::SimpleCopyIn, backend::SimpleCopyInInternalMessage;
332        crate::server_session::ServerCopyIn<crate::server_session::CopyExtended> => backend::ExtendedCopyIn, backend::ExtendedCopyInInternalMessage;
333        crate::server_session::ServerCopyInDone<crate::server_session::CopySimple> => backend::SimpleCopyInDone, backend::SimpleCopyInDoneInternalMessage;
334        crate::server_session::ServerCopyInDone<crate::server_session::CopyExtended> => backend::ExtendedCopyInDone, backend::ExtendedCopyInDoneInternalMessage;
335        crate::server_session::ServerCopyInFailed<crate::server_session::CopySimple> => backend::SimpleCopyInFailed, backend::SimpleCopyInFailedInternalMessage;
336        crate::server_session::ServerCopyInFailed<crate::server_session::CopyExtended> => backend::ExtendedCopyInFailed, backend::ExtendedCopyInFailedInternalMessage;
337        crate::server_session::ServerCopyOut<crate::server_session::CopySimple> => backend::SimpleCopyOut, backend::SimpleCopyOutInternalMessage;
338        crate::server_session::ServerCopyOut<crate::server_session::CopyExtended> => backend::ExtendedCopyOut, backend::ExtendedCopyOutInternalMessage;
339        crate::server_session::ServerCopyOutDone<crate::server_session::CopySimple> => backend::SimpleCopyOutDone, backend::SimpleCopyOutDoneInternalMessage;
340        crate::server_session::ServerCopyOutDone<crate::server_session::CopyExtended> => backend::ExtendedCopyOutDone, backend::ExtendedCopyOutDoneInternalMessage;
341        crate::server_session::ServerCopyBoth<crate::server_session::CopySimple, crate::server_session::BothOpen> => backend::SimpleCopyBoth, backend::SimpleCopyBothInternalMessage;
342        crate::server_session::ServerCopyBoth<crate::server_session::CopyExtended, crate::server_session::BothOpen> => backend::ExtendedCopyBoth, backend::ExtendedCopyBothInternalMessage;
343        crate::server_session::ServerCopyBoth<crate::server_session::CopySimple, crate::server_session::BothClientDone> => backend::SimpleCopyBothClientDone, backend::SimpleCopyBothClientDoneInternalMessage;
344        crate::server_session::ServerCopyBoth<crate::server_session::CopyExtended, crate::server_session::BothClientDone> => backend::ExtendedCopyBothClientDone, backend::ExtendedCopyBothClientDoneInternalMessage;
345        crate::server_session::ServerCopyBoth<crate::server_session::CopySimple, crate::server_session::BothServerDone> => backend::SimpleCopyBothServerDone, backend::SimpleCopyBothServerDoneInternalMessage;
346        crate::server_session::ServerCopyBoth<crate::server_session::CopyExtended, crate::server_session::BothServerDone> => backend::ExtendedCopyBothServerDone, backend::ExtendedCopyBothServerDoneInternalMessage;
347        crate::server_session::ServerCopyBoth<crate::server_session::CopySimple, crate::server_session::BothDone> => backend::SimpleCopyBothDone, backend::SimpleCopyBothDoneInternalMessage;
348        crate::server_session::ServerCopyBoth<crate::server_session::CopyExtended, crate::server_session::BothDone> => backend::ExtendedCopyBothDone, backend::ExtendedCopyBothDoneInternalMessage;
349        crate::server_session::ServerCopyBothFailed<crate::server_session::CopySimple> => backend::SimpleCopyBothFailed, backend::SimpleCopyBothFailedInternalMessage;
350        crate::server_session::ServerCopyBothFailed<crate::server_session::CopyExtended> => backend::ExtendedCopyBothFailed, backend::ExtendedCopyBothFailedInternalMessage;
351    );
352
353    impl TypedPhase<ServerRole, BackendMessage> for crate::auth::Ready {
354        type ProtocolPhase = frontend::Ready;
355        type Message = TypedBackendMessage<frontend::ReadyExternalMessage>;
356    }
357
358    impl TypedPhase<ClientRole, FrontendMessage> for crate::auth::Ready {
359        type ProtocolPhase = backend::Ready;
360        type Message = backend::ReadyExternalMessage;
361    }
362
363    impl TypedPhase<ClientRole, PreStartupMessage> for crate::pre_startup::PreStartup {
364        type ProtocolPhase = server_pre_startup::PreStartup;
365        type Message = server_pre_startup::PreStartupExternalMessage;
366    }
367
368    impl TypedPhase<ServerRole, EncryptionReply> for crate::pre_startup::AwaitingSslReply {
369        type ProtocolPhase = pre_startup::AwaitingSslReply;
370        type Message = pre_startup::AwaitingSslReplyExternalMessage;
371    }
372
373    impl TypedPhase<ServerRole, EncryptionReply> for crate::pre_startup::AwaitingGssReply {
374        type ProtocolPhase = pre_startup::AwaitingGssReply;
375        type Message = pre_startup::AwaitingGssReplyExternalMessage;
376    }
377
378    macro_rules! typed_backend_phase {
379        ($connection:path => $protocol:path, $message:path) => {
380            impl TypedPhase<ServerRole, BackendMessage> for $connection {
381                type ProtocolPhase = $protocol;
382                type Message = TypedBackendMessage<$message>;
383            }
384        };
385    }
386
387    typed_backend_phase!(crate::auth::Auth => authentication::Auth, authentication::AuthExternalMessage);
388    typed_backend_phase!(crate::auth::TokenChallenge => authentication::TokenChallenge, authentication::TokenChallengeExternalMessage);
389    typed_backend_phase!(crate::auth::Sasl => authentication::Sasl, authentication::SaslExternalMessage);
390    typed_backend_phase!(crate::auth::AwaitingAuthOk => authentication::AwaitingAuthOk, authentication::AwaitingAuthOkExternalMessage);
391    typed_backend_phase!(crate::auth::AwaitingStartupReady => authentication::AwaitingStartupReady, authentication::AwaitingStartupReadyExternalMessage);
392    typed_backend_phase!(crate::session::SimpleQuery => frontend::Simple, frontend::SimpleExternalMessage);
393    typed_backend_phase!(crate::session::FunctionCalling => frontend::FunctionCalling, frontend::FunctionCallingExternalMessage);
394    typed_backend_phase!(crate::session::Building => frontend::Building, frontend::BuildingExternalMessage);
395    typed_backend_phase!(crate::session::BoundBuilding => frontend::BoundBuilding, frontend::BoundBuildingExternalMessage);
396    typed_backend_phase!(crate::session::AwaitingReady => frontend::AwaitingReady, frontend::AwaitingReadyExternalMessage);
397    typed_backend_phase!(crate::session::CopyIn => frontend::CopyIn, frontend::CopyInExternalMessage);
398    typed_backend_phase!(crate::session::CopyOut => frontend::CopyOut, frontend::CopyOutExternalMessage);
399    typed_backend_phase!(crate::session::CopyBoth => frontend::CopyBoth, frontend::CopyBothExternalMessage);
400    typed_backend_phase!(crate::session::CopyBothClientDone => frontend::CopyBothClientDone, frontend::CopyBothClientDoneExternalMessage);
401    typed_backend_phase!(crate::session::CopyBothServerDone => frontend::CopyBothServerDone, frontend::CopyBothServerDoneExternalMessage);
402    typed_backend_phase!(crate::session::Draining => frontend::Draining, frontend::DrainingExternalMessage);
403    typed_backend_phase!(crate::session::Resetting => frontend::Resetting, frontend::ResettingExternalMessage);
404    typed_backend_phase!(crate::session::ResetComplete => frontend::ResetComplete, frontend::ResetCompleteExternalMessage);
405
406    macro_rules! typed_frontend_phase {
407        ($connection:ty => $protocol:path, $message:path) => {
408            impl TypedPhase<ClientRole, FrontendMessage> for $connection {
409                type ProtocolPhase = $protocol;
410                type Message = $message;
411            }
412        };
413    }
414
415    typed_frontend_phase!(crate::server_auth::ServerAuth => server_authentication::Auth, server_authentication::AuthExternalMessage);
416    typed_frontend_phase!(crate::server_auth::ServerPassword => server_authentication::PasswordResponse, server_authentication::PasswordResponseExternalMessage);
417    typed_frontend_phase!(crate::server_auth::ServerSaslInitial => server_authentication::SaslInitial, server_authentication::SaslInitialExternalMessage);
418    typed_frontend_phase!(crate::server_auth::ServerSaslResponse => server_authentication::SaslResponse, server_authentication::SaslResponseExternalMessage);
419    typed_frontend_phase!(crate::server_auth::ServerAuthResponse => server_authentication::TokenResponse, server_authentication::TokenResponseExternalMessage);
420    typed_frontend_phase!(crate::server_auth::ServerStartupReady => server_authentication::StartupReady, server_authentication::StartupReadyExternalMessage);
421    typed_frontend_phase!(crate::server_session::ServerBuilding => backend::Building, backend::BuildingExternalMessage);
422    typed_frontend_phase!(crate::server_session::ServerExtendedError => backend::ExtendedError, backend::ExtendedErrorExternalMessage);
423    typed_frontend_phase!(crate::server_session::ServerCopyIn<crate::server_session::CopySimple> => backend::SimpleCopyIn, backend::SimpleCopyInExternalMessage);
424    typed_frontend_phase!(crate::server_session::ServerCopyIn<crate::server_session::CopyExtended> => backend::ExtendedCopyIn, backend::ExtendedCopyInExternalMessage);
425    typed_frontend_phase!(crate::server_session::ServerCopyBoth<crate::server_session::CopySimple, crate::server_session::BothOpen> => backend::SimpleCopyBoth, backend::SimpleCopyBothExternalMessage);
426    typed_frontend_phase!(crate::server_session::ServerCopyBoth<crate::server_session::CopyExtended, crate::server_session::BothOpen> => backend::ExtendedCopyBoth, backend::ExtendedCopyBothExternalMessage);
427    typed_frontend_phase!(crate::server_session::ServerCopyBoth<crate::server_session::CopySimple, crate::server_session::BothServerDone> => backend::SimpleCopyBothServerDone, backend::SimpleCopyBothServerDoneExternalMessage);
428    typed_frontend_phase!(crate::server_session::ServerCopyBoth<crate::server_session::CopyExtended, crate::server_session::BothServerDone> => backend::ExtendedCopyBothServerDone, backend::ExtendedCopyBothServerDoneExternalMessage);
429}
430
431/// Async middleware whose role, protocol phase, and legal message set are type indexed.
432///
433/// `Message` should be a phase-specific message type generated by
434/// [`pg_proto_fsm::protocol`]. Such values can only be obtained after a decoded
435/// wire message has been projected into a legal transition for `Phase`, so an
436/// implementation cannot return a replacement from another role or phase.
437#[allow(async_fn_in_trait)]
438pub trait TypedMiddleware<Role, Phase, Message, State> {
439    /// An error which prevents the message from continuing through the chain.
440    type Error;
441
442    /// Observes, mutates, or replaces one phase-legal message and may await while
443    /// borrowing both the handler and caller-defined state.
444    ///
445    /// # Errors
446    ///
447    /// Returns a policy-defined error to stop message processing.
448    async fn intercept_typed(
449        &mut self,
450        state: &mut State,
451        message: Message,
452    ) -> Result<Message, Self::Error>;
453}
454
455/// Adapts one direction-wide wire middleware to every generated typed phase.
456///
457/// Messages returned by the wrapped middleware are re-projected into the same
458/// phase-specific `Message` type. This provides a pass-through default for
459/// policies which inspect only selected wire families; a replacement which is
460/// illegal in the inferred phase is returned as an error.
461pub struct WireAdapter<Wire, Handler> {
462    handler: Handler,
463    _wire: PhantomData<fn(Wire) -> Wire>,
464}
465
466impl<Wire, Handler> WireAdapter<Wire, Handler> {
467    /// Wraps direction-wide wire middleware for use at typed interception points.
468    pub const fn new(handler: Handler) -> Self {
469        Self {
470            handler,
471            _wire: PhantomData,
472        }
473    }
474
475    /// Returns the wrapped wire middleware.
476    pub fn into_inner(self) -> Handler {
477        self.handler
478    }
479}
480
481/// Failure from direction-wide middleware adapted to a typed phase.
482#[derive(Clone, Debug, Eq, PartialEq)]
483pub enum WireAdapterError<Error, Wire> {
484    /// The wrapped middleware rejected the message according to its policy.
485    Middleware(Error),
486    /// The wrapped middleware returned a wire message illegal in the typed phase.
487    IllegalReplacement(Wire),
488}
489
490impl<Role, Phase, Message, State, Wire, Handler> TypedMiddleware<Role, Phase, Message, State>
491    for WireAdapter<Wire, Handler>
492where
493    Message: Into<Wire> + TryFrom<Wire, Error = Wire>,
494    Handler: MessageMiddleware<Wire, State>,
495{
496    type Error = WireAdapterError<Handler::Error, Wire>;
497
498    async fn intercept_typed(
499        &mut self,
500        state: &mut State,
501        message: Message,
502    ) -> Result<Message, Self::Error> {
503        let message = self
504            .handler
505            .intercept(state, message.into())
506            .await
507            .map_err(WireAdapterError::Middleware)?;
508        Message::try_from(message).map_err(WireAdapterError::IllegalReplacement)
509    }
510}
511
512impl<Role, Phase, Message, State, Error, F> TypedMiddleware<Role, Phase, Message, State> for F
513where
514    F: for<'a> AsyncFnMut(&'a mut State, Message) -> Result<Message, Error>,
515{
516    type Error = Error;
517
518    async fn intercept_typed(
519        &mut self,
520        state: &mut State,
521        message: Message,
522    ) -> Result<Message, Self::Error> {
523        self(state, message).await
524    }
525}
526
527/// Adds composition to every sized middleware implementation.
528pub trait MessageMiddlewareExt: Sized {
529    /// Runs this value followed by `next` whenever both implement middleware for
530    /// the intercepted message and state types.
531    fn then<Next>(self, next: Next) -> Then<Self, Next> {
532        Then {
533            first: self,
534            second: next,
535        }
536    }
537}
538
539impl<Handler> MessageMiddlewareExt for Handler {}
540
541impl<Message, State, Error, F> MessageMiddleware<Message, State> for F
542where
543    F: for<'a> AsyncFnMut(&'a mut State, Message) -> Result<Message, Error>,
544{
545    type Error = Error;
546
547    async fn intercept(
548        &mut self,
549        state: &mut State,
550        message: Message,
551    ) -> Result<Message, Self::Error> {
552        self(state, message).await
553    }
554}
555
556/// Middleware which returns every message unchanged.
557#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
558pub struct Identity;
559
560impl<Message, State> MessageMiddleware<Message, State> for Identity {
561    type Error = Infallible;
562
563    async fn intercept(
564        &mut self,
565        _state: &mut State,
566        message: Message,
567    ) -> Result<Message, Self::Error> {
568        Ok(message)
569    }
570}
571
572impl<Role, Phase, Message, State> TypedMiddleware<Role, Phase, Message, State> for Identity {
573    type Error = Infallible;
574
575    async fn intercept_typed(
576        &mut self,
577        _state: &mut State,
578        message: Message,
579    ) -> Result<Message, Self::Error> {
580        Ok(message)
581    }
582}
583
584/// Two middleware stages evaluated from `first` to `second`.
585#[derive(Clone, Copy, Debug, Eq, PartialEq)]
586pub struct Then<First, Second> {
587    first: First,
588    second: Second,
589}
590
591impl<First, Second> Then<First, Second> {
592    pub(crate) fn parts_mut(&mut self) -> (&mut First, &mut Second) {
593        (&mut self.first, &mut self.second)
594    }
595}
596
597impl<Message, State, First, Second> MessageMiddleware<Message, State> for Then<First, Second>
598where
599    First: MessageMiddleware<Message, State>,
600    Second: MessageMiddleware<Message, State>,
601{
602    type Error = ChainError<First::Error, Second::Error>;
603
604    async fn intercept(
605        &mut self,
606        state: &mut State,
607        message: Message,
608    ) -> Result<Message, Self::Error> {
609        let message = self
610            .first
611            .intercept(state, message)
612            .await
613            .map_err(ChainError::First)?;
614        self.second
615            .intercept(state, message)
616            .await
617            .map_err(ChainError::Second)
618    }
619}
620
621impl<Role, Phase, Message, State, First, Second> TypedMiddleware<Role, Phase, Message, State>
622    for Then<First, Second>
623where
624    First: TypedMiddleware<Role, Phase, Message, State>,
625    Second: TypedMiddleware<Role, Phase, Message, State>,
626{
627    type Error = ChainError<First::Error, Second::Error>;
628
629    async fn intercept_typed(
630        &mut self,
631        state: &mut State,
632        message: Message,
633    ) -> Result<Message, Self::Error> {
634        let message = self
635            .first
636            .intercept_typed(state, message)
637            .await
638            .map_err(ChainError::First)?;
639        self.second
640            .intercept_typed(state, message)
641            .await
642            .map_err(ChainError::Second)
643    }
644}
645
646/// Identifies which stage of a two-part middleware chain failed.
647#[derive(Clone, Copy, Debug, Eq, PartialEq)]
648pub enum ChainError<First, Second> {
649    /// The first stage rejected the message.
650    First(First),
651    /// The second stage rejected the message.
652    Second(Second),
653}
654
655/// Failure while applying or validating middleware output.
656#[derive(Clone, Copy, Debug, Eq, PartialEq)]
657pub enum InterceptError<Error, Message> {
658    /// Middleware rejected the message according to its own policy.
659    Middleware(Error),
660    /// Middleware returned a message which is illegal in the supplied state.
661    Invalid(Message),
662}
663
664/// I/O or interception failure while receiving a middleware-checked message.
665#[derive(Debug)]
666pub enum ReceiveError<Error, Message> {
667    /// Reading or decoding the message failed.
668    Io(io::Error),
669    /// Middleware rejected the message or produced an illegal replacement.
670    Intercept(InterceptError<Error, Message>),
671}
672
673/// Failure while receiving through compile-time phase-checked middleware.
674#[derive(Debug)]
675pub enum TypedReceiveError<Error, Wire> {
676    /// Reading or decoding the message failed.
677    Io(io::Error),
678    /// The peer sent a decoded message which is illegal in the connection phase.
679    Illegal(Wire),
680    /// Middleware rejected the phase-legal message according to its policy.
681    Middleware(Error),
682    /// Middleware produced a phase-legal value with an invalid wire shape.
683    InvalidWire(Wire),
684}
685
686/// Owns user state and middleware as one reusable interception unit.
687#[derive(Clone, Copy, Debug, Eq, PartialEq)]
688pub struct Middleware<State, Handler> {
689    state: State,
690    handler: Handler,
691}
692
693impl<State, Handler> Middleware<State, Handler> {
694    /// Creates middleware with its connection- or application-local state.
695    pub const fn new(state: State, handler: Handler) -> Self {
696        Self { state, handler }
697    }
698
699    /// Borrows the accumulated user state.
700    pub const fn state(&self) -> &State {
701        &self.state
702    }
703
704    /// Mutably borrows the accumulated user state.
705    pub const fn state_mut(&mut self) -> &mut State {
706        &mut self.state
707    }
708
709    /// Borrows the middleware implementation.
710    pub const fn handler(&self) -> &Handler {
711        &self.handler
712    }
713
714    /// Mutably borrows the middleware implementation.
715    pub const fn handler_mut(&mut self) -> &mut Handler {
716        &mut self.handler
717    }
718
719    /// Separates the accumulated state from its middleware implementation.
720    pub fn into_parts(self) -> (State, Handler) {
721        (self.state, self.handler)
722    }
723
724    pub(crate) fn parts_mut(&mut self) -> (&mut State, &mut Handler) {
725        (&mut self.state, &mut self.handler)
726    }
727
728    /// Intercepts one owned message.
729    ///
730    /// # Errors
731    ///
732    /// Returns the middleware's policy-defined error.
733    pub async fn intercept<Message>(&mut self, message: Message) -> Result<Message, Handler::Error>
734    where
735        Handler: MessageMiddleware<Message, State>,
736    {
737        self.handler.intercept(&mut self.state, message).await
738    }
739
740    /// Intercepts a message whose role and legal protocol phase are type indexed.
741    ///
742    /// This operation performs no dynamic protocol-state check: `Role`, `Phase`,
743    /// and the generated `Message` type are selected together by the typed caller.
744    /// Wire-shape validation remains a separate runtime boundary after converting
745    /// the result back into its decoded wire representation.
746    ///
747    /// # Errors
748    ///
749    /// Returns the middleware's policy-defined error.
750    pub async fn intercept_typed<Role, Phase, Message>(
751        &mut self,
752        message: Message,
753    ) -> Result<Message, Handler::Error>
754    where
755        Handler: TypedMiddleware<Role, Phase, Message, State>,
756    {
757        self.handler.intercept_typed(&mut self.state, message).await
758    }
759
760    /// Intercepts a message and checks the result against `protocol_state` at runtime.
761    ///
762    /// The compiler enforces the message direction and requires `ProtocolState`
763    /// to implement [`AcceptsMessage`] for that message type. The replacement's
764    /// concrete variant and the supplied generated [`crate::grammar`] runtime
765    /// state value are dynamic, however, so protocol legality and wire
766    /// reconstructability are checked at runtime after the complete middleware
767    /// chain. Call this immediately before projecting and advancing the same
768    /// protocol state.
769    ///
770    /// # Errors
771    ///
772    /// Returns a middleware policy error, or the unchanged replacement when it
773    /// is not legal in `protocol_state`.
774    pub async fn intercept_checked<Message, ProtocolState>(
775        &mut self,
776        protocol_state: &ProtocolState,
777        message: Message,
778    ) -> Result<Message, InterceptError<Handler::Error, Message>>
779    where
780        Message: ReconstructableMessage,
781        Handler: MessageMiddleware<Message, State>,
782        ProtocolState: AcceptsMessage<Message>,
783    {
784        let message = self
785            .intercept(message)
786            .await
787            .map_err(InterceptError::Middleware)?;
788        if message.is_reconstructable() && protocol_state.accepts(&message) {
789            Ok(message)
790        } else {
791            Err(InterceptError::Invalid(message))
792        }
793    }
794}
795
796impl<Transport, Phase, Cleanliness> crate::Conn<Transport, Phase, Cleanliness> {
797    /// Intercepts one locally generated message indexed by this connection phase.
798    ///
799    /// The returned generated enum may select a different legal transition in
800    /// the same phase. Match it and apply the corresponding existing typestate
801    /// operation before encoding or forwarding the value.
802    ///
803    /// # Errors
804    ///
805    /// Returns a middleware policy error or an invalid replacement wire shape.
806    pub async fn intercept_outbound_typed<Role, Wire, State, Handler>(
807        &self,
808        middleware: &mut Middleware<State, Handler>,
809        message: <Phase as TypedOutboundPhase<Role, Wire>>::Message,
810    ) -> Result<
811        <Phase as TypedOutboundPhase<Role, Wire>>::Message,
812        TypedReceiveError<Handler::Error, Wire>,
813    >
814    where
815        Phase: TypedOutboundPhase<Role, Wire>,
816        Wire: ReconstructableMessage,
817        Handler: TypedMiddleware<
818                Role,
819                <Phase as TypedOutboundPhase<Role, Wire>>::ProtocolPhase,
820                <Phase as TypedOutboundPhase<Role, Wire>>::Message,
821                State,
822            >,
823    {
824        let message = middleware
825            .intercept_typed::<Role, <Phase as TypedOutboundPhase<Role, Wire>>::ProtocolPhase, _>(
826                message,
827            )
828            .await
829            .map_err(TypedReceiveError::Middleware)?;
830        if message.as_ref().is_reconstructable() {
831            Ok(message)
832        } else {
833            Err(TypedReceiveError::InvalidWire(message.into()))
834        }
835    }
836}
837
838#[cfg(test)]
839mod tests {
840    use std::convert::Infallible;
841
842    use bytes::Bytes;
843
844    use super::{
845        AcceptsMessage as _, ChainError, ClientRole, Identity, InterceptError,
846        MessageMiddlewareExt as _, Middleware, ServerRole, TypedOutboundPhase, TypedPhase,
847        WireAdapter,
848    };
849    use crate::{
850        Conn,
851        codec::{BackendMessage, FrontendMessage, Parse},
852        grammar::{
853            backend, pre_startup as pre_startup_grammar, server_authentication, server_pre_startup,
854        },
855        pre_startup::{EncryptionReply, PreStartupMessage},
856    };
857
858    #[tokio::test]
859    async fn identity_is_a_no_op() {
860        let mut middleware = Middleware::new((), Identity);
861        assert_eq!(
862            middleware.intercept(String::from("message")).await,
863            Ok(String::from("message"))
864        );
865    }
866
867    #[tokio::test]
868    async fn closure_can_replace_message_and_accumulate_state() {
869        let mut middleware = Middleware::new(
870            Vec::new(),
871            async |seen: &mut Vec<String>, message: String| {
872                seen.push(message.clone());
873                Ok::<_, &'static str>(message.to_uppercase())
874            },
875        );
876
877        assert_eq!(
878            middleware.intercept(String::from("hello")).await,
879            Ok(String::from("HELLO"))
880        );
881        assert_eq!(middleware.state(), &[String::from("hello")]);
882    }
883
884    #[tokio::test]
885    async fn connection_phase_indexes_locally_generated_typed_middleware() {
886        let conn = Conn::new(());
887        let message =
888            pre_startup_grammar::PreStartupInternalMessage::try_from(PreStartupMessage::SslRequest)
889                .expect("SSLRequest is legal before startup");
890        let mut middleware = Middleware::new(
891            Vec::new(),
892            async |seen: &mut Vec<&'static str>,
893                   _message: pre_startup_grammar::PreStartupInternalMessage| {
894                seen.push("outbound");
895                Ok::<_, Infallible>(
896                    pre_startup_grammar::PreStartupInternalMessage::try_from(
897                        PreStartupMessage::GssEncRequest,
898                    )
899                    .expect("GSSENCRequest is another legal pre-startup choice"),
900                )
901            },
902        );
903
904        let output = conn
905            .intercept_outbound_typed::<ClientRole, PreStartupMessage, _, _>(
906                &mut middleware,
907                message,
908            )
909            .await
910            .expect("replacement remains legal in the connection phase");
911
912        assert!(matches!(output.as_ref(), PreStartupMessage::GssEncRequest));
913        assert_eq!(middleware.state(), &["outbound"]);
914        conn.into_transport();
915    }
916
917    #[tokio::test]
918    async fn middleware_can_borrow_user_state_across_await() {
919        let handler = async |steps: &mut Vec<&'static str>, message: String| {
920            steps.push("before");
921            tokio::task::yield_now().await;
922            steps.push("after");
923            Ok::<_, Infallible>(message)
924        };
925        let mut middleware = Middleware::new(Vec::new(), handler);
926
927        assert_eq!(
928            middleware.intercept(String::from("message")).await,
929            Ok(String::from("message"))
930        );
931        assert_eq!(middleware.state(), &["before", "after"]);
932    }
933
934    #[tokio::test]
935    async fn typed_closure_replaces_only_within_its_role_and_phase() {
936        let handler = async |seen: &mut usize, _message: backend::ReadyExternalMessage| {
937            *seen += 1;
938            backend::ReadyExternalMessage::try_from(FrontendMessage::Terminate)
939                .map_err(|_| "terminate must be legal while ready")
940        };
941        let mut middleware = Middleware::new(0, handler);
942        let Ok(input) = backend::ReadyExternalMessage::try_from(FrontendMessage::Query(
943            Bytes::from_static(b"select 1"),
944        )) else {
945            panic!("query must be legal while ready");
946        };
947
948        let output = middleware
949            .intercept_typed::<ClientRole, backend::Ready, _>(input)
950            .await
951            .expect("middleware accepts the message");
952
953        assert_eq!(output.event(), backend::Event::Terminate);
954        assert!(matches!(output.into_wire(), FrontendMessage::Terminate));
955        assert_eq!(*middleware.state(), 1);
956    }
957
958    #[tokio::test]
959    async fn typed_chain_is_ordered_and_threads_shared_state() {
960        let first = async |order: &mut Vec<&'static str>,
961                           message: backend::ReadyExternalMessage| {
962            order.push("first");
963            Ok::<_, Infallible>(message)
964        };
965        let second = async |order: &mut Vec<&'static str>,
966                            message: backend::ReadyExternalMessage| {
967            order.push("second");
968            Ok::<_, Infallible>(message)
969        };
970        let mut middleware = Middleware::new(Vec::new(), first.then(second));
971        let Ok(input) = backend::ReadyExternalMessage::try_from(FrontendMessage::Terminate) else {
972            panic!("terminate must be legal while ready");
973        };
974
975        let output = middleware
976            .intercept_typed::<ClientRole, backend::Ready, _>(input)
977            .await
978            .expect("both typed stages accept the message");
979
980        assert_eq!(output.event(), backend::Event::Terminate);
981        assert_eq!(middleware.state(), &["first", "second"]);
982    }
983
984    #[tokio::test]
985    async fn wire_adapter_passes_unhandled_families_through_multiple_phases() {
986        let handler = async |seen: &mut usize, message: FrontendMessage| {
987            *seen += 1;
988            Ok::<_, Infallible>(message)
989        };
990        let mut middleware = Middleware::new(0, WireAdapter::new(handler));
991
992        let Ok(ready) = backend::ReadyExternalMessage::try_from(FrontendMessage::Terminate) else {
993            panic!("terminate must be legal while ready");
994        };
995        middleware
996            .intercept_typed::<ClientRole, backend::Ready, _>(ready)
997            .await
998            .expect("ready pass-through");
999
1000        let Ok(building) = backend::BuildingExternalMessage::try_from(FrontendMessage::Sync) else {
1001            panic!("sync must be legal while building");
1002        };
1003        middleware
1004            .intercept_typed::<ClientRole, backend::Building, _>(building)
1005            .await
1006            .expect("building pass-through");
1007
1008        assert_eq!(*middleware.state(), 2);
1009    }
1010
1011    #[tokio::test]
1012    async fn chain_passes_replacement_to_next_stage_in_order() {
1013        let first = async |order: &mut Vec<&'static str>, mut message: String| {
1014            order.push("first");
1015            message.push('1');
1016            Ok::<_, &'static str>(message)
1017        };
1018        let second = async |order: &mut Vec<&'static str>, mut message: String| {
1019            order.push("second");
1020            message.push('2');
1021            Ok::<_, u8>(message)
1022        };
1023        let mut middleware = Middleware::new(Vec::new(), first.then(second));
1024
1025        assert_eq!(
1026            middleware.intercept(String::from("m")).await,
1027            Ok(String::from("m12"))
1028        );
1029        assert_eq!(middleware.state(), &["first", "second"]);
1030    }
1031
1032    #[tokio::test]
1033    async fn chain_stops_after_first_error() {
1034        let first = async |calls: &mut usize, _message: String| {
1035            *calls += 1;
1036            Err::<String, _>("rejected")
1037        };
1038        let second = async |calls: &mut usize, message: String| {
1039            *calls += 1;
1040            Ok::<_, u8>(message)
1041        };
1042        let mut middleware = Middleware::new(0, first.then(second));
1043
1044        assert_eq!(
1045            middleware.intercept(String::from("message")).await,
1046            Err(ChainError::First("rejected"))
1047        );
1048        assert_eq!(*middleware.state(), 1);
1049    }
1050
1051    #[tokio::test]
1052    async fn checked_interception_accepts_a_legal_replacement() {
1053        let mut middleware =
1054            Middleware::new((), async |_state: &mut (), _message: FrontendMessage| {
1055                Ok::<_, Infallible>(FrontendMessage::Terminate)
1056            });
1057
1058        assert_eq!(
1059            middleware
1060                .intercept_checked(
1061                    &backend::RuntimeState::Ready,
1062                    FrontendMessage::Query(Bytes::from_static(b"select 1")),
1063                )
1064                .await,
1065            Ok(FrontendMessage::Terminate)
1066        );
1067    }
1068
1069    #[tokio::test]
1070    async fn checked_interception_returns_an_illegal_replacement() {
1071        let replacement = FrontendMessage::Parse(Parse {
1072            statement: Bytes::new(),
1073            query: Bytes::from_static(b"select 2"),
1074            parameter_types: Vec::new(),
1075        });
1076        let expected = replacement.clone();
1077        let mut middleware = Middleware::new(
1078            (),
1079            async move |_state: &mut (), _message: FrontendMessage| {
1080                Ok::<_, Infallible>(replacement.clone())
1081            },
1082        );
1083
1084        assert_eq!(
1085            middleware
1086                .intercept_checked(
1087                    &backend::RuntimeState::Simple,
1088                    FrontendMessage::Query(Bytes::from_static(b"select 1")),
1089                )
1090                .await,
1091            Err(InterceptError::Invalid(expected))
1092        );
1093    }
1094
1095    #[test]
1096    fn generated_states_cover_authentication_extended_query_copy_and_replication() {
1097        let password = FrontendMessage::PasswordResponse(Bytes::from_static(b"secret"));
1098        assert!(server_authentication::RuntimeState::PasswordResponse.accepts(&password));
1099        assert!(
1100            !server_authentication::RuntimeState::PasswordResponse
1101                .accepts(&FrontendMessage::Query(Bytes::from_static(b"select 1")))
1102        );
1103
1104        let parse = FrontendMessage::Parse(Parse {
1105            statement: Bytes::from_static(b"statement"),
1106            query: Bytes::from_static(b"select 1"),
1107            parameter_types: Vec::new(),
1108        });
1109        assert!(backend::RuntimeState::Building.accepts(&parse));
1110        assert!(
1111            !backend::RuntimeState::Building
1112                .accepts(&FrontendMessage::Query(Bytes::from_static(b"select 1")))
1113        );
1114        assert!(backend::RuntimeState::ExtendedError.accepts(&parse));
1115        assert!(backend::RuntimeState::ExtendedError.accepts(&FrontendMessage::Sync));
1116
1117        let copy = FrontendMessage::CopyData(Bytes::from_static(b"data"));
1118        assert!(backend::RuntimeState::SimpleCopyIn.accepts(&copy));
1119        assert!(backend::RuntimeState::ExtendedCopyBoth.accepts(&copy));
1120        assert!(
1121            !backend::RuntimeState::ExtendedCopyBoth
1122                .accepts(&FrontendMessage::Query(Bytes::from_static(b"select 1")))
1123        );
1124
1125        assert!(
1126            server_pre_startup::RuntimeState::PreStartup.accepts(&PreStartupMessage::SslRequest)
1127        );
1128        assert!(
1129            !server_pre_startup::RuntimeState::SslDecision.accepts(&PreStartupMessage::SslRequest)
1130        );
1131    }
1132
1133    #[test]
1134    #[allow(clippy::too_many_lines)]
1135    fn grammar_catalogue_covers_inbound_and_outbound_typestate_indices() {
1136        fn inbound<Connection, Role, Wire>()
1137        where
1138            Connection: TypedPhase<Role, Wire>,
1139        {
1140        }
1141        fn outbound<Connection, Role, Wire>()
1142        where
1143            Connection: TypedOutboundPhase<Role, Wire>,
1144        {
1145        }
1146
1147        inbound::<crate::auth::Ready, ServerRole, BackendMessage>();
1148        inbound::<crate::auth::Ready, ClientRole, FrontendMessage>();
1149        outbound::<crate::auth::Ready, ClientRole, FrontendMessage>();
1150        inbound::<crate::pre_startup::PreStartup, ClientRole, PreStartupMessage>();
1151        outbound::<crate::pre_startup::PreStartup, ClientRole, PreStartupMessage>();
1152        inbound::<crate::server_auth::ServerAuth, ClientRole, FrontendMessage>();
1153        outbound::<crate::server_auth::ServerAuth, ServerRole, BackendMessage>();
1154        inbound::<crate::session::CopyBoth, ServerRole, BackendMessage>();
1155        outbound::<crate::session::CopyBoth, ClientRole, FrontendMessage>();
1156
1157        inbound::<crate::auth::Auth, ServerRole, BackendMessage>();
1158        inbound::<crate::auth::TokenChallenge, ServerRole, BackendMessage>();
1159        inbound::<crate::auth::Sasl, ServerRole, BackendMessage>();
1160        inbound::<crate::auth::AwaitingAuthOk, ServerRole, BackendMessage>();
1161        inbound::<crate::auth::AwaitingStartupReady, ServerRole, BackendMessage>();
1162        inbound::<crate::session::SimpleQuery, ServerRole, BackendMessage>();
1163        inbound::<crate::session::FunctionCalling, ServerRole, BackendMessage>();
1164        inbound::<crate::session::Building, ServerRole, BackendMessage>();
1165        inbound::<crate::session::BoundBuilding, ServerRole, BackendMessage>();
1166        inbound::<crate::session::AwaitingReady, ServerRole, BackendMessage>();
1167        inbound::<crate::session::CopyIn, ServerRole, BackendMessage>();
1168        inbound::<crate::session::CopyOut, ServerRole, BackendMessage>();
1169        inbound::<crate::session::CopyBothClientDone, ServerRole, BackendMessage>();
1170        inbound::<crate::session::CopyBothServerDone, ServerRole, BackendMessage>();
1171        inbound::<crate::session::Draining, ServerRole, BackendMessage>();
1172        inbound::<crate::session::Resetting, ServerRole, BackendMessage>();
1173        inbound::<crate::session::ResetComplete, ServerRole, BackendMessage>();
1174
1175        outbound::<crate::auth::PasswordResponse, ClientRole, FrontendMessage>();
1176        outbound::<crate::auth::TokenResponse, ClientRole, FrontendMessage>();
1177        outbound::<crate::auth::SaslInitial, ClientRole, FrontendMessage>();
1178        outbound::<crate::auth::SaslChallenge, ClientRole, FrontendMessage>();
1179        outbound::<crate::session::Building, ClientRole, FrontendMessage>();
1180        outbound::<crate::session::BoundBuilding, ClientRole, FrontendMessage>();
1181        outbound::<crate::session::CopyIn, ClientRole, FrontendMessage>();
1182        outbound::<crate::session::CopyBothServerDone, ClientRole, FrontendMessage>();
1183        outbound::<crate::pre_startup::ServerSslDecision, ServerRole, EncryptionReply>();
1184        outbound::<crate::pre_startup::ServerGssDecision, ServerRole, EncryptionReply>();
1185
1186        inbound::<crate::pre_startup::AwaitingSslReply, ServerRole, EncryptionReply>();
1187        inbound::<crate::pre_startup::AwaitingGssReply, ServerRole, EncryptionReply>();
1188        inbound::<crate::server_auth::ServerPassword, ClientRole, FrontendMessage>();
1189        inbound::<crate::server_auth::ServerSaslInitial, ClientRole, FrontendMessage>();
1190        inbound::<crate::server_auth::ServerSaslResponse, ClientRole, FrontendMessage>();
1191        inbound::<crate::server_auth::ServerAuthResponse, ClientRole, FrontendMessage>();
1192        inbound::<crate::server_auth::ServerStartupReady, ClientRole, FrontendMessage>();
1193        inbound::<crate::server_session::ServerBuilding, ClientRole, FrontendMessage>();
1194        inbound::<crate::server_session::ServerExtendedError, ClientRole, FrontendMessage>();
1195        inbound::<
1196            crate::server_session::ServerCopyIn<crate::server_session::CopySimple>,
1197            ClientRole,
1198            FrontendMessage,
1199        >();
1200        inbound::<
1201            crate::server_session::ServerCopyIn<crate::server_session::CopyExtended>,
1202            ClientRole,
1203            FrontendMessage,
1204        >();
1205        inbound::<
1206            crate::server_session::ServerCopyBoth<
1207                crate::server_session::CopySimple,
1208                crate::server_session::BothOpen,
1209            >,
1210            ClientRole,
1211            FrontendMessage,
1212        >();
1213        inbound::<
1214            crate::server_session::ServerCopyBoth<
1215                crate::server_session::CopyExtended,
1216                crate::server_session::BothOpen,
1217            >,
1218            ClientRole,
1219            FrontendMessage,
1220        >();
1221        inbound::<
1222            crate::server_session::ServerCopyBoth<
1223                crate::server_session::CopySimple,
1224                crate::server_session::BothServerDone,
1225            >,
1226            ClientRole,
1227            FrontendMessage,
1228        >();
1229        inbound::<
1230            crate::server_session::ServerCopyBoth<
1231                crate::server_session::CopyExtended,
1232                crate::server_session::BothServerDone,
1233            >,
1234            ClientRole,
1235            FrontendMessage,
1236        >();
1237
1238        outbound::<crate::server_auth::ServerStartupRejected, ServerRole, BackendMessage>();
1239        outbound::<crate::server_auth::ServerPassword, ServerRole, BackendMessage>();
1240        outbound::<crate::server_auth::ServerSaslInitial, ServerRole, BackendMessage>();
1241        outbound::<crate::server_auth::ServerSasl, ServerRole, BackendMessage>();
1242        outbound::<crate::server_auth::ServerSaslResponse, ServerRole, BackendMessage>();
1243        outbound::<crate::server_auth::ServerAuthResponse, ServerRole, BackendMessage>();
1244        outbound::<crate::server_auth::ServerAuthPolicy, ServerRole, BackendMessage>();
1245        outbound::<crate::server_auth::ServerStartupReady, ServerRole, BackendMessage>();
1246        outbound::<crate::server_session::ServerSimpleQuery, ServerRole, BackendMessage>();
1247        outbound::<crate::server_session::ServerSimpleError, ServerRole, BackendMessage>();
1248        outbound::<crate::server_session::ServerFunctionCall, ServerRole, BackendMessage>();
1249        outbound::<crate::server_session::ServerFunctionCallDone, ServerRole, BackendMessage>();
1250        outbound::<crate::server_session::ServerFunctionCallError, ServerRole, BackendMessage>();
1251        outbound::<crate::server_session::ServerParse, ServerRole, BackendMessage>();
1252        outbound::<crate::server_session::ServerBind, ServerRole, BackendMessage>();
1253        outbound::<crate::server_session::ServerDescribe, ServerRole, BackendMessage>();
1254        outbound::<crate::server_session::ServerExecute, ServerRole, BackendMessage>();
1255        outbound::<crate::server_session::ServerClose, ServerRole, BackendMessage>();
1256        outbound::<crate::server_session::ServerSync, ServerRole, BackendMessage>();
1257        outbound::<crate::server_session::ServerBuilding, ServerRole, BackendMessage>();
1258        outbound::<crate::server_session::ServerExtendedError, ServerRole, BackendMessage>();
1259        outbound::<
1260            crate::server_session::ServerCopyIn<crate::server_session::CopySimple>,
1261            ServerRole,
1262            BackendMessage,
1263        >();
1264        outbound::<
1265            crate::server_session::ServerCopyIn<crate::server_session::CopyExtended>,
1266            ServerRole,
1267            BackendMessage,
1268        >();
1269        outbound::<
1270            crate::server_session::ServerCopyInDone<crate::server_session::CopySimple>,
1271            ServerRole,
1272            BackendMessage,
1273        >();
1274        outbound::<
1275            crate::server_session::ServerCopyInDone<crate::server_session::CopyExtended>,
1276            ServerRole,
1277            BackendMessage,
1278        >();
1279        outbound::<
1280            crate::server_session::ServerCopyInFailed<crate::server_session::CopySimple>,
1281            ServerRole,
1282            BackendMessage,
1283        >();
1284        outbound::<
1285            crate::server_session::ServerCopyInFailed<crate::server_session::CopyExtended>,
1286            ServerRole,
1287            BackendMessage,
1288        >();
1289        outbound::<
1290            crate::server_session::ServerCopyOut<crate::server_session::CopySimple>,
1291            ServerRole,
1292            BackendMessage,
1293        >();
1294        outbound::<
1295            crate::server_session::ServerCopyOut<crate::server_session::CopyExtended>,
1296            ServerRole,
1297            BackendMessage,
1298        >();
1299        outbound::<
1300            crate::server_session::ServerCopyOutDone<crate::server_session::CopySimple>,
1301            ServerRole,
1302            BackendMessage,
1303        >();
1304        outbound::<
1305            crate::server_session::ServerCopyOutDone<crate::server_session::CopyExtended>,
1306            ServerRole,
1307            BackendMessage,
1308        >();
1309        outbound::<
1310            crate::server_session::ServerCopyBoth<
1311                crate::server_session::CopySimple,
1312                crate::server_session::BothOpen,
1313            >,
1314            ServerRole,
1315            BackendMessage,
1316        >();
1317        outbound::<
1318            crate::server_session::ServerCopyBoth<
1319                crate::server_session::CopyExtended,
1320                crate::server_session::BothOpen,
1321            >,
1322            ServerRole,
1323            BackendMessage,
1324        >();
1325        outbound::<
1326            crate::server_session::ServerCopyBoth<
1327                crate::server_session::CopySimple,
1328                crate::server_session::BothClientDone,
1329            >,
1330            ServerRole,
1331            BackendMessage,
1332        >();
1333        outbound::<
1334            crate::server_session::ServerCopyBoth<
1335                crate::server_session::CopyExtended,
1336                crate::server_session::BothClientDone,
1337            >,
1338            ServerRole,
1339            BackendMessage,
1340        >();
1341        outbound::<
1342            crate::server_session::ServerCopyBoth<
1343                crate::server_session::CopySimple,
1344                crate::server_session::BothServerDone,
1345            >,
1346            ServerRole,
1347            BackendMessage,
1348        >();
1349        outbound::<
1350            crate::server_session::ServerCopyBoth<
1351                crate::server_session::CopyExtended,
1352                crate::server_session::BothServerDone,
1353            >,
1354            ServerRole,
1355            BackendMessage,
1356        >();
1357        outbound::<
1358            crate::server_session::ServerCopyBoth<
1359                crate::server_session::CopySimple,
1360                crate::server_session::BothDone,
1361            >,
1362            ServerRole,
1363            BackendMessage,
1364        >();
1365        outbound::<
1366            crate::server_session::ServerCopyBoth<
1367                crate::server_session::CopyExtended,
1368                crate::server_session::BothDone,
1369            >,
1370            ServerRole,
1371            BackendMessage,
1372        >();
1373        outbound::<
1374            crate::server_session::ServerCopyBothFailed<crate::server_session::CopySimple>,
1375            ServerRole,
1376            BackendMessage,
1377        >();
1378        outbound::<
1379            crate::server_session::ServerCopyBothFailed<crate::server_session::CopyExtended>,
1380            ServerRole,
1381            BackendMessage,
1382        >();
1383    }
1384
1385    #[tokio::test]
1386    async fn checked_interception_rejects_an_unencodable_message() {
1387        let invalid = FrontendMessage::Parse(Parse {
1388            statement: Bytes::from_static(b"invalid\0name"),
1389            query: Bytes::from_static(b"select 1"),
1390            parameter_types: Vec::new(),
1391        });
1392        let expected = invalid.clone();
1393        let mut middleware = Middleware::new((), async move |_state: &mut (), _message| {
1394            Ok::<_, Infallible>(invalid.clone())
1395        });
1396
1397        assert_eq!(
1398            middleware
1399                .intercept_checked(&backend::RuntimeState::Ready, FrontendMessage::Terminate)
1400                .await,
1401            Err(InterceptError::Invalid(expected))
1402        );
1403    }
1404}
pg_proto/middleware.rs

pg_proto/
middleware.rs
