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
244impl TypedPhase<ServerRole, BackendMessage> for crate::auth::Ready {
245    type ProtocolPhase = frontend::Ready;
246    type Message = TypedBackendMessage<frontend::ReadyExternalMessage>;
247}
248
249impl TypedPhase<ClientRole, FrontendMessage> for crate::auth::Ready {
250    type ProtocolPhase = backend::Ready;
251    type Message = backend::ReadyExternalMessage;
252}
253
254impl TypedPhase<ClientRole, PreStartupMessage> for crate::pre_startup::PreStartup {
255    type ProtocolPhase = server_pre_startup::PreStartup;
256    type Message = server_pre_startup::PreStartupExternalMessage;
257}
258
259impl TypedPhase<ServerRole, EncryptionReply> for crate::pre_startup::AwaitingSslReply {
260    type ProtocolPhase = pre_startup::AwaitingSslReply;
261    type Message = pre_startup::AwaitingSslReplyExternalMessage;
262}
263
264impl TypedPhase<ServerRole, EncryptionReply> for crate::pre_startup::AwaitingGssReply {
265    type ProtocolPhase = pre_startup::AwaitingGssReply;
266    type Message = pre_startup::AwaitingGssReplyExternalMessage;
267}
268
269macro_rules! typed_backend_phase {
270    ($connection:path => $protocol:path, $message:path) => {
271        impl TypedPhase<ServerRole, BackendMessage> for $connection {
272            type ProtocolPhase = $protocol;
273            type Message = TypedBackendMessage<$message>;
274        }
275    };
276}
277
278typed_backend_phase!(crate::auth::Auth => authentication::Auth, authentication::AuthExternalMessage);
279typed_backend_phase!(crate::auth::TokenChallenge => authentication::TokenChallenge, authentication::TokenChallengeExternalMessage);
280typed_backend_phase!(crate::auth::Sasl => authentication::Sasl, authentication::SaslExternalMessage);
281typed_backend_phase!(crate::auth::AwaitingAuthOk => authentication::AwaitingAuthOk, authentication::AwaitingAuthOkExternalMessage);
282typed_backend_phase!(crate::auth::AwaitingStartupReady => authentication::AwaitingStartupReady, authentication::AwaitingStartupReadyExternalMessage);
283typed_backend_phase!(crate::session::SimpleQuery => frontend::Simple, frontend::SimpleExternalMessage);
284typed_backend_phase!(crate::session::FunctionCalling => frontend::FunctionCalling, frontend::FunctionCallingExternalMessage);
285typed_backend_phase!(crate::session::Building => frontend::Building, frontend::BuildingExternalMessage);
286typed_backend_phase!(crate::session::BoundBuilding => frontend::BoundBuilding, frontend::BoundBuildingExternalMessage);
287typed_backend_phase!(crate::session::AwaitingReady => frontend::AwaitingReady, frontend::AwaitingReadyExternalMessage);
288typed_backend_phase!(crate::session::CopyIn => frontend::CopyIn, frontend::CopyInExternalMessage);
289typed_backend_phase!(crate::session::CopyOut => frontend::CopyOut, frontend::CopyOutExternalMessage);
290typed_backend_phase!(crate::session::CopyBoth => frontend::CopyBoth, frontend::CopyBothExternalMessage);
291typed_backend_phase!(crate::session::CopyBothClientDone => frontend::CopyBothClientDone, frontend::CopyBothClientDoneExternalMessage);
292typed_backend_phase!(crate::session::CopyBothServerDone => frontend::CopyBothServerDone, frontend::CopyBothServerDoneExternalMessage);
293typed_backend_phase!(crate::session::Draining => frontend::Draining, frontend::DrainingExternalMessage);
294typed_backend_phase!(crate::session::Resetting => frontend::Resetting, frontend::ResettingExternalMessage);
295typed_backend_phase!(crate::session::ResetComplete => frontend::ResetComplete, frontend::ResetCompleteExternalMessage);
296
297macro_rules! typed_frontend_phase {
298    ($connection:ty => $protocol:path, $message:path) => {
299        impl TypedPhase<ClientRole, FrontendMessage> for $connection {
300            type ProtocolPhase = $protocol;
301            type Message = $message;
302        }
303    };
304}
305
306typed_frontend_phase!(crate::server_auth::ServerAuth => server_authentication::Auth, server_authentication::AuthExternalMessage);
307typed_frontend_phase!(crate::server_auth::ServerPassword => server_authentication::PasswordResponse, server_authentication::PasswordResponseExternalMessage);
308typed_frontend_phase!(crate::server_auth::ServerSaslInitial => server_authentication::SaslInitial, server_authentication::SaslInitialExternalMessage);
309typed_frontend_phase!(crate::server_auth::ServerSasl => server_authentication::SaslResponse, server_authentication::SaslResponseExternalMessage);
310typed_frontend_phase!(crate::server_auth::ServerAuthResponse => server_authentication::TokenResponse, server_authentication::TokenResponseExternalMessage);
311typed_frontend_phase!(crate::server_auth::ServerStartupReady => server_authentication::StartupReady, server_authentication::StartupReadyExternalMessage);
312typed_frontend_phase!(crate::server_session::ServerBuilding => backend::Building, backend::BuildingExternalMessage);
313typed_frontend_phase!(crate::server_session::ServerExtendedError => backend::ExtendedError, backend::ExtendedErrorExternalMessage);
314typed_frontend_phase!(crate::server_session::ServerCopyIn<crate::server_session::CopySimple> => backend::SimpleCopyIn, backend::SimpleCopyInExternalMessage);
315typed_frontend_phase!(crate::server_session::ServerCopyIn<crate::server_session::CopyExtended> => backend::ExtendedCopyIn, backend::ExtendedCopyInExternalMessage);
316typed_frontend_phase!(crate::server_session::ServerCopyBoth<crate::server_session::CopySimple, crate::server_session::BothOpen> => backend::SimpleCopyBoth, backend::SimpleCopyBothExternalMessage);
317typed_frontend_phase!(crate::server_session::ServerCopyBoth<crate::server_session::CopyExtended, crate::server_session::BothOpen> => backend::ExtendedCopyBoth, backend::ExtendedCopyBothExternalMessage);
318typed_frontend_phase!(crate::server_session::ServerCopyBoth<crate::server_session::CopySimple, crate::server_session::BothServerDone> => backend::SimpleCopyBothServerDone, backend::SimpleCopyBothServerDoneExternalMessage);
319typed_frontend_phase!(crate::server_session::ServerCopyBoth<crate::server_session::CopyExtended, crate::server_session::BothServerDone> => backend::ExtendedCopyBothServerDone, backend::ExtendedCopyBothServerDoneExternalMessage);
320
321/// Async middleware whose role, protocol phase, and legal message set are type indexed.
322///
323/// `Message` should be a phase-specific message type generated by
324/// [`pg_proto_fsm::protocol`]. Such values can only be obtained after a decoded
325/// wire message has been projected into a legal transition for `Phase`, so an
326/// implementation cannot return a replacement from another role or phase.
327#[allow(async_fn_in_trait)]
328pub trait TypedMiddleware<Role, Phase, Message, State> {
329    /// An error which prevents the message from continuing through the chain.
330    type Error;
331
332    /// Observes, mutates, or replaces one phase-legal message and may await while
333    /// borrowing both the handler and caller-defined state.
334    ///
335    /// # Errors
336    ///
337    /// Returns a policy-defined error to stop message processing.
338    async fn intercept_typed(
339        &mut self,
340        state: &mut State,
341        message: Message,
342    ) -> Result<Message, Self::Error>;
343}
344
345/// Adapts one direction-wide wire middleware to every generated typed phase.
346///
347/// Messages returned by the wrapped middleware are re-projected into the same
348/// phase-specific `Message` type. This provides a pass-through default for
349/// policies which inspect only selected wire families; a replacement which is
350/// illegal in the inferred phase is returned as an error.
351pub struct WireAdapter<Wire, Handler> {
352    handler: Handler,
353    _wire: PhantomData<fn(Wire) -> Wire>,
354}
355
356impl<Wire, Handler> WireAdapter<Wire, Handler> {
357    /// Wraps direction-wide wire middleware for use at typed interception points.
358    pub const fn new(handler: Handler) -> Self {
359        Self {
360            handler,
361            _wire: PhantomData,
362        }
363    }
364
365    /// Returns the wrapped wire middleware.
366    pub fn into_inner(self) -> Handler {
367        self.handler
368    }
369}
370
371/// Failure from direction-wide middleware adapted to a typed phase.
372#[derive(Clone, Debug, Eq, PartialEq)]
373pub enum WireAdapterError<Error, Wire> {
374    /// The wrapped middleware rejected the message according to its policy.
375    Middleware(Error),
376    /// The wrapped middleware returned a wire message illegal in the typed phase.
377    IllegalReplacement(Wire),
378}
379
380impl<Role, Phase, Message, State, Wire, Handler> TypedMiddleware<Role, Phase, Message, State>
381    for WireAdapter<Wire, Handler>
382where
383    Message: Into<Wire> + TryFrom<Wire, Error = Wire>,
384    Handler: MessageMiddleware<Wire, State>,
385{
386    type Error = WireAdapterError<Handler::Error, Wire>;
387
388    async fn intercept_typed(
389        &mut self,
390        state: &mut State,
391        message: Message,
392    ) -> Result<Message, Self::Error> {
393        let message = self
394            .handler
395            .intercept(state, message.into())
396            .await
397            .map_err(WireAdapterError::Middleware)?;
398        Message::try_from(message).map_err(WireAdapterError::IllegalReplacement)
399    }
400}
401
402impl<Role, Phase, Message, State, Error, F> TypedMiddleware<Role, Phase, Message, State> for F
403where
404    F: for<'a> AsyncFnMut(&'a mut State, Message) -> Result<Message, Error>,
405{
406    type Error = Error;
407
408    async fn intercept_typed(
409        &mut self,
410        state: &mut State,
411        message: Message,
412    ) -> Result<Message, Self::Error> {
413        self(state, message).await
414    }
415}
416
417/// Adds composition to every sized middleware implementation.
418pub trait MessageMiddlewareExt: Sized {
419    /// Runs this value followed by `next` whenever both implement middleware for
420    /// the intercepted message and state types.
421    fn then<Next>(self, next: Next) -> Then<Self, Next> {
422        Then {
423            first: self,
424            second: next,
425        }
426    }
427}
428
429impl<Handler> MessageMiddlewareExt for Handler {}
430
431impl<Message, State, Error, F> MessageMiddleware<Message, State> for F
432where
433    F: for<'a> AsyncFnMut(&'a mut State, Message) -> Result<Message, Error>,
434{
435    type Error = Error;
436
437    async fn intercept(
438        &mut self,
439        state: &mut State,
440        message: Message,
441    ) -> Result<Message, Self::Error> {
442        self(state, message).await
443    }
444}
445
446/// Middleware which returns every message unchanged.
447#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
448pub struct Identity;
449
450impl<Message, State> MessageMiddleware<Message, State> for Identity {
451    type Error = Infallible;
452
453    async fn intercept(
454        &mut self,
455        _state: &mut State,
456        message: Message,
457    ) -> Result<Message, Self::Error> {
458        Ok(message)
459    }
460}
461
462impl<Role, Phase, Message, State> TypedMiddleware<Role, Phase, Message, State> for Identity {
463    type Error = Infallible;
464
465    async fn intercept_typed(
466        &mut self,
467        _state: &mut State,
468        message: Message,
469    ) -> Result<Message, Self::Error> {
470        Ok(message)
471    }
472}
473
474/// Two middleware stages evaluated from `first` to `second`.
475#[derive(Clone, Copy, Debug, Eq, PartialEq)]
476pub struct Then<First, Second> {
477    first: First,
478    second: Second,
479}
480
481impl<Message, State, First, Second> MessageMiddleware<Message, State> for Then<First, Second>
482where
483    First: MessageMiddleware<Message, State>,
484    Second: MessageMiddleware<Message, State>,
485{
486    type Error = ChainError<First::Error, Second::Error>;
487
488    async fn intercept(
489        &mut self,
490        state: &mut State,
491        message: Message,
492    ) -> Result<Message, Self::Error> {
493        let message = self
494            .first
495            .intercept(state, message)
496            .await
497            .map_err(ChainError::First)?;
498        self.second
499            .intercept(state, message)
500            .await
501            .map_err(ChainError::Second)
502    }
503}
504
505impl<Role, Phase, Message, State, First, Second> TypedMiddleware<Role, Phase, Message, State>
506    for Then<First, Second>
507where
508    First: TypedMiddleware<Role, Phase, Message, State>,
509    Second: TypedMiddleware<Role, Phase, Message, State>,
510{
511    type Error = ChainError<First::Error, Second::Error>;
512
513    async fn intercept_typed(
514        &mut self,
515        state: &mut State,
516        message: Message,
517    ) -> Result<Message, Self::Error> {
518        let message = self
519            .first
520            .intercept_typed(state, message)
521            .await
522            .map_err(ChainError::First)?;
523        self.second
524            .intercept_typed(state, message)
525            .await
526            .map_err(ChainError::Second)
527    }
528}
529
530/// Identifies which stage of a two-part middleware chain failed.
531#[derive(Clone, Copy, Debug, Eq, PartialEq)]
532pub enum ChainError<First, Second> {
533    /// The first stage rejected the message.
534    First(First),
535    /// The second stage rejected the message.
536    Second(Second),
537}
538
539/// Failure while applying or validating middleware output.
540#[derive(Clone, Copy, Debug, Eq, PartialEq)]
541pub enum InterceptError<Error, Message> {
542    /// Middleware rejected the message according to its own policy.
543    Middleware(Error),
544    /// Middleware returned a message which is illegal in the supplied state.
545    Invalid(Message),
546}
547
548/// I/O or interception failure while receiving a middleware-checked message.
549#[derive(Debug)]
550pub enum ReceiveError<Error, Message> {
551    /// Reading or decoding the message failed.
552    Io(io::Error),
553    /// Middleware rejected the message or produced an illegal replacement.
554    Intercept(InterceptError<Error, Message>),
555}
556
557/// Failure while receiving through compile-time phase-checked middleware.
558#[derive(Debug)]
559pub enum TypedReceiveError<Error, Wire> {
560    /// Reading or decoding the message failed.
561    Io(io::Error),
562    /// The peer sent a decoded message which is illegal in the connection phase.
563    Illegal(Wire),
564    /// Middleware rejected the phase-legal message according to its policy.
565    Middleware(Error),
566    /// Middleware produced a phase-legal value with an invalid wire shape.
567    InvalidWire(Wire),
568}
569
570/// Owns user state and middleware as one reusable interception unit.
571#[derive(Clone, Copy, Debug, Eq, PartialEq)]
572pub struct Middleware<State, Handler> {
573    state: State,
574    handler: Handler,
575}
576
577impl<State, Handler> Middleware<State, Handler> {
578    /// Creates middleware with its connection- or application-local state.
579    pub const fn new(state: State, handler: Handler) -> Self {
580        Self { state, handler }
581    }
582
583    /// Borrows the accumulated user state.
584    pub const fn state(&self) -> &State {
585        &self.state
586    }
587
588    /// Mutably borrows the accumulated user state.
589    pub const fn state_mut(&mut self) -> &mut State {
590        &mut self.state
591    }
592
593    /// Borrows the middleware implementation.
594    pub const fn handler(&self) -> &Handler {
595        &self.handler
596    }
597
598    /// Mutably borrows the middleware implementation.
599    pub const fn handler_mut(&mut self) -> &mut Handler {
600        &mut self.handler
601    }
602
603    /// Separates the accumulated state from its middleware implementation.
604    pub fn into_parts(self) -> (State, Handler) {
605        (self.state, self.handler)
606    }
607
608    /// Intercepts one owned message.
609    ///
610    /// # Errors
611    ///
612    /// Returns the middleware's policy-defined error.
613    pub async fn intercept<Message>(&mut self, message: Message) -> Result<Message, Handler::Error>
614    where
615        Handler: MessageMiddleware<Message, State>,
616    {
617        self.handler.intercept(&mut self.state, message).await
618    }
619
620    /// Intercepts a message whose role and legal protocol phase are type indexed.
621    ///
622    /// This operation performs no dynamic protocol-state check: `Role`, `Phase`,
623    /// and the generated `Message` type are selected together by the typed caller.
624    /// Wire-shape validation remains a separate runtime boundary after converting
625    /// the result back into its decoded wire representation.
626    ///
627    /// # Errors
628    ///
629    /// Returns the middleware's policy-defined error.
630    pub async fn intercept_typed<Role, Phase, Message>(
631        &mut self,
632        message: Message,
633    ) -> Result<Message, Handler::Error>
634    where
635        Handler: TypedMiddleware<Role, Phase, Message, State>,
636    {
637        self.handler.intercept_typed(&mut self.state, message).await
638    }
639
640    /// Intercepts a message and checks the result against `protocol_state` at runtime.
641    ///
642    /// The compiler enforces the message direction and requires `ProtocolState`
643    /// to implement [`AcceptsMessage`] for that message type. The replacement's
644    /// concrete variant and the supplied generated [`crate::grammar`] runtime
645    /// state value are dynamic, however, so protocol legality and wire
646    /// reconstructability are checked at runtime after the complete middleware
647    /// chain. Call this immediately before projecting and advancing the same
648    /// protocol state.
649    ///
650    /// # Errors
651    ///
652    /// Returns a middleware policy error, or the unchanged replacement when it
653    /// is not legal in `protocol_state`.
654    pub async fn intercept_checked<Message, ProtocolState>(
655        &mut self,
656        protocol_state: &ProtocolState,
657        message: Message,
658    ) -> Result<Message, InterceptError<Handler::Error, Message>>
659    where
660        Message: ReconstructableMessage,
661        Handler: MessageMiddleware<Message, State>,
662        ProtocolState: AcceptsMessage<Message>,
663    {
664        let message = self
665            .intercept(message)
666            .await
667            .map_err(InterceptError::Middleware)?;
668        if message.is_reconstructable() && protocol_state.accepts(&message) {
669            Ok(message)
670        } else {
671            Err(InterceptError::Invalid(message))
672        }
673    }
674}
675
676#[cfg(test)]
677mod tests {
678    use std::convert::Infallible;
679
680    use bytes::Bytes;
681
682    use super::{
683        AcceptsMessage as _, ChainError, ClientRole, Identity, InterceptError,
684        MessageMiddlewareExt as _, Middleware, WireAdapter,
685    };
686    use crate::{
687        codec::{FrontendMessage, Parse},
688        grammar::{backend, server_authentication, server_pre_startup},
689        pre_startup::PreStartupMessage,
690    };
691
692    #[tokio::test]
693    async fn identity_is_a_no_op() {
694        let mut middleware = Middleware::new((), Identity);
695        assert_eq!(
696            middleware.intercept(String::from("message")).await,
697            Ok(String::from("message"))
698        );
699    }
700
701    #[tokio::test]
702    async fn closure_can_replace_message_and_accumulate_state() {
703        let mut middleware = Middleware::new(
704            Vec::new(),
705            async |seen: &mut Vec<String>, message: String| {
706                seen.push(message.clone());
707                Ok::<_, &'static str>(message.to_uppercase())
708            },
709        );
710
711        assert_eq!(
712            middleware.intercept(String::from("hello")).await,
713            Ok(String::from("HELLO"))
714        );
715        assert_eq!(middleware.state(), &[String::from("hello")]);
716    }
717
718    #[tokio::test]
719    async fn middleware_can_borrow_user_state_across_await() {
720        let handler = async |steps: &mut Vec<&'static str>, message: String| {
721            steps.push("before");
722            tokio::task::yield_now().await;
723            steps.push("after");
724            Ok::<_, Infallible>(message)
725        };
726        let mut middleware = Middleware::new(Vec::new(), handler);
727
728        assert_eq!(
729            middleware.intercept(String::from("message")).await,
730            Ok(String::from("message"))
731        );
732        assert_eq!(middleware.state(), &["before", "after"]);
733    }
734
735    #[tokio::test]
736    async fn typed_closure_replaces_only_within_its_role_and_phase() {
737        let handler = async |seen: &mut usize, _message: backend::ReadyExternalMessage| {
738            *seen += 1;
739            backend::ReadyExternalMessage::try_from(FrontendMessage::Terminate)
740                .map_err(|_| "terminate must be legal while ready")
741        };
742        let mut middleware = Middleware::new(0, handler);
743        let Ok(input) = backend::ReadyExternalMessage::try_from(FrontendMessage::Query(
744            Bytes::from_static(b"select 1"),
745        )) else {
746            panic!("query must be legal while ready");
747        };
748
749        let output = middleware
750            .intercept_typed::<ClientRole, backend::Ready, _>(input)
751            .await
752            .expect("middleware accepts the message");
753
754        assert_eq!(output.event(), backend::Event::Terminate);
755        assert!(matches!(output.into_wire(), FrontendMessage::Terminate));
756        assert_eq!(*middleware.state(), 1);
757    }
758
759    #[tokio::test]
760    async fn typed_chain_is_ordered_and_threads_shared_state() {
761        let first = async |order: &mut Vec<&'static str>,
762                           message: backend::ReadyExternalMessage| {
763            order.push("first");
764            Ok::<_, Infallible>(message)
765        };
766        let second = async |order: &mut Vec<&'static str>,
767                            message: backend::ReadyExternalMessage| {
768            order.push("second");
769            Ok::<_, Infallible>(message)
770        };
771        let mut middleware = Middleware::new(Vec::new(), first.then(second));
772        let Ok(input) = backend::ReadyExternalMessage::try_from(FrontendMessage::Terminate) else {
773            panic!("terminate must be legal while ready");
774        };
775
776        let output = middleware
777            .intercept_typed::<ClientRole, backend::Ready, _>(input)
778            .await
779            .expect("both typed stages accept the message");
780
781        assert_eq!(output.event(), backend::Event::Terminate);
782        assert_eq!(middleware.state(), &["first", "second"]);
783    }
784
785    #[tokio::test]
786    async fn wire_adapter_passes_unhandled_families_through_multiple_phases() {
787        let handler = async |seen: &mut usize, message: FrontendMessage| {
788            *seen += 1;
789            Ok::<_, Infallible>(message)
790        };
791        let mut middleware = Middleware::new(0, WireAdapter::new(handler));
792
793        let Ok(ready) = backend::ReadyExternalMessage::try_from(FrontendMessage::Terminate) else {
794            panic!("terminate must be legal while ready");
795        };
796        middleware
797            .intercept_typed::<ClientRole, backend::Ready, _>(ready)
798            .await
799            .expect("ready pass-through");
800
801        let Ok(building) = backend::BuildingExternalMessage::try_from(FrontendMessage::Sync) else {
802            panic!("sync must be legal while building");
803        };
804        middleware
805            .intercept_typed::<ClientRole, backend::Building, _>(building)
806            .await
807            .expect("building pass-through");
808
809        assert_eq!(*middleware.state(), 2);
810    }
811
812    #[tokio::test]
813    async fn chain_passes_replacement_to_next_stage_in_order() {
814        let first = async |order: &mut Vec<&'static str>, mut message: String| {
815            order.push("first");
816            message.push('1');
817            Ok::<_, &'static str>(message)
818        };
819        let second = async |order: &mut Vec<&'static str>, mut message: String| {
820            order.push("second");
821            message.push('2');
822            Ok::<_, u8>(message)
823        };
824        let mut middleware = Middleware::new(Vec::new(), first.then(second));
825
826        assert_eq!(
827            middleware.intercept(String::from("m")).await,
828            Ok(String::from("m12"))
829        );
830        assert_eq!(middleware.state(), &["first", "second"]);
831    }
832
833    #[tokio::test]
834    async fn chain_stops_after_first_error() {
835        let first = async |calls: &mut usize, _message: String| {
836            *calls += 1;
837            Err::<String, _>("rejected")
838        };
839        let second = async |calls: &mut usize, message: String| {
840            *calls += 1;
841            Ok::<_, u8>(message)
842        };
843        let mut middleware = Middleware::new(0, first.then(second));
844
845        assert_eq!(
846            middleware.intercept(String::from("message")).await,
847            Err(ChainError::First("rejected"))
848        );
849        assert_eq!(*middleware.state(), 1);
850    }
851
852    #[tokio::test]
853    async fn checked_interception_accepts_a_legal_replacement() {
854        let mut middleware =
855            Middleware::new((), async |_state: &mut (), _message: FrontendMessage| {
856                Ok::<_, Infallible>(FrontendMessage::Terminate)
857            });
858
859        assert_eq!(
860            middleware
861                .intercept_checked(
862                    &backend::RuntimeState::Ready,
863                    FrontendMessage::Query(Bytes::from_static(b"select 1")),
864                )
865                .await,
866            Ok(FrontendMessage::Terminate)
867        );
868    }
869
870    #[tokio::test]
871    async fn checked_interception_returns_an_illegal_replacement() {
872        let replacement = FrontendMessage::Parse(Parse {
873            statement: Bytes::new(),
874            query: Bytes::from_static(b"select 2"),
875            parameter_types: Vec::new(),
876        });
877        let expected = replacement.clone();
878        let mut middleware = Middleware::new(
879            (),
880            async move |_state: &mut (), _message: FrontendMessage| {
881                Ok::<_, Infallible>(replacement.clone())
882            },
883        );
884
885        assert_eq!(
886            middleware
887                .intercept_checked(
888                    &backend::RuntimeState::Simple,
889                    FrontendMessage::Query(Bytes::from_static(b"select 1")),
890                )
891                .await,
892            Err(InterceptError::Invalid(expected))
893        );
894    }
895
896    #[test]
897    fn generated_states_cover_authentication_extended_query_copy_and_replication() {
898        let password = FrontendMessage::PasswordResponse(Bytes::from_static(b"secret"));
899        assert!(server_authentication::RuntimeState::PasswordResponse.accepts(&password));
900        assert!(
901            !server_authentication::RuntimeState::PasswordResponse
902                .accepts(&FrontendMessage::Query(Bytes::from_static(b"select 1")))
903        );
904
905        let parse = FrontendMessage::Parse(Parse {
906            statement: Bytes::from_static(b"statement"),
907            query: Bytes::from_static(b"select 1"),
908            parameter_types: Vec::new(),
909        });
910        assert!(backend::RuntimeState::Building.accepts(&parse));
911        assert!(
912            !backend::RuntimeState::Building
913                .accepts(&FrontendMessage::Query(Bytes::from_static(b"select 1")))
914        );
915        assert!(backend::RuntimeState::ExtendedError.accepts(&parse));
916        assert!(backend::RuntimeState::ExtendedError.accepts(&FrontendMessage::Sync));
917
918        let copy = FrontendMessage::CopyData(Bytes::from_static(b"data"));
919        assert!(backend::RuntimeState::SimpleCopyIn.accepts(&copy));
920        assert!(backend::RuntimeState::ExtendedCopyBoth.accepts(&copy));
921        assert!(
922            !backend::RuntimeState::ExtendedCopyBoth
923                .accepts(&FrontendMessage::Query(Bytes::from_static(b"select 1")))
924        );
925
926        assert!(
927            server_pre_startup::RuntimeState::PreStartup.accepts(&PreStartupMessage::SslRequest)
928        );
929        assert!(
930            !server_pre_startup::RuntimeState::SslDecision.accepts(&PreStartupMessage::SslRequest)
931        );
932    }
933
934    #[tokio::test]
935    async fn checked_interception_rejects_an_unencodable_message() {
936        let invalid = FrontendMessage::Parse(Parse {
937            statement: Bytes::from_static(b"invalid\0name"),
938            query: Bytes::from_static(b"select 1"),
939            parameter_types: Vec::new(),
940        });
941        let expected = invalid.clone();
942        let mut middleware = Middleware::new((), async move |_state: &mut (), _message| {
943            Ok::<_, Infallible>(invalid.clone())
944        });
945
946        assert_eq!(
947            middleware
948                .intercept_checked(&backend::RuntimeState::Ready, FrontendMessage::Terminate)
949                .await,
950            Err(InterceptError::Invalid(expected))
951        );
952    }
953}
pg_proto/middleware.rs

pg_proto/
middleware.rs
