mz_pgwire_common/
codec.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! Encoding/decoding of messages in pgwire. See "[Frontend/Backend Protocol:
11//! Message Formats][1]" in the PostgreSQL reference for the specification.
12//!
13//! See the [crate docs](crate) for higher level concerns.
14//!
15//! [1]: https://www.postgresql.org/docs/11/protocol-message-formats.html
16
17use std::collections::BTreeMap;
18use std::error::Error;
19use std::{fmt, str};
20
21use byteorder::{ByteOrder, NetworkEndian};
22use bytes::{BufMut, BytesMut};
23use mz_ore::cast::{CastFrom, u64_to_usize};
24use mz_ore::netio::{self};
25use tokio::io::{self, AsyncRead, AsyncReadExt};
26
27use crate::FrontendMessage;
28use crate::format::Format;
29use crate::message::{FrontendStartupMessage, VERSION_CANCEL, VERSION_GSSENC, VERSION_SSL};
30
31pub const REJECT_ENCRYPTION: u8 = b'N';
32pub const ACCEPT_SSL_ENCRYPTION: u8 = b'S';
33
34/// Maximum allowed size for a request.
35pub const MAX_REQUEST_SIZE: usize = u64_to_usize(2 * bytesize::MB);
36
37#[derive(Debug)]
38pub enum CodecError {
39    StringNoTerminator,
40}
41
42impl Error for CodecError {}
43
44impl fmt::Display for CodecError {
45    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
46        f.write_str(match self {
47            CodecError::StringNoTerminator => "The string does not have a terminator",
48        })
49    }
50}
51
52pub trait Pgbuf: BufMut {
53    fn put_string(&mut self, s: &str);
54    fn put_length_i16(&mut self, len: usize) -> Result<(), io::Error>;
55    fn put_length_u16(&mut self, len: usize) -> Result<(), io::Error>;
56    fn put_format_i8(&mut self, format: Format);
57    fn put_format_i16(&mut self, format: Format);
58}
59
60impl<B: BufMut> Pgbuf for B {
61    fn put_string(&mut self, s: &str) {
62        self.put(s.as_bytes());
63        self.put_u8(b'\0');
64    }
65
66    fn put_length_i16(&mut self, len: usize) -> Result<(), io::Error> {
67        let len = i16::try_from(len).map_err(|_| {
68            io::Error::new(io::ErrorKind::InvalidData, "length does not fit in an i16")
69        })?;
70        self.put_i16(len);
71        Ok(())
72    }
73
74    /// Writes a count field as unsigned, so it may exceed 32767. The protocol
75    /// calls these fields `Int16`, but PostgreSQL clients decode them as
76    /// unsigned.
77    fn put_length_u16(&mut self, len: usize) -> Result<(), io::Error> {
78        let len = u16::try_from(len).map_err(|_| {
79            io::Error::new(io::ErrorKind::InvalidData, "length does not fit in a u16")
80        })?;
81        self.put_u16(len);
82        Ok(())
83    }
84
85    fn put_format_i8(&mut self, format: Format) {
86        self.put_i8(format.into())
87    }
88
89    fn put_format_i16(&mut self, format: Format) {
90        self.put_i8(0);
91        self.put_format_i8(format);
92    }
93}
94
95pub async fn decode_startup<A>(mut conn: A) -> Result<Option<FrontendStartupMessage>, io::Error>
96where
97    A: AsyncRead + Unpin,
98{
99    let mut frame_len = [0; 4];
100    let nread = netio::read_exact_or_eof(&mut conn, &mut frame_len).await?;
101    match nread {
102        // Complete frame length. Continue.
103        4 => (),
104        // Connection closed cleanly. Indicate that the startup sequence has
105        // been terminated by the client.
106        0 => return Ok(None),
107        // Partial frame length. Likely a client bug or network glitch, so
108        // surface the unexpected EOF.
109        _ => return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "early eof")),
110    };
111    let frame_len = parse_frame_len(&frame_len)?;
112
113    let mut buf = BytesMut::new();
114    buf.resize(frame_len, b'0');
115    conn.read_exact(&mut buf).await?;
116
117    let mut buf = Cursor::new(&buf);
118    let version = buf.read_i32()?;
119    let message = match version {
120        VERSION_CANCEL => FrontendStartupMessage::CancelRequest {
121            conn_id: buf.read_u32()?,
122            secret_key: buf.read_u32()?,
123        },
124        VERSION_SSL => FrontendStartupMessage::SslRequest,
125        VERSION_GSSENC => FrontendStartupMessage::GssEncRequest,
126        _ => {
127            let mut params = BTreeMap::new();
128            while buf.peek_byte()? != 0 {
129                let name = buf.read_cstr()?.to_owned();
130                let value = buf.read_cstr()?.to_owned();
131                params.insert(name, value);
132            }
133            FrontendStartupMessage::Startup { version, params }
134        }
135    };
136    Ok(Some(message))
137}
138
139impl FrontendStartupMessage {
140    /// Encodes self into dst.
141    pub fn encode(&self, dst: &mut BytesMut) -> Result<(), io::Error> {
142        // Write message length placeholder. The true length is filled in later.
143        let base = dst.len();
144        dst.put_u32(0);
145
146        // Write message contents.
147        match self {
148            FrontendStartupMessage::Startup { version, params } => {
149                dst.put_i32(*version);
150                for (k, v) in params {
151                    dst.put_string(k);
152                    dst.put_string(v);
153                }
154                dst.put_i8(0);
155            }
156            FrontendStartupMessage::CancelRequest {
157                conn_id,
158                secret_key,
159            } => {
160                dst.put_i32(VERSION_CANCEL);
161                dst.put_u32(*conn_id);
162                dst.put_u32(*secret_key);
163            }
164            FrontendStartupMessage::SslRequest {} => dst.put_i32(VERSION_SSL),
165            FrontendStartupMessage::GssEncRequest => panic!("unsupported"),
166        }
167
168        let len = dst.len() - base;
169
170        // Overwrite length placeholder with true length.
171        let len = i32::try_from(len).map_err(|_| {
172            io::Error::new(
173                io::ErrorKind::InvalidData,
174                "length of encoded message does not fit into an i32",
175            )
176        })?;
177        dst[base..base + 4].copy_from_slice(&len.to_be_bytes());
178
179        Ok(())
180    }
181}
182
183impl FrontendMessage {
184    /// Encodes self into dst.
185    pub fn encode(&self, dst: &mut BytesMut) -> Result<(), io::Error> {
186        // Write type byte.
187        let byte = match self {
188            FrontendMessage::Password { .. } => b'p',
189            _ => panic!("unsupported"),
190        };
191        dst.put_u8(byte);
192
193        // Write message length placeholder. The true length is filled in later.
194        let base = dst.len();
195        dst.put_u32(0);
196
197        // Write message contents.
198        match self {
199            FrontendMessage::Password { password } => {
200                dst.put_string(password);
201            }
202            _ => panic!("unsupported"),
203        }
204
205        let len = dst.len() - base;
206
207        // Overwrite length placeholder with true length.
208        let len = i32::try_from(len).map_err(|_| {
209            io::Error::new(
210                io::ErrorKind::InvalidData,
211                "length of encoded message does not fit into an i32",
212            )
213        })?;
214        dst[base..base + 4].copy_from_slice(&len.to_be_bytes());
215
216        Ok(())
217    }
218}
219
220#[derive(Debug)]
221pub enum DecodeState {
222    Head,
223    Data(u8, usize),
224}
225
226pub fn parse_frame_len(src: &[u8]) -> Result<usize, io::Error> {
227    let n = usize::cast_from(NetworkEndian::read_u32(src));
228    if n > netio::MAX_FRAME_SIZE {
229        return Err(io::Error::new(
230            io::ErrorKind::InvalidData,
231            netio::FrameTooBig,
232        ));
233    } else if n < 4 {
234        return Err(io::Error::new(
235            io::ErrorKind::InvalidInput,
236            "invalid frame length",
237        ));
238    }
239    Ok(n - 4)
240}
241
242/// Decodes data within pgwire messages.
243///
244/// The API provided is very similar to [`bytes::Buf`], but operations return
245/// errors rather than panicking. This is important for safety, as we don't want
246/// to crash if the user sends us malformed pgwire messages.
247///
248/// There are also some special-purpose methods, like [`Cursor::read_cstr`],
249/// that are specific to pgwire messages.
250#[derive(Debug)]
251pub struct Cursor<'a> {
252    buf: &'a [u8],
253}
254
255impl<'a> Cursor<'a> {
256    /// Constructs a new `Cursor` from a byte slice. The cursor will begin
257    /// decoding from the beginning of the slice.
258    pub fn new(buf: &'a [u8]) -> Cursor<'a> {
259        Cursor { buf }
260    }
261
262    /// Returns the next byte without advancing the cursor.
263    pub fn peek_byte(&self) -> Result<u8, io::Error> {
264        self.buf
265            .get(0)
266            .copied()
267            .ok_or_else(|| input_err("No byte to read"))
268    }
269
270    /// Returns the next byte, advancing the cursor by one byte.
271    pub fn read_byte(&mut self) -> Result<u8, io::Error> {
272        let byte = self.peek_byte()?;
273        self.advance(1);
274        Ok(byte)
275    }
276
277    /// Returns the next null-terminated string. The null character is not
278    /// included the returned string. The cursor is advanced past the null-
279    /// terminated string.
280    ///
281    /// If there is no null byte remaining in the string, returns
282    /// `CodecError::StringNoTerminator`. If the string is not valid UTF-8,
283    /// returns an `io::Error` with an error kind of
284    /// `io::ErrorKind::InvalidInput`.
285    ///
286    /// NOTE(benesch): it is possible that returning a string here is wrong, and
287    /// we should be returning bytes, so that we can support messages that are
288    /// not UTF-8 encoded. At the moment, we've not discovered a need for this,
289    /// though, and using proper strings is convenient.
290    pub fn read_cstr(&mut self) -> Result<&'a str, io::Error> {
291        if let Some(pos) = self.buf.iter().position(|b| *b == 0) {
292            let val = std::str::from_utf8(&self.buf[..pos]).map_err(input_err)?;
293            self.advance(pos + 1);
294            Ok(val)
295        } else {
296            Err(input_err(CodecError::StringNoTerminator))
297        }
298    }
299
300    /// Reads the next 32-bit signed integer, advancing the cursor by four
301    /// bytes.
302    pub fn read_i32(&mut self) -> Result<i32, io::Error> {
303        if self.buf.len() < 4 {
304            return Err(input_err("not enough buffer for an Int32"));
305        }
306        let val = NetworkEndian::read_i32(self.buf);
307        self.advance(4);
308        Ok(val)
309    }
310
311    /// Reads the next 16-bit unsigned integer, advancing the cursor by two
312    /// bytes.
313    pub fn read_u16(&mut self) -> Result<u16, io::Error> {
314        if self.buf.len() < 2 {
315            return Err(input_err("not enough buffer for an Int16"));
316        }
317        let val = NetworkEndian::read_u16(self.buf);
318        self.advance(2);
319        Ok(val)
320    }
321
322    /// Reads the next 32-bit unsigned integer, advancing the cursor by four
323    /// bytes.
324    pub fn read_u32(&mut self) -> Result<u32, io::Error> {
325        if self.buf.len() < 4 {
326            return Err(input_err("not enough buffer for an Int32"));
327        }
328        let val = NetworkEndian::read_u32(self.buf);
329        self.advance(4);
330        Ok(val)
331    }
332
333    /// Reads the next 16-bit format code, advancing the cursor by two bytes.
334    pub fn read_format(&mut self) -> Result<Format, io::Error> {
335        Format::try_from(self.read_u16()?)
336    }
337
338    /// Advances the cursor by `n` bytes.
339    pub fn advance(&mut self, n: usize) {
340        self.buf = &self.buf[n..]
341    }
342}
343
344/// Constructs an error indicating that the client has violated the pgwire
345/// protocol.
346pub fn input_err(source: impl Into<Box<dyn std::error::Error + Send + Sync>>) -> io::Error {
347    io::Error::new(io::ErrorKind::InvalidInput, source.into())
348}
mz_pgwire_common/codec.rs

mz_pgwire_common/
codec.rs
