zero_postgres/protocol/frontend/
startup.rs

1//! Startup and termination messages.
2
3use crate::protocol::codec::MessageBuilder;
4
5/// Protocol version 3.0 (0x00030000)
6pub const PROTOCOL_VERSION_3_0: i32 = 196608;
7
8/// Protocol version 3.2 (0x00030002)
9pub const PROTOCOL_VERSION_3_2: i32 = 196610;
10
11/// SSL request code
12pub const SSL_REQUEST_CODE: i32 = 80877103;
13
14/// GSSAPI encryption request code
15pub const GSSENC_REQUEST_CODE: i32 = 80877104;
16
17/// Cancel request code
18pub const CANCEL_REQUEST_CODE: i32 = 80877102;
19
20/// Write an SSLRequest message.
21///
22/// This is sent before StartupMessage to request TLS encryption.
23/// Server responds with single byte: 'S' (accepted) or 'N' (rejected).
24pub fn write_ssl_request(buf: &mut Vec<u8>) {
25    let mut msg = MessageBuilder::new_startup(buf);
26    msg.write_i32(SSL_REQUEST_CODE);
27    msg.finish();
28}
29
30/// Write a GSSENCRequest message.
31///
32/// This is sent before StartupMessage to request GSSAPI encryption.
33/// Server responds with single byte: 'G' (accepted) or 'N' (rejected).
34pub fn write_gssenc_request(buf: &mut Vec<u8>) {
35    let mut msg = MessageBuilder::new_startup(buf);
36    msg.write_i32(GSSENC_REQUEST_CODE);
37    msg.finish();
38}
39
40/// Write a StartupMessage.
41///
42/// Parameters is a list of (name, value) pairs.
43/// Required: "user" - database username
44/// Optional: "database", "options", "replication", "client_encoding", etc.
45pub fn write_startup(buf: &mut Vec<u8>, params: &[(&str, &str)]) {
46    write_startup_with_version(buf, PROTOCOL_VERSION_3_2, params);
47}
48
49/// Write a StartupMessage with a specific protocol version.
50pub fn write_startup_with_version(buf: &mut Vec<u8>, version: i32, params: &[(&str, &str)]) {
51    let mut msg = MessageBuilder::new_startup(buf);
52    msg.write_i32(version);
53
54    for (name, value) in params {
55        msg.write_cstr(name);
56        msg.write_cstr(value);
57    }
58
59    // Terminator
60    msg.write_u8(0);
61    msg.finish();
62}
63
64/// Write a CancelRequest message.
65///
66/// This is sent on a NEW connection to cancel a query running on another connection.
67/// The server closes the connection immediately with no response.
68///
69/// In protocol 3.2, the secret key is variable-length (4-256 bytes).
70pub fn write_cancel_request(buf: &mut Vec<u8>, pid: u32, secret_key: &[u8]) {
71    let mut msg = MessageBuilder::new_startup(buf);
72    msg.write_i32(CANCEL_REQUEST_CODE);
73    msg.write_i32(pid as i32);
74    msg.write_bytes(secret_key);
75    msg.finish();
76}
77
78/// Write a Terminate message.
79///
80/// Sent to cleanly close the connection.
81pub fn write_terminate(buf: &mut Vec<u8>) {
82    let msg = MessageBuilder::new(buf, super::msg_type::TERMINATE);
83    msg.finish();
84}
85
86#[cfg(test)]
87mod tests {
88    use super::*;
89
90    #[test]
91    fn test_ssl_request() {
92        let mut buf = Vec::new();
93        write_ssl_request(&mut buf);
94
95        assert_eq!(buf.len(), 8);
96        assert_eq!(&buf[0..4], &8_i32.to_be_bytes());
97        assert_eq!(&buf[4..8], &SSL_REQUEST_CODE.to_be_bytes());
98    }
99
100    #[test]
101    fn test_startup() {
102        let mut buf = Vec::new();
103        write_startup(&mut buf, &[("user", "postgres"), ("database", "test")]);
104
105        // Check length is at start
106        let len = i32::from_be_bytes([buf[0], buf[1], buf[2], buf[3]]);
107        assert_eq!(len as usize, buf.len());
108
109        // Check protocol version
110        let version = i32::from_be_bytes([buf[4], buf[5], buf[6], buf[7]]);
111        assert_eq!(version, PROTOCOL_VERSION_3_2);
112    }
113
114    #[test]
115    fn test_terminate() {
116        let mut buf = Vec::new();
117        write_terminate(&mut buf);
118
119        assert_eq!(buf.len(), 5);
120        assert_eq!(buf[0], b'X');
121        assert_eq!(&buf[1..5], &4_i32.to_be_bytes());
122    }
123}