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

pg_proto/
transport.rs
