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

pg_proto/
pre_startup.rs
