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

pg_proto/
middleware.rs
