pg_proto/
transport.rs

1//! Buffered, cancellation-safe outbound transport.
2
3use std::{collections::BTreeMap, io, sync::Arc};
4
5use bytes::{Buf, Bytes, BytesMut};
6use rustls::{
7    ClientConfig, ServerConfig,
8    pki_types::{CertificateDer, ServerName},
9};
10use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
11use tokio_util::codec::{Decoder, Encoder};
12
13use crate::{
14    Conn,
15    auth::TlsServerEndPoint,
16    codec::{Backend, BackendMessage, Direction, Frame, Frontend, FrontendMessage, PgCodec},
17    demux::{
18        CancelKey, Demux, Notification, OrderedAsyncEvent, ParameterStatus, SessionItem,
19        TaggedNotice,
20    },
21    middleware::{
22        AcceptsMessage, ClientRole, MessageMiddleware, Middleware, ReceiveError,
23        ReconstructableMessage as _, ServerRole, TypedMiddleware, TypedPhase, TypedReceiveError,
24    },
25    pre_startup::{
26        AwaitingSslReply, DEFAULT_MAX_PRE_STARTUP_PACKET_LEN, EncryptionReply, Negotiation,
27        PreStartup, PreStartupMessage, ServerSslDecision, SslMode, SslModeNegotiation,
28        TlsHandshake, decode_pre_startup_with_limit, gssenc_request_packet, ssl_request_packet,
29    },
30    tls::{ClientTls, ServerTls},
31};
32
33/// Transport wrapper which retains bytes until each write has completed.
34#[derive(Debug)]
35pub struct Buffered<S, D = Backend> {
36    io: S,
37    outbound: BytesMut,
38    inbound: BytesMut,
39    inbound_codec: PgCodec<D>,
40    max_pre_startup_packet_len: usize,
41    demux: Demux,
42}
43
44impl<S> Buffered<S, Backend> {
45    /// Wraps an upstream-facing transport which receives backend messages.
46    pub fn new(io: S) -> Self {
47        Self {
48            io,
49            outbound: BytesMut::new(),
50            inbound: BytesMut::new(),
51            inbound_codec: PgCodec::default(),
52            max_pre_startup_packet_len: DEFAULT_MAX_PRE_STARTUP_PACKET_LEN,
53            demux: Demux::default(),
54        }
55    }
56
57    /// Creates a backend-facing transport with a bounded tagged-frame size.
58    ///
59    /// # Errors
60    ///
61    /// Returns an error when the limit is outside `PostgreSQL`'s frame range.
62    pub fn with_max_frame_len(io: S, max_frame_len: usize) -> io::Result<Self> {
63        Ok(Self {
64            io,
65            outbound: BytesMut::new(),
66            inbound: BytesMut::new(),
67            inbound_codec: PgCodec::with_max_frame_len(max_frame_len)?,
68            max_pre_startup_packet_len: DEFAULT_MAX_PRE_STARTUP_PACKET_LEN,
69            demux: Demux::default(),
70        })
71    }
72}
73
74impl<S> Buffered<S, Frontend> {
75    /// Wraps a client-facing transport which receives frontend messages.
76    pub fn new_frontend(io: S) -> Self {
77        Self {
78            io,
79            outbound: BytesMut::new(),
80            inbound: BytesMut::new(),
81            inbound_codec: PgCodec::default(),
82            max_pre_startup_packet_len: DEFAULT_MAX_PRE_STARTUP_PACKET_LEN,
83            demux: Demux::default(),
84        }
85    }
86
87    /// Creates a frontend-facing transport with a bounded tagged-frame size.
88    ///
89    /// # Errors
90    ///
91    /// Returns an error when the limit is outside `PostgreSQL`'s frame range.
92    pub fn with_max_frame_len_frontend(io: S, max_frame_len: usize) -> io::Result<Self> {
93        Self::with_limits_frontend(io, max_frame_len, DEFAULT_MAX_PRE_STARTUP_PACKET_LEN)
94    }
95
96    /// Creates a frontend-facing transport with bounded tagged and pre-startup packets.
97    ///
98    /// # Errors
99    ///
100    /// Returns an error when either limit is outside `PostgreSQL`'s framing range.
101    pub fn with_limits_frontend(
102        io: S,
103        max_frame_len: usize,
104        max_pre_startup_packet_len: usize,
105    ) -> io::Result<Self> {
106        if !(8..=i32::MAX as usize).contains(&max_pre_startup_packet_len) {
107            return Err(io::Error::new(
108                io::ErrorKind::InvalidInput,
109                "pre-startup packet limit must be between 8 and i32::MAX bytes",
110            ));
111        }
112        Ok(Self {
113            io,
114            outbound: BytesMut::new(),
115            inbound: BytesMut::new(),
116            inbound_codec: PgCodec::with_max_frame_len(max_frame_len)?,
117            max_pre_startup_packet_len,
118            demux: Demux::default(),
119        })
120    }
121}
122
123impl<S, D> Buffered<S, D> {
124    /// Encodes a frame synchronously into the outbound buffer.
125    ///
126    /// # Errors
127    ///
128    /// Returns an error when the frame is too large to encode.
129    pub fn push(&mut self, frame: Frame) -> io::Result<()> {
130        self.inbound_codec.encode(frame, &mut self.outbound)
131    }
132
133    #[must_use]
134    /// Returns encoded bytes which have not yet been fully written.
135    pub fn pending(&self) -> &[u8] {
136        &self.outbound
137    }
138
139    /// Removes buffering and returns the underlying I/O transport.
140    pub fn into_inner(self) -> S {
141        self.io
142    }
143
144    /// Borrows the underlying I/O transport without disturbing codec buffers.
145    pub const fn get_ref(&self) -> &S {
146        &self.io
147    }
148
149    /// Mutably borrows the underlying I/O transport without disturbing codec buffers.
150    pub const fn get_mut(&mut self) -> &mut S {
151        &mut self.io
152    }
153
154    fn push_raw(&mut self, bytes: &[u8]) {
155        self.outbound.extend_from_slice(bytes);
156    }
157
158    #[must_use]
159    /// Returns the backend asynchronous-message demultiplexer.
160    pub const fn demux(&self) -> &Demux {
161        &self.demux
162    }
163
164    /// Returns mutable access to the backend asynchronous-message demultiplexer.
165    pub const fn demux_mut(&mut self) -> &mut Demux {
166        &mut self.demux
167    }
168}
169
170impl<S, D> Buffered<S, D>
171where
172    S: AsyncRead + AsyncWrite + Unpin,
173{
174    async fn connect_tls(
175        self,
176        server_name: ServerName<'static>,
177        config: Arc<ClientConfig>,
178    ) -> io::Result<Buffered<ClientTls<S>, D>> {
179        if !self.outbound.is_empty() || !self.inbound.is_empty() {
180            return Err(io::Error::new(
181                io::ErrorKind::InvalidInput,
182                "TLS upgrade requires empty plaintext buffers",
183            ));
184        }
185        Ok(Buffered {
186            io: crate::tls::connect(self.io, server_name, config).await?,
187            outbound: self.outbound,
188            inbound: self.inbound,
189            inbound_codec: self.inbound_codec,
190            max_pre_startup_packet_len: self.max_pre_startup_packet_len,
191            demux: self.demux,
192        })
193    }
194
195    async fn accept_tls(
196        self,
197        config: Arc<ServerConfig>,
198        leaf_certificate: CertificateDer<'static>,
199    ) -> io::Result<Buffered<ServerTls<S>, D>> {
200        if !self.outbound.is_empty() || !self.inbound.is_empty() {
201            return Err(io::Error::new(
202                io::ErrorKind::InvalidInput,
203                "TLS upgrade requires empty plaintext buffers",
204            ));
205        }
206        Ok(Buffered {
207            io: crate::tls::accept(self.io, config, &leaf_certificate).await?,
208            outbound: self.outbound,
209            inbound: self.inbound,
210            inbound_codec: self.inbound_codec,
211            max_pre_startup_packet_len: self.max_pre_startup_packet_len,
212            demux: self.demux,
213        })
214    }
215}
216
217impl<S: TlsServerEndPoint, D> TlsServerEndPoint for Buffered<S, D> {
218    fn tls_server_end_point(&self) -> &[u8] {
219        self.io.tls_server_end_point()
220    }
221}
222
223impl<S: AsyncWrite + Unpin, D> Buffered<S, D> {
224    /// Writes all buffered bytes without consuming the connection.
225    ///
226    /// Completed partial writes are removed immediately. If this future is
227    /// cancelled, the connection remains owned by the caller and all unwritten
228    /// bytes remain buffered for the next call.
229    ///
230    /// # Errors
231    ///
232    /// Returns the underlying transport's write error or `WriteZero`.
233    pub async fn flush(&mut self) -> io::Result<()> {
234        while !self.outbound.is_empty() {
235            let written = self.io.write(&self.outbound).await?;
236            if written == 0 {
237                return Err(io::Error::new(
238                    io::ErrorKind::WriteZero,
239                    "transport wrote zero buffered bytes",
240                ));
241            }
242            self.outbound.advance(written);
243        }
244        self.io.flush().await
245    }
246}
247
248impl<S: AsyncRead + Unpin, D: Direction> Buffered<S, D> {
249    /// Receives one typed message in this transport's inbound direction.
250    ///
251    /// # Errors
252    ///
253    /// Returns decoding and underlying transport read errors, or `UnexpectedEof`.
254    pub async fn receive_wire(&mut self) -> io::Result<D::Message> {
255        loop {
256            if let Some(message) = self.inbound_codec.decode(&mut self.inbound)? {
257                return Ok(message);
258            }
259            if self.io.read_buf(&mut self.inbound).await? == 0 {
260                return Err(io::Error::new(
261                    io::ErrorKind::UnexpectedEof,
262                    "peer closed with no complete message",
263                ));
264            }
265        }
266    }
267}
268
269impl<S: AsyncRead + Unpin> Buffered<S, Backend> {
270    async fn receive_encryption_reply(&mut self) -> io::Result<EncryptionReply> {
271        let byte = self.io.read_u8().await?;
272        EncryptionReply::try_from(byte)
273            .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "invalid encryption reply"))
274    }
275}
276
277impl<S: AsyncRead + Unpin> Buffered<S, Frontend> {
278    /// Receives one raw first packet before tagged frontend framing begins.
279    ///
280    /// # Errors
281    ///
282    /// Returns malformed pre-startup data and underlying transport read errors.
283    pub async fn receive_pre_startup(&mut self) -> io::Result<PreStartupMessage> {
284        loop {
285            if let Some(message) =
286                decode_pre_startup_with_limit(&mut self.inbound, self.max_pre_startup_packet_len)?
287            {
288                return Ok(message);
289            }
290            if self.io.read_buf(&mut self.inbound).await? == 0 {
291                return Err(io::Error::new(
292                    io::ErrorKind::UnexpectedEof,
293                    "client closed with no complete pre-startup packet",
294                ));
295            }
296        }
297    }
298}
299
300impl<S: AsyncRead + Unpin> Buffered<S, Backend> {
301    /// Receives one decoded backend message while retaining partial input.
302    ///
303    /// # Errors
304    ///
305    /// Returns decoding and underlying transport read errors, or `UnexpectedEof`.
306    pub async fn receive_backend(&mut self) -> io::Result<BackendMessage> {
307        self.receive_wire().await
308    }
309
310    /// Receives the next protocol-advancing message through the async demux.
311    ///
312    /// # Errors
313    ///
314    /// Returns decoding and underlying transport read errors, or `UnexpectedEof`.
315    pub async fn receive_session(&mut self) -> io::Result<SessionItem> {
316        loop {
317            let message = self.receive_backend().await?;
318            if let Some(item) = self.project_backend(message) {
319                return Ok(item);
320            }
321        }
322    }
323
324    /// Projects an inspected or modified backend message into the session stream.
325    pub fn project_backend(&mut self, message: BackendMessage) -> Option<SessionItem> {
326        self.demux.route(message)
327    }
328}
329
330impl<S, D, Phase, Cleanliness> Conn<Buffered<S, D>, Phase, Cleanliness> {
331    /// Adds an already-typed message to this connection's outbound buffer.
332    ///
333    /// # Errors
334    ///
335    /// Returns an error when the frame is too large to encode.
336    pub fn push_frame(&mut self, frame: Frame) -> io::Result<()> {
337        self.transport_mut().push(frame)
338    }
339
340    #[must_use]
341    /// Returns encoded output which has not yet been flushed.
342    pub fn pending_output(&self) -> &[u8] {
343        self.transport().pending()
344    }
345}
346
347impl<S, Cleanliness> Conn<Buffered<S, Backend>, PreStartup, Cleanliness> {
348    /// Buffers an `SSLRequest` and enters the raw single-byte reply phase.
349    pub fn request_ssl(mut self) -> Conn<Buffered<S, Backend>, AwaitingSslReply, Cleanliness> {
350        self.transport_mut().push_raw(&ssl_request_packet());
351        self.transition()
352    }
353
354    /// Buffers a `GSSENCRequest` and enters the raw single-byte reply phase.
355    pub fn request_gss(
356        mut self,
357    ) -> Conn<Buffered<S, Backend>, crate::pre_startup::AwaitingGssReply, Cleanliness> {
358        self.transport_mut().push_raw(&gssenc_request_packet());
359        self.transition()
360    }
361}
362
363impl<S, Cleanliness> Conn<Buffered<S, Frontend>, ServerSslDecision, Cleanliness> {
364    /// Buffers the server's raw `S` response and enters the TLS handshake phase.
365    pub fn approve_ssl(mut self) -> Conn<Buffered<S, Frontend>, TlsHandshake, Cleanliness> {
366        self.transport_mut().push_raw(b"S");
367        self.transition()
368    }
369
370    /// Buffers the server's raw `N` response and returns to pre-startup choice.
371    pub fn decline_ssl(mut self) -> Conn<Buffered<S, Frontend>, PreStartup, Cleanliness> {
372        self.transport_mut().push_raw(b"N");
373        self.transition()
374    }
375
376    /// Buffers the historical raw `E` response and terminates negotiation.
377    pub fn reject_ssl_with_legacy_error(
378        mut self,
379    ) -> Conn<Buffered<S, Frontend>, crate::pre_startup::Terminated, Cleanliness> {
380        self.transport_mut().push_raw(b"E");
381        self.transition()
382    }
383}
384
385impl<S, Cleanliness>
386    Conn<Buffered<S, Frontend>, crate::pre_startup::ServerGssDecision, Cleanliness>
387{
388    /// Buffers the server's raw `S` response and enters the GSS handshake phase.
389    pub fn approve_gss(
390        mut self,
391    ) -> Conn<Buffered<S, Frontend>, crate::pre_startup::GssHandshake, Cleanliness> {
392        self.transport_mut().push_raw(b"S");
393        self.transition()
394    }
395
396    /// Buffers the server's raw `N` response and returns to pre-startup choice.
397    pub fn decline_gss(mut self) -> Conn<Buffered<S, Frontend>, PreStartup, Cleanliness> {
398        self.transport_mut().push_raw(b"N");
399        self.transition()
400    }
401
402    /// Buffers the historical raw `E` response and terminates negotiation.
403    pub fn reject_gss_with_legacy_error(
404        mut self,
405    ) -> Conn<Buffered<S, Frontend>, crate::pre_startup::Terminated, Cleanliness> {
406        self.transport_mut().push_raw(b"E");
407        self.transition()
408    }
409}
410
411impl<S: AsyncRead + Unpin, Cleanliness> Conn<Buffered<S, Backend>, AwaitingSslReply, Cleanliness> {
412    /// Receives and projects the server's raw SSL decision byte.
413    ///
414    /// # Errors
415    ///
416    /// Returns an I/O error or rejects a byte other than `S`, `N`, or `E`.
417    pub async fn receive_ssl_reply(
418        mut self,
419    ) -> io::Result<Negotiation<Buffered<S, Backend>, TlsHandshake, Cleanliness>> {
420        let reply = self.transport_mut().receive_encryption_reply().await?;
421        Ok(match reply {
422            EncryptionReply::Accepted => Negotiation::Accepted(self.transition()),
423            EncryptionReply::Rejected => Negotiation::Rejected(self.transition()),
424            EncryptionReply::LegacyError => Negotiation::LegacyError(self.transition()),
425        })
426    }
427
428    /// Receives the server decision and enforces the selected plaintext fallback policy.
429    ///
430    /// # Errors
431    ///
432    /// Returns an I/O error or rejects a byte other than `S`, `N`, or `E`.
433    pub async fn receive_ssl_reply_for_mode(
434        mut self,
435        mode: SslMode,
436    ) -> io::Result<SslModeNegotiation<Buffered<S, Backend>, Cleanliness>> {
437        let reply = self.transport_mut().receive_encryption_reply().await?;
438        Ok(self.apply_ssl_reply(reply, mode))
439    }
440}
441
442impl<S: AsyncRead + Unpin, Cleanliness>
443    Conn<Buffered<S, Backend>, crate::pre_startup::AwaitingGssReply, Cleanliness>
444{
445    /// Receives and projects the server's raw GSSENC decision byte.
446    ///
447    /// # Errors
448    ///
449    /// Returns an I/O error or rejects a byte other than `S`, `N`, or `E`.
450    pub async fn receive_gss_reply(
451        mut self,
452    ) -> io::Result<Negotiation<Buffered<S, Backend>, crate::pre_startup::GssHandshake, Cleanliness>>
453    {
454        let reply = self.transport_mut().receive_encryption_reply().await?;
455        Ok(match reply {
456            EncryptionReply::Accepted => Negotiation::Accepted(self.transition()),
457            EncryptionReply::Rejected => Negotiation::Rejected(self.transition()),
458            EncryptionReply::LegacyError => Negotiation::LegacyError(self.transition()),
459        })
460    }
461}
462
463impl<S, Cleanliness> Conn<Buffered<S, Backend>, TlsHandshake, Cleanliness>
464where
465    S: AsyncRead + AsyncWrite + Unpin,
466{
467    /// Completes a client-side TLS handshake and changes the transport type.
468    ///
469    /// # Errors
470    ///
471    /// Returns a TLS handshake, certificate, channel-binding, or buffer-state error.
472    pub async fn connect_tls(
473        self,
474        server_name: ServerName<'static>,
475        config: Arc<ClientConfig>,
476    ) -> io::Result<Conn<Buffered<ClientTls<S>, Backend>, PreStartup, Cleanliness>> {
477        let transport = self.into_transport();
478        Ok(Conn::new(transport.connect_tls(server_name, config).await?)
479            .transition::<PreStartup, Cleanliness>())
480    }
481}
482
483impl<S, Cleanliness> Conn<Buffered<S, Frontend>, TlsHandshake, Cleanliness>
484where
485    S: AsyncRead + AsyncWrite + Unpin,
486{
487    /// Completes a server-side TLS handshake and changes the transport type.
488    ///
489    /// # Errors
490    ///
491    /// Returns a TLS handshake, certificate, channel-binding, or buffer-state error.
492    pub async fn accept_tls(
493        self,
494        config: Arc<ServerConfig>,
495        leaf_certificate: CertificateDer<'static>,
496    ) -> io::Result<Conn<Buffered<ServerTls<S>, Frontend>, PreStartup, Cleanliness>> {
497        let transport = self.into_transport();
498        Ok(
499            Conn::new(transport.accept_tls(config, leaf_certificate).await?)
500                .transition::<PreStartup, Cleanliness>(),
501        )
502    }
503}
504
505impl<S, D, Cleanliness> Conn<Buffered<S, D>, crate::pre_startup::Startup, Cleanliness> {
506    /// Buffers the raw, untagged startup packet before normal framing begins.
507    pub fn push_startup_packet(&mut self, packet: &[u8]) {
508        self.transport_mut().outbound.extend_from_slice(packet);
509    }
510}
511
512impl<S: AsyncWrite + Unpin, D, Phase, Cleanliness> Conn<Buffered<S, D>, Phase, Cleanliness> {
513    /// Flushes buffered output while retaining ownership of the typed connection.
514    ///
515    /// # Errors
516    ///
517    /// Returns an error from the underlying transport.
518    pub async fn flush(&mut self) -> io::Result<()> {
519        self.transport_mut().flush().await
520    }
521}
522
523impl<S: AsyncRead + Unpin, Phase, Cleanliness> Conn<Buffered<S, Backend>, Phase, Cleanliness> {
524    /// Receives one backend message before demultiplexing or state advancement.
525    /// This is the interception point for proxy policy and message rewriting.
526    ///
527    /// # Errors
528    ///
529    /// Returns decoding and underlying transport read errors, or `UnexpectedEof`.
530    pub async fn receive_backend_wire(&mut self) -> io::Result<BackendMessage> {
531        self.transport_mut().receive_backend().await
532    }
533
534    /// Receives one backend message through middleware indexed by this connection phase.
535    ///
536    /// Unlike [`Self::receive_backend_wire_with_middleware`], callers do not pass
537    /// a runtime protocol state. `Phase` selects the generated legal message set
538    /// and the server sender role at compile time.
539    ///
540    /// # Errors
541    ///
542    /// Returns an I/O or decoding error, an illegal peer message, a middleware
543    /// policy error, or a phase-legal replacement with an invalid wire shape.
544    pub async fn receive_backend_typed<State, Handler>(
545        &mut self,
546        middleware: &mut Middleware<State, Handler>,
547    ) -> Result<
548        <Phase as TypedPhase<ServerRole, BackendMessage>>::Message,
549        TypedReceiveError<Handler::Error, BackendMessage>,
550    >
551    where
552        Phase: TypedPhase<ServerRole, BackendMessage>,
553        Handler: TypedMiddleware<
554                ServerRole,
555                <Phase as TypedPhase<ServerRole, BackendMessage>>::ProtocolPhase,
556                <Phase as TypedPhase<ServerRole, BackendMessage>>::Message,
557                State,
558            >,
559    {
560        let message = self
561            .receive_backend_wire()
562            .await
563            .map_err(TypedReceiveError::Io)?;
564        let message = <Phase as TypedPhase<ServerRole, BackendMessage>>::Message::try_from(message)
565            .map_err(TypedReceiveError::Illegal)?;
566        let message = middleware
567            .intercept_typed::<
568                ServerRole,
569                <Phase as TypedPhase<ServerRole, BackendMessage>>::ProtocolPhase,
570                _,
571            >(message)
572            .map_err(TypedReceiveError::Middleware)?;
573        if message.as_ref().is_reconstructable() {
574            Ok(message)
575        } else {
576            Err(TypedReceiveError::InvalidWire(message.into()))
577        }
578    }
579
580    /// Receives typed backend traffic until one protocol-advancing item remains.
581    ///
582    /// Asynchronous messages pass through the same middleware, are recorded by
583    /// the demultiplexer in wire order, and leave `Phase` unchanged.
584    ///
585    /// # Errors
586    ///
587    /// Returns the same failures as [`Self::receive_backend_typed`].
588    pub async fn receive_typed<State, Handler>(
589        &mut self,
590        middleware: &mut Middleware<State, Handler>,
591    ) -> Result<SessionItem, TypedReceiveError<Handler::Error, BackendMessage>>
592    where
593        Phase: TypedPhase<ServerRole, BackendMessage>,
594        Handler: TypedMiddleware<
595                ServerRole,
596                <Phase as TypedPhase<ServerRole, BackendMessage>>::ProtocolPhase,
597                <Phase as TypedPhase<ServerRole, BackendMessage>>::Message,
598                State,
599            >,
600    {
601        loop {
602            let message = self.receive_backend_typed(middleware).await?;
603            if let Some(item) = self.project_backend(message.into()) {
604                return Ok(item);
605            }
606        }
607    }
608
609    /// Receives one SSL or GSSENC decision through phase-typed middleware.
610    ///
611    /// `Phase` must be an encryption-reply phase; callers then consume the
612    /// connection with its existing typed `receive_reply` projection.
613    ///
614    /// # Errors
615    ///
616    /// Returns an I/O error, illegal decision, middleware policy error, or an
617    /// invalid replacement wire shape.
618    pub async fn receive_encryption_reply_typed<State, Handler>(
619        &mut self,
620        middleware: &mut Middleware<State, Handler>,
621    ) -> Result<
622        <Phase as TypedPhase<ServerRole, EncryptionReply>>::Message,
623        TypedReceiveError<Handler::Error, EncryptionReply>,
624    >
625    where
626        Phase: TypedPhase<ServerRole, EncryptionReply>,
627        Handler: TypedMiddleware<
628                ServerRole,
629                <Phase as TypedPhase<ServerRole, EncryptionReply>>::ProtocolPhase,
630                <Phase as TypedPhase<ServerRole, EncryptionReply>>::Message,
631                State,
632            >,
633    {
634        let message = self
635            .transport_mut()
636            .receive_encryption_reply()
637            .await
638            .map_err(TypedReceiveError::Io)?;
639        let message =
640            <Phase as TypedPhase<ServerRole, EncryptionReply>>::Message::try_from(message)
641                .map_err(TypedReceiveError::Illegal)?;
642        let message = middleware
643            .intercept_typed::<
644                ServerRole,
645                <Phase as TypedPhase<ServerRole, EncryptionReply>>::ProtocolPhase,
646                _,
647            >(message)
648            .map_err(TypedReceiveError::Middleware)?;
649        if message.as_ref().is_reconstructable() {
650            Ok(message)
651        } else {
652            Err(TypedReceiveError::InvalidWire(message.into()))
653        }
654    }
655
656    /// Receives, intercepts, and validates one backend message before projection.
657    ///
658    /// # Errors
659    ///
660    /// Returns an I/O, decoding, middleware-policy, or state-validation error.
661    pub async fn receive_backend_wire_with_middleware<State, Handler, ProtocolState>(
662        &mut self,
663        middleware: &mut Middleware<State, Handler>,
664        protocol_state: &ProtocolState,
665    ) -> Result<BackendMessage, ReceiveError<Handler::Error, BackendMessage>>
666    where
667        Handler: MessageMiddleware<BackendMessage, State>,
668        ProtocolState: AcceptsMessage<BackendMessage>,
669    {
670        let message = self
671            .receive_backend_wire()
672            .await
673            .map_err(ReceiveError::Io)?;
674        middleware
675            .intercept_checked(protocol_state, message)
676            .map_err(ReceiveError::Intercept)
677    }
678
679    /// Projects an inspected or modified message into the filtered session stream.
680    pub fn project_backend(&mut self, message: BackendMessage) -> Option<SessionItem> {
681        self.transport_mut().project_backend(message)
682    }
683
684    /// Receives the next message in the filtered session projection.
685    ///
686    /// # Errors
687    ///
688    /// Returns decoding and underlying transport read errors, or `UnexpectedEof`.
689    pub async fn receive(&mut self) -> io::Result<SessionItem> {
690        self.transport_mut().receive_session().await
691    }
692
693    /// Receives backend messages through middleware before demultiplexing.
694    ///
695    /// Asynchronous messages are intercepted and then recorded by the demux;
696    /// this method continues until a protocol-advancing item is available.
697    ///
698    /// # Errors
699    ///
700    /// Returns an I/O, decoding, middleware-policy, or state-validation error.
701    pub async fn receive_with_middleware<State, Handler, ProtocolState>(
702        &mut self,
703        middleware: &mut Middleware<State, Handler>,
704        protocol_state: &ProtocolState,
705    ) -> Result<SessionItem, ReceiveError<Handler::Error, BackendMessage>>
706    where
707        Handler: MessageMiddleware<BackendMessage, State>,
708        ProtocolState: AcceptsMessage<BackendMessage>,
709    {
710        loop {
711            let message = self
712                .receive_backend_wire_with_middleware(middleware, protocol_state)
713                .await?;
714            if let Some(item) = self.project_backend(message) {
715                return Ok(item);
716            }
717        }
718    }
719
720    #[must_use]
721    /// Returns the latest upstream cancellation key observed during startup.
722    pub fn cancel_key(&self) -> Option<&CancelKey> {
723        self.transport().demux().cancel_key()
724    }
725
726    /// Returns the latest backend parameter values observed by the demux.
727    #[must_use]
728    pub fn parameters(&self) -> &BTreeMap<Bytes, Bytes> {
729        self.transport().demux().parameters()
730    }
731
732    /// Returns whether current parameters differ from the startup baseline.
733    #[must_use]
734    pub fn parameters_changed(&self) -> bool {
735        self.transport().demux().parameters_changed()
736    }
737
738    /// Returns the latest transaction status observed in `ReadyForQuery`.
739    #[must_use]
740    pub fn transaction_status(&self) -> Option<crate::codec::TransactionStatus> {
741        self.transport().demux().transaction_status()
742    }
743
744    /// Removes the oldest queued asynchronous notification.
745    pub fn pop_notification(&mut self) -> Option<Notification> {
746        self.transport_mut().demux_mut().pop_notification()
747    }
748
749    /// Removes the next tagged notice for prompt forwarding to the client.
750    pub fn pop_notice(&mut self) -> Option<TaggedNotice> {
751        self.transport_mut().demux_mut().pop_notice()
752    }
753
754    /// Removes the next ordered parameter update for forwarding to the client.
755    pub fn pop_parameter_status(&mut self) -> Option<ParameterStatus> {
756        self.transport_mut().demux_mut().pop_parameter_status()
757    }
758
759    /// Removes the next independent backend event in original wire order.
760    pub fn pop_async_event(&mut self) -> Option<OrderedAsyncEvent> {
761        self.transport_mut().demux_mut().pop_async_event()
762    }
763}
764
765impl<S: AsyncRead + Unpin, Phase, Cleanliness> Conn<Buffered<S, Frontend>, Phase, Cleanliness> {
766    /// Receives one frontend message before any server-role state advancement.
767    ///
768    /// # Errors
769    ///
770    /// Returns decoding and underlying transport read errors, or `UnexpectedEof`.
771    pub async fn receive_frontend_wire(&mut self) -> io::Result<FrontendMessage> {
772        self.transport_mut().receive_wire().await
773    }
774
775    /// Receives one frontend message through middleware indexed by this connection phase.
776    ///
777    /// `Phase` selects the generated legal message set and the client sender role
778    /// at compile time, so no runtime protocol-state argument is accepted.
779    ///
780    /// # Errors
781    ///
782    /// Returns an I/O or decoding error, an illegal peer message, a middleware
783    /// policy error, or a phase-legal replacement with an invalid wire shape.
784    pub async fn receive_frontend_typed<State, Handler>(
785        &mut self,
786        middleware: &mut Middleware<State, Handler>,
787    ) -> Result<
788        <Phase as TypedPhase<ClientRole, FrontendMessage>>::Message,
789        TypedReceiveError<Handler::Error, FrontendMessage>,
790    >
791    where
792        Phase: TypedPhase<ClientRole, FrontendMessage>,
793        Handler: TypedMiddleware<
794                ClientRole,
795                <Phase as TypedPhase<ClientRole, FrontendMessage>>::ProtocolPhase,
796                <Phase as TypedPhase<ClientRole, FrontendMessage>>::Message,
797                State,
798            >,
799    {
800        let message = self
801            .receive_frontend_wire()
802            .await
803            .map_err(TypedReceiveError::Io)?;
804        let message =
805            <Phase as TypedPhase<ClientRole, FrontendMessage>>::Message::try_from(message)
806                .map_err(TypedReceiveError::Illegal)?;
807        let message = middleware
808            .intercept_typed::<
809                ClientRole,
810                <Phase as TypedPhase<ClientRole, FrontendMessage>>::ProtocolPhase,
811                _,
812            >(message)
813            .map_err(TypedReceiveError::Middleware)?;
814        if message.as_ref().is_reconstructable() {
815            Ok(message)
816        } else {
817            Err(TypedReceiveError::InvalidWire(message.into()))
818        }
819    }
820
821    /// Receives, intercepts, and validates one frontend message before projection.
822    ///
823    /// # Errors
824    ///
825    /// Returns an I/O, decoding, middleware-policy, or state-validation error.
826    pub async fn receive_frontend_wire_with_middleware<State, Handler, ProtocolState>(
827        &mut self,
828        middleware: &mut Middleware<State, Handler>,
829        protocol_state: &ProtocolState,
830    ) -> Result<FrontendMessage, ReceiveError<Handler::Error, FrontendMessage>>
831    where
832        Handler: MessageMiddleware<FrontendMessage, State>,
833        ProtocolState: AcceptsMessage<FrontendMessage>,
834    {
835        let message = self
836            .receive_frontend_wire()
837            .await
838            .map_err(ReceiveError::Io)?;
839        middleware
840            .intercept_checked(protocol_state, message)
841            .map_err(ReceiveError::Intercept)
842    }
843}
844
845impl<S: AsyncRead + Unpin, Cleanliness> Conn<Buffered<S, Frontend>, PreStartup, Cleanliness> {
846    /// Receives a raw pre-startup packet before server-role state projection.
847    ///
848    /// # Errors
849    ///
850    /// Returns malformed pre-startup data and underlying transport read errors.
851    pub async fn receive_pre_startup_wire(&mut self) -> io::Result<PreStartupMessage> {
852        self.transport_mut().receive_pre_startup().await
853    }
854
855    /// Receives a client pre-startup packet through phase-typed middleware.
856    ///
857    /// # Errors
858    ///
859    /// Returns an I/O or decoding error, an illegal pre-startup packet, a
860    /// middleware policy error, or an invalid replacement wire shape.
861    pub async fn receive_pre_startup_typed<State, Handler>(
862        &mut self,
863        middleware: &mut Middleware<State, Handler>,
864    ) -> Result<
865        <PreStartup as TypedPhase<ClientRole, PreStartupMessage>>::Message,
866        TypedReceiveError<Handler::Error, PreStartupMessage>,
867    >
868    where
869        Handler: TypedMiddleware<
870                ClientRole,
871                <PreStartup as TypedPhase<ClientRole, PreStartupMessage>>::ProtocolPhase,
872                <PreStartup as TypedPhase<ClientRole, PreStartupMessage>>::Message,
873                State,
874            >,
875    {
876        let message = self
877            .receive_pre_startup_wire()
878            .await
879            .map_err(TypedReceiveError::Io)?;
880        let message =
881            <PreStartup as TypedPhase<ClientRole, PreStartupMessage>>::Message::try_from(message)
882                .map_err(TypedReceiveError::Illegal)?;
883        let message = middleware
884            .intercept_typed::<
885                ClientRole,
886                <PreStartup as TypedPhase<ClientRole, PreStartupMessage>>::ProtocolPhase,
887                _,
888            >(message)
889            .map_err(TypedReceiveError::Middleware)?;
890        if message.as_ref().is_reconstructable() {
891            Ok(message)
892        } else {
893            Err(TypedReceiveError::InvalidWire(message.into()))
894        }
895    }
896
897    /// Receives, intercepts, and validates one untagged pre-startup message.
898    ///
899    /// # Errors
900    ///
901    /// Returns an I/O, decoding, middleware-policy, or state-validation error.
902    pub async fn receive_pre_startup_wire_with_middleware<State, Handler, ProtocolState>(
903        &mut self,
904        middleware: &mut Middleware<State, Handler>,
905        protocol_state: &ProtocolState,
906    ) -> Result<PreStartupMessage, ReceiveError<Handler::Error, PreStartupMessage>>
907    where
908        Handler: MessageMiddleware<PreStartupMessage, State>,
909        ProtocolState: AcceptsMessage<PreStartupMessage>,
910    {
911        let message = self
912            .receive_pre_startup_wire()
913            .await
914            .map_err(ReceiveError::Io)?;
915        middleware
916            .intercept_checked(protocol_state, message)
917            .map_err(ReceiveError::Intercept)
918    }
919}
920
921#[cfg(test)]
922mod tests {
923    use std::{
924        convert::Infallible,
925        future::Future,
926        pin::Pin,
927        task::{Context, Poll},
928    };
929
930    use bytes::Bytes;
931    use tokio::io::AsyncWrite;
932
933    use super::*;
934    use crate::{
935        grammar::{backend, frontend, server_pre_startup},
936        middleware::{InterceptError, Middleware, ReceiveError},
937    };
938
939    #[derive(Debug, Default)]
940    struct ShortWriter {
941        output: Vec<u8>,
942    }
943
944    impl AsyncWrite for ShortWriter {
945        fn poll_write(
946            mut self: Pin<&mut Self>,
947            _cx: &mut Context<'_>,
948            buffer: &[u8],
949        ) -> Poll<io::Result<usize>> {
950            let written = buffer.len().min(2);
951            self.output.extend_from_slice(&buffer[..written]);
952            Poll::Ready(Ok(written))
953        }
954
955        fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
956            Poll::Ready(Ok(()))
957        }
958
959        fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
960            Poll::Ready(Ok(()))
961        }
962    }
963
964    #[tokio::test]
965    async fn flush_handles_partial_writes_without_losing_bytes() {
966        let frame = Frame {
967            tag: b'S',
968            body: Bytes::new(),
969        };
970        let mut transport = Buffered::new(ShortWriter::default());
971        transport.push(frame).expect("encodable frame");
972        assert_eq!(transport.pending(), &[b'S', 0, 0, 0, 4]);
973        transport.flush().await.expect("writable transport");
974        assert!(transport.pending().is_empty());
975        assert_eq!(transport.into_inner().output, [b'S', 0, 0, 0, 4]);
976    }
977
978    #[test]
979    fn buffered_transport_enforces_its_frame_limit_on_output() {
980        let mut transport = Buffered::<_, Backend>::with_max_frame_len((), 9).unwrap();
981        let error = transport
982            .push(Frame {
983                tag: b'Q',
984                body: Bytes::from_static(b"12345"),
985            })
986            .unwrap_err();
987        assert_eq!(error.kind(), io::ErrorKind::InvalidInput);
988        assert!(transport.pending().is_empty());
989    }
990
991    #[test]
992    fn cancelling_flush_retains_unwritten_bytes() {
993        #[derive(Debug, Default)]
994        struct PausingWriter {
995            output: Vec<u8>,
996            blocked: bool,
997        }
998
999        impl AsyncWrite for PausingWriter {
1000            fn poll_write(
1001                mut self: Pin<&mut Self>,
1002                _cx: &mut Context<'_>,
1003                buffer: &[u8],
1004            ) -> Poll<io::Result<usize>> {
1005                if self.blocked {
1006                    return Poll::Pending;
1007                }
1008                let written = buffer.len().min(2);
1009                self.output.extend_from_slice(&buffer[..written]);
1010                self.blocked = true;
1011                Poll::Ready(Ok(written))
1012            }
1013
1014            fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
1015                Poll::Ready(Ok(()))
1016            }
1017
1018            fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
1019                Poll::Ready(Ok(()))
1020            }
1021        }
1022
1023        let mut transport = Buffered::new(PausingWriter::default());
1024        transport
1025            .push(Frame {
1026                tag: b'S',
1027                body: Bytes::new(),
1028            })
1029            .expect("encodable frame");
1030
1031        let mut flush = Box::pin(transport.flush());
1032        let waker = std::task::Waker::noop();
1033        let mut context = Context::from_waker(waker);
1034        assert!(flush.as_mut().poll(&mut context).is_pending());
1035        drop(flush);
1036
1037        assert_eq!(transport.pending(), &[0, 0, 4]);
1038        assert_eq!(transport.io.output, [b'S', 0]);
1039    }
1040
1041    #[tokio::test]
1042    async fn receive_filters_parameter_status_before_session_message() {
1043        let (client, mut server) = tokio::io::duplex(256);
1044        let mut wire = BytesMut::new();
1045        let mut encoder = PgCodec::<Backend>::default();
1046        encoder
1047            .encode(
1048                Frame {
1049                    tag: b'S',
1050                    body: Bytes::from_static(b"client_encoding\0UTF8\0"),
1051                },
1052                &mut wire,
1053            )
1054            .expect("encodable ParameterStatus");
1055        encoder
1056            .encode(
1057                Frame {
1058                    tag: b'Z',
1059                    body: Bytes::from_static(b"I"),
1060                },
1061                &mut wire,
1062            )
1063            .expect("encodable ReadyForQuery");
1064        server.write_all(&wire).await.expect("writable test peer");
1065
1066        let mut transport = Buffered::new(client);
1067        assert_eq!(
1068            transport.receive_session().await.expect("valid messages"),
1069            SessionItem::ReadyForQuery {
1070                status: crate::codec::TransactionStatus::Idle,
1071                parameters_changed: false,
1072            }
1073        );
1074        assert_eq!(
1075            transport
1076                .demux()
1077                .parameters()
1078                .get(&Bytes::from_static(b"client_encoding")),
1079            Some(&Bytes::from_static(b"UTF8"))
1080        );
1081        let conn: Conn<_, crate::auth::Ready> = Conn::new(transport).transition();
1082        assert_eq!(
1083            conn.parameters().get(b"client_encoding".as_slice()),
1084            Some(&Bytes::from_static(b"UTF8"))
1085        );
1086        assert!(!conn.parameters_changed());
1087        assert_eq!(
1088            conn.transaction_status(),
1089            Some(crate::codec::TransactionStatus::Idle)
1090        );
1091        conn.into_transport();
1092    }
1093
1094    #[tokio::test]
1095    async fn wire_message_can_be_modified_before_projection() {
1096        let (client, mut server) = tokio::io::duplex(128);
1097        let original = BackendMessage::ParameterStatus {
1098            name: Bytes::from_static(b"application_name"),
1099            value: Bytes::from_static(b"upstream"),
1100        };
1101        let mut bytes = BytesMut::new();
1102        PgCodec::<Backend>::default()
1103            .encode(
1104                original.to_frame().expect("reconstructable message"),
1105                &mut bytes,
1106            )
1107            .expect("encodable message");
1108        server.write_all(&bytes).await.expect("writable test peer");
1109
1110        let mut transport = Buffered::new(client);
1111        let mut message = transport
1112            .receive_backend()
1113            .await
1114            .expect("decodable message");
1115        let BackendMessage::ParameterStatus { value, .. } = &mut message else {
1116            panic!("unexpected message")
1117        };
1118        *value = Bytes::from_static(b"proxy");
1119        assert!(transport.project_backend(message).is_none());
1120        assert_eq!(
1121            transport
1122                .demux()
1123                .parameters()
1124                .get(&Bytes::from_static(b"application_name")),
1125            Some(&Bytes::from_static(b"proxy"))
1126        );
1127    }
1128
1129    #[tokio::test]
1130    async fn typed_middleware_accepts_async_traffic_without_advancing_ready() {
1131        let (client, mut server) = tokio::io::duplex(128);
1132        let original = BackendMessage::ParameterStatus {
1133            name: Bytes::from_static(b"application_name"),
1134            value: Bytes::from_static(b"upstream"),
1135        };
1136        let mut bytes = BytesMut::new();
1137        PgCodec::<Backend>::default()
1138            .encode(
1139                original.to_frame().expect("reconstructable message"),
1140                &mut bytes,
1141            )
1142            .expect("encodable message");
1143        server.write_all(&bytes).await.expect("writable test peer");
1144
1145        let transport = Buffered::new(client);
1146        let mut conn: Conn<_, crate::auth::Ready> = Conn::new(transport).transition();
1147        let mut middleware = Middleware::new(0_usize, |seen: &mut usize, _message| {
1148            *seen += 1;
1149            let replacement = BackendMessage::ParameterStatus {
1150                name: Bytes::from_static(b"application_name"),
1151                value: Bytes::from_static(b"proxy"),
1152            };
1153            match crate::middleware::TypedBackendMessage::try_from(replacement) {
1154                Ok(replacement) => Ok::<_, Infallible>(replacement),
1155                Err(message) => panic!("parameter status must be asynchronous: {message:?}"),
1156            }
1157        });
1158
1159        let message = conn
1160            .receive_backend_typed(&mut middleware)
1161            .await
1162            .expect("typed asynchronous message");
1163        assert!(conn.project_backend(message.into()).is_none());
1164        assert_eq!(*middleware.state(), 1);
1165        assert_eq!(
1166            conn.parameters().get(b"application_name".as_slice()),
1167            Some(&Bytes::from_static(b"proxy"))
1168        );
1169        conn.into_transport();
1170    }
1171
1172    #[tokio::test]
1173    async fn typed_receive_projects_into_the_existing_next_connection_enum() {
1174        let (client, mut server) = tokio::io::duplex(128);
1175        let ready = BackendMessage::ReadyForQuery(crate::codec::TransactionStatus::Idle);
1176        let mut bytes = BytesMut::new();
1177        PgCodec::<Backend>::default()
1178            .encode(
1179                ready.to_frame().expect("reconstructable message"),
1180                &mut bytes,
1181            )
1182            .expect("encodable message");
1183        server.write_all(&bytes).await.expect("writable test peer");
1184
1185        let transport = Buffered::new(client);
1186        let conn: Conn<_, crate::auth::Ready> = Conn::new(transport).transition();
1187        let (mut query, _) = conn
1188            .push_stateless_query(b"select 1")
1189            .expect("encodable query");
1190        let mut middleware = Middleware::new((), crate::middleware::Identity);
1191        let item = query
1192            .receive_typed(&mut middleware)
1193            .await
1194            .expect("phase-legal ready message");
1195
1196        let transition = query.offer(item).expect("typed next-state projection");
1197        let crate::session::SimpleTransition::Ready(crate::session::ReadyState::Clean(ready)) =
1198            transition
1199        else {
1200            panic!("idle readiness must return a clean ready connection");
1201        };
1202        ready.into_transport();
1203    }
1204
1205    #[tokio::test]
1206    async fn typed_receive_keeps_wire_shape_validation_at_runtime() {
1207        let (client, mut peer) = tokio::io::duplex(256);
1208        let query = FrontendMessage::Query(Bytes::from_static(b"select 1"));
1209        let mut bytes = BytesMut::new();
1210        PgCodec::<Frontend>::default()
1211            .encode(query.to_frame().expect("reconstructable query"), &mut bytes)
1212            .expect("encodable query");
1213        peer.write_all(&bytes).await.expect("writable test peer");
1214
1215        let transport = Buffered::<_, Frontend>::new_frontend(client);
1216        let mut conn: Conn<_, crate::auth::Ready> = Conn::new(transport).transition();
1217        let invalid = FrontendMessage::Parse(crate::codec::Parse {
1218            statement: Bytes::from_static(b"invalid\0statement"),
1219            query: Bytes::from_static(b"select 2"),
1220            parameter_types: Vec::new(),
1221        });
1222        let mut middleware = Middleware::new((), move |_state: &mut (), _message| {
1223            let invalid = invalid.clone();
1224            match backend::ReadyExternalMessage::try_from(invalid) {
1225                Ok(invalid) => Ok::<_, Infallible>(invalid),
1226                Err(message) => panic!("parse must be protocol-legal while ready: {message:?}"),
1227            }
1228        });
1229
1230        let result = conn.receive_frontend_typed(&mut middleware).await;
1231        assert!(matches!(
1232            result,
1233            Err(TypedReceiveError::InvalidWire(FrontendMessage::Parse(_)))
1234        ));
1235        conn.into_transport();
1236    }
1237
1238    #[tokio::test]
1239    async fn middleware_rewrites_backend_before_demux_bookkeeping() {
1240        let (client, mut server) = tokio::io::duplex(256);
1241        let mut wire = BytesMut::new();
1242        let mut encoder = PgCodec::<Backend>::default();
1243        for message in [
1244            BackendMessage::BackendKeyData {
1245                process_id: 7,
1246                secret_key: Bytes::from_static(b"old!"),
1247            },
1248            BackendMessage::ParameterStatus {
1249                name: Bytes::from_static(b"application_name"),
1250                value: Bytes::from_static(b"upstream"),
1251            },
1252            BackendMessage::ReadyForQuery(crate::codec::TransactionStatus::Idle),
1253        ] {
1254            encoder
1255                .encode(
1256                    message.to_frame().expect("reconstructable message"),
1257                    &mut wire,
1258                )
1259                .expect("encodable message");
1260        }
1261        server.write_all(&wire).await.expect("writable test peer");
1262
1263        let transport = Buffered::new(client);
1264        let mut conn: Conn<_, crate::auth::Ready> = Conn::new(transport).transition();
1265        let mut middleware = Middleware::new(0_usize, |seen: &mut usize, mut message| {
1266            *seen += 1;
1267            if let BackendMessage::ParameterStatus { value, .. } = &mut message {
1268                *value = Bytes::from_static(b"proxy");
1269            }
1270            if let BackendMessage::BackendKeyData {
1271                process_id,
1272                secret_key,
1273            } = &mut message
1274            {
1275                *process_id = 9;
1276                *secret_key = Bytes::from_static(b"new!");
1277            }
1278            if let BackendMessage::ReadyForQuery(status) = &mut message {
1279                *status = crate::codec::TransactionStatus::InTransaction;
1280            }
1281            Ok::<_, Infallible>(message)
1282        });
1283
1284        assert!(matches!(
1285            conn.receive_with_middleware(&mut middleware, &frontend::RuntimeState::Simple)
1286                .await,
1287            Ok(SessionItem::Message(BackendMessage::BackendKeyData { .. }))
1288        ));
1289        assert!(matches!(
1290            conn.receive_with_middleware(&mut middleware, &frontend::RuntimeState::Simple)
1291                .await,
1292            Ok(SessionItem::ReadyForQuery { .. })
1293        ));
1294        assert_eq!(*middleware.state(), 3);
1295        assert_eq!(
1296            conn.parameters().get(b"application_name".as_slice()),
1297            Some(&Bytes::from_static(b"proxy"))
1298        );
1299        assert_eq!(
1300            conn.cancel_key(),
1301            Some(&CancelKey {
1302                process_id: 9,
1303                secret_key: Bytes::from_static(b"new!"),
1304            })
1305        );
1306        assert_eq!(
1307            conn.transaction_status(),
1308            Some(crate::codec::TransactionStatus::InTransaction)
1309        );
1310        conn.into_transport();
1311    }
1312
1313    #[tokio::test]
1314    async fn frontend_middleware_returns_illegal_replacement_before_projection() {
1315        let (proxy, mut client) = tokio::io::duplex(128);
1316        let original = FrontendMessage::CopyData(Bytes::from_static(b"row"));
1317        let mut bytes = BytesMut::new();
1318        PgCodec::<Frontend>::default()
1319            .encode(
1320                original.to_frame().expect("reconstructable message"),
1321                &mut bytes,
1322            )
1323            .expect("encodable message");
1324        client.write_all(&bytes).await.expect("writable client");
1325
1326        let mut conn = Conn::new(Buffered::<_, Frontend>::new_frontend(proxy));
1327        let mut middleware = Middleware::new((), |_state: &mut (), _message| {
1328            Ok::<_, Infallible>(FrontendMessage::Query(Bytes::from_static(b"select 1")))
1329        });
1330        let result = conn
1331            .receive_frontend_wire_with_middleware(
1332                &mut middleware,
1333                &backend::RuntimeState::SimpleCopyIn,
1334            )
1335            .await;
1336
1337        assert!(matches!(
1338            result,
1339            Err(ReceiveError::Intercept(InterceptError::Invalid(
1340                FrontendMessage::Query(_)
1341            )))
1342        ));
1343        conn.into_transport();
1344    }
1345
1346    #[tokio::test]
1347    async fn middleware_replacement_reaches_the_forwarded_peer() {
1348        let (client_side, mut client) = tokio::io::duplex(128);
1349        let original = FrontendMessage::Query(Bytes::from_static(b"select plaintext"));
1350        let mut bytes = BytesMut::new();
1351        PgCodec::<Frontend>::default()
1352            .encode(
1353                original.to_frame().expect("reconstructable message"),
1354                &mut bytes,
1355            )
1356            .expect("encodable message");
1357        client.write_all(&bytes).await.expect("writable client");
1358
1359        let mut downstream = Conn::new(Buffered::<_, Frontend>::new_frontend(client_side));
1360        let mut middleware = Middleware::new((), |_state: &mut (), _message| {
1361            Ok::<_, Infallible>(FrontendMessage::Query(Bytes::from_static(
1362                b"select encrypted",
1363            )))
1364        });
1365        let rewritten = downstream
1366            .receive_frontend_wire_with_middleware(&mut middleware, &backend::RuntimeState::Ready)
1367            .await
1368            .expect("legal rewritten query");
1369
1370        let (upstream_side, mut server) = tokio::io::duplex(128);
1371        let mut upstream = Buffered::<_, Backend>::new(upstream_side);
1372        upstream
1373            .push(rewritten.to_frame().expect("reconstructable replacement"))
1374            .expect("encodable replacement");
1375        upstream.flush().await.expect("writable upstream");
1376
1377        let mut received = BytesMut::new();
1378        server
1379            .read_buf(&mut received)
1380            .await
1381            .expect("readable upstream peer");
1382        assert_eq!(
1383            PgCodec::<Frontend>::default()
1384                .decode(&mut received)
1385                .expect("decodable frame"),
1386            Some(FrontendMessage::Query(Bytes::from_static(
1387                b"select encrypted"
1388            )))
1389        );
1390        downstream.into_transport();
1391    }
1392
1393    #[tokio::test]
1394    async fn pre_startup_middleware_can_replace_with_another_legal_choice() {
1395        let (proxy, mut client) = tokio::io::duplex(128);
1396        client
1397            .write_all(
1398                &PreStartupMessage::SslRequest
1399                    .to_packet()
1400                    .expect("encodable SSLRequest"),
1401            )
1402            .await
1403            .expect("writable client");
1404
1405        let mut conn = Conn::new(Buffered::<_, Frontend>::new_frontend(proxy));
1406        let replacement = PreStartupMessage::CancelRequest {
1407            process_id: 42,
1408            secret_key: Bytes::from_static(b"key!"),
1409        };
1410        let expected = replacement.clone();
1411        let mut middleware = Middleware::new((), move |_state: &mut (), _message| {
1412            Ok::<_, Infallible>(replacement.clone())
1413        });
1414
1415        assert_eq!(
1416            conn.receive_pre_startup_wire_with_middleware(
1417                &mut middleware,
1418                &server_pre_startup::RuntimeState::PreStartup,
1419            )
1420            .await
1421            .expect("legal replacement"),
1422            expected
1423        );
1424        conn.into_transport();
1425    }
1426
1427    #[tokio::test]
1428    async fn client_facing_transport_intercepts_typed_frontend_messages() {
1429        let (proxy, mut client) = tokio::io::duplex(128);
1430        let message = FrontendMessage::Query(Bytes::from_static(b"select plaintext"));
1431        let mut bytes = BytesMut::new();
1432        PgCodec::<Frontend>::default()
1433            .encode(
1434                message.to_frame().expect("reconstructable Query"),
1435                &mut bytes,
1436            )
1437            .expect("encodable Query");
1438        client.write_all(&bytes).await.expect("writable client");
1439
1440        let mut transport = Buffered::<_, Frontend>::new_frontend(proxy);
1441        let mut intercepted = transport.receive_wire().await.expect("decodable Query");
1442        let FrontendMessage::Query(query) = &mut intercepted else {
1443            panic!("unexpected frontend message")
1444        };
1445        *query = Bytes::from_static(b"select encrypted");
1446        assert_eq!(
1447            intercepted,
1448            FrontendMessage::Query(Bytes::from_static(b"select encrypted"))
1449        );
1450    }
1451
1452    #[tokio::test]
1453    async fn client_facing_transport_projects_repeated_pre_startup_choice() {
1454        let (proxy, mut client) = tokio::io::duplex(256);
1455        let ssl = PreStartupMessage::SslRequest
1456            .to_packet()
1457            .expect("encodable SSLRequest");
1458        let startup = PreStartupMessage::Startup(crate::startup::StartupMessage {
1459            version: crate::startup::ProtocolVersion::V3_2,
1460            parameters: std::collections::BTreeMap::from([(
1461                Bytes::from_static(b"user"),
1462                Bytes::from_static(b"postgres"),
1463            )]),
1464        });
1465        let startup_packet = startup.to_packet().expect("encodable StartupMessage");
1466        client.write_all(&ssl).await.expect("writable client");
1467        client
1468            .write_all(&startup_packet)
1469            .await
1470            .expect("writable client");
1471
1472        let mut conn = Conn::new(Buffered::<_, Frontend>::new_frontend(proxy));
1473        let ssl = conn
1474            .receive_pre_startup_wire()
1475            .await
1476            .expect("decodable SSLRequest");
1477        let crate::pre_startup::PreStartupOffer::Ssl(decision) = conn.offer_pre_startup(ssl) else {
1478            panic!("unexpected pre-startup branch")
1479        };
1480        let (mut conn, reply) = decision.reject_ssl();
1481        assert_eq!(reply, b'N');
1482        let message = conn
1483            .receive_pre_startup_wire()
1484            .await
1485            .expect("decodable StartupMessage");
1486        assert_eq!(message, startup);
1487        let crate::pre_startup::PreStartupOffer::Startup { conn, .. } =
1488            conn.offer_pre_startup(message)
1489        else {
1490            panic!("unexpected pre-startup branch")
1491        };
1492        let _transport = conn.into_transport();
1493    }
1494
1495    #[tokio::test]
1496    async fn client_facing_transport_applies_its_pre_startup_limit() {
1497        let (proxy, mut client) = tokio::io::duplex(32);
1498        client
1499            .write_all(&17_u32.to_be_bytes())
1500            .await
1501            .expect("writable client");
1502
1503        let mut transport =
1504            Buffered::<_, Frontend>::with_limits_frontend(proxy, 64, 16).expect("valid limits");
1505        let error = transport
1506            .receive_pre_startup()
1507            .await
1508            .expect_err("declared packet exceeds the configured limit");
1509
1510        assert_eq!(error.kind(), io::ErrorKind::InvalidData);
1511    }
1512
1513    #[tokio::test]
1514    async fn upstream_transport_negotiates_raw_gssenc_reply() {
1515        let (proxy, mut server) = tokio::io::duplex(32);
1516        let mut pending = Conn::new(Buffered::new(proxy)).request_gss();
1517        pending.flush().await.expect("GSSENCRequest is writable");
1518
1519        let mut request = [0_u8; 8];
1520        server
1521            .read_exact(&mut request)
1522            .await
1523            .expect("server receives GSSENCRequest");
1524        assert_eq!(request, gssenc_request_packet());
1525        server
1526            .write_all(b"N")
1527            .await
1528            .expect("server writes decision");
1529
1530        let Negotiation::Rejected(plaintext) = pending
1531            .receive_gss_reply()
1532            .await
1533            .expect("valid GSSENC decision")
1534        else {
1535            panic!("expected plaintext fallback")
1536        };
1537        plaintext.into_transport();
1538    }
1539
1540    #[test]
1541    fn client_facing_transport_buffers_raw_gssenc_decision() {
1542        let conn = Conn::new(Buffered::<_, Frontend>::new_frontend(()));
1543        let crate::pre_startup::PreStartupOffer::Gss(decision) =
1544            conn.offer_pre_startup(PreStartupMessage::GssEncRequest)
1545        else {
1546            panic!("expected GSSENC decision")
1547        };
1548
1549        let handshake = decision.approve_gss();
1550        assert_eq!(handshake.pending_output(), b"S");
1551        handshake.into_transport();
1552
1553        let conn = Conn::new(Buffered::<_, Frontend>::new_frontend(()));
1554        let crate::pre_startup::PreStartupOffer::Gss(decision) =
1555            conn.offer_pre_startup(PreStartupMessage::GssEncRequest)
1556        else {
1557            panic!("expected GSSENC decision")
1558        };
1559        let terminated = decision.reject_gss_with_legacy_error();
1560        assert_eq!(terminated.pending_output(), b"E");
1561        terminated.into_transport();
1562    }
1563
1564    #[test]
1565    fn client_facing_transport_buffers_legacy_ssl_error() {
1566        let conn = Conn::new(Buffered::<_, Frontend>::new_frontend(()));
1567        let crate::pre_startup::PreStartupOffer::Ssl(decision) =
1568            conn.offer_pre_startup(PreStartupMessage::SslRequest)
1569        else {
1570            panic!("expected SSL decision")
1571        };
1572
1573        let terminated = decision.reject_ssl_with_legacy_error();
1574        assert_eq!(terminated.pending_output(), b"E");
1575        terminated.into_transport();
1576    }
1577}
pg_proto/transport.rs

pg_proto/
transport.rs
