pg_proto/
pre_startup.rs

1//! `PostgreSQL`'s untagged pre-startup protocol.
2//!
3//! These packets precede normal message framing. Encryption acceptance is a raw
4//! byte, and a successful negotiation changes the connection's transport type.
5
6use crate::{Conn, Pristine, startup::StartupMessage};
7use bytes::{BufMut as _, Bytes, BytesMut};
8
9const SSL_REQUEST_CODE: u32 = 80_877_103;
10const GSSENC_REQUEST_CODE: u32 = 80_877_104;
11const CANCEL_REQUEST_CODE: u32 = 80_877_102;
12
13/// `PostgreSQL`'s maximum accepted untagged startup packet size.
14pub const DEFAULT_MAX_PRE_STARTUP_PACKET_LEN: usize = 10_000;
15
16/// Connection awaits the client's first untagged startup-family packet.
17#[derive(Debug)]
18pub enum PreStartup {}
19
20/// A decoded `StartupMessage` is ready for protocol validation.
21#[derive(Debug)]
22pub enum Startup {}
23
24/// Client role awaits the server's raw SSL decision byte.
25#[derive(Debug)]
26pub enum AwaitingSslReply {}
27
28/// Client role awaits the server's raw GSS encryption decision byte.
29#[derive(Debug)]
30pub enum AwaitingGssReply {}
31
32/// SSL was accepted and the transport must complete a TLS handshake.
33#[derive(Debug)]
34pub enum TlsHandshake {}
35
36/// GSS encryption was accepted and the transport must complete its handshake.
37#[derive(Debug)]
38pub enum GssHandshake {}
39
40/// Pre-startup processing terminated without entering a normal session.
41#[derive(Debug)]
42pub enum Terminated {}
43
44/// Server role must accept or reject a client's `SSLRequest`.
45#[derive(Debug)]
46pub enum ServerSslDecision {}
47
48/// Server role must accept or reject a client's `GSSENCRequest`.
49#[derive(Debug)]
50pub enum ServerGssDecision {}
51
52/// Server-role projection of the client's first-packet external choice.
53#[derive(Debug)]
54pub enum PreStartupOffer<S, C = Pristine> {
55    /// Client requested TLS negotiation.
56    Ssl(Conn<S, ServerSslDecision, C>),
57    /// Client requested GSS encryption negotiation.
58    Gss(Conn<S, ServerGssDecision, C>),
59    /// Client sent an out-of-band cancellation request.
60    Cancel {
61        /// Terminal connection carrying the received transport.
62        conn: Conn<S, Terminated, C>,
63        /// Target backend process ID.
64        process_id: u32,
65        /// Target backend secret key.
66        secret_key: Bytes,
67    },
68    /// Client supplied normal startup parameters.
69    Startup {
70        /// Connection ready for protocol validation and authentication.
71        conn: Conn<S, Startup, C>,
72        /// Decoded startup version and parameters.
73        message: StartupMessage,
74    },
75}
76
77/// The server's single-byte answer to an SSL request.
78#[derive(Clone, Copy, Debug, Eq, PartialEq)]
79pub enum EncryptionReply {
80    /// Server sent `S` and requires an in-place encryption handshake.
81    Accepted,
82    /// Server sent `N` and declined encryption.
83    Rejected,
84    /// Historical server sent `E` and terminated the connection.
85    LegacyError,
86}
87
88impl EncryptionReply {
89    /// Returns `PostgreSQL`'s raw one-byte wire representation.
90    #[must_use]
91    pub const fn as_byte(self) -> u8 {
92        match self {
93            Self::Accepted => b'S',
94            Self::Rejected => b'N',
95            Self::LegacyError => b'E',
96        }
97    }
98}
99
100/// The external choice occupying a new connection's untagged first packet.
101#[derive(Clone, Debug, Eq, PartialEq)]
102pub enum PreStartupMessage {
103    /// Raw SSL negotiation request code 80877103.
104    SslRequest,
105    /// Raw GSS encryption request code 80877104.
106    GssEncRequest,
107    /// Out-of-band cancellation packet.
108    CancelRequest {
109        /// Target backend process ID.
110        process_id: u32,
111        /// Target backend secret key.
112        secret_key: Bytes,
113    },
114    /// Normal protocol startup packet.
115    Startup(StartupMessage),
116}
117
118impl PreStartupMessage {
119    /// Reconstructs the complete raw packet.
120    ///
121    /// # Errors
122    ///
123    /// Returns an error for an invalid cancellation key or startup message.
124    pub fn to_packet(&self) -> std::io::Result<Bytes> {
125        match self {
126            Self::SslRequest => Ok(Bytes::copy_from_slice(&request_packet(SSL_REQUEST_CODE))),
127            Self::GssEncRequest => Ok(Bytes::copy_from_slice(&request_packet(GSSENC_REQUEST_CODE))),
128            Self::CancelRequest {
129                process_id,
130                secret_key,
131            } => cancel_packet(*process_id, secret_key),
132            Self::Startup(message) => message.encode(),
133        }
134    }
135}
136
137/// Incrementally decodes one raw pre-startup packet without consuming partial input.
138///
139/// # Errors
140///
141/// Returns an error for invalid lengths, special-request shapes, or startup data.
142pub fn decode_pre_startup(input: &mut BytesMut) -> std::io::Result<Option<PreStartupMessage>> {
143    decode_pre_startup_with_limit(input, DEFAULT_MAX_PRE_STARTUP_PACKET_LEN)
144}
145
146/// Incrementally decodes one raw pre-startup packet with an allocation bound.
147///
148/// The declared length is checked before reserving space for the remainder of
149/// a partial packet.
150///
151/// # Errors
152///
153/// Returns an error for an invalid limit, an oversized packet, invalid lengths,
154/// special-request shapes, or startup data.
155pub fn decode_pre_startup_with_limit(
156    input: &mut BytesMut,
157    max_packet_len: usize,
158) -> std::io::Result<Option<PreStartupMessage>> {
159    if !(8..=i32::MAX as usize).contains(&max_packet_len) {
160        return Err(invalid(
161            "pre-startup packet limit must be between 8 and i32::MAX bytes",
162        ));
163    }
164    if input.len() < 4 {
165        input.reserve(4 - input.len());
166        return Ok(None);
167    }
168    let length = usize::try_from(u32::from_be_bytes([input[0], input[1], input[2], input[3]]))
169        .map_err(|_| invalid("pre-startup packet length overflow"))?;
170    if length < 8 {
171        return Err(invalid("pre-startup packet is shorter than 8 bytes"));
172    }
173    if length > i32::MAX as usize {
174        return Err(invalid("pre-startup packet length exceeds i32::MAX"));
175    }
176    if length > max_packet_len {
177        return Err(invalid("pre-startup packet exceeds configured limit"));
178    }
179    if input.len() < length {
180        input.reserve(length - input.len());
181        return Ok(None);
182    }
183    let packet = input.split_to(length).freeze();
184    let code = u32::from_be_bytes([packet[4], packet[5], packet[6], packet[7]]);
185    match code {
186        SSL_REQUEST_CODE if length == 8 => Ok(Some(PreStartupMessage::SslRequest)),
187        GSSENC_REQUEST_CODE if length == 8 => Ok(Some(PreStartupMessage::GssEncRequest)),
188        CANCEL_REQUEST_CODE => {
189            if !(16..=268).contains(&length) {
190                return Err(invalid("invalid CancelRequest length"));
191            }
192            let process_id = u32::from_be_bytes([packet[8], packet[9], packet[10], packet[11]]);
193            let secret_key = packet.slice(12..);
194            Ok(Some(PreStartupMessage::CancelRequest {
195                process_id,
196                secret_key,
197            }))
198        }
199        _ => StartupMessage::decode(packet)
200            .map(PreStartupMessage::Startup)
201            .map(Some),
202    }
203}
204
205impl TryFrom<u8> for EncryptionReply {
206    type Error = InvalidEncryptionReply;
207
208    fn try_from(value: u8) -> Result<Self, Self::Error> {
209        match value {
210            b'S' => Ok(Self::Accepted),
211            b'N' => Ok(Self::Rejected),
212            b'E' => Ok(Self::LegacyError),
213            byte => Err(InvalidEncryptionReply(byte)),
214        }
215    }
216}
217
218/// A raw encryption decision byte other than `S`, `N`, or legacy `E`.
219#[derive(Clone, Copy, Debug, Eq, PartialEq)]
220pub struct InvalidEncryptionReply(
221    /// Invalid byte received from the server.
222    pub u8,
223);
224
225/// libpq-compatible TLS negotiation policy.
226#[derive(Clone, Copy, Debug, Eq, PartialEq)]
227pub enum SslMode {
228    /// Never request TLS.
229    Disable,
230    /// Try plaintext first and retry with TLS if plaintext fails.
231    Allow,
232    /// Request TLS first but permit plaintext fallback.
233    Prefer,
234    /// Require encryption without certificate verification.
235    Require,
236    /// Require TLS and validate the certificate chain.
237    VerifyCa,
238    /// Require TLS and validate both chain and server hostname.
239    VerifyFull,
240}
241
242/// Peer-certificate checks required by an SSL mode.
243#[derive(Clone, Copy, Debug, Eq, PartialEq)]
244pub enum CertificateVerification {
245    /// Do not authenticate the peer certificate.
246    None,
247    /// Validate the certificate chain against configured roots.
248    CertificateAuthority,
249    /// Validate the certificate chain and requested hostname.
250    CertificateAuthorityAndHost,
251}
252
253/// Actions needed to apply an [`SslMode`] across connection attempts.
254#[derive(Clone, Copy, Debug, Eq, PartialEq)]
255pub struct SslStrategy {
256    /// Whether the first connection should send `SSLRequest`.
257    pub request_on_first_connection: bool,
258    /// Whether plaintext failure should open a new TLS-first connection.
259    pub retry_with_ssl_after_plaintext_failure: bool,
260    /// Whether an `N` response permits continuation in plaintext.
261    pub allow_server_rejection: bool,
262    /// Certificate verification required after TLS acceptance.
263    pub verification: CertificateVerification,
264}
265
266impl SslMode {
267    /// Converts libpq-compatible mode semantics into connection actions.
268    #[must_use]
269    pub const fn strategy(self) -> SslStrategy {
270        match self {
271            Self::Disable => SslStrategy {
272                request_on_first_connection: false,
273                retry_with_ssl_after_plaintext_failure: false,
274                allow_server_rejection: true,
275                verification: CertificateVerification::None,
276            },
277            Self::Allow => SslStrategy {
278                request_on_first_connection: false,
279                retry_with_ssl_after_plaintext_failure: true,
280                allow_server_rejection: true,
281                verification: CertificateVerification::None,
282            },
283            Self::Prefer => SslStrategy {
284                request_on_first_connection: true,
285                retry_with_ssl_after_plaintext_failure: false,
286                allow_server_rejection: true,
287                verification: CertificateVerification::None,
288            },
289            Self::Require => SslStrategy {
290                request_on_first_connection: true,
291                retry_with_ssl_after_plaintext_failure: false,
292                allow_server_rejection: false,
293                verification: CertificateVerification::None,
294            },
295            Self::VerifyCa => SslStrategy {
296                request_on_first_connection: true,
297                retry_with_ssl_after_plaintext_failure: false,
298                allow_server_rejection: false,
299                verification: CertificateVerification::CertificateAuthority,
300            },
301            Self::VerifyFull => SslStrategy {
302                request_on_first_connection: true,
303                retry_with_ssl_after_plaintext_failure: false,
304                allow_server_rejection: false,
305                verification: CertificateVerification::CertificateAuthorityAndHost,
306            },
307        }
308    }
309}
310
311/// A reply whose branches deliberately have different typestates.
312#[derive(Debug)]
313pub enum Negotiation<S, Handshake, C = Pristine> {
314    /// Encryption accepted; complete the transport handshake next.
315    Accepted(Conn<S, Handshake, C>),
316    /// Encryption rejected; plaintext pre-startup choice resumes.
317    Rejected(Conn<S, PreStartup, C>),
318    /// Historical `E` response terminated negotiation.
319    LegacyError(Conn<S, Terminated, C>),
320}
321
322/// SSL negotiation after applying plaintext-fallback policy.
323#[derive(Debug)]
324pub enum SslModeNegotiation<S, C = Pristine> {
325    /// TLS was accepted and must be handshaken.
326    Accepted(Conn<S, TlsHandshake, C>),
327    /// Mode permits continuation in plaintext.
328    Plaintext(Conn<S, PreStartup, C>),
329    /// Server rejected TLS required by the configured mode.
330    RequiredRejected {
331        /// Terminal connection retaining the transport.
332        conn: Conn<S, Terminated, C>,
333        /// Mode whose requirement could not be met.
334        mode: SslMode,
335    },
336    /// Historical server error terminated negotiation.
337    LegacyError(Conn<S, Terminated, C>),
338}
339
340impl<S> Conn<S, PreStartup, Pristine> {
341    /// Encodes `SSLRequest` and enters the raw-reply phase.
342    pub fn ssl_request(self) -> (Conn<S, AwaitingSslReply>, [u8; 8]) {
343        (self.transition(), ssl_request_packet())
344    }
345
346    /// Encodes `GSSENCRequest` and enters the raw-reply phase.
347    pub fn gssenc_request(self) -> (Conn<S, AwaitingGssReply>, [u8; 8]) {
348        (self.transition(), gssenc_request_packet())
349    }
350
351    /// Encodes and enters the startup phase.
352    ///
353    /// # Errors
354    ///
355    /// Returns an error when the startup parameters cannot be encoded.
356    pub fn startup(
357        self,
358        message: &StartupMessage,
359    ) -> std::io::Result<(Conn<S, Startup>, bytes::Bytes)> {
360        Ok((self.transition(), message.encode()?))
361    }
362
363    /// Encodes a version 3.0 or 3.2 out-of-band cancellation request.
364    ///
365    /// # Errors
366    ///
367    /// Returns an error unless the cancellation key is between 4 and 256 bytes.
368    pub fn cancel_request(
369        self,
370        process_id: u32,
371        secret_key: &[u8],
372    ) -> std::io::Result<(Conn<S, Terminated>, bytes::Bytes)> {
373        if !(4..=256).contains(&secret_key.len()) {
374            return Err(std::io::Error::new(
375                std::io::ErrorKind::InvalidInput,
376                "cancellation key length is outside 4..=256",
377            ));
378        }
379        Ok((self.transition(), cancel_packet(process_id, secret_key)?))
380    }
381}
382
383impl<S, C> Conn<S, PreStartup, C> {
384    /// Projects an inspected client pre-startup packet into the server role.
385    pub fn offer_pre_startup(self, message: PreStartupMessage) -> PreStartupOffer<S, C> {
386        match message {
387            PreStartupMessage::SslRequest => PreStartupOffer::Ssl(self.transition()),
388            PreStartupMessage::GssEncRequest => PreStartupOffer::Gss(self.transition()),
389            PreStartupMessage::CancelRequest {
390                process_id,
391                secret_key,
392            } => PreStartupOffer::Cancel {
393                conn: self.transition(),
394                process_id,
395                secret_key,
396            },
397            PreStartupMessage::Startup(message) => PreStartupOffer::Startup {
398                conn: self.transition(),
399                message,
400            },
401        }
402    }
403}
404
405impl<S, C> Conn<S, ServerSslDecision, C> {
406    /// Rejects SSL and returns to the pre-startup choice on the same transport.
407    pub fn reject_ssl(self) -> (Conn<S, PreStartup, C>, u8) {
408        (self.transition(), EncryptionReply::Rejected.as_byte())
409    }
410
411    /// Accepts SSL and requires the transport handshake before startup is legal.
412    pub fn accept_ssl(self) -> (Conn<S, TlsHandshake, C>, u8) {
413        (self.transition(), EncryptionReply::Accepted.as_byte())
414    }
415
416    /// Emits the historical raw `E` response and terminates negotiation.
417    pub fn legacy_ssl_error(self) -> (Conn<S, Terminated, C>, u8) {
418        (self.transition(), EncryptionReply::LegacyError.as_byte())
419    }
420}
421
422impl<S, C> Conn<S, ServerGssDecision, C> {
423    /// Sends `N` and returns to plaintext pre-startup choice.
424    pub fn reject_gss(self) -> (Conn<S, PreStartup, C>, u8) {
425        (self.transition(), EncryptionReply::Rejected.as_byte())
426    }
427
428    /// Sends `S` and requires a server-side GSS transport handshake.
429    pub fn accept_gss(self) -> (Conn<S, GssHandshake, C>, u8) {
430        (self.transition(), EncryptionReply::Accepted.as_byte())
431    }
432
433    /// Emits the historical raw `E` response and terminates negotiation.
434    pub fn legacy_gss_error(self) -> (Conn<S, Terminated, C>, u8) {
435        (self.transition(), EncryptionReply::LegacyError.as_byte())
436    }
437}
438
439impl<S, C> Conn<S, AwaitingSslReply, C> {
440    /// Resolves the raw SSL response byte.
441    ///
442    /// A pending negotiation cannot send a startup message:
443    ///
444    /// ```compile_fail
445    /// use pg_proto::Conn;
446    /// let (pending, _) = Conn::new(()).ssl_request();
447    /// let _ = pending.startup();
448    /// ```
449    pub fn receive_reply(self, reply: EncryptionReply) -> Negotiation<S, TlsHandshake, C> {
450        match reply {
451            EncryptionReply::Accepted => Negotiation::Accepted(self.transition()),
452            EncryptionReply::Rejected => Negotiation::Rejected(self.transition()),
453            EncryptionReply::LegacyError => Negotiation::LegacyError(self.transition()),
454        }
455    }
456
457    /// Resolves the SSL response while enforcing an [`SslMode`]'s fallback rule.
458    pub fn apply_ssl_reply(
459        self,
460        reply: EncryptionReply,
461        mode: SslMode,
462    ) -> SslModeNegotiation<S, C> {
463        match reply {
464            EncryptionReply::Accepted => SslModeNegotiation::Accepted(self.transition()),
465            EncryptionReply::Rejected if mode.strategy().allow_server_rejection => {
466                SslModeNegotiation::Plaintext(self.transition())
467            }
468            EncryptionReply::Rejected => SslModeNegotiation::RequiredRejected {
469                conn: self.transition(),
470                mode,
471            },
472            EncryptionReply::LegacyError => SslModeNegotiation::LegacyError(self.transition()),
473        }
474    }
475}
476
477impl<S, C> Conn<S, AwaitingGssReply, C> {
478    /// Applies the server's GSS encryption decision byte.
479    pub fn receive_reply(self, reply: EncryptionReply) -> Negotiation<S, GssHandshake, C> {
480        match reply {
481            EncryptionReply::Accepted => Negotiation::Accepted(self.transition()),
482            EncryptionReply::Rejected => Negotiation::Rejected(self.transition()),
483            EncryptionReply::LegacyError => Negotiation::LegacyError(self.transition()),
484        }
485    }
486}
487
488impl<S, C> Conn<S, TlsHandshake, C> {
489    /// Records a completed in-place TLS upgrade, changing the transport type.
490    pub fn finish_tls<Tls>(self, upgrade: impl FnOnce(S) -> Tls) -> Conn<Tls, PreStartup, C> {
491        self.map_transport(upgrade).transition()
492    }
493}
494
495impl<S, C> Conn<S, TlsHandshake, C> {
496    /// Records a server-side TLS upgrade while preserving cleanliness.
497    pub fn finish_server_tls<Tls>(
498        self,
499        upgrade: impl FnOnce(S) -> Tls,
500    ) -> Conn<Tls, PreStartup, C> {
501        self.map_transport(upgrade).transition()
502    }
503}
504
505impl<S> Conn<S, GssHandshake, Pristine> {
506    /// Records a completed in-place GSS encryption upgrade.
507    pub fn finish_gss<Gss>(self, upgrade: impl FnOnce(S) -> Gss) -> Conn<Gss, PreStartup> {
508        Conn::new(upgrade(self.into_transport()))
509    }
510}
511
512impl<S, C> Conn<S, GssHandshake, C> {
513    /// Completes a server-side GSS upgrade and returns to encrypted pre-startup.
514    pub fn finish_server_gss<Gss>(
515        self,
516        upgrade: impl FnOnce(S) -> Gss,
517    ) -> Conn<Gss, PreStartup, C> {
518        self.map_transport(upgrade).transition()
519    }
520}
521
522pub(crate) const fn ssl_request_packet() -> [u8; 8] {
523    request_packet(SSL_REQUEST_CODE)
524}
525
526pub(crate) const fn gssenc_request_packet() -> [u8; 8] {
527    request_packet(GSSENC_REQUEST_CODE)
528}
529
530const fn request_packet(code: u32) -> [u8; 8] {
531    let length = 8_u32.to_be_bytes();
532    let code = code.to_be_bytes();
533    [
534        length[0], length[1], length[2], length[3], code[0], code[1], code[2], code[3],
535    ]
536}
537
538fn cancel_packet(process_id: u32, secret_key: &[u8]) -> std::io::Result<Bytes> {
539    if !(4..=256).contains(&secret_key.len()) {
540        return Err(std::io::Error::new(
541            std::io::ErrorKind::InvalidInput,
542            "cancellation key length is outside 4..=256",
543        ));
544    }
545    let key_length =
546        u32::try_from(secret_key.len()).map_err(|_| invalid("cancellation key length overflow"))?;
547    let length = 12 + key_length;
548    let mut packet = BytesMut::with_capacity(12 + secret_key.len());
549    packet.put_u32(length);
550    packet.put_u32(CANCEL_REQUEST_CODE);
551    packet.put_u32(process_id);
552    packet.extend_from_slice(secret_key);
553    Ok(packet.freeze())
554}
555
556fn invalid(message: &'static str) -> std::io::Error {
557    std::io::Error::new(std::io::ErrorKind::InvalidData, message)
558}
559
560#[cfg(test)]
561mod tests {
562    use super::*;
563
564    #[test]
565    fn encodes_special_requests_in_network_byte_order() {
566        let (pending, ssl) = Conn::new(()).ssl_request();
567        assert_eq!(ssl, [0, 0, 0, 8, 4, 210, 22, 47]);
568        pending.into_transport();
569
570        let (terminated, cancel) = Conn::new(())
571            .cancel_request(0x0102_0304, &[5, 6, 7, 8])
572            .expect("valid protocol 3.0 cancellation key");
573        assert_eq!(
574            &cancel[..],
575            [0, 0, 0, 16, 4, 210, 22, 46, 1, 2, 3, 4, 5, 6, 7, 8]
576        );
577        terminated.into_transport();
578    }
579
580    #[test]
581    fn tls_upgrade_changes_the_transport_type() {
582        struct Tcp;
583        struct Tls;
584
585        let (pending, _) = Conn::new(Tcp).ssl_request();
586        let Negotiation::Accepted(handshake) = pending.receive_reply(EncryptionReply::Accepted)
587        else {
588            panic!("unexpected negotiation branch")
589        };
590        let upgraded: Conn<Tls, PreStartup> = handshake.finish_tls(|Tcp| Tls);
591        let message = StartupMessage {
592            version: crate::startup::ProtocolVersion::V3_0,
593            parameters: std::collections::BTreeMap::new(),
594        };
595        let (startup, _) = upgraded.startup(&message).expect("valid startup message");
596        let _transport = startup.into_transport();
597    }
598
599    #[test]
600    fn sslmode_allow_starts_plaintext_then_reconnects_with_tls() {
601        assert_eq!(
602            SslMode::Allow.strategy(),
603            SslStrategy {
604                request_on_first_connection: false,
605                retry_with_ssl_after_plaintext_failure: true,
606                allow_server_rejection: true,
607                verification: CertificateVerification::None,
608            }
609        );
610    }
611
612    #[test]
613    fn verify_full_requires_tls_ca_and_hostname() {
614        assert_eq!(
615            SslMode::VerifyFull.strategy(),
616            SslStrategy {
617                request_on_first_connection: true,
618                retry_with_ssl_after_plaintext_failure: false,
619                allow_server_rejection: false,
620                verification: CertificateVerification::CertificateAuthorityAndHost,
621            }
622        );
623    }
624
625    #[test]
626    fn sslmode_rejection_is_plaintext_only_when_policy_allows_it() {
627        let (pending, _) = Conn::new(()).ssl_request();
628        let SslModeNegotiation::Plaintext(plaintext) =
629            pending.apply_ssl_reply(EncryptionReply::Rejected, SslMode::Prefer)
630        else {
631            panic!("prefer should permit a plaintext fallback")
632        };
633        plaintext.into_transport();
634
635        let (pending, _) = Conn::new(()).ssl_request();
636        let SslModeNegotiation::RequiredRejected { conn, mode } =
637            pending.apply_ssl_reply(EncryptionReply::Rejected, SslMode::VerifyFull)
638        else {
639            panic!("verify-full must reject a server without TLS")
640        };
641        assert_eq!(mode, SslMode::VerifyFull);
642        conn.into_transport();
643    }
644
645    #[test]
646    fn encryption_negotiation_and_upgrade_preserve_cleanliness() {
647        fn require_dirty<S>(conn: Conn<S, PreStartup, crate::Dirty>) {
648            conn.into_transport();
649        }
650
651        let pending: Conn<(), AwaitingSslReply, crate::Dirty> = Conn::new(()).transition();
652        let Negotiation::Accepted(handshake) = pending.receive_reply(EncryptionReply::Accepted)
653        else {
654            panic!("expected the TLS handshake branch")
655        };
656        let upgraded = handshake.finish_tls(|()| 42_u8);
657
658        require_dirty(upgraded);
659    }
660
661    #[test]
662    fn incrementally_decodes_each_pre_startup_branch() {
663        let messages = [
664            PreStartupMessage::SslRequest,
665            PreStartupMessage::GssEncRequest,
666            PreStartupMessage::CancelRequest {
667                process_id: 42,
668                secret_key: Bytes::from_static(&[7; 32]),
669            },
670            PreStartupMessage::Startup(StartupMessage {
671                version: crate::startup::ProtocolVersion::V3_2,
672                parameters: std::collections::BTreeMap::from([(
673                    Bytes::from_static(b"user"),
674                    Bytes::from_static(b"postgres"),
675                )]),
676            }),
677        ];
678
679        for message in messages {
680            let packet = message.to_packet().expect("encodable pre-startup message");
681            let mut input = BytesMut::from(&packet[..3]);
682            assert_eq!(
683                decode_pre_startup(&mut input).expect("partial input is valid"),
684                None
685            );
686            input.extend_from_slice(&packet[3..]);
687            assert_eq!(
688                decode_pre_startup(&mut input).expect("complete packet is valid"),
689                Some(message)
690            );
691            assert!(input.is_empty());
692        }
693    }
694
695    #[test]
696    fn rejects_oversized_pre_startup_before_reserving_body() {
697        let mut input = BytesMut::from(&10_001_u32.to_be_bytes()[..]);
698        let capacity = input.capacity();
699
700        let error = decode_pre_startup(&mut input).expect_err("packet exceeds the default limit");
701
702        assert_eq!(error.kind(), std::io::ErrorKind::InvalidData);
703        assert_eq!(input.len(), 4);
704        assert_eq!(input.capacity(), capacity);
705    }
706
707    #[test]
708    fn validates_custom_pre_startup_limit() {
709        let packet = PreStartupMessage::SslRequest
710            .to_packet()
711            .expect("SSL request is encodable");
712
713        assert!(decode_pre_startup_with_limit(&mut BytesMut::from(&packet[..]), 7).is_err());
714        assert!(
715            decode_pre_startup_with_limit(&mut BytesMut::from(&packet[..]), i32::MAX as usize + 1)
716                .is_err()
717        );
718    }
719}
pg_proto/pre_startup.rs

pg_proto/
pre_startup.rs
