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/// Marker selecting associations for messages received from a peer.
233pub enum Inbound {}
234
235/// Marker selecting associations for messages sent by the local role.
236pub enum Outbound {}
237
238/// Associates a connection typestate with a generated protocol phase and legal message set.
239///
240/// Grammar declarations generate all implementations. The direction, sender role, and decoded
241/// wire type remain explicit indices so callers can use one interface for inbound and outbound
242/// capabilities without exposing generated association machinery.
243pub trait PhaseAssociation<Direction, Role, Wire>:
244    phase_association_seal::Sealed<Direction, Role, Wire>
245{
246    /// Generated grammar phase corresponding to the connection typestate.
247    type ProtocolPhase;
248    /// Opaque set of messages legal for this direction, role, and phase.
249    type Message: AsRef<Wire> + TryFrom<Wire, Error = Wire> + Into<Wire>;
250}
251
252/// Seals [`PhaseAssociation`] while remaining nameable by generated grammar implementations.
253#[doc(hidden)]
254pub(crate) mod phase_association_seal {
255    /// Implemented alongside each generated phase association.
256    pub trait Sealed<Direction, Role, Wire> {}
257}
258
259/// Async middleware whose role, protocol phase, and legal message set are type indexed.
260///
261/// `Message` should be a phase-specific message type generated by
262/// [`pg_proto_fsm::protocol`]. Such values can only be obtained after a decoded
263/// wire message has been projected into a legal transition for `Phase`, so an
264/// implementation cannot return a replacement from another role or phase.
265#[allow(async_fn_in_trait)]
266pub trait TypedMiddleware<Role, Phase, Message, State> {
267    /// An error which prevents the message from continuing through the chain.
268    type Error;
269
270    /// Observes, mutates, or replaces one phase-legal message and may await while
271    /// borrowing both the handler and caller-defined state.
272    ///
273    /// # Errors
274    ///
275    /// Returns a policy-defined error to stop message processing.
276    async fn intercept_typed(
277        &mut self,
278        state: &mut State,
279        message: Message,
280    ) -> Result<Message, Self::Error>;
281}
282
283/// Adapts one direction-wide wire middleware to every generated typed phase.
284///
285/// Messages returned by the wrapped middleware are re-projected into the same
286/// phase-specific `Message` type. This provides a pass-through default for
287/// policies which inspect only selected wire families; a replacement which is
288/// illegal in the inferred phase is returned as an error.
289pub struct WireAdapter<Wire, Handler> {
290    handler: Handler,
291    _wire: PhantomData<fn(Wire) -> Wire>,
292}
293
294impl<Wire, Handler> WireAdapter<Wire, Handler> {
295    /// Wraps direction-wide wire middleware for use at typed interception points.
296    pub const fn new(handler: Handler) -> Self {
297        Self {
298            handler,
299            _wire: PhantomData,
300        }
301    }
302
303    /// Returns the wrapped wire middleware.
304    pub fn into_inner(self) -> Handler {
305        self.handler
306    }
307}
308
309/// Failure from direction-wide middleware adapted to a typed phase.
310#[derive(Clone, Debug, Eq, PartialEq)]
311pub enum WireAdapterError<Error, Wire> {
312    /// The wrapped middleware rejected the message according to its policy.
313    Middleware(Error),
314    /// The wrapped middleware returned a wire message illegal in the typed phase.
315    IllegalReplacement(Wire),
316}
317
318impl<Role, Phase, Message, State, Wire, Handler> TypedMiddleware<Role, Phase, Message, State>
319    for WireAdapter<Wire, Handler>
320where
321    Message: Into<Wire> + TryFrom<Wire, Error = Wire>,
322    Handler: MessageMiddleware<Wire, State>,
323{
324    type Error = WireAdapterError<Handler::Error, Wire>;
325
326    async fn intercept_typed(
327        &mut self,
328        state: &mut State,
329        message: Message,
330    ) -> Result<Message, Self::Error> {
331        let message = self
332            .handler
333            .intercept(state, message.into())
334            .await
335            .map_err(WireAdapterError::Middleware)?;
336        Message::try_from(message).map_err(WireAdapterError::IllegalReplacement)
337    }
338}
339
340impl<Role, Phase, Message, State, Error, F> TypedMiddleware<Role, Phase, Message, State> for F
341where
342    F: for<'a> AsyncFnMut(&'a mut State, Message) -> Result<Message, Error>,
343{
344    type Error = Error;
345
346    async fn intercept_typed(
347        &mut self,
348        state: &mut State,
349        message: Message,
350    ) -> Result<Message, Self::Error> {
351        self(state, message).await
352    }
353}
354
355/// Adds composition to every sized middleware implementation.
356pub trait MessageMiddlewareExt: Sized {
357    /// Runs this value followed by `next` whenever both implement middleware for
358    /// the intercepted message and state types.
359    fn then<Next>(self, next: Next) -> Then<Self, Next> {
360        Then {
361            first: self,
362            second: next,
363        }
364    }
365}
366
367impl<Handler> MessageMiddlewareExt for Handler {}
368
369impl<Message, State, Error, F> MessageMiddleware<Message, State> for F
370where
371    F: for<'a> AsyncFnMut(&'a mut State, Message) -> Result<Message, Error>,
372{
373    type Error = Error;
374
375    async fn intercept(
376        &mut self,
377        state: &mut State,
378        message: Message,
379    ) -> Result<Message, Self::Error> {
380        self(state, message).await
381    }
382}
383
384/// Middleware which returns every message unchanged.
385#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
386pub struct Identity;
387
388impl<Message, State> MessageMiddleware<Message, State> for Identity {
389    type Error = Infallible;
390
391    async fn intercept(
392        &mut self,
393        _state: &mut State,
394        message: Message,
395    ) -> Result<Message, Self::Error> {
396        Ok(message)
397    }
398}
399
400impl<Role, Phase, Message, State> TypedMiddleware<Role, Phase, Message, State> for Identity {
401    type Error = Infallible;
402
403    async fn intercept_typed(
404        &mut self,
405        _state: &mut State,
406        message: Message,
407    ) -> Result<Message, Self::Error> {
408        Ok(message)
409    }
410}
411
412/// Two middleware stages evaluated from `first` to `second`.
413#[derive(Clone, Copy, Debug, Eq, PartialEq)]
414pub struct Then<First, Second> {
415    first: First,
416    second: Second,
417}
418
419impl<First, Second> Then<First, Second> {
420    pub(crate) fn parts_mut(&mut self) -> (&mut First, &mut Second) {
421        (&mut self.first, &mut self.second)
422    }
423}
424
425impl<Message, State, First, Second> MessageMiddleware<Message, State> for Then<First, Second>
426where
427    First: MessageMiddleware<Message, State>,
428    Second: MessageMiddleware<Message, State>,
429{
430    type Error = ChainError<First::Error, Second::Error>;
431
432    async fn intercept(
433        &mut self,
434        state: &mut State,
435        message: Message,
436    ) -> Result<Message, Self::Error> {
437        let message = self
438            .first
439            .intercept(state, message)
440            .await
441            .map_err(ChainError::First)?;
442        self.second
443            .intercept(state, message)
444            .await
445            .map_err(ChainError::Second)
446    }
447}
448
449impl<Role, Phase, Message, State, First, Second> TypedMiddleware<Role, Phase, Message, State>
450    for Then<First, Second>
451where
452    First: TypedMiddleware<Role, Phase, Message, State>,
453    Second: TypedMiddleware<Role, Phase, Message, State>,
454{
455    type Error = ChainError<First::Error, Second::Error>;
456
457    async fn intercept_typed(
458        &mut self,
459        state: &mut State,
460        message: Message,
461    ) -> Result<Message, Self::Error> {
462        let message = self
463            .first
464            .intercept_typed(state, message)
465            .await
466            .map_err(ChainError::First)?;
467        self.second
468            .intercept_typed(state, message)
469            .await
470            .map_err(ChainError::Second)
471    }
472}
473
474/// Identifies which stage of a two-part middleware chain failed.
475#[derive(Clone, Copy, Debug, Eq, PartialEq)]
476pub enum ChainError<First, Second> {
477    /// The first stage rejected the message.
478    First(First),
479    /// The second stage rejected the message.
480    Second(Second),
481}
482
483/// Failure while applying or validating middleware output.
484#[derive(Clone, Copy, Debug, Eq, PartialEq)]
485pub enum InterceptError<Error, Message> {
486    /// Middleware rejected the message according to its own policy.
487    Middleware(Error),
488    /// Middleware returned a message which is illegal in the supplied state.
489    Invalid(Message),
490}
491
492/// I/O or interception failure while receiving a middleware-checked message.
493#[derive(Debug)]
494pub enum ReceiveError<Error, Message> {
495    /// Reading or decoding the message failed.
496    Io(io::Error),
497    /// Middleware rejected the message or produced an illegal replacement.
498    Intercept(InterceptError<Error, Message>),
499}
500
501/// Failure while receiving through compile-time phase-checked middleware.
502#[derive(Debug)]
503pub enum TypedReceiveError<Error, Wire> {
504    /// Reading or decoding the message failed.
505    Io(io::Error),
506    /// The peer sent a decoded message which is illegal in the connection phase.
507    Illegal(Wire),
508    /// Middleware rejected the phase-legal message according to its policy.
509    Middleware(Error),
510    /// Middleware produced a phase-legal value with an invalid wire shape.
511    InvalidWire(Wire),
512}
513
514/// Owns user state and middleware as one reusable interception unit.
515#[derive(Clone, Copy, Debug, Eq, PartialEq)]
516pub struct Middleware<State, Handler> {
517    state: State,
518    handler: Handler,
519}
520
521impl<State, Handler> Middleware<State, Handler> {
522    /// Creates middleware with its connection- or application-local state.
523    pub const fn new(state: State, handler: Handler) -> Self {
524        Self { state, handler }
525    }
526
527    /// Borrows the accumulated user state.
528    pub const fn state(&self) -> &State {
529        &self.state
530    }
531
532    /// Mutably borrows the accumulated user state.
533    pub const fn state_mut(&mut self) -> &mut State {
534        &mut self.state
535    }
536
537    /// Borrows the middleware implementation.
538    pub const fn handler(&self) -> &Handler {
539        &self.handler
540    }
541
542    /// Mutably borrows the middleware implementation.
543    pub const fn handler_mut(&mut self) -> &mut Handler {
544        &mut self.handler
545    }
546
547    /// Separates the accumulated state from its middleware implementation.
548    pub fn into_parts(self) -> (State, Handler) {
549        (self.state, self.handler)
550    }
551
552    pub(crate) fn parts_mut(&mut self) -> (&mut State, &mut Handler) {
553        (&mut self.state, &mut self.handler)
554    }
555
556    /// Intercepts one owned message.
557    ///
558    /// # Errors
559    ///
560    /// Returns the middleware's policy-defined error.
561    pub async fn intercept<Message>(&mut self, message: Message) -> Result<Message, Handler::Error>
562    where
563        Handler: MessageMiddleware<Message, State>,
564    {
565        self.handler.intercept(&mut self.state, message).await
566    }
567
568    /// Intercepts a message whose role and legal protocol phase are type indexed.
569    ///
570    /// This operation performs no dynamic protocol-state check: `Role`, `Phase`,
571    /// and the generated `Message` type are selected together by the typed caller.
572    /// Wire-shape validation remains a separate runtime boundary after converting
573    /// the result back into its decoded wire representation.
574    ///
575    /// # Errors
576    ///
577    /// Returns the middleware's policy-defined error.
578    pub async fn intercept_typed<Role, Phase, Message>(
579        &mut self,
580        message: Message,
581    ) -> Result<Message, Handler::Error>
582    where
583        Handler: TypedMiddleware<Role, Phase, Message, State>,
584    {
585        self.handler.intercept_typed(&mut self.state, message).await
586    }
587
588    /// Intercepts a message and checks the result against `protocol_state` at runtime.
589    ///
590    /// The compiler enforces the message direction and requires `ProtocolState`
591    /// to implement [`AcceptsMessage`] for that message type. The replacement's
592    /// concrete variant and the supplied generated [`crate::grammar`] runtime
593    /// state value are dynamic, however, so protocol legality and wire
594    /// reconstructability are checked at runtime after the complete middleware
595    /// chain. Call this immediately before projecting and advancing the same
596    /// protocol state.
597    ///
598    /// # Errors
599    ///
600    /// Returns a middleware policy error, or the unchanged replacement when it
601    /// is not legal in `protocol_state`.
602    pub async fn intercept_checked<Message, ProtocolState>(
603        &mut self,
604        protocol_state: &ProtocolState,
605        message: Message,
606    ) -> Result<Message, InterceptError<Handler::Error, Message>>
607    where
608        Message: ReconstructableMessage,
609        Handler: MessageMiddleware<Message, State>,
610        ProtocolState: AcceptsMessage<Message>,
611    {
612        let message = self
613            .intercept(message)
614            .await
615            .map_err(InterceptError::Middleware)?;
616        if message.is_reconstructable() && protocol_state.accepts(&message) {
617            Ok(message)
618        } else {
619            Err(InterceptError::Invalid(message))
620        }
621    }
622}
623
624impl<Transport, Phase, Cleanliness> crate::Conn<Transport, Phase, Cleanliness> {
625    /// Intercepts one locally generated message indexed by this connection phase.
626    ///
627    /// The returned generated enum may select a different legal transition in
628    /// the same phase. Match it and apply the corresponding existing typestate
629    /// operation before encoding or forwarding the value.
630    ///
631    /// # Errors
632    ///
633    /// Returns a middleware policy error or an invalid replacement wire shape.
634    pub async fn intercept_outbound_typed<Role, Wire, State, Handler>(
635        &self,
636        middleware: &mut Middleware<State, Handler>,
637        message: <Phase as PhaseAssociation<Outbound, Role, Wire>>::Message,
638    ) -> Result<
639        <Phase as PhaseAssociation<Outbound, Role, Wire>>::Message,
640        TypedReceiveError<Handler::Error, Wire>,
641    >
642    where
643        Phase: PhaseAssociation<Outbound, Role, Wire>,
644        Wire: ReconstructableMessage,
645        Handler: TypedMiddleware<
646                Role,
647                <Phase as PhaseAssociation<Outbound, Role, Wire>>::ProtocolPhase,
648                <Phase as PhaseAssociation<Outbound, Role, Wire>>::Message,
649                State,
650            >,
651    {
652        let message = middleware
653            .intercept_typed::<Role, <Phase as PhaseAssociation<Outbound, Role, Wire>>::ProtocolPhase, _>(
654                message,
655            )
656            .await
657            .map_err(TypedReceiveError::Middleware)?;
658        if message.as_ref().is_reconstructable() {
659            Ok(message)
660        } else {
661            Err(TypedReceiveError::InvalidWire(message.into()))
662        }
663    }
664}
665
666#[cfg(test)]
667mod tests {
668    use std::convert::Infallible;
669
670    use bytes::Bytes;
671
672    use super::{
673        AcceptsMessage as _, ChainError, ClientRole, Identity, Inbound, InterceptError,
674        MessageMiddlewareExt as _, Middleware, Outbound, PhaseAssociation, ServerRole, WireAdapter,
675    };
676    use crate::{
677        Conn,
678        codec::{BackendMessage, FrontendMessage, Parse},
679        grammar::{
680            authentication, backend, frontend, pre_startup as pre_startup_grammar,
681            server_authentication, server_pre_startup,
682        },
683        middleware::TypedBackendMessage,
684        pre_startup::{EncryptionReply, PreStartupMessage},
685    };
686
687    #[tokio::test]
688    async fn identity_is_a_no_op() {
689        let mut middleware = Middleware::new((), Identity);
690        assert_eq!(
691            middleware.intercept(String::from("message")).await,
692            Ok(String::from("message"))
693        );
694    }
695
696    #[tokio::test]
697    async fn closure_can_replace_message_and_accumulate_state() {
698        let mut middleware = Middleware::new(
699            Vec::new(),
700            async |seen: &mut Vec<String>, message: String| {
701                seen.push(message.clone());
702                Ok::<_, &'static str>(message.to_uppercase())
703            },
704        );
705
706        assert_eq!(
707            middleware.intercept(String::from("hello")).await,
708            Ok(String::from("HELLO"))
709        );
710        assert_eq!(middleware.state(), &[String::from("hello")]);
711    }
712
713    #[tokio::test]
714    async fn connection_phase_indexes_locally_generated_typed_middleware() {
715        let conn = Conn::new(());
716        let message =
717            pre_startup_grammar::PreStartupInternalMessage::try_from(PreStartupMessage::SslRequest)
718                .expect("SSLRequest is legal before startup");
719        let mut middleware = Middleware::new(
720            Vec::new(),
721            async |seen: &mut Vec<&'static str>,
722                   _message: pre_startup_grammar::PreStartupInternalMessage| {
723                seen.push("outbound");
724                Ok::<_, Infallible>(
725                    pre_startup_grammar::PreStartupInternalMessage::try_from(
726                        PreStartupMessage::GssEncRequest,
727                    )
728                    .expect("GSSENCRequest is another legal pre-startup choice"),
729                )
730            },
731        );
732
733        let output = conn
734            .intercept_outbound_typed::<ClientRole, PreStartupMessage, _, _>(
735                &mut middleware,
736                message,
737            )
738            .await
739            .expect("replacement remains legal in the connection phase");
740
741        assert!(matches!(output.as_ref(), PreStartupMessage::GssEncRequest));
742        assert_eq!(middleware.state(), &["outbound"]);
743        conn.into_transport();
744    }
745
746    #[tokio::test]
747    async fn middleware_can_borrow_user_state_across_await() {
748        let handler = async |steps: &mut Vec<&'static str>, message: String| {
749            steps.push("before");
750            tokio::task::yield_now().await;
751            steps.push("after");
752            Ok::<_, Infallible>(message)
753        };
754        let mut middleware = Middleware::new(Vec::new(), handler);
755
756        assert_eq!(
757            middleware.intercept(String::from("message")).await,
758            Ok(String::from("message"))
759        );
760        assert_eq!(middleware.state(), &["before", "after"]);
761    }
762
763    #[tokio::test]
764    async fn typed_closure_replaces_only_within_its_role_and_phase() {
765        let handler = async |seen: &mut usize, _message: backend::ReadyExternalMessage| {
766            *seen += 1;
767            backend::ReadyExternalMessage::try_from(FrontendMessage::Terminate)
768                .map_err(|_| "terminate must be legal while ready")
769        };
770        let mut middleware = Middleware::new(0, handler);
771        let Ok(input) = backend::ReadyExternalMessage::try_from(FrontendMessage::Query(
772            Bytes::from_static(b"select 1"),
773        )) else {
774            panic!("query must be legal while ready");
775        };
776
777        let output = middleware
778            .intercept_typed::<ClientRole, backend::Ready, _>(input)
779            .await
780            .expect("middleware accepts the message");
781
782        assert_eq!(output.event(), backend::Event::Terminate);
783        assert!(matches!(output.into_wire(), FrontendMessage::Terminate));
784        assert_eq!(*middleware.state(), 1);
785    }
786
787    #[tokio::test]
788    async fn typed_chain_is_ordered_and_threads_shared_state() {
789        let first = async |order: &mut Vec<&'static str>,
790                           message: backend::ReadyExternalMessage| {
791            order.push("first");
792            Ok::<_, Infallible>(message)
793        };
794        let second = async |order: &mut Vec<&'static str>,
795                            message: backend::ReadyExternalMessage| {
796            order.push("second");
797            Ok::<_, Infallible>(message)
798        };
799        let mut middleware = Middleware::new(Vec::new(), first.then(second));
800        let Ok(input) = backend::ReadyExternalMessage::try_from(FrontendMessage::Terminate) else {
801            panic!("terminate must be legal while ready");
802        };
803
804        let output = middleware
805            .intercept_typed::<ClientRole, backend::Ready, _>(input)
806            .await
807            .expect("both typed stages accept the message");
808
809        assert_eq!(output.event(), backend::Event::Terminate);
810        assert_eq!(middleware.state(), &["first", "second"]);
811    }
812
813    #[tokio::test]
814    async fn wire_adapter_passes_unhandled_families_through_multiple_phases() {
815        let handler = async |seen: &mut usize, message: FrontendMessage| {
816            *seen += 1;
817            Ok::<_, Infallible>(message)
818        };
819        let mut middleware = Middleware::new(0, WireAdapter::new(handler));
820
821        let Ok(ready) = backend::ReadyExternalMessage::try_from(FrontendMessage::Terminate) else {
822            panic!("terminate must be legal while ready");
823        };
824        middleware
825            .intercept_typed::<ClientRole, backend::Ready, _>(ready)
826            .await
827            .expect("ready pass-through");
828
829        let Ok(building) = backend::BuildingExternalMessage::try_from(FrontendMessage::Sync) else {
830            panic!("sync must be legal while building");
831        };
832        middleware
833            .intercept_typed::<ClientRole, backend::Building, _>(building)
834            .await
835            .expect("building pass-through");
836
837        assert_eq!(*middleware.state(), 2);
838    }
839
840    #[tokio::test]
841    async fn chain_passes_replacement_to_next_stage_in_order() {
842        let first = async |order: &mut Vec<&'static str>, mut message: String| {
843            order.push("first");
844            message.push('1');
845            Ok::<_, &'static str>(message)
846        };
847        let second = async |order: &mut Vec<&'static str>, mut message: String| {
848            order.push("second");
849            message.push('2');
850            Ok::<_, u8>(message)
851        };
852        let mut middleware = Middleware::new(Vec::new(), first.then(second));
853
854        assert_eq!(
855            middleware.intercept(String::from("m")).await,
856            Ok(String::from("m12"))
857        );
858        assert_eq!(middleware.state(), &["first", "second"]);
859    }
860
861    #[tokio::test]
862    async fn chain_stops_after_first_error() {
863        let first = async |calls: &mut usize, _message: String| {
864            *calls += 1;
865            Err::<String, _>("rejected")
866        };
867        let second = async |calls: &mut usize, message: String| {
868            *calls += 1;
869            Ok::<_, u8>(message)
870        };
871        let mut middleware = Middleware::new(0, first.then(second));
872
873        assert_eq!(
874            middleware.intercept(String::from("message")).await,
875            Err(ChainError::First("rejected"))
876        );
877        assert_eq!(*middleware.state(), 1);
878    }
879
880    #[tokio::test]
881    async fn checked_interception_accepts_a_legal_replacement() {
882        let mut middleware =
883            Middleware::new((), async |_state: &mut (), _message: FrontendMessage| {
884                Ok::<_, Infallible>(FrontendMessage::Terminate)
885            });
886
887        assert_eq!(
888            middleware
889                .intercept_checked(
890                    &backend::RuntimeState::Ready,
891                    FrontendMessage::Query(Bytes::from_static(b"select 1")),
892                )
893                .await,
894            Ok(FrontendMessage::Terminate)
895        );
896    }
897
898    #[tokio::test]
899    async fn checked_interception_returns_an_illegal_replacement() {
900        let replacement = FrontendMessage::Parse(Parse {
901            statement: Bytes::new(),
902            query: Bytes::from_static(b"select 2"),
903            parameter_types: Vec::new(),
904        });
905        let expected = replacement.clone();
906        let mut middleware = Middleware::new(
907            (),
908            async move |_state: &mut (), _message: FrontendMessage| {
909                Ok::<_, Infallible>(replacement.clone())
910            },
911        );
912
913        assert_eq!(
914            middleware
915                .intercept_checked(
916                    &backend::RuntimeState::Simple,
917                    FrontendMessage::Query(Bytes::from_static(b"select 1")),
918                )
919                .await,
920            Err(InterceptError::Invalid(expected))
921        );
922    }
923
924    #[test]
925    fn generated_states_cover_authentication_extended_query_copy_and_replication() {
926        let password = FrontendMessage::PasswordResponse(Bytes::from_static(b"secret"));
927        assert!(server_authentication::RuntimeState::PasswordResponse.accepts(&password));
928        assert!(
929            !server_authentication::RuntimeState::PasswordResponse
930                .accepts(&FrontendMessage::Query(Bytes::from_static(b"select 1")))
931        );
932
933        let parse = FrontendMessage::Parse(Parse {
934            statement: Bytes::from_static(b"statement"),
935            query: Bytes::from_static(b"select 1"),
936            parameter_types: Vec::new(),
937        });
938        assert!(backend::RuntimeState::Building.accepts(&parse));
939        assert!(
940            !backend::RuntimeState::Building
941                .accepts(&FrontendMessage::Query(Bytes::from_static(b"select 1")))
942        );
943        assert!(backend::RuntimeState::ExtendedError.accepts(&parse));
944        assert!(backend::RuntimeState::ExtendedError.accepts(&FrontendMessage::Sync));
945
946        let copy = FrontendMessage::CopyData(Bytes::from_static(b"data"));
947        assert!(backend::RuntimeState::SimpleCopyIn.accepts(&copy));
948        assert!(backend::RuntimeState::ExtendedCopyBoth.accepts(&copy));
949        assert!(
950            !backend::RuntimeState::ExtendedCopyBoth
951                .accepts(&FrontendMessage::Query(Bytes::from_static(b"select 1")))
952        );
953
954        assert!(
955            server_pre_startup::RuntimeState::PreStartup.accepts(&PreStartupMessage::SslRequest)
956        );
957        assert!(
958            !server_pre_startup::RuntimeState::SslDecision.accepts(&PreStartupMessage::SslRequest)
959        );
960    }
961
962    #[test]
963    #[allow(clippy::too_many_lines)]
964    fn grammar_associations_cover_inbound_and_outbound_typestate_indices() {
965        fn inbound<Connection, Role, Wire>()
966        where
967            Connection: PhaseAssociation<Inbound, Role, Wire>,
968        {
969        }
970        fn outbound<Connection, Role, Wire>()
971        where
972            Connection: PhaseAssociation<Outbound, Role, Wire>,
973        {
974        }
975        fn exact<Connection, Direction, Role, Wire, ProtocolPhase, Message>()
976        where
977            Connection: PhaseAssociation<
978                    Direction,
979                    Role,
980                    Wire,
981                    ProtocolPhase = ProtocolPhase,
982                    Message = Message,
983                >,
984            Message: AsRef<Wire> + TryFrom<Wire, Error = Wire> + Into<Wire>,
985        {
986        }
987
988        exact::<
989            crate::auth::Ready,
990            Inbound,
991            ServerRole,
992            BackendMessage,
993            frontend::Ready,
994            TypedBackendMessage<frontend::ReadyExternalMessage>,
995        >();
996        exact::<
997            crate::auth::Ready,
998            Outbound,
999            ClientRole,
1000            FrontendMessage,
1001            frontend::Ready,
1002            frontend::ReadyInternalMessage,
1003        >();
1004        exact::<
1005            crate::auth::Ready,
1006            Inbound,
1007            ClientRole,
1008            FrontendMessage,
1009            backend::Ready,
1010            backend::ReadyExternalMessage,
1011        >();
1012        exact::<
1013            crate::server_session::ServerFunctionCallDone,
1014            Outbound,
1015            ServerRole,
1016            BackendMessage,
1017            backend::FunctionReady,
1018            TypedBackendMessage<backend::FunctionReadyInternalMessage>,
1019        >();
1020        exact::<
1021            crate::pre_startup::AwaitingSslReply,
1022            Inbound,
1023            ServerRole,
1024            EncryptionReply,
1025            pre_startup_grammar::AwaitingSslReply,
1026            pre_startup_grammar::AwaitingSslReplyExternalMessage,
1027        >();
1028        exact::<
1029            crate::pre_startup::PreStartup,
1030            Outbound,
1031            ClientRole,
1032            PreStartupMessage,
1033            pre_startup_grammar::PreStartup,
1034            pre_startup_grammar::PreStartupInternalMessage,
1035        >();
1036        exact::<
1037            crate::pre_startup::PreStartup,
1038            Inbound,
1039            ClientRole,
1040            PreStartupMessage,
1041            server_pre_startup::PreStartup,
1042            server_pre_startup::PreStartupExternalMessage,
1043        >();
1044        exact::<
1045            crate::pre_startup::ServerSslDecision,
1046            Outbound,
1047            ServerRole,
1048            EncryptionReply,
1049            server_pre_startup::SslDecision,
1050            server_pre_startup::SslDecisionInternalMessage,
1051        >();
1052        exact::<
1053            crate::auth::Auth,
1054            Inbound,
1055            ServerRole,
1056            BackendMessage,
1057            authentication::Auth,
1058            TypedBackendMessage<authentication::AuthExternalMessage>,
1059        >();
1060        exact::<
1061            crate::auth::PasswordResponse,
1062            Outbound,
1063            ClientRole,
1064            FrontendMessage,
1065            authentication::PasswordResponse,
1066            authentication::PasswordResponseInternalMessage,
1067        >();
1068        exact::<
1069            crate::server_auth::ServerAuth,
1070            Inbound,
1071            ClientRole,
1072            FrontendMessage,
1073            server_authentication::Auth,
1074            server_authentication::AuthExternalMessage,
1075        >();
1076        exact::<
1077            crate::server_auth::ServerAuth,
1078            Outbound,
1079            ServerRole,
1080            BackendMessage,
1081            server_authentication::Auth,
1082            TypedBackendMessage<server_authentication::AuthInternalMessage>,
1083        >();
1084
1085        inbound::<crate::auth::Ready, ServerRole, BackendMessage>();
1086        inbound::<crate::auth::Ready, ClientRole, FrontendMessage>();
1087        outbound::<crate::auth::Ready, ClientRole, FrontendMessage>();
1088        inbound::<crate::pre_startup::PreStartup, ClientRole, PreStartupMessage>();
1089        outbound::<crate::pre_startup::PreStartup, ClientRole, PreStartupMessage>();
1090        inbound::<crate::server_auth::ServerAuth, ClientRole, FrontendMessage>();
1091        outbound::<crate::server_auth::ServerAuth, ServerRole, BackendMessage>();
1092        inbound::<crate::session::CopyBoth, ServerRole, BackendMessage>();
1093        outbound::<crate::session::CopyBoth, ClientRole, FrontendMessage>();
1094
1095        inbound::<crate::auth::Auth, ServerRole, BackendMessage>();
1096        inbound::<crate::auth::TokenChallenge, ServerRole, BackendMessage>();
1097        inbound::<crate::auth::Sasl, ServerRole, BackendMessage>();
1098        inbound::<crate::auth::AwaitingAuthOk, ServerRole, BackendMessage>();
1099        inbound::<crate::auth::AwaitingStartupReady, ServerRole, BackendMessage>();
1100        inbound::<crate::session::SimpleQuery, ServerRole, BackendMessage>();
1101        inbound::<crate::session::FunctionCalling, ServerRole, BackendMessage>();
1102        inbound::<crate::session::Building, ServerRole, BackendMessage>();
1103        inbound::<crate::session::BoundBuilding, ServerRole, BackendMessage>();
1104        inbound::<crate::session::AwaitingReady, ServerRole, BackendMessage>();
1105        inbound::<crate::session::CopyIn, ServerRole, BackendMessage>();
1106        inbound::<crate::session::CopyOut, ServerRole, BackendMessage>();
1107        inbound::<crate::session::CopyBothClientDone, ServerRole, BackendMessage>();
1108        inbound::<crate::session::CopyBothServerDone, ServerRole, BackendMessage>();
1109        inbound::<crate::session::Draining, ServerRole, BackendMessage>();
1110        inbound::<crate::session::Resetting, ServerRole, BackendMessage>();
1111        inbound::<crate::session::ResetComplete, ServerRole, BackendMessage>();
1112
1113        outbound::<crate::auth::PasswordResponse, ClientRole, FrontendMessage>();
1114        outbound::<crate::auth::TokenResponse, ClientRole, FrontendMessage>();
1115        outbound::<crate::auth::SaslInitial, ClientRole, FrontendMessage>();
1116        outbound::<crate::auth::SaslChallenge, ClientRole, FrontendMessage>();
1117        outbound::<crate::session::Building, ClientRole, FrontendMessage>();
1118        outbound::<crate::session::BoundBuilding, ClientRole, FrontendMessage>();
1119        outbound::<crate::session::CopyIn, ClientRole, FrontendMessage>();
1120        outbound::<crate::session::CopyBothServerDone, ClientRole, FrontendMessage>();
1121        outbound::<crate::pre_startup::ServerSslDecision, ServerRole, EncryptionReply>();
1122        outbound::<crate::pre_startup::ServerGssDecision, ServerRole, EncryptionReply>();
1123
1124        inbound::<crate::pre_startup::AwaitingSslReply, ServerRole, EncryptionReply>();
1125        inbound::<crate::pre_startup::AwaitingGssReply, ServerRole, EncryptionReply>();
1126        inbound::<crate::server_auth::ServerPassword, ClientRole, FrontendMessage>();
1127        inbound::<crate::server_auth::ServerSaslInitial, ClientRole, FrontendMessage>();
1128        inbound::<crate::server_auth::ServerSaslResponse, ClientRole, FrontendMessage>();
1129        inbound::<crate::server_auth::ServerAuthResponse, ClientRole, FrontendMessage>();
1130        inbound::<crate::server_auth::ServerStartupReady, ClientRole, FrontendMessage>();
1131        inbound::<crate::server_session::ServerBuilding, ClientRole, FrontendMessage>();
1132        inbound::<crate::server_session::ServerExtendedError, ClientRole, FrontendMessage>();
1133        inbound::<
1134            crate::server_session::ServerCopyIn<crate::server_session::CopySimple>,
1135            ClientRole,
1136            FrontendMessage,
1137        >();
1138        inbound::<
1139            crate::server_session::ServerCopyIn<crate::server_session::CopyExtended>,
1140            ClientRole,
1141            FrontendMessage,
1142        >();
1143        inbound::<
1144            crate::server_session::ServerCopyBoth<
1145                crate::server_session::CopySimple,
1146                crate::server_session::BothOpen,
1147            >,
1148            ClientRole,
1149            FrontendMessage,
1150        >();
1151        inbound::<
1152            crate::server_session::ServerCopyBoth<
1153                crate::server_session::CopyExtended,
1154                crate::server_session::BothOpen,
1155            >,
1156            ClientRole,
1157            FrontendMessage,
1158        >();
1159        inbound::<
1160            crate::server_session::ServerCopyBoth<
1161                crate::server_session::CopySimple,
1162                crate::server_session::BothServerDone,
1163            >,
1164            ClientRole,
1165            FrontendMessage,
1166        >();
1167        inbound::<
1168            crate::server_session::ServerCopyBoth<
1169                crate::server_session::CopyExtended,
1170                crate::server_session::BothServerDone,
1171            >,
1172            ClientRole,
1173            FrontendMessage,
1174        >();
1175
1176        outbound::<crate::server_auth::ServerStartupRejected, ServerRole, BackendMessage>();
1177        outbound::<crate::server_auth::ServerPassword, ServerRole, BackendMessage>();
1178        outbound::<crate::server_auth::ServerSaslInitial, ServerRole, BackendMessage>();
1179        outbound::<crate::server_auth::ServerSasl, ServerRole, BackendMessage>();
1180        outbound::<crate::server_auth::ServerSaslResponse, ServerRole, BackendMessage>();
1181        outbound::<crate::server_auth::ServerAuthResponse, ServerRole, BackendMessage>();
1182        outbound::<crate::server_auth::ServerAuthPolicy, ServerRole, BackendMessage>();
1183        outbound::<crate::server_auth::ServerStartupReady, ServerRole, BackendMessage>();
1184        outbound::<crate::server_session::ServerSimpleQuery, ServerRole, BackendMessage>();
1185        outbound::<crate::server_session::ServerSimpleError, ServerRole, BackendMessage>();
1186        outbound::<crate::server_session::ServerFunctionCall, ServerRole, BackendMessage>();
1187        outbound::<crate::server_session::ServerFunctionCallDone, ServerRole, BackendMessage>();
1188        outbound::<crate::server_session::ServerFunctionCallError, ServerRole, BackendMessage>();
1189        outbound::<crate::server_session::ServerParse, ServerRole, BackendMessage>();
1190        outbound::<crate::server_session::ServerBind, ServerRole, BackendMessage>();
1191        outbound::<crate::server_session::ServerDescribe, ServerRole, BackendMessage>();
1192        outbound::<crate::server_session::ServerExecute, ServerRole, BackendMessage>();
1193        outbound::<crate::server_session::ServerClose, ServerRole, BackendMessage>();
1194        outbound::<crate::server_session::ServerSync, ServerRole, BackendMessage>();
1195        outbound::<crate::server_session::ServerBuilding, ServerRole, BackendMessage>();
1196        outbound::<crate::server_session::ServerExtendedError, ServerRole, BackendMessage>();
1197        outbound::<
1198            crate::server_session::ServerCopyIn<crate::server_session::CopySimple>,
1199            ServerRole,
1200            BackendMessage,
1201        >();
1202        outbound::<
1203            crate::server_session::ServerCopyIn<crate::server_session::CopyExtended>,
1204            ServerRole,
1205            BackendMessage,
1206        >();
1207        outbound::<
1208            crate::server_session::ServerCopyInDone<crate::server_session::CopySimple>,
1209            ServerRole,
1210            BackendMessage,
1211        >();
1212        outbound::<
1213            crate::server_session::ServerCopyInDone<crate::server_session::CopyExtended>,
1214            ServerRole,
1215            BackendMessage,
1216        >();
1217        outbound::<
1218            crate::server_session::ServerCopyInFailed<crate::server_session::CopySimple>,
1219            ServerRole,
1220            BackendMessage,
1221        >();
1222        outbound::<
1223            crate::server_session::ServerCopyInFailed<crate::server_session::CopyExtended>,
1224            ServerRole,
1225            BackendMessage,
1226        >();
1227        outbound::<
1228            crate::server_session::ServerCopyOut<crate::server_session::CopySimple>,
1229            ServerRole,
1230            BackendMessage,
1231        >();
1232        outbound::<
1233            crate::server_session::ServerCopyOut<crate::server_session::CopyExtended>,
1234            ServerRole,
1235            BackendMessage,
1236        >();
1237        outbound::<
1238            crate::server_session::ServerCopyOutDone<crate::server_session::CopySimple>,
1239            ServerRole,
1240            BackendMessage,
1241        >();
1242        outbound::<
1243            crate::server_session::ServerCopyOutDone<crate::server_session::CopyExtended>,
1244            ServerRole,
1245            BackendMessage,
1246        >();
1247        outbound::<
1248            crate::server_session::ServerCopyBoth<
1249                crate::server_session::CopySimple,
1250                crate::server_session::BothOpen,
1251            >,
1252            ServerRole,
1253            BackendMessage,
1254        >();
1255        outbound::<
1256            crate::server_session::ServerCopyBoth<
1257                crate::server_session::CopyExtended,
1258                crate::server_session::BothOpen,
1259            >,
1260            ServerRole,
1261            BackendMessage,
1262        >();
1263        outbound::<
1264            crate::server_session::ServerCopyBoth<
1265                crate::server_session::CopySimple,
1266                crate::server_session::BothClientDone,
1267            >,
1268            ServerRole,
1269            BackendMessage,
1270        >();
1271        outbound::<
1272            crate::server_session::ServerCopyBoth<
1273                crate::server_session::CopyExtended,
1274                crate::server_session::BothClientDone,
1275            >,
1276            ServerRole,
1277            BackendMessage,
1278        >();
1279        outbound::<
1280            crate::server_session::ServerCopyBoth<
1281                crate::server_session::CopySimple,
1282                crate::server_session::BothServerDone,
1283            >,
1284            ServerRole,
1285            BackendMessage,
1286        >();
1287        outbound::<
1288            crate::server_session::ServerCopyBoth<
1289                crate::server_session::CopyExtended,
1290                crate::server_session::BothServerDone,
1291            >,
1292            ServerRole,
1293            BackendMessage,
1294        >();
1295        outbound::<
1296            crate::server_session::ServerCopyBoth<
1297                crate::server_session::CopySimple,
1298                crate::server_session::BothDone,
1299            >,
1300            ServerRole,
1301            BackendMessage,
1302        >();
1303        outbound::<
1304            crate::server_session::ServerCopyBoth<
1305                crate::server_session::CopyExtended,
1306                crate::server_session::BothDone,
1307            >,
1308            ServerRole,
1309            BackendMessage,
1310        >();
1311        outbound::<
1312            crate::server_session::ServerCopyBothFailed<crate::server_session::CopySimple>,
1313            ServerRole,
1314            BackendMessage,
1315        >();
1316        outbound::<
1317            crate::server_session::ServerCopyBothFailed<crate::server_session::CopyExtended>,
1318            ServerRole,
1319            BackendMessage,
1320        >();
1321    }
1322
1323    #[tokio::test]
1324    async fn checked_interception_rejects_an_unencodable_message() {
1325        let invalid = FrontendMessage::Parse(Parse {
1326            statement: Bytes::from_static(b"invalid\0name"),
1327            query: Bytes::from_static(b"select 1"),
1328            parameter_types: Vec::new(),
1329        });
1330        let expected = invalid.clone();
1331        let mut middleware = Middleware::new((), async move |_state: &mut (), _message| {
1332            Ok::<_, Infallible>(invalid.clone())
1333        });
1334
1335        assert_eq!(
1336            middleware
1337                .intercept_checked(&backend::RuntimeState::Ready, FrontendMessage::Terminate)
1338                .await,
1339            Err(InterceptError::Invalid(expected))
1340        );
1341    }
1342}
pg_proto/middleware.rs

pg_proto/
middleware.rs
