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    pre_startup::{
22        AwaitingSslReply, DEFAULT_MAX_PRE_STARTUP_PACKET_LEN, EncryptionReply, Negotiation,
23        PreStartup, PreStartupMessage, ServerSslDecision, SslMode, SslModeNegotiation,
24        TlsHandshake, decode_pre_startup_with_limit, gssenc_request_packet, ssl_request_packet,
25    },
26    tls::{ClientTls, ServerTls},
27};
28
29/// Transport wrapper which retains bytes until each write has completed.
30#[derive(Debug)]
31pub struct Buffered<S, D = Backend> {
32    io: S,
33    outbound: BytesMut,
34    inbound: BytesMut,
35    inbound_codec: PgCodec<D>,
36    max_pre_startup_packet_len: usize,
37    demux: Demux,
38}
39
40impl<S> Buffered<S, Backend> {
41    /// Wraps an upstream-facing transport which receives backend messages.
42    pub fn new(io: S) -> Self {
43        Self {
44            io,
45            outbound: BytesMut::new(),
46            inbound: BytesMut::new(),
47            inbound_codec: PgCodec::default(),
48            max_pre_startup_packet_len: DEFAULT_MAX_PRE_STARTUP_PACKET_LEN,
49            demux: Demux::default(),
50        }
51    }
52
53    /// Creates a backend-facing transport with a bounded tagged-frame size.
54    ///
55    /// # Errors
56    ///
57    /// Returns an error when the limit is outside `PostgreSQL`'s frame range.
58    pub fn with_max_frame_len(io: S, max_frame_len: usize) -> io::Result<Self> {
59        Ok(Self {
60            io,
61            outbound: BytesMut::new(),
62            inbound: BytesMut::new(),
63            inbound_codec: PgCodec::with_max_frame_len(max_frame_len)?,
64            max_pre_startup_packet_len: DEFAULT_MAX_PRE_STARTUP_PACKET_LEN,
65            demux: Demux::default(),
66        })
67    }
68}
69
70impl<S> Buffered<S, Frontend> {
71    /// Wraps a client-facing transport which receives frontend messages.
72    pub fn new_frontend(io: S) -> Self {
73        Self {
74            io,
75            outbound: BytesMut::new(),
76            inbound: BytesMut::new(),
77            inbound_codec: PgCodec::default(),
78            max_pre_startup_packet_len: DEFAULT_MAX_PRE_STARTUP_PACKET_LEN,
79            demux: Demux::default(),
80        }
81    }
82
83    /// Creates a frontend-facing transport with a bounded tagged-frame size.
84    ///
85    /// # Errors
86    ///
87    /// Returns an error when the limit is outside `PostgreSQL`'s frame range.
88    pub fn with_max_frame_len_frontend(io: S, max_frame_len: usize) -> io::Result<Self> {
89        Self::with_limits_frontend(io, max_frame_len, DEFAULT_MAX_PRE_STARTUP_PACKET_LEN)
90    }
91
92    /// Creates a frontend-facing transport with bounded tagged and pre-startup packets.
93    ///
94    /// # Errors
95    ///
96    /// Returns an error when either limit is outside `PostgreSQL`'s framing range.
97    pub fn with_limits_frontend(
98        io: S,
99        max_frame_len: usize,
100        max_pre_startup_packet_len: usize,
101    ) -> io::Result<Self> {
102        if !(8..=i32::MAX as usize).contains(&max_pre_startup_packet_len) {
103            return Err(io::Error::new(
104                io::ErrorKind::InvalidInput,
105                "pre-startup packet limit must be between 8 and i32::MAX bytes",
106            ));
107        }
108        Ok(Self {
109            io,
110            outbound: BytesMut::new(),
111            inbound: BytesMut::new(),
112            inbound_codec: PgCodec::with_max_frame_len(max_frame_len)?,
113            max_pre_startup_packet_len,
114            demux: Demux::default(),
115        })
116    }
117}
118
119impl<S, D> Buffered<S, D> {
120    /// Encodes a frame synchronously into the outbound buffer.
121    ///
122    /// # Errors
123    ///
124    /// Returns an error when the frame is too large to encode.
125    pub fn push(&mut self, frame: Frame) -> io::Result<()> {
126        self.inbound_codec.encode(frame, &mut self.outbound)
127    }
128
129    #[must_use]
130    /// Returns encoded bytes which have not yet been fully written.
131    pub fn pending(&self) -> &[u8] {
132        &self.outbound
133    }
134
135    /// Removes buffering and returns the underlying I/O transport.
136    pub fn into_inner(self) -> S {
137        self.io
138    }
139
140    fn push_raw(&mut self, bytes: &[u8]) {
141        self.outbound.extend_from_slice(bytes);
142    }
143
144    #[must_use]
145    /// Returns the backend asynchronous-message demultiplexer.
146    pub const fn demux(&self) -> &Demux {
147        &self.demux
148    }
149
150    /// Returns mutable access to the backend asynchronous-message demultiplexer.
151    pub const fn demux_mut(&mut self) -> &mut Demux {
152        &mut self.demux
153    }
154}
155
156impl<S, D> Buffered<S, D>
157where
158    S: AsyncRead + AsyncWrite + Unpin,
159{
160    async fn connect_tls(
161        self,
162        server_name: ServerName<'static>,
163        config: Arc<ClientConfig>,
164    ) -> io::Result<Buffered<ClientTls<S>, D>> {
165        if !self.outbound.is_empty() || !self.inbound.is_empty() {
166            return Err(io::Error::new(
167                io::ErrorKind::InvalidInput,
168                "TLS upgrade requires empty plaintext buffers",
169            ));
170        }
171        Ok(Buffered {
172            io: crate::tls::connect(self.io, server_name, config).await?,
173            outbound: self.outbound,
174            inbound: self.inbound,
175            inbound_codec: self.inbound_codec,
176            max_pre_startup_packet_len: self.max_pre_startup_packet_len,
177            demux: self.demux,
178        })
179    }
180
181    async fn accept_tls(
182        self,
183        config: Arc<ServerConfig>,
184        leaf_certificate: CertificateDer<'static>,
185    ) -> io::Result<Buffered<ServerTls<S>, D>> {
186        if !self.outbound.is_empty() || !self.inbound.is_empty() {
187            return Err(io::Error::new(
188                io::ErrorKind::InvalidInput,
189                "TLS upgrade requires empty plaintext buffers",
190            ));
191        }
192        Ok(Buffered {
193            io: crate::tls::accept(self.io, config, &leaf_certificate).await?,
194            outbound: self.outbound,
195            inbound: self.inbound,
196            inbound_codec: self.inbound_codec,
197            max_pre_startup_packet_len: self.max_pre_startup_packet_len,
198            demux: self.demux,
199        })
200    }
201}
202
203impl<S: TlsServerEndPoint, D> TlsServerEndPoint for Buffered<S, D> {
204    fn tls_server_end_point(&self) -> &[u8] {
205        self.io.tls_server_end_point()
206    }
207}
208
209impl<S: AsyncWrite + Unpin, D> Buffered<S, D> {
210    /// Writes all buffered bytes without consuming the connection.
211    ///
212    /// Completed partial writes are removed immediately. If this future is
213    /// cancelled, the connection remains owned by the caller and all unwritten
214    /// bytes remain buffered for the next call.
215    ///
216    /// # Errors
217    ///
218    /// Returns the underlying transport's write error or `WriteZero`.
219    pub async fn flush(&mut self) -> io::Result<()> {
220        while !self.outbound.is_empty() {
221            let written = self.io.write(&self.outbound).await?;
222            if written == 0 {
223                return Err(io::Error::new(
224                    io::ErrorKind::WriteZero,
225                    "transport wrote zero buffered bytes",
226                ));
227            }
228            self.outbound.advance(written);
229        }
230        self.io.flush().await
231    }
232}
233
234impl<S: AsyncRead + Unpin, D: Direction> Buffered<S, D> {
235    /// Receives one typed message in this transport's inbound direction.
236    ///
237    /// # Errors
238    ///
239    /// Returns decoding and underlying transport read errors, or `UnexpectedEof`.
240    pub async fn receive_wire(&mut self) -> io::Result<D::Message> {
241        loop {
242            if let Some(message) = self.inbound_codec.decode(&mut self.inbound)? {
243                return Ok(message);
244            }
245            if self.io.read_buf(&mut self.inbound).await? == 0 {
246                return Err(io::Error::new(
247                    io::ErrorKind::UnexpectedEof,
248                    "peer closed with no complete message",
249                ));
250            }
251        }
252    }
253}
254
255impl<S: AsyncRead + Unpin> Buffered<S, Backend> {
256    async fn receive_encryption_reply(&mut self) -> io::Result<EncryptionReply> {
257        let byte = self.io.read_u8().await?;
258        EncryptionReply::try_from(byte)
259            .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "invalid encryption reply"))
260    }
261}
262
263impl<S: AsyncRead + Unpin> Buffered<S, Frontend> {
264    /// Receives one raw first packet before tagged frontend framing begins.
265    ///
266    /// # Errors
267    ///
268    /// Returns malformed pre-startup data and underlying transport read errors.
269    pub async fn receive_pre_startup(&mut self) -> io::Result<PreStartupMessage> {
270        loop {
271            if let Some(message) =
272                decode_pre_startup_with_limit(&mut self.inbound, self.max_pre_startup_packet_len)?
273            {
274                return Ok(message);
275            }
276            if self.io.read_buf(&mut self.inbound).await? == 0 {
277                return Err(io::Error::new(
278                    io::ErrorKind::UnexpectedEof,
279                    "client closed with no complete pre-startup packet",
280                ));
281            }
282        }
283    }
284}
285
286impl<S: AsyncRead + Unpin> Buffered<S, Backend> {
287    /// Receives one decoded backend message while retaining partial input.
288    ///
289    /// # Errors
290    ///
291    /// Returns decoding and underlying transport read errors, or `UnexpectedEof`.
292    pub async fn receive_backend(&mut self) -> io::Result<BackendMessage> {
293        self.receive_wire().await
294    }
295
296    /// Receives the next protocol-advancing message through the async demux.
297    ///
298    /// # Errors
299    ///
300    /// Returns decoding and underlying transport read errors, or `UnexpectedEof`.
301    pub async fn receive_session(&mut self) -> io::Result<SessionItem> {
302        loop {
303            let message = self.receive_backend().await?;
304            if let Some(item) = self.project_backend(message) {
305                return Ok(item);
306            }
307        }
308    }
309
310    /// Projects an inspected or modified backend message into the session stream.
311    pub fn project_backend(&mut self, message: BackendMessage) -> Option<SessionItem> {
312        self.demux.route(message)
313    }
314}
315
316impl<S, D, Phase, Cleanliness> Conn<Buffered<S, D>, Phase, Cleanliness> {
317    /// Adds an already-typed message to this connection's outbound buffer.
318    ///
319    /// # Errors
320    ///
321    /// Returns an error when the frame is too large to encode.
322    pub fn push_frame(&mut self, frame: Frame) -> io::Result<()> {
323        self.transport_mut().push(frame)
324    }
325
326    #[must_use]
327    /// Returns encoded output which has not yet been flushed.
328    pub fn pending_output(&self) -> &[u8] {
329        self.transport().pending()
330    }
331}
332
333impl<S, Cleanliness> Conn<Buffered<S, Backend>, PreStartup, Cleanliness> {
334    /// Buffers an `SSLRequest` and enters the raw single-byte reply phase.
335    pub fn request_ssl(mut self) -> Conn<Buffered<S, Backend>, AwaitingSslReply, Cleanliness> {
336        self.transport_mut().push_raw(&ssl_request_packet());
337        self.transition()
338    }
339
340    /// Buffers a `GSSENCRequest` and enters the raw single-byte reply phase.
341    pub fn request_gss(
342        mut self,
343    ) -> Conn<Buffered<S, Backend>, crate::pre_startup::AwaitingGssReply, Cleanliness> {
344        self.transport_mut().push_raw(&gssenc_request_packet());
345        self.transition()
346    }
347}
348
349impl<S, Cleanliness> Conn<Buffered<S, Frontend>, ServerSslDecision, Cleanliness> {
350    /// Buffers the server's raw `S` response and enters the TLS handshake phase.
351    pub fn approve_ssl(mut self) -> Conn<Buffered<S, Frontend>, TlsHandshake, Cleanliness> {
352        self.transport_mut().push_raw(b"S");
353        self.transition()
354    }
355
356    /// Buffers the server's raw `N` response and returns to pre-startup choice.
357    pub fn decline_ssl(mut self) -> Conn<Buffered<S, Frontend>, PreStartup, Cleanliness> {
358        self.transport_mut().push_raw(b"N");
359        self.transition()
360    }
361
362    /// Buffers the historical raw `E` response and terminates negotiation.
363    pub fn reject_ssl_with_legacy_error(
364        mut self,
365    ) -> Conn<Buffered<S, Frontend>, crate::pre_startup::Terminated, Cleanliness> {
366        self.transport_mut().push_raw(b"E");
367        self.transition()
368    }
369}
370
371impl<S, Cleanliness>
372    Conn<Buffered<S, Frontend>, crate::pre_startup::ServerGssDecision, Cleanliness>
373{
374    /// Buffers the server's raw `S` response and enters the GSS handshake phase.
375    pub fn approve_gss(
376        mut self,
377    ) -> Conn<Buffered<S, Frontend>, crate::pre_startup::GssHandshake, Cleanliness> {
378        self.transport_mut().push_raw(b"S");
379        self.transition()
380    }
381
382    /// Buffers the server's raw `N` response and returns to pre-startup choice.
383    pub fn decline_gss(mut self) -> Conn<Buffered<S, Frontend>, PreStartup, Cleanliness> {
384        self.transport_mut().push_raw(b"N");
385        self.transition()
386    }
387
388    /// Buffers the historical raw `E` response and terminates negotiation.
389    pub fn reject_gss_with_legacy_error(
390        mut self,
391    ) -> Conn<Buffered<S, Frontend>, crate::pre_startup::Terminated, Cleanliness> {
392        self.transport_mut().push_raw(b"E");
393        self.transition()
394    }
395}
396
397impl<S: AsyncRead + Unpin, Cleanliness> Conn<Buffered<S, Backend>, AwaitingSslReply, Cleanliness> {
398    /// Receives and projects the server's raw SSL decision byte.
399    ///
400    /// # Errors
401    ///
402    /// Returns an I/O error or rejects a byte other than `S`, `N`, or `E`.
403    pub async fn receive_ssl_reply(
404        mut self,
405    ) -> io::Result<Negotiation<Buffered<S, Backend>, TlsHandshake, Cleanliness>> {
406        let reply = self.transport_mut().receive_encryption_reply().await?;
407        Ok(match reply {
408            EncryptionReply::Accepted => Negotiation::Accepted(self.transition()),
409            EncryptionReply::Rejected => Negotiation::Rejected(self.transition()),
410            EncryptionReply::LegacyError => Negotiation::LegacyError(self.transition()),
411        })
412    }
413
414    /// Receives the server decision and enforces the selected plaintext fallback policy.
415    ///
416    /// # Errors
417    ///
418    /// Returns an I/O error or rejects a byte other than `S`, `N`, or `E`.
419    pub async fn receive_ssl_reply_for_mode(
420        mut self,
421        mode: SslMode,
422    ) -> io::Result<SslModeNegotiation<Buffered<S, Backend>, Cleanliness>> {
423        let reply = self.transport_mut().receive_encryption_reply().await?;
424        Ok(self.apply_ssl_reply(reply, mode))
425    }
426}
427
428impl<S: AsyncRead + Unpin, Cleanliness>
429    Conn<Buffered<S, Backend>, crate::pre_startup::AwaitingGssReply, Cleanliness>
430{
431    /// Receives and projects the server's raw GSSENC decision byte.
432    ///
433    /// # Errors
434    ///
435    /// Returns an I/O error or rejects a byte other than `S`, `N`, or `E`.
436    pub async fn receive_gss_reply(
437        mut self,
438    ) -> io::Result<Negotiation<Buffered<S, Backend>, crate::pre_startup::GssHandshake, Cleanliness>>
439    {
440        let reply = self.transport_mut().receive_encryption_reply().await?;
441        Ok(match reply {
442            EncryptionReply::Accepted => Negotiation::Accepted(self.transition()),
443            EncryptionReply::Rejected => Negotiation::Rejected(self.transition()),
444            EncryptionReply::LegacyError => Negotiation::LegacyError(self.transition()),
445        })
446    }
447}
448
449impl<S, Cleanliness> Conn<Buffered<S, Backend>, TlsHandshake, Cleanliness>
450where
451    S: AsyncRead + AsyncWrite + Unpin,
452{
453    /// Completes a client-side TLS handshake and changes the transport type.
454    ///
455    /// # Errors
456    ///
457    /// Returns a TLS handshake, certificate, channel-binding, or buffer-state error.
458    pub async fn connect_tls(
459        self,
460        server_name: ServerName<'static>,
461        config: Arc<ClientConfig>,
462    ) -> io::Result<Conn<Buffered<ClientTls<S>, Backend>, PreStartup, Cleanliness>> {
463        let transport = self.into_transport();
464        Ok(Conn::new(transport.connect_tls(server_name, config).await?)
465            .transition::<PreStartup, Cleanliness>())
466    }
467}
468
469impl<S, Cleanliness> Conn<Buffered<S, Frontend>, TlsHandshake, Cleanliness>
470where
471    S: AsyncRead + AsyncWrite + Unpin,
472{
473    /// Completes a server-side TLS handshake and changes the transport type.
474    ///
475    /// # Errors
476    ///
477    /// Returns a TLS handshake, certificate, channel-binding, or buffer-state error.
478    pub async fn accept_tls(
479        self,
480        config: Arc<ServerConfig>,
481        leaf_certificate: CertificateDer<'static>,
482    ) -> io::Result<Conn<Buffered<ServerTls<S>, Frontend>, PreStartup, Cleanliness>> {
483        let transport = self.into_transport();
484        Ok(
485            Conn::new(transport.accept_tls(config, leaf_certificate).await?)
486                .transition::<PreStartup, Cleanliness>(),
487        )
488    }
489}
490
491impl<S, D, Cleanliness> Conn<Buffered<S, D>, crate::pre_startup::Startup, Cleanliness> {
492    /// Buffers the raw, untagged startup packet before normal framing begins.
493    pub fn push_startup_packet(&mut self, packet: &[u8]) {
494        self.transport_mut().outbound.extend_from_slice(packet);
495    }
496}
497
498impl<S: AsyncWrite + Unpin, D, Phase, Cleanliness> Conn<Buffered<S, D>, Phase, Cleanliness> {
499    /// Flushes buffered output while retaining ownership of the typed connection.
500    ///
501    /// # Errors
502    ///
503    /// Returns an error from the underlying transport.
504    pub async fn flush(&mut self) -> io::Result<()> {
505        self.transport_mut().flush().await
506    }
507}
508
509impl<S: AsyncRead + Unpin, Phase, Cleanliness> Conn<Buffered<S, Backend>, Phase, Cleanliness> {
510    /// Receives one backend message before demultiplexing or state advancement.
511    /// This is the interception point for proxy policy and message rewriting.
512    ///
513    /// # Errors
514    ///
515    /// Returns decoding and underlying transport read errors, or `UnexpectedEof`.
516    pub async fn receive_backend_wire(&mut self) -> io::Result<BackendMessage> {
517        self.transport_mut().receive_backend().await
518    }
519
520    /// Projects an inspected or modified message into the filtered session stream.
521    pub fn project_backend(&mut self, message: BackendMessage) -> Option<SessionItem> {
522        self.transport_mut().project_backend(message)
523    }
524
525    /// Receives the next message in the filtered session projection.
526    ///
527    /// # Errors
528    ///
529    /// Returns decoding and underlying transport read errors, or `UnexpectedEof`.
530    pub async fn receive(&mut self) -> io::Result<SessionItem> {
531        self.transport_mut().receive_session().await
532    }
533
534    #[must_use]
535    /// Returns the latest upstream cancellation key observed during startup.
536    pub fn cancel_key(&self) -> Option<&CancelKey> {
537        self.transport().demux().cancel_key()
538    }
539
540    /// Returns the latest backend parameter values observed by the demux.
541    #[must_use]
542    pub fn parameters(&self) -> &BTreeMap<Bytes, Bytes> {
543        self.transport().demux().parameters()
544    }
545
546    /// Returns whether current parameters differ from the startup baseline.
547    #[must_use]
548    pub fn parameters_changed(&self) -> bool {
549        self.transport().demux().parameters_changed()
550    }
551
552    /// Returns the latest transaction status observed in `ReadyForQuery`.
553    #[must_use]
554    pub fn transaction_status(&self) -> Option<crate::codec::TransactionStatus> {
555        self.transport().demux().transaction_status()
556    }
557
558    /// Removes the oldest queued asynchronous notification.
559    pub fn pop_notification(&mut self) -> Option<Notification> {
560        self.transport_mut().demux_mut().pop_notification()
561    }
562
563    /// Removes the next tagged notice for prompt forwarding to the client.
564    pub fn pop_notice(&mut self) -> Option<TaggedNotice> {
565        self.transport_mut().demux_mut().pop_notice()
566    }
567
568    /// Removes the next ordered parameter update for forwarding to the client.
569    pub fn pop_parameter_status(&mut self) -> Option<ParameterStatus> {
570        self.transport_mut().demux_mut().pop_parameter_status()
571    }
572
573    /// Removes the next independent backend event in original wire order.
574    pub fn pop_async_event(&mut self) -> Option<OrderedAsyncEvent> {
575        self.transport_mut().demux_mut().pop_async_event()
576    }
577}
578
579impl<S: AsyncRead + Unpin, Phase, Cleanliness> Conn<Buffered<S, Frontend>, Phase, Cleanliness> {
580    /// Receives one frontend message before any server-role state advancement.
581    ///
582    /// # Errors
583    ///
584    /// Returns decoding and underlying transport read errors, or `UnexpectedEof`.
585    pub async fn receive_frontend_wire(&mut self) -> io::Result<FrontendMessage> {
586        self.transport_mut().receive_wire().await
587    }
588}
589
590impl<S: AsyncRead + Unpin, Cleanliness> Conn<Buffered<S, Frontend>, PreStartup, Cleanliness> {
591    /// Receives a raw pre-startup packet before server-role state projection.
592    ///
593    /// # Errors
594    ///
595    /// Returns malformed pre-startup data and underlying transport read errors.
596    pub async fn receive_pre_startup_wire(&mut self) -> io::Result<PreStartupMessage> {
597        self.transport_mut().receive_pre_startup().await
598    }
599}
600
601#[cfg(test)]
602mod tests {
603    use std::{
604        future::Future,
605        pin::Pin,
606        task::{Context, Poll},
607    };
608
609    use bytes::Bytes;
610    use tokio::io::AsyncWrite;
611
612    use super::*;
613
614    #[derive(Debug, Default)]
615    struct ShortWriter {
616        output: Vec<u8>,
617    }
618
619    impl AsyncWrite for ShortWriter {
620        fn poll_write(
621            mut self: Pin<&mut Self>,
622            _cx: &mut Context<'_>,
623            buffer: &[u8],
624        ) -> Poll<io::Result<usize>> {
625            let written = buffer.len().min(2);
626            self.output.extend_from_slice(&buffer[..written]);
627            Poll::Ready(Ok(written))
628        }
629
630        fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
631            Poll::Ready(Ok(()))
632        }
633
634        fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
635            Poll::Ready(Ok(()))
636        }
637    }
638
639    #[tokio::test]
640    async fn flush_handles_partial_writes_without_losing_bytes() {
641        let frame = Frame {
642            tag: b'S',
643            body: Bytes::new(),
644        };
645        let mut transport = Buffered::new(ShortWriter::default());
646        transport.push(frame).expect("encodable frame");
647        assert_eq!(transport.pending(), &[b'S', 0, 0, 0, 4]);
648        transport.flush().await.expect("writable transport");
649        assert!(transport.pending().is_empty());
650        assert_eq!(transport.into_inner().output, [b'S', 0, 0, 0, 4]);
651    }
652
653    #[test]
654    fn buffered_transport_enforces_its_frame_limit_on_output() {
655        let mut transport = Buffered::<_, Backend>::with_max_frame_len((), 9).unwrap();
656        let error = transport
657            .push(Frame {
658                tag: b'Q',
659                body: Bytes::from_static(b"12345"),
660            })
661            .unwrap_err();
662        assert_eq!(error.kind(), io::ErrorKind::InvalidInput);
663        assert!(transport.pending().is_empty());
664    }
665
666    #[test]
667    fn cancelling_flush_retains_unwritten_bytes() {
668        #[derive(Debug, Default)]
669        struct PausingWriter {
670            output: Vec<u8>,
671            blocked: bool,
672        }
673
674        impl AsyncWrite for PausingWriter {
675            fn poll_write(
676                mut self: Pin<&mut Self>,
677                _cx: &mut Context<'_>,
678                buffer: &[u8],
679            ) -> Poll<io::Result<usize>> {
680                if self.blocked {
681                    return Poll::Pending;
682                }
683                let written = buffer.len().min(2);
684                self.output.extend_from_slice(&buffer[..written]);
685                self.blocked = true;
686                Poll::Ready(Ok(written))
687            }
688
689            fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
690                Poll::Ready(Ok(()))
691            }
692
693            fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
694                Poll::Ready(Ok(()))
695            }
696        }
697
698        let mut transport = Buffered::new(PausingWriter::default());
699        transport
700            .push(Frame {
701                tag: b'S',
702                body: Bytes::new(),
703            })
704            .expect("encodable frame");
705
706        let mut flush = Box::pin(transport.flush());
707        let waker = std::task::Waker::noop();
708        let mut context = Context::from_waker(waker);
709        assert!(flush.as_mut().poll(&mut context).is_pending());
710        drop(flush);
711
712        assert_eq!(transport.pending(), &[0, 0, 4]);
713        assert_eq!(transport.io.output, [b'S', 0]);
714    }
715
716    #[tokio::test]
717    async fn receive_filters_parameter_status_before_session_message() {
718        let (client, mut server) = tokio::io::duplex(256);
719        let mut wire = BytesMut::new();
720        let mut encoder = PgCodec::<Backend>::default();
721        encoder
722            .encode(
723                Frame {
724                    tag: b'S',
725                    body: Bytes::from_static(b"client_encoding\0UTF8\0"),
726                },
727                &mut wire,
728            )
729            .expect("encodable ParameterStatus");
730        encoder
731            .encode(
732                Frame {
733                    tag: b'Z',
734                    body: Bytes::from_static(b"I"),
735                },
736                &mut wire,
737            )
738            .expect("encodable ReadyForQuery");
739        server.write_all(&wire).await.expect("writable test peer");
740
741        let mut transport = Buffered::new(client);
742        assert_eq!(
743            transport.receive_session().await.expect("valid messages"),
744            SessionItem::ReadyForQuery {
745                status: crate::codec::TransactionStatus::Idle,
746                parameters_changed: false,
747            }
748        );
749        assert_eq!(
750            transport
751                .demux()
752                .parameters()
753                .get(&Bytes::from_static(b"client_encoding")),
754            Some(&Bytes::from_static(b"UTF8"))
755        );
756        let conn: Conn<_, crate::auth::Ready> = Conn::new(transport).transition();
757        assert_eq!(
758            conn.parameters().get(b"client_encoding".as_slice()),
759            Some(&Bytes::from_static(b"UTF8"))
760        );
761        assert!(!conn.parameters_changed());
762        assert_eq!(
763            conn.transaction_status(),
764            Some(crate::codec::TransactionStatus::Idle)
765        );
766        conn.into_transport();
767    }
768
769    #[tokio::test]
770    async fn wire_message_can_be_modified_before_projection() {
771        let (client, mut server) = tokio::io::duplex(128);
772        let original = BackendMessage::ParameterStatus {
773            name: Bytes::from_static(b"application_name"),
774            value: Bytes::from_static(b"upstream"),
775        };
776        let mut bytes = BytesMut::new();
777        PgCodec::<Backend>::default()
778            .encode(
779                original.to_frame().expect("reconstructable message"),
780                &mut bytes,
781            )
782            .expect("encodable message");
783        server.write_all(&bytes).await.expect("writable test peer");
784
785        let mut transport = Buffered::new(client);
786        let mut message = transport
787            .receive_backend()
788            .await
789            .expect("decodable message");
790        let BackendMessage::ParameterStatus { value, .. } = &mut message else {
791            panic!("unexpected message")
792        };
793        *value = Bytes::from_static(b"proxy");
794        assert!(transport.project_backend(message).is_none());
795        assert_eq!(
796            transport
797                .demux()
798                .parameters()
799                .get(&Bytes::from_static(b"application_name")),
800            Some(&Bytes::from_static(b"proxy"))
801        );
802    }
803
804    #[tokio::test]
805    async fn client_facing_transport_intercepts_typed_frontend_messages() {
806        let (proxy, mut client) = tokio::io::duplex(128);
807        let message = FrontendMessage::Query(Bytes::from_static(b"select plaintext"));
808        let mut bytes = BytesMut::new();
809        PgCodec::<Frontend>::default()
810            .encode(
811                message.to_frame().expect("reconstructable Query"),
812                &mut bytes,
813            )
814            .expect("encodable Query");
815        client.write_all(&bytes).await.expect("writable client");
816
817        let mut transport = Buffered::<_, Frontend>::new_frontend(proxy);
818        let mut intercepted = transport.receive_wire().await.expect("decodable Query");
819        let FrontendMessage::Query(query) = &mut intercepted else {
820            panic!("unexpected frontend message")
821        };
822        *query = Bytes::from_static(b"select encrypted");
823        assert_eq!(
824            intercepted,
825            FrontendMessage::Query(Bytes::from_static(b"select encrypted"))
826        );
827    }
828
829    #[tokio::test]
830    async fn client_facing_transport_projects_repeated_pre_startup_choice() {
831        let (proxy, mut client) = tokio::io::duplex(256);
832        let ssl = PreStartupMessage::SslRequest
833            .to_packet()
834            .expect("encodable SSLRequest");
835        let startup = PreStartupMessage::Startup(crate::startup::StartupMessage {
836            version: crate::startup::ProtocolVersion::V3_2,
837            parameters: std::collections::BTreeMap::from([(
838                Bytes::from_static(b"user"),
839                Bytes::from_static(b"postgres"),
840            )]),
841        });
842        let startup_packet = startup.to_packet().expect("encodable StartupMessage");
843        client.write_all(&ssl).await.expect("writable client");
844        client
845            .write_all(&startup_packet)
846            .await
847            .expect("writable client");
848
849        let mut conn = Conn::new(Buffered::<_, Frontend>::new_frontend(proxy));
850        let ssl = conn
851            .receive_pre_startup_wire()
852            .await
853            .expect("decodable SSLRequest");
854        let crate::pre_startup::PreStartupOffer::Ssl(decision) = conn.offer_pre_startup(ssl) else {
855            panic!("unexpected pre-startup branch")
856        };
857        let (mut conn, reply) = decision.reject_ssl();
858        assert_eq!(reply, b'N');
859        let message = conn
860            .receive_pre_startup_wire()
861            .await
862            .expect("decodable StartupMessage");
863        assert_eq!(message, startup);
864        let crate::pre_startup::PreStartupOffer::Startup { conn, .. } =
865            conn.offer_pre_startup(message)
866        else {
867            panic!("unexpected pre-startup branch")
868        };
869        let _transport = conn.into_transport();
870    }
871
872    #[tokio::test]
873    async fn client_facing_transport_applies_its_pre_startup_limit() {
874        let (proxy, mut client) = tokio::io::duplex(32);
875        client
876            .write_all(&17_u32.to_be_bytes())
877            .await
878            .expect("writable client");
879
880        let mut transport =
881            Buffered::<_, Frontend>::with_limits_frontend(proxy, 64, 16).expect("valid limits");
882        let error = transport
883            .receive_pre_startup()
884            .await
885            .expect_err("declared packet exceeds the configured limit");
886
887        assert_eq!(error.kind(), io::ErrorKind::InvalidData);
888    }
889
890    #[tokio::test]
891    async fn upstream_transport_negotiates_raw_gssenc_reply() {
892        let (proxy, mut server) = tokio::io::duplex(32);
893        let mut pending = Conn::new(Buffered::new(proxy)).request_gss();
894        pending.flush().await.expect("GSSENCRequest is writable");
895
896        let mut request = [0_u8; 8];
897        server
898            .read_exact(&mut request)
899            .await
900            .expect("server receives GSSENCRequest");
901        assert_eq!(request, gssenc_request_packet());
902        server
903            .write_all(b"N")
904            .await
905            .expect("server writes decision");
906
907        let Negotiation::Rejected(plaintext) = pending
908            .receive_gss_reply()
909            .await
910            .expect("valid GSSENC decision")
911        else {
912            panic!("expected plaintext fallback")
913        };
914        plaintext.into_transport();
915    }
916
917    #[test]
918    fn client_facing_transport_buffers_raw_gssenc_decision() {
919        let conn = Conn::new(Buffered::<_, Frontend>::new_frontend(()));
920        let crate::pre_startup::PreStartupOffer::Gss(decision) =
921            conn.offer_pre_startup(PreStartupMessage::GssEncRequest)
922        else {
923            panic!("expected GSSENC decision")
924        };
925
926        let handshake = decision.approve_gss();
927        assert_eq!(handshake.pending_output(), b"S");
928        handshake.into_transport();
929
930        let conn = Conn::new(Buffered::<_, Frontend>::new_frontend(()));
931        let crate::pre_startup::PreStartupOffer::Gss(decision) =
932            conn.offer_pre_startup(PreStartupMessage::GssEncRequest)
933        else {
934            panic!("expected GSSENC decision")
935        };
936        let terminated = decision.reject_gss_with_legacy_error();
937        assert_eq!(terminated.pending_output(), b"E");
938        terminated.into_transport();
939    }
940
941    #[test]
942    fn client_facing_transport_buffers_legacy_ssl_error() {
943        let conn = Conn::new(Buffered::<_, Frontend>::new_frontend(()));
944        let crate::pre_startup::PreStartupOffer::Ssl(decision) =
945            conn.offer_pre_startup(PreStartupMessage::SslRequest)
946        else {
947            panic!("expected SSL decision")
948        };
949
950        let terminated = decision.reject_ssl_with_legacy_error();
951        assert_eq!(terminated.pending_output(), b"E");
952        terminated.into_transport();
953    }
954}
pg_proto/transport.rs

pg_proto/
transport.rs
