pg_proto/
startup.rs

1//! Untagged startup messages and protocol-version negotiation.
2
3use std::{collections::BTreeMap, io};
4
5use bytes::{Buf, BufMut, Bytes, BytesMut};
6
7/// A frontend startup message, retained as bytes for lossless proxy forwarding.
8#[derive(Clone, Debug, Eq, PartialEq)]
9pub struct StartupMessage {
10    /// Requested frontend/backend protocol version.
11    pub version: ProtocolVersion,
12    /// Startup parameter names and values, including user and database.
13    pub parameters: BTreeMap<Bytes, Bytes>,
14}
15
16impl StartupMessage {
17    /// Encodes the untagged startup packet.
18    ///
19    /// # Errors
20    ///
21    /// Returns an error for embedded NUL bytes or a packet larger than `i32::MAX`.
22    pub fn encode(&self) -> io::Result<Bytes> {
23        let mut output = BytesMut::new();
24        output.extend_from_slice(&[0; 4]);
25        output.put_u16(self.version.major);
26        output.put_u16(self.version.minor);
27        for (name, value) in &self.parameters {
28            put_cstr(name, &mut output)?;
29            put_cstr(value, &mut output)?;
30        }
31        output.put_u8(0);
32        let length = i32::try_from(output.len())
33            .map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "startup packet too large"))?;
34        output[..4].copy_from_slice(&length.to_be_bytes());
35        Ok(output.freeze())
36    }
37
38    /// Decodes one complete untagged startup packet.
39    ///
40    /// # Errors
41    ///
42    /// Returns an error for an invalid length, malformed strings, duplicate
43    /// parameters, or trailing bytes.
44    pub fn decode(mut packet: Bytes) -> io::Result<Self> {
45        if packet.len() < 8 {
46            return Err(invalid("startup packet is shorter than 8 bytes"));
47        }
48        let declared = usize::try_from(packet.get_u32())
49            .map_err(|_| invalid("startup packet length overflow"))?;
50        if declared != packet.len() + 4 {
51            return Err(invalid("startup packet length does not match its bytes"));
52        }
53        let version = ProtocolVersion {
54            major: packet.get_u16(),
55            minor: packet.get_u16(),
56        };
57        let mut parameters = BTreeMap::new();
58        loop {
59            if packet.is_empty() {
60                return Err(invalid("startup parameters have no final terminator"));
61            }
62            if packet[0] == 0 {
63                packet.advance(1);
64                break;
65            }
66            let name = take_cstr(&mut packet)?;
67            let value = take_cstr(&mut packet)?;
68            if parameters.insert(name, value).is_some() {
69                return Err(invalid("duplicate startup parameter"));
70            }
71        }
72        if !packet.is_empty() {
73            return Err(invalid("startup packet has trailing bytes"));
74        }
75        Ok(Self {
76            version,
77            parameters,
78        })
79    }
80}
81
82/// `PostgreSQL` protocol version, including supported 3.x minor versions.
83#[derive(Clone, Copy, Debug, Eq, PartialEq)]
84pub struct ProtocolVersion {
85    /// Protocol major version.
86    pub major: u16,
87    /// Protocol minor version.
88    pub minor: u16,
89}
90
91impl ProtocolVersion {
92    /// `PostgreSQL` protocol 3.0.
93    pub const V3_0: Self = Self { major: 3, minor: 0 };
94    /// `PostgreSQL` protocol 3.1.
95    pub const V3_1: Self = Self { major: 3, minor: 1 };
96    /// `PostgreSQL` protocol 3.2.
97    pub const V3_2: Self = Self { major: 3, minor: 2 };
98}
99
100fn put_cstr(value: &[u8], output: &mut BytesMut) -> io::Result<()> {
101    if value.contains(&0) {
102        return Err(io::Error::new(
103            io::ErrorKind::InvalidInput,
104            "startup parameter contains a NUL byte",
105        ));
106    }
107    output.extend_from_slice(value);
108    output.put_u8(0);
109    Ok(())
110}
111
112fn take_cstr(input: &mut Bytes) -> io::Result<Bytes> {
113    let end = input
114        .iter()
115        .position(|byte| *byte == 0)
116        .ok_or_else(|| invalid("unterminated startup parameter"))?;
117    let value = input.split_to(end);
118    input.advance(1);
119    Ok(value)
120}
121
122fn invalid(message: &'static str) -> io::Error {
123    io::Error::new(io::ErrorKind::InvalidData, message)
124}
125
126#[cfg(test)]
127mod tests {
128    use super::*;
129
130    #[test]
131    fn startup_packet_has_raw_length_and_version() {
132        let message = StartupMessage {
133            version: ProtocolVersion::V3_2,
134            parameters: BTreeMap::from([
135                (Bytes::from_static(b"database"), Bytes::from_static(b"db")),
136                (Bytes::from_static(b"user"), Bytes::from_static(b"alice")),
137            ]),
138        };
139        let encoded = message.encode().expect("valid startup message");
140        assert_eq!(u32::from_be_bytes(encoded[..4].try_into().unwrap()), 32);
141        assert_eq!(&encoded[4..8], &[0, 3, 0, 2]);
142        assert_eq!(encoded.last(), Some(&0));
143        assert_eq!(
144            StartupMessage::decode(encoded).expect("decodable startup message"),
145            message
146        );
147    }
148}
pg_proto/startup.rs

pg_proto/
startup.rs
