postgres_protocol/message/
startup.rs

1#![allow(missing_docs)]
2
3use std::collections::HashMap;
4use std::io;
5use std::iter::FromIterator;
6
7use crate::message::shared::Buffer;
8use byteorder::{BigEndian, ReadBytesExt};
9use bytes::{BufMut, Bytes, BytesMut};
10
11#[derive(Debug)]
12pub enum StartupRequest {
13    Startup(Bytes),
14    Cancel(CancelData),
15    SSLRequest,
16    GSSEncRequest,
17    Password(Bytes),
18}
19
20#[derive(Debug)]
21pub struct CancelData {
22    pub process_id: i32,
23    pub secret_key: i32,
24}
25
26impl StartupRequest {
27    #[inline]
28    pub fn parse_without_tag(buf: &mut BytesMut) -> io::Result<Option<StartupRequest>> {
29        if buf.len() < 4 {
30            let to_read = 4 - buf.len();
31            buf.reserve(to_read);
32            return Ok(None);
33        }
34
35        let len = (&buf[0..4]).read_u32::<BigEndian>().unwrap() as usize;
36
37        // defined in pg source
38        const MIN_STARTUP_LEN: usize = 4;
39        const MAX_STARTUP_LEN: usize = 10000;
40
41        if !(MIN_STARTUP_LEN..=MAX_STARTUP_LEN).contains(&len) {
42            return Err(io::Error::new(
43                io::ErrorKind::InvalidInput,
44                "invalid message length: parsing u32",
45            ));
46        }
47
48        // Read the message contents.
49        if buf.len() < len {
50            let to_read = len - buf.len();
51            buf.reserve(to_read);
52            return Ok(None);
53        }
54
55        let mut buf = Buffer {
56            bytes: buf.split_to(len).freeze(),
57            idx: 4,
58        };
59
60        let code = buf.read_u32::<BigEndian>()?;
61        let message = match code {
62            // Startup message
63            196_608 => {
64                let data = buf.read_all();
65                StartupRequest::Startup(data)
66            }
67
68            // Cancel Request
69            80_877_102 => {
70                let process_id = buf.read_i32::<BigEndian>()?;
71                let secret_key = buf.read_i32::<BigEndian>()?;
72                StartupRequest::Cancel(CancelData {
73                    process_id,
74                    secret_key,
75                })
76            }
77
78            // SSL Request
79            80_877_103 => StartupRequest::SSLRequest,
80
81            // GSS Encode Request
82            80_877_104 => StartupRequest::GSSEncRequest,
83
84            _ => {
85                return Err(io::Error::new(
86                    io::ErrorKind::InvalidInput,
87                    format!("unknown startup message code `{}`", code),
88                ));
89            }
90        };
91
92        if !buf.is_empty() {
93            return Err(io::Error::new(
94                io::ErrorKind::InvalidInput,
95                "invalid message length: expected buffer to be empty",
96            ));
97        }
98
99        Ok(Some(message))
100    }
101
102    #[inline]
103    pub fn parse_with_tag(buf: &mut BytesMut) -> io::Result<Option<StartupRequest>> {
104        if buf.len() < 5 {
105            let to_read = 5 - buf.len();
106            buf.reserve(to_read);
107            return Ok(None);
108        }
109
110        let len = (&buf[1..5]).read_u32::<BigEndian>().unwrap() as usize;
111
112        // defined in pg source
113        const MIN_STARTUP_LEN: usize = 4;
114        const MAX_STARTUP_LEN: usize = 10000;
115
116        if !(MIN_STARTUP_LEN..=MAX_STARTUP_LEN).contains(&len) {
117            return Err(io::Error::new(
118                io::ErrorKind::InvalidInput,
119                "invalid message length: parsing u32",
120            ));
121        }
122
123        // Read the message contents.
124        if buf.len() < len {
125            let to_read = len - buf.len();
126            buf.reserve(to_read);
127            return Ok(None);
128        }
129
130        let tag = buf[0] as char;
131        let mut buf = Buffer {
132            bytes: buf.split_to(len).freeze(),
133            idx: 5,
134        };
135
136        let message = match tag {
137            // PasswordMessage
138            'p' => {
139                let passwd = buf.read_cstr()?;
140                StartupRequest::Password(passwd)
141            }
142
143            _ => {
144                return Err(io::Error::new(
145                    io::ErrorKind::InvalidInput,
146                    format!("unknown startup message tag '{}'", tag),
147                ));
148            }
149        };
150
151        if !buf.is_empty() {
152            return Err(io::Error::new(
153                io::ErrorKind::InvalidInput,
154                "invalid message length: expected buffer to be empty",
155            ));
156        }
157
158        Ok(Some(message))
159    }
160}
161
162#[derive(Debug)]
163pub struct StartupData {
164    pub parameters: HashMap<String, Bytes>,
165}
166
167impl StartupData {
168    pub fn parse(bytes: Bytes) -> io::Result<StartupData> {
169        let mut buf = Buffer { bytes, idx: 0 };
170
171        let pairs = std::iter::from_fn(move || {
172            let key = buf.read_cstr().ok()?;
173            if key.is_empty() {
174                return None;
175            }
176            let value = buf.read_cstr().ok()?;
177            let key_str = String::from_utf8(key.to_vec()).ok()?;
178            Some((key_str, value))
179        });
180
181        let parameters = HashMap::from_iter(pairs);
182        Ok(StartupData { parameters })
183    }
184}
185
186#[derive(Debug)]
187pub enum StartupResponse {
188    AuthenticationOk,
189    AuthenticationMD5Password { salt: [u8; 4] },
190    SSLResponse(bool),
191    GSSEncResponse(bool),
192    ErrorResponse { code: &'static str, message: String },
193    ParameterStatus { key: String, value: Bytes },
194    BackendKeyData { process_id: i32, secret_key: i32 },
195    ReadyForQuery,
196}
197
198impl StartupResponse {
199    pub fn encode(&self, dst: &mut BytesMut) -> Result<(), io::Error> {
200        match self {
201            StartupResponse::AuthenticationOk => {
202                dst.reserve(1 + 8);
203                dst.put_u8(b'R');
204                dst.put_u32(8);
205                dst.put_u32(0); // auth ok
206            }
207            StartupResponse::AuthenticationMD5Password { salt } => {
208                dst.reserve(1 + 12);
209                dst.put_u8(b'R');
210                dst.put_u32(12);
211                dst.put_slice(&salt[..]);
212                dst.put_u32(0); // salt
213            }
214            StartupResponse::SSLResponse(ok) => {
215                dst.reserve(1);
216                dst.put_u8(if *ok { b'S' } else { b'N' });
217            }
218            StartupResponse::GSSEncResponse(ok) => {
219                dst.reserve(1);
220                dst.put_u8(if *ok { b'G' } else { b'N' });
221            }
222            StartupResponse::ErrorResponse { code, message } => {
223                let code_bytes = code.as_bytes();
224                let err_bytes = message.as_bytes();
225                let len =
226                    4 + (1 + 5 + 1) + (1 + code_bytes.len() + 1) + (1 + err_bytes.len() + 1) + 1;
227                dst.reserve(1 + len);
228                dst.put_u8(b'E');
229                dst.put_u32(len as u32);
230
231                dst.put_u8(b'S');
232                dst.put_slice(b"FATAL");
233                dst.put_u8(0);
234
235                // Clients reject an ErrorResponse without a SQLSTATE.
236                dst.put_u8(b'C');
237                dst.put_slice(code_bytes);
238                dst.put_u8(0);
239
240                dst.put_u8(b'M');
241                dst.put_slice(err_bytes);
242                dst.put_u8(0);
243                dst.put_u8(0);
244            }
245
246            StartupResponse::ParameterStatus { key, value } => {
247                let key_bytes = key.as_bytes();
248                let value_bytes = value.as_ref();
249
250                let len = 4 + (key_bytes.len() + 1) + (value_bytes.len() + 1);
251                dst.reserve(1 + len);
252                dst.put_u8(b'S');
253                dst.put_u32(len as u32);
254
255                dst.put_slice(key_bytes);
256                dst.put_u8(0);
257                dst.put_slice(value_bytes);
258                dst.put_u8(0);
259            }
260
261            StartupResponse::BackendKeyData {
262                process_id,
263                secret_key,
264            } => {
265                dst.reserve(1 + 12);
266                dst.put_u8(b'K');
267                dst.put_u32(12);
268                dst.put_i32(*process_id);
269                dst.put_i32(*secret_key);
270            }
271
272            StartupResponse::ReadyForQuery => {
273                dst.reserve(1 + 5);
274                dst.put_u8(b'Z');
275                dst.put_u32(5);
276                dst.put_u8(b'I');
277            }
278        }
279        Ok(())
280    }
281}
282
283#[cfg(test)]
284mod tests {
285    use fallible_iterator::FallibleIterator;
286
287    use super::*;
288    use crate::message::backend::Message;
289
290    #[test]
291    fn error_response_is_parseable_by_clients() {
292        let mut buf = BytesMut::new();
293        StartupResponse::ErrorResponse {
294            code: "3D000",
295            message: "connection rejected: unknown database".to_string(),
296        }
297        .encode(&mut buf)
298        .unwrap();
299
300        let body = match Message::parse(&mut buf).unwrap() {
301            Some(Message::ErrorResponse(body)) => body,
302            _ => panic!("expected an ErrorResponse"),
303        };
304        assert!(buf.is_empty());
305
306        let fields: Vec<(u8, Vec<u8>)> = body
307            .fields()
308            .map(|f| Ok((f.type_(), f.value_bytes().to_vec())))
309            .collect()
310            .unwrap();
311        assert_eq!(
312            fields,
313            vec![
314                (b'S', b"FATAL".to_vec()),
315                (b'C', b"3D000".to_vec()),
316                (b'M', b"connection rejected: unknown database".to_vec()),
317            ]
318        );
319    }
320}
postgres_protocol/message/startup.rs

postgres_protocol/message/
startup.rs
