mz_balancerd/
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
10use async_trait::async_trait;
11use bytes::{Buf, BufMut, BytesMut};
12use bytesize::ByteSize;
13use futures::{SinkExt, TryStreamExt, sink};
14use mz_ore::cast::CastFrom;
15use mz_ore::future::OreSinkExt;
16use mz_ore::netio::AsyncReady;
17use mz_pgwire_common::{
18    Conn, Cursor, DecodeState, ErrorResponse, FrontendMessage, MAX_PREAUTH_FRAME_SIZE,
19    MAX_REQUEST_SIZE, Pgbuf, parse_frame_len,
20};
21use tokio::io::{self, AsyncRead, AsyncWrite, Interest, Ready};
22use tokio_util::codec::{Decoder, Encoder, Framed};
23
24/// Internal representation of a backend [message].
25///
26/// [message]: https://www.postgresql.org/docs/11/protocol-message-formats.html
27#[derive(Debug)]
28pub enum BackendMessage {
29    AuthenticationCleartextPassword,
30    ErrorResponse(ErrorResponse),
31}
32
33impl From<ErrorResponse> for BackendMessage {
34    fn from(err: ErrorResponse) -> BackendMessage {
35        BackendMessage::ErrorResponse(err)
36    }
37}
38
39/// A connection that manages the encoding and decoding of pgwire frames.
40///
41/// This decodes at most one frame per connection, the client's credential, and
42/// is bounded by [`MAX_PREAUTH_FRAME_SIZE`] throughout. Once the destination is
43/// resolved the connection is spliced and the remaining bytes are proxied
44/// without being framed.
45pub struct FramedConn<A> {
46    inner: sink::Buffer<Framed<Conn<A>, Codec>, BackendMessage>,
47}
48
49impl<A> FramedConn<A>
50where
51    A: AsyncRead + AsyncWrite + Unpin,
52{
53    /// Constructs a new framed connection.
54    ///
55    /// The underlying connection, `inner`, is expected to be something like a
56    /// TCP stream. Anything that implements [`AsyncRead`] and [`AsyncWrite`]
57    /// will do.
58    pub fn new(inner: Conn<A>) -> FramedConn<A> {
59        FramedConn {
60            inner: Framed::new(inner, Codec::new()).buffer(32),
61        }
62    }
63
64    /// Reads and decodes one frontend message from the client.
65    ///
66    /// Blocks until the client sends a complete message. If the client
67    /// terminates the stream, returns `None`. Returns an error if the client
68    /// sends a malformed message or if the connection underlying is broken.
69    ///
70    /// # Cancel safety
71    ///
72    /// This method is cancel safe. The returned future only holds onto a
73    /// reference to thea underlying stream, so dropping it will never lose a
74    /// value.
75    ///
76    /// <https://docs.rs/tokio-stream/latest/tokio_stream/trait.StreamExt.html#cancel-safety-1>
77    pub async fn recv(&mut self) -> Result<Option<FrontendMessage>, io::Error> {
78        let message = self.inner.try_next().await?;
79        Ok(message)
80    }
81
82    /// Encodes and sends one backend message to the client.
83    ///
84    /// Note that the connection is not flushed after calling this method. You
85    /// must call [`FramedConn::flush`] explicitly. Returns an error if the
86    /// underlying connection is broken.
87    ///
88    /// Please use `StateMachine::send` instead if calling from `StateMachine`,
89    /// as it applies session-based filters before calling this method.
90    pub async fn send<M>(&mut self, message: M) -> Result<(), io::Error>
91    where
92        M: Into<BackendMessage>,
93    {
94        let message = message.into();
95        self.inner.enqueue(message).await
96    }
97
98    /// Flushes all outstanding messages.
99    pub async fn flush(&mut self) -> Result<(), io::Error> {
100        self.inner.flush().await
101    }
102}
103
104impl<A> FramedConn<A>
105where
106    A: AsyncRead + AsyncWrite + Unpin,
107{
108    pub fn inner(&self) -> &Conn<A> {
109        self.inner.get_ref().get_ref()
110    }
111    pub fn inner_mut(&mut self) -> &mut Conn<A> {
112        self.inner.get_mut().get_mut()
113    }
114
115    /// Releases the framing buffers, returning the connection and any bytes
116    /// read from it but not yet decoded.
117    ///
118    /// Messages sent but not flushed are lost, so callers must
119    /// [`FramedConn::flush`] first.
120    pub fn into_parts(self) -> (Conn<A>, BytesMut) {
121        let parts = self.inner.into_inner().into_parts();
122        debug_assert!(parts.write_buf.is_empty(), "unflushed messages");
123        (parts.io, parts.read_buf)
124    }
125}
126
127#[async_trait]
128impl<A> AsyncReady for FramedConn<A>
129where
130    A: AsyncRead + AsyncWrite + AsyncReady + Send + Sync + Unpin,
131{
132    async fn ready(&self, interest: Interest) -> io::Result<Ready> {
133        self.inner.get_ref().get_ref().ready(interest).await
134    }
135}
136
137struct Codec {
138    decode_state: DecodeState,
139}
140
141impl Codec {
142    /// Creates a new `Codec`.
143    pub fn new() -> Codec {
144        Codec {
145            decode_state: DecodeState::Head,
146        }
147    }
148}
149
150impl Default for Codec {
151    fn default() -> Codec {
152        Codec::new()
153    }
154}
155
156impl Encoder<BackendMessage> for Codec {
157    type Error = io::Error;
158
159    /// Encode a backend message into `dst`.
160    /// If this function returns an error result, `dst` is left unmodified.
161    fn encode(&mut self, msg: BackendMessage, dst: &mut BytesMut) -> Result<(), io::Error> {
162        // Record the starting position so we can truncate on error.
163        // This prevents partial messages from being left in the buffer,
164        // which could be sent to the client and cause "lost synchronization" errors.
165        let start = dst.len();
166        match self.encode_inner(msg, dst) {
167            Ok(()) => Ok(()),
168            Err(e) => {
169                dst.truncate(start);
170                Err(e)
171            }
172        }
173    }
174}
175
176impl Codec {
177    /// This is the meat of the encoding logic. It's a separate function so that errors returned by
178    /// `?` can be handled in the outer `encode` function.
179    fn encode_inner(&self, msg: BackendMessage, dst: &mut BytesMut) -> Result<(), io::Error> {
180        // Write type byte.
181        let byte = match &msg {
182            BackendMessage::AuthenticationCleartextPassword => b'R',
183            BackendMessage::ErrorResponse(r) => {
184                if r.severity.is_error() {
185                    b'E'
186                } else {
187                    b'N'
188                }
189            }
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 msg {
199            BackendMessage::AuthenticationCleartextPassword => {
200                dst.put_u32(3);
201            }
202            BackendMessage::ErrorResponse(ErrorResponse {
203                severity,
204                code,
205                message,
206                detail,
207                hint,
208                position,
209            }) => {
210                dst.put_u8(b'S');
211                dst.put_string(severity.as_str());
212                dst.put_u8(b'C');
213                dst.put_string(code.code());
214                dst.put_u8(b'M');
215                dst.put_string(&message);
216                if let Some(detail) = &detail {
217                    dst.put_u8(b'D');
218                    dst.put_string(detail);
219                }
220                if let Some(hint) = &hint {
221                    dst.put_u8(b'H');
222                    dst.put_string(hint);
223                }
224                if let Some(position) = &position {
225                    dst.put_u8(b'P');
226                    dst.put_string(&position.to_string());
227                }
228                dst.put_u8(b'\0');
229            }
230        }
231
232        let len = dst.len() - base;
233
234        // Overwrite length placeholder with true length.
235        let len = i32::try_from(len).map_err(|_| {
236            io::Error::new(
237                io::ErrorKind::InvalidData,
238                "length of encoded message does not fit into an i32",
239            )
240        })?;
241        dst[base..base + 4].copy_from_slice(&len.to_be_bytes());
242
243        Ok(())
244    }
245}
246
247impl Decoder for Codec {
248    type Item = FrontendMessage;
249    type Error = io::Error;
250
251    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<FrontendMessage>, io::Error> {
252        if src.len() > MAX_REQUEST_SIZE {
253            return Err(io::Error::new(
254                io::ErrorKind::InvalidData,
255                format!(
256                    "request larger than {}",
257                    ByteSize::b(u64::cast_from(MAX_REQUEST_SIZE))
258                ),
259            ));
260        }
261        loop {
262            match self.decode_state {
263                DecodeState::Head => {
264                    if src.len() < 5 {
265                        return Ok(None);
266                    }
267                    let msg_type = src[0];
268                    let frame_len = parse_frame_len(&src[1..], MAX_PREAUTH_FRAME_SIZE)?;
269                    src.advance(5);
270                    src.reserve(frame_len);
271                    self.decode_state = DecodeState::Data(msg_type, frame_len);
272                }
273
274                DecodeState::Data(msg_type, frame_len) => {
275                    if src.len() < frame_len {
276                        return Ok(None);
277                    }
278                    let buf = src.split_to(frame_len).freeze();
279                    let buf = Cursor::new(&buf);
280                    let msg = match msg_type {
281                        // Termination.
282                        b'X' => decode_terminate(buf)?,
283
284                        // Authentication.
285                        b'p' => decode_password(buf)?,
286
287                        // Invalid.
288                        _ => {
289                            return Err(io::Error::new(
290                                io::ErrorKind::InvalidData,
291                                format!("unknown message type {}", msg_type),
292                            ));
293                        }
294                    };
295                    src.reserve(5);
296                    self.decode_state = DecodeState::Head;
297                    return Ok(Some(msg));
298                }
299            }
300        }
301    }
302}
303
304fn decode_terminate(mut _buf: Cursor) -> Result<FrontendMessage, io::Error> {
305    // Nothing more to decode.
306    Ok(FrontendMessage::Terminate)
307}
308
309fn decode_password(mut buf: Cursor) -> Result<FrontendMessage, io::Error> {
310    Ok(FrontendMessage::Password {
311        password: buf.read_cstr()?.to_owned(),
312    })
313}
mz_balancerd/codec.rs

mz_balancerd/
codec.rs
