Skip to main content

pg_proto/
runtime_middleware.rs

1//! Contextual, per-connection middleware used by the builder facade.
2
3use crate::codec::{BackendMessage, FrontendMessage};
4
5/// Creates one isolated handler synchronously for a new connection.
6pub trait MiddlewareFactory<Context> {
7    /// Handler owned by that connection.
8    type Handler;
9    /// Creates the handler. Factories are deliberately infallible.
10    fn create(&self, context: &Context) -> Self::Handler;
11}
12
13impl<Context, Handler, Factory> MiddlewareFactory<Context> for Factory
14where
15    Factory: Fn(&Context) -> Handler,
16{
17    type Handler = Handler;
18    fn create(&self, context: &Context) -> Handler {
19        self(context)
20    }
21}
22
23/// Default factory and handler: messages pass through unchanged.
24#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
25pub struct IdentityMiddleware;
26
27impl<Context> MiddlewareFactory<Context> for IdentityMiddleware {
28    type Handler = Self;
29    fn create(&self, _context: &Context) -> Self {
30        *self
31    }
32}
33
34/// Two factories/handlers composed in builder declaration order.
35#[derive(Clone, Debug, Eq, PartialEq)]
36pub struct MiddlewareChain<First, Second>(pub First, pub Second);
37
38impl<Context, First, Second> MiddlewareFactory<Context> for MiddlewareChain<First, Second>
39where
40    First: MiddlewareFactory<Context>,
41    Second: MiddlewareFactory<Context>,
42{
43    type Handler = MiddlewareChain<First::Handler, Second::Handler>;
44    fn create(&self, context: &Context) -> Self::Handler {
45        MiddlewareChain(self.0.create(context), self.1.create(context))
46    }
47}
48
49/// Middleware for a PostgreSQL client role. Implement only the directions used.
50pub trait ClientMiddleware<State, Context> {
51    /// Intercepts a client-originated pre-startup negotiation packet.
52    fn pre_startup(
53        &mut self,
54        _context: &Context,
55        _state: &mut State,
56        message: crate::pre_startup::PreStartupMessage,
57    ) -> crate::pre_startup::PreStartupMessage {
58        message
59    }
60    /// Intercepts the owned startup message before it is sent.
61    fn startup(
62        &mut self,
63        _context: &Context,
64        _state: &mut State,
65        message: crate::startup::StartupMessage,
66    ) -> crate::startup::StartupMessage {
67        message
68    }
69    /// Intercepts one message sent by the client.
70    fn frontend(
71        &mut self,
72        _context: &Context,
73        _state: &mut State,
74        message: FrontendMessage,
75    ) -> FrontendMessage {
76        message
77    }
78    /// Intercepts one message received from the server.
79    fn backend(
80        &mut self,
81        _context: &Context,
82        _state: &mut State,
83        message: BackendMessage,
84    ) -> BackendMessage {
85        message
86    }
87}
88
89impl<State, Context> ClientMiddleware<State, Context> for IdentityMiddleware {}
90
91impl<State, Context, First, Second> ClientMiddleware<State, Context>
92    for MiddlewareChain<First, Second>
93where
94    First: ClientMiddleware<State, Context>,
95    Second: ClientMiddleware<State, Context>,
96{
97    fn pre_startup(
98        &mut self,
99        context: &Context,
100        state: &mut State,
101        message: crate::pre_startup::PreStartupMessage,
102    ) -> crate::pre_startup::PreStartupMessage {
103        let message = self.0.pre_startup(context, state, message);
104        self.1.pre_startup(context, state, message)
105    }
106    fn startup(
107        &mut self,
108        context: &Context,
109        state: &mut State,
110        message: crate::startup::StartupMessage,
111    ) -> crate::startup::StartupMessage {
112        let message = self.0.startup(context, state, message);
113        self.1.startup(context, state, message)
114    }
115    fn frontend(
116        &mut self,
117        context: &Context,
118        state: &mut State,
119        message: FrontendMessage,
120    ) -> FrontendMessage {
121        let message = self.0.frontend(context, state, message);
122        self.1.frontend(context, state, message)
123    }
124    fn backend(
125        &mut self,
126        context: &Context,
127        state: &mut State,
128        message: BackendMessage,
129    ) -> BackendMessage {
130        let message = self.0.backend(context, state, message);
131        self.1.backend(context, state, message)
132    }
133}
134
135/// Middleware for a PostgreSQL server role. Implement only the directions used.
136pub trait ServerMiddleware<State, Context> {
137    /// Intercepts one pre-startup packet before protocol dispatch.
138    fn pre_startup(
139        &mut self,
140        _context: &Context,
141        _state: &mut State,
142        message: crate::pre_startup::PreStartupMessage,
143    ) -> crate::pre_startup::PreStartupMessage {
144        message
145    }
146    /// Intercepts the owned startup message before authentication.
147    fn startup(
148        &mut self,
149        _context: &Context,
150        _state: &mut State,
151        message: crate::startup::StartupMessage,
152    ) -> crate::startup::StartupMessage {
153        message
154    }
155    /// Intercepts an out-of-band cancellation request before it is returned.
156    fn cancellation(
157        &mut self,
158        _context: &Context,
159        _state: &mut State,
160        request: crate::CancellationRequest,
161    ) -> crate::CancellationRequest {
162        request
163    }
164    /// Intercepts one message received from the client.
165    fn frontend(
166        &mut self,
167        _context: &Context,
168        _state: &mut State,
169        message: FrontendMessage,
170    ) -> FrontendMessage {
171        message
172    }
173    /// Intercepts one message sent by the server.
174    fn backend(
175        &mut self,
176        _context: &Context,
177        _state: &mut State,
178        message: BackendMessage,
179    ) -> BackendMessage {
180        message
181    }
182}
183
184impl<State, Context> ServerMiddleware<State, Context> for IdentityMiddleware {}
185
186impl<State, Context, First, Second> ServerMiddleware<State, Context>
187    for MiddlewareChain<First, Second>
188where
189    First: ServerMiddleware<State, Context>,
190    Second: ServerMiddleware<State, Context>,
191{
192    fn pre_startup(
193        &mut self,
194        context: &Context,
195        state: &mut State,
196        message: crate::pre_startup::PreStartupMessage,
197    ) -> crate::pre_startup::PreStartupMessage {
198        let message = self.0.pre_startup(context, state, message);
199        self.1.pre_startup(context, state, message)
200    }
201    fn startup(
202        &mut self,
203        context: &Context,
204        state: &mut State,
205        message: crate::startup::StartupMessage,
206    ) -> crate::startup::StartupMessage {
207        let message = self.0.startup(context, state, message);
208        self.1.startup(context, state, message)
209    }
210    fn cancellation(
211        &mut self,
212        context: &Context,
213        state: &mut State,
214        request: crate::CancellationRequest,
215    ) -> crate::CancellationRequest {
216        let request = self.0.cancellation(context, state, request);
217        self.1.cancellation(context, state, request)
218    }
219    fn frontend(
220        &mut self,
221        context: &Context,
222        state: &mut State,
223        message: FrontendMessage,
224    ) -> FrontendMessage {
225        let message = self.0.frontend(context, state, message);
226        self.1.frontend(context, state, message)
227    }
228    fn backend(
229        &mut self,
230        context: &Context,
231        state: &mut State,
232        message: BackendMessage,
233    ) -> BackendMessage {
234        let message = self.0.backend(context, state, message);
235        self.1.backend(context, state, message)
236    }
237}