mitosis/router/
wire.rs

1//! Postgres wire protocol pieces the router needs: the untagged first message a client sends,
2//! tagged messages, and the few messages the router builds itself. Pure parsing and building.
3
4use std::io::{Error, ErrorKind, Result};
5use tokio::io::{AsyncRead, AsyncReadExt};
6
7/// Longest first message accepted (a startup message with its parameters).
8pub const MAX_FIRST: usize = 10_000;
9/// Longest tagged message `read_message` accepts; only the startup exchange goes through it.
10pub const MAX_MESSAGE: usize = 1 << 24;
11
12const SSL_CODE: u32 = 80_877_103;
13const GSSENC_CODE: u32 = 80_877_104;
14const CANCEL_CODE: u32 = 80_877_102;
15/// First byte of a TLS ClientHello (`sslnegotiation=direct`).
16const TLS_HANDSHAKE: u8 = 0x16;
17
18#[derive(Debug, PartialEq, Eq)]
19pub enum FirstMessage {
20    Ssl,
21    GssEnc,
22    Cancel {
23        pid: i32,
24        key: Vec<u8>,
25    },
26    Startup {
27        major: u16,
28        minor: u16,
29        params: Vec<(String, String)>,
30    },
31    DirectTls,
32}
33
34fn bad(msg: impl Into<String>) -> Error {
35    Error::new(ErrorKind::InvalidData, msg.into())
36}
37
38/// Reads one length-prefixed first message (startup, SSL/GSS request or cancel).
39pub async fn read_first<R: AsyncRead + Unpin>(r: &mut R) -> Result<FirstMessage> {
40    let mut head = [0u8; 4];
41    r.read_exact(&mut head[..1]).await?;
42    // A real length would have to be over 350 MB to start with this byte.
43    if head[0] == TLS_HANDSHAKE {
44        return Ok(FirstMessage::DirectTls);
45    }
46    r.read_exact(&mut head[1..]).await?;
47    let len = u32::from_be_bytes(head) as usize;
48    if !(8..=MAX_FIRST).contains(&len) {
49        return Err(bad(format!("first message length {len} out of range")));
50    }
51    let mut body = vec![0u8; len - 4];
52    r.read_exact(&mut body).await?;
53    parse_first(&body)
54}
55
56/// `body` is the message after its length.
57fn parse_first(body: &[u8]) -> Result<FirstMessage> {
58    let code = u32::from_be_bytes(body[..4].try_into().expect("4 bytes"));
59    let rest = &body[4..];
60    match code {
61        SSL_CODE if rest.is_empty() => Ok(FirstMessage::Ssl),
62        GSSENC_CODE if rest.is_empty() => Ok(FirstMessage::GssEnc),
63        CANCEL_CODE if rest.len() >= 4 => {
64            let pid = i32::from_be_bytes(rest[..4].try_into().expect("4 bytes"));
65            let key = rest[4..].to_vec();
66            if !(4..=256).contains(&key.len()) {
67                return Err(bad("cancel key length out of range"));
68            }
69            Ok(FirstMessage::Cancel { pid, key })
70        }
71        _ => {
72            let (major, minor) = ((code >> 16) as u16, (code & 0xffff) as u16);
73            if major == 1234 {
74                return Err(bad(format!("unknown request code {code}")));
75            }
76            Ok(FirstMessage::Startup {
77                major,
78                minor,
79                params: parse_params(rest)?,
80            })
81        }
82    }
83}
84
85/// `name\0value\0 ... \0`.
86fn parse_params(mut b: &[u8]) -> Result<Vec<(String, String)>> {
87    let cstr = |b: &mut &[u8]| -> Result<String> {
88        let end = b
89            .iter()
90            .position(|&c| c == 0)
91            .ok_or_else(|| bad("startup parameters are not terminated"))?;
92        let s = String::from_utf8(b[..end].to_vec())
93            .map_err(|_| bad("startup parameter is not UTF-8"))?;
94        *b = &b[end + 1..];
95        Ok(s)
96    };
97    let mut params = Vec::new();
98    loop {
99        let name = cstr(&mut b)?;
100        if name.is_empty() {
101            return if b.is_empty() {
102                Ok(params)
103            } else {
104                Err(bad("bytes after the startup parameters"))
105            };
106        }
107        let value = cstr(&mut b)?;
108        params.push((name, value));
109    }
110}
111
112fn with_len(mut b: Vec<u8>, at: usize) -> Vec<u8> {
113    let len = u32::try_from(b.len() - at).expect("message under 4 GB");
114    b[at..at + 4].copy_from_slice(&len.to_be_bytes());
115    b
116}
117
118/// A startup message for protocol `major.minor` with `params`.
119pub fn startup_bytes(major: u16, minor: u16, params: &[(String, String)]) -> Vec<u8> {
120    let mut b = vec![0u8; 4];
121    b.extend_from_slice(&((u32::from(major) << 16) | u32::from(minor)).to_be_bytes());
122    for (k, v) in params {
123        b.extend_from_slice(k.as_bytes());
124        b.push(0);
125        b.extend_from_slice(v.as_bytes());
126        b.push(0);
127    }
128    b.push(0);
129    with_len(b, 0)
130}
131
132/// A cancel request; its length follows the key (4 bytes in 3.0, up to 256 in 3.2).
133pub fn cancel_bytes(pid: i32, key: &[u8]) -> Vec<u8> {
134    let mut b = vec![0u8; 4];
135    b.extend_from_slice(&CANCEL_CODE.to_be_bytes());
136    b.extend_from_slice(&pid.to_be_bytes());
137    b.extend_from_slice(key);
138    with_len(b, 0)
139}
140
141/// A tagged message: `tag`, length, `body`.
142pub fn message(tag: u8, body: &[u8]) -> Vec<u8> {
143    let mut b = vec![tag, 0, 0, 0, 0];
144    b.extend_from_slice(body);
145    with_len(b, 1)
146}
147
148/// Reads one tagged message; returns its tag and body.
149pub async fn read_message<R: AsyncRead + Unpin>(r: &mut R) -> Result<(u8, Vec<u8>)> {
150    let mut head = [0u8; 5];
151    r.read_exact(&mut head).await?;
152    let len = u32::from_be_bytes(head[1..].try_into().expect("4 bytes")) as usize;
153    if !(4..=MAX_MESSAGE).contains(&len) {
154        return Err(bad(format!("message length {len} out of range")));
155    }
156    let mut body = vec![0u8; len - 4];
157    r.read_exact(&mut body).await?;
158    Ok((head[0], body))
159}
160
161/// A FATAL `ErrorResponse` with SQLSTATE `sqlstate`.
162pub fn error_response(sqlstate: &str, message_text: &str, hint: Option<&str>) -> Vec<u8> {
163    let mut body = Vec::new();
164    let mut field = |code: u8, value: &str| {
165        body.push(code);
166        // A NUL would end the field early.
167        body.extend(value.bytes().filter(|&c| c != 0));
168        body.push(0);
169    };
170    field(b'S', "FATAL");
171    field(b'V', "FATAL");
172    field(b'C', sqlstate);
173    field(b'M', message_text);
174    if let Some(h) = hint {
175        field(b'H', h);
176    }
177    body.push(0);
178    message(b'E', &body)
179}
180
181/// The process id and secret key of a `BackendKeyData` body.
182pub fn parse_backend_key(body: &[u8]) -> Option<(i32, Vec<u8>)> {
183    let pid = i32::from_be_bytes(body.get(..4)?.try_into().ok()?);
184    let key = &body[4..];
185    (4..=256).contains(&key.len()).then(|| (pid, key.to_vec()))
186}
187
188#[cfg(test)]
189mod tests {
190    use super::*;
191
192    async fn first(b: &[u8]) -> Result<FirstMessage> {
193        read_first(&mut &b[..]).await
194    }
195
196    fn params() -> Vec<(String, String)> {
197        vec![
198            ("user".into(), "postgres".into()),
199            ("database".into(), "feat".into()),
200        ]
201    }
202
203    #[tokio::test]
204    async fn parses_startup_30_and_32() {
205        for minor in [0, 2] {
206            let b = startup_bytes(3, minor, &params());
207            assert_eq!(
208                u32::from_be_bytes(b[..4].try_into().unwrap()) as usize,
209                b.len()
210            );
211            assert_eq!(
212                first(&b).await.unwrap(),
213                FirstMessage::Startup {
214                    major: 3,
215                    minor,
216                    params: params()
217                }
218            );
219        }
220        // Unterminated parameters.
221        let mut b = startup_bytes(3, 0, &params());
222        b.pop();
223        let len = (b.len() as u32).to_be_bytes();
224        b[..4].copy_from_slice(&len);
225        assert!(first(&b).await.is_err());
226    }
227
228    #[tokio::test]
229    async fn parses_ssl_gss_cancel_3_0_and_3_2() {
230        let req = |code: u32| {
231            let mut b = 8u32.to_be_bytes().to_vec();
232            b.extend_from_slice(&code.to_be_bytes());
233            b
234        };
235        assert_eq!(first(&req(SSL_CODE)).await.unwrap(), FirstMessage::Ssl);
236        assert_eq!(
237            first(&req(GSSENC_CODE)).await.unwrap(),
238            FirstMessage::GssEnc
239        );
240        for key in [vec![1, 2, 3, 4], (0..32).collect::<Vec<u8>>()] {
241            let b = cancel_bytes(77, &key);
242            assert_eq!(b.len(), 12 + key.len());
243            assert_eq!(
244                first(&b).await.unwrap(),
245                FirstMessage::Cancel {
246                    pid: 77,
247                    key: key.clone()
248                }
249            );
250        }
251        assert!(first(&cancel_bytes(1, &[1, 2])).await.is_err());
252        assert_eq!(
253            first(&[0x16, 3, 1, 0, 0]).await.unwrap(),
254            FirstMessage::DirectTls
255        );
256        assert!(first(&req((1234 << 16) | 9999)).await.is_err());
257    }
258
259    #[tokio::test]
260    async fn rejects_oversized_first_message() {
261        let mut b = 10_001u32.to_be_bytes().to_vec();
262        b.extend(vec![0u8; 10_000]);
263        assert!(first(&b).await.is_err());
264        assert!(first(&4u32.to_be_bytes()).await.is_err());
265        // Exactly the limit is fine when well formed.
266        let big = vec![("x".to_string(), "y".repeat(MAX_FIRST - 12))];
267        let b = startup_bytes(3, 0, &big);
268        assert_eq!(b.len(), MAX_FIRST);
269        assert!(first(&b).await.is_ok());
270    }
271
272    #[test]
273    fn backend_key_4_and_32_bytes() {
274        for n in [4usize, 32, 256] {
275            let mut body = 42i32.to_be_bytes().to_vec();
276            body.extend(vec![7u8; n]);
277            assert_eq!(parse_backend_key(&body), Some((42, vec![7u8; n])));
278        }
279        for n in [0usize, 3, 257] {
280            let mut body = 42i32.to_be_bytes().to_vec();
281            body.extend(vec![7u8; n]);
282            assert_eq!(parse_backend_key(&body), None);
283        }
284        assert_eq!(parse_backend_key(&[1, 2]), None);
285    }
286
287    #[tokio::test]
288    async fn error_response_shape() {
289        let b = error_response("3D000", "no clone named \"x\"", Some("connect as y"));
290        let (tag, body) = read_message(&mut &b[..]).await.unwrap();
291        assert_eq!(tag, b'E');
292        let mut fields = Vec::new();
293        let mut rest = &body[..];
294        while rest[0] != 0 {
295            let end = rest.iter().position(|&c| c == 0).unwrap();
296            fields.push((rest[0], String::from_utf8(rest[1..end].to_vec()).unwrap()));
297            rest = &rest[end + 1..];
298        }
299        assert_eq!(rest, [0]);
300        assert_eq!(
301            fields,
302            vec![
303                (b'S', "FATAL".into()),
304                (b'V', "FATAL".into()),
305                (b'C', "3D000".into()),
306                (b'M', "no clone named \"x\"".into()),
307                (b'H', "connect as y".into()),
308            ]
309        );
310        let b = error_response("57P03", "m", None);
311        assert!(!b.contains(&b'H'));
312    }
313}
mitosis/router/wire.rs

mitosis/router/
wire.rs
