mz_pgwire/
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::net::IpAddr;
18
19use async_trait::async_trait;
20use bytes::{Buf, BufMut, BytesMut};
21use futures::{SinkExt, TryStreamExt, sink};
22use itertools::Itertools;
23use mz_adapter_types::connection::ConnectionId;
24use mz_ore::future::OreSinkExt;
25use mz_ore::netio::AsyncReady;
26use mz_pgwire_common::{
27    ChannelBinding, Conn, Cursor, DecodeState, ErrorResponse, FrontendMessage, GS2Header, Pgbuf,
28    SASLClientFinalResponse, SASLInitialResponse, input_err, parse_frame_len,
29};
30use tokio::io::{self, AsyncRead, AsyncWrite, Interest, Ready};
31use tokio::time::{self, Duration};
32use tokio_util::codec::{Decoder, Encoder, Framed};
33use tracing::trace;
34
35use crate::message::{BackendMessage, BackendMessageKind, SASLServerFinalMessageKinds};
36
37/// A connection that manages the encoding and decoding of pgwire frames.
38pub struct FramedConn<A> {
39    conn_id: ConnectionId,
40    peer_addr: Option<IpAddr>,
41    inner: sink::Buffer<Framed<Conn<A>, Codec>, BackendMessage>,
42}
43
44impl<A> FramedConn<A>
45where
46    A: AsyncRead + AsyncWrite + Unpin,
47{
48    /// Constructs a new framed connection.
49    ///
50    /// The underlying connection, `inner`, is expected to be something like a
51    /// TCP stream. Anything that implements [`AsyncRead`] and [`AsyncWrite`]
52    /// will do.
53    ///
54    /// The supplied `conn_id` is used to identify the connection in logging
55    /// messages.
56    pub fn new(conn_id: ConnectionId, peer_addr: Option<IpAddr>, inner: Conn<A>) -> FramedConn<A> {
57        FramedConn {
58            conn_id,
59            peer_addr,
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        match &message {
80            Some(message) => trace!("cid={} recv_name={}", self.conn_id, message.name()),
81            None => trace!("cid={} recv=<eof>", self.conn_id),
82        }
83        Ok(message)
84    }
85
86    /// Encodes and sends one backend message to the client.
87    ///
88    /// Note that the connection is not flushed after calling this method. You
89    /// must call [`FramedConn::flush`] explicitly. Returns an error if the
90    /// underlying connection is broken.
91    ///
92    /// Please use `StateMachine::send` instead if calling from `StateMachine`,
93    /// as it applies session-based filters before calling this method.
94    pub async fn send<M>(&mut self, message: M) -> Result<(), io::Error>
95    where
96        M: Into<BackendMessage>,
97    {
98        let message = message.into();
99        trace!(
100            "cid={} send={:?}",
101            self.conn_id,
102            BackendMessageKind::from(&message)
103        );
104        self.inner.enqueue(message).await
105    }
106
107    /// Encodes and sends the backend messages in the `messages` iterator to the
108    /// client.
109    ///
110    /// As with [`FramedConn::send`], the connection is not flushed after
111    /// calling this method. You must call [`FramedConn::flush`] explicitly.
112    /// Returns an error if the underlying connection is broken.
113    pub async fn send_all(
114        &mut self,
115        messages: impl IntoIterator<Item = BackendMessage>,
116    ) -> Result<(), io::Error> {
117        // N.B. we intentionally don't use `self.conn.send_all` here to avoid
118        // flushing the sink unnecessarily.
119        for m in messages {
120            self.send(m).await?;
121        }
122        Ok(())
123    }
124
125    /// Flushes all outstanding messages.
126    pub async fn flush(&mut self) -> Result<(), io::Error> {
127        self.inner.flush().await
128    }
129
130    /// Injects state that affects how certain backend messages are encoded.
131    ///
132    /// Specifically, the encoding of `BackendMessage::DataRow` depends upon the
133    /// types of the datums in the row. To avoid including the same type
134    /// information in each message, we use this side channel to install the
135    /// type information in the codec before sending any data row messages. This
136    /// violates the abstraction boundary a bit but results in much better
137    /// performance.
138    pub fn set_encode_state(
139        &mut self,
140        encode_state: Vec<(mz_pgrepr::Type, mz_pgwire_common::Format)>,
141        text_settings: mz_pgrepr::TextEncodeSettings,
142    ) {
143        let codec = self.inner.get_mut().codec_mut();
144        codec.encode_state = encode_state;
145        codec.text_settings = text_settings;
146    }
147
148    /// Waits for the connection to be closed.
149    ///
150    /// Returns a "connection closed" error when the connection is closed. If
151    /// another error occurs before the connection is closed, that error is
152    /// returned instead.
153    ///
154    /// Use this method when you have an unbounded stream of data to forward to
155    /// the connection and the protocol does not require the client to
156    /// periodically acknowledge receipt. If you don't call this method to
157    /// periodically check if the connection has closed, you may not notice that
158    /// the client has gone away for an unboundedly long amount of time; usually
159    /// not until the stream of data produces its next message and you attempt
160    /// to write the data to the connection.
161    pub async fn wait_closed(&self) -> io::Error
162    where
163        A: AsyncReady + Send + Sync,
164    {
165        loop {
166            time::sleep(Duration::from_secs(1)).await;
167
168            match self.ready(Interest::READABLE | Interest::WRITABLE).await {
169                Ok(ready) if ready.is_read_closed() || ready.is_write_closed() => {
170                    return io::Error::new(io::ErrorKind::Other, "connection closed");
171                }
172                Ok(_) => (),
173                Err(err) => return err,
174            }
175        }
176    }
177
178    /// Returns the ID associated with this connection.
179    pub fn conn_id(&self) -> &ConnectionId {
180        &self.conn_id
181    }
182
183    /// Returns the peer address of the connection.
184    pub fn peer_addr(&self) -> &Option<IpAddr> {
185        &self.peer_addr
186    }
187}
188
189impl<A> FramedConn<A>
190where
191    A: AsyncRead + AsyncWrite + Unpin,
192{
193    pub fn inner(&self) -> &Conn<A> {
194        self.inner.get_ref().get_ref()
195    }
196}
197
198#[async_trait]
199impl<A> AsyncReady for FramedConn<A>
200where
201    A: AsyncRead + AsyncWrite + AsyncReady + Send + Sync + Unpin,
202{
203    async fn ready(&self, interest: Interest) -> io::Result<Ready> {
204        self.inner.get_ref().get_ref().ready(interest).await
205    }
206}
207
208pub struct Codec {
209    decode_state: DecodeState,
210    encode_state: Vec<(mz_pgrepr::Type, mz_pgwire_common::Format)>,
211    /// The session's text encoding settings when `encode_state` was installed.
212    text_settings: mz_pgrepr::TextEncodeSettings,
213}
214
215impl Codec {
216    /// Creates a new `Codec`.
217    pub fn new() -> Codec {
218        Codec {
219            decode_state: DecodeState::Head,
220            encode_state: vec![],
221            text_settings: mz_pgrepr::TextEncodeSettings::STABLE,
222        }
223    }
224}
225
226impl Default for Codec {
227    fn default() -> Codec {
228        Codec::new()
229    }
230}
231
232impl Encoder<BackendMessage> for Codec {
233    type Error = io::Error;
234
235    /// Encode a backend message into `dst`.
236    /// If this function returns an error result, `dst` is left unmodified.
237    fn encode(&mut self, msg: BackendMessage, dst: &mut BytesMut) -> Result<(), io::Error> {
238        // Record the starting position so we can truncate on error.
239        // This prevents partial messages from being left in the buffer,
240        // which could be sent to the client and cause "lost synchronization" errors.
241        let start = dst.len();
242        match self.encode_inner(msg, dst) {
243            Ok(()) => Ok(()),
244            Err(e) => {
245                dst.truncate(start);
246                Err(e)
247            }
248        }
249    }
250}
251
252impl Codec {
253    /// This is the meat of the encoding logic. It's a separate function so that errors returned by
254    /// `?` can be handled in the outer `encode` function.
255    fn encode_inner(&self, msg: BackendMessage, dst: &mut BytesMut) -> Result<(), io::Error> {
256        // Write type byte.
257        let byte = match &msg {
258            BackendMessage::AuthenticationOk => b'R',
259            BackendMessage::AuthenticationCleartextPassword
260            | BackendMessage::AuthenticationSASL
261            | BackendMessage::AuthenticationSASLContinue(_)
262            | BackendMessage::AuthenticationSASLFinal(_) => b'R',
263            BackendMessage::RowDescription(_) => b'T',
264            BackendMessage::DataRow(_) => b'D',
265            BackendMessage::CommandComplete { .. } => b'C',
266            BackendMessage::EmptyQueryResponse => b'I',
267            BackendMessage::ReadyForQuery(_) => b'Z',
268            BackendMessage::NoData => b'n',
269            BackendMessage::ParameterStatus(_, _) => b'S',
270            BackendMessage::PortalSuspended => b's',
271            BackendMessage::BackendKeyData { .. } => b'K',
272            BackendMessage::ParameterDescription(_) => b't',
273            BackendMessage::ParseComplete => b'1',
274            BackendMessage::BindComplete => b'2',
275            BackendMessage::CloseComplete => b'3',
276            BackendMessage::ErrorResponse(r) => {
277                if r.severity.is_error() {
278                    b'E'
279                } else {
280                    b'N'
281                }
282            }
283            BackendMessage::CopyInResponse { .. } => b'G',
284            BackendMessage::CopyOutResponse { .. } => b'H',
285            BackendMessage::CopyData(_) => b'd',
286            BackendMessage::CopyDone => b'c',
287        };
288        dst.put_u8(byte);
289
290        // Write message length placeholder. The true length is filled in later.
291        let base = dst.len();
292        dst.put_u32(0);
293
294        // Write message contents.
295        match msg {
296            BackendMessage::CopyInResponse {
297                overall_format,
298                column_formats,
299            }
300            | BackendMessage::CopyOutResponse {
301                overall_format,
302                column_formats,
303            } => {
304                dst.put_format_i8(overall_format);
305                if column_formats.len() > usize::try_from(i16::MAX).expect("i16::MAX is positive") {
306                    return Err(io::Error::new(
307                        io::ErrorKind::InvalidData,
308                        format!(
309                            "{} columns in COPY response, which exceeds {}",
310                            column_formats.len(),
311                            i16::MAX
312                        ),
313                    ));
314                }
315                dst.put_length_i16(column_formats.len())?;
316                for format in column_formats {
317                    dst.put_format_i16(format);
318                }
319            }
320            BackendMessage::CopyData(data) => {
321                dst.put_slice(&data);
322            }
323            BackendMessage::CopyDone => (),
324            BackendMessage::AuthenticationOk => {
325                dst.put_u32(0);
326            }
327            BackendMessage::AuthenticationCleartextPassword => {
328                dst.put_u32(3);
329            }
330            BackendMessage::AuthenticationSASL => {
331                dst.put_u32(10);
332                dst.put_string("SCRAM-SHA-256");
333                dst.put_u8(b'\0');
334            }
335            BackendMessage::AuthenticationSASLContinue(data) => {
336                dst.put_u32(11);
337                let data = format!(
338                    "r={},s={},i={}",
339                    data.nonce, data.salt, data.iteration_count
340                );
341                dst.put_slice(data.as_bytes());
342            }
343            BackendMessage::AuthenticationSASLFinal(data) => {
344                dst.put_u32(12);
345                let res = match data.kind {
346                    SASLServerFinalMessageKinds::Verifier(verifier) => {
347                        format!("v={}", verifier)
348                    }
349                };
350                dst.put_slice(res.as_bytes());
351                if !data.extensions.is_empty() {
352                    dst.put_slice(b",");
353                    dst.put_slice(data.extensions.join(",").as_bytes());
354                }
355            }
356            BackendMessage::RowDescription(fields) => {
357                if fields.len() > usize::try_from(i16::MAX).expect("i16::MAX is positive") {
358                    return Err(io::Error::new(
359                        io::ErrorKind::InvalidData,
360                        format!(
361                            "{} fields in row description, which exceeds {}",
362                            fields.len(),
363                            i16::MAX
364                        ),
365                    ));
366                }
367                dst.put_length_i16(fields.len())?;
368                for f in &fields {
369                    dst.put_string(&f.name.to_string());
370                    dst.put_u32(f.table_id);
371                    dst.put_u16(f.column_id);
372                    dst.put_u32(f.type_oid);
373                    dst.put_i16(f.type_len);
374                    dst.put_i32(f.type_mod);
375                    // TODO: make the format correct
376                    dst.put_format_i16(f.format);
377                }
378            }
379            BackendMessage::DataRow(fields) => {
380                if fields.len() > usize::try_from(i16::MAX).expect("i16::MAX is positive") {
381                    return Err(io::Error::new(
382                        io::ErrorKind::InvalidData,
383                        format!(
384                            "{} fields in data row, which exceeds {}",
385                            fields.len(),
386                            i16::MAX
387                        ),
388                    ));
389                }
390                dst.put_length_i16(fields.len())?;
391                for (f, (ty, format)) in fields.iter().zip_eq(&self.encode_state) {
392                    if let Some(f) = f {
393                        let base = dst.len();
394                        dst.put_u32(0);
395                        f.encode(ty, *format, dst, self.text_settings)?;
396                        let len = dst.len() - base - 4;
397                        let len = i32::try_from(len).map_err(|_| {
398                            io::Error::new(
399                                io::ErrorKind::InvalidData,
400                                "length of encoded data row field does not fit into an i32",
401                            )
402                        })?;
403                        dst[base..base + 4].copy_from_slice(&len.to_be_bytes());
404                    } else {
405                        dst.put_i32(-1);
406                    }
407                }
408            }
409            BackendMessage::CommandComplete { tag } => {
410                dst.put_string(&tag);
411            }
412            BackendMessage::ParseComplete => (),
413            BackendMessage::BindComplete => (),
414            BackendMessage::CloseComplete => (),
415            BackendMessage::EmptyQueryResponse => (),
416            BackendMessage::ReadyForQuery(status) => {
417                dst.put_u8(status.into());
418            }
419            BackendMessage::ParameterStatus(name, value) => {
420                dst.put_string(name);
421                dst.put_string(&value);
422            }
423            BackendMessage::PortalSuspended => (),
424            BackendMessage::NoData => (),
425            BackendMessage::BackendKeyData {
426                conn_id,
427                secret_key,
428            } => {
429                dst.put_u32(conn_id);
430                dst.put_u32(secret_key);
431            }
432            BackendMessage::ParameterDescription(params) => {
433                if params.len() > usize::try_from(i16::MAX).expect("i16::MAX is positive") {
434                    return Err(io::Error::new(
435                        io::ErrorKind::InvalidData,
436                        format!(
437                            "{} params in parameter description, which exceeds {}",
438                            params.len(),
439                            i16::MAX
440                        ),
441                    ));
442                }
443                dst.put_length_i16(params.len())?;
444                for param in params {
445                    dst.put_u32(param.oid());
446                }
447            }
448            BackendMessage::ErrorResponse(ErrorResponse {
449                severity,
450                code,
451                message,
452                detail,
453                hint,
454                position,
455            }) => {
456                dst.put_u8(b'S');
457                dst.put_string(severity.as_str());
458                dst.put_u8(b'C');
459                dst.put_string(code.code());
460                dst.put_u8(b'M');
461                dst.put_string(&message);
462                if let Some(detail) = &detail {
463                    dst.put_u8(b'D');
464                    dst.put_string(detail);
465                }
466                if let Some(hint) = &hint {
467                    dst.put_u8(b'H');
468                    dst.put_string(hint);
469                }
470                if let Some(position) = &position {
471                    dst.put_u8(b'P');
472                    dst.put_string(&position.to_string());
473                }
474                dst.put_u8(b'\0');
475            }
476        }
477
478        let len = dst.len() - base;
479
480        // Overwrite length placeholder with true length.
481        let len = i32::try_from(len).map_err(|_| {
482            io::Error::new(
483                io::ErrorKind::InvalidData,
484                "length of encoded message does not fit into an i32",
485            )
486        })?;
487        dst[base..base + 4].copy_from_slice(&len.to_be_bytes());
488
489        Ok(())
490    }
491}
492
493impl Decoder for Codec {
494    type Item = FrontendMessage;
495    type Error = io::Error;
496
497    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<FrontendMessage>, io::Error> {
498        loop {
499            match self.decode_state {
500                DecodeState::Head => {
501                    if src.len() < 5 {
502                        return Ok(None);
503                    }
504                    let msg_type = src[0];
505                    let frame_len = parse_frame_len(&src[1..])?;
506                    src.advance(5);
507                    src.reserve(frame_len);
508                    self.decode_state = DecodeState::Data(msg_type, frame_len);
509                }
510
511                DecodeState::Data(msg_type, frame_len) => {
512                    if src.len() < frame_len {
513                        return Ok(None);
514                    }
515                    let buf = src.split_to(frame_len).freeze();
516                    let buf = Cursor::new(&buf);
517                    let msg = match msg_type {
518                        // Simple query flow.
519                        b'Q' => decode_query(buf)?,
520
521                        // Extended query flow.
522                        b'P' => decode_parse(buf)?,
523                        b'D' => decode_describe(buf)?,
524                        b'B' => decode_bind(buf)?,
525                        b'E' => decode_execute(buf)?,
526                        b'H' => decode_flush(buf)?,
527                        b'S' => decode_sync(buf)?,
528                        b'C' => decode_close(buf)?,
529
530                        // Termination.
531                        b'X' => decode_terminate(buf)?,
532
533                        // Authentication.
534                        b'p' => decode_auth(buf)?,
535
536                        // Copy from flow.
537                        b'f' => decode_copy_fail(buf)?,
538                        b'd' => decode_copy_data(buf, frame_len)?,
539                        b'c' => decode_copy_done(buf)?,
540
541                        // Invalid.
542                        _ => {
543                            return Err(io::Error::new(
544                                io::ErrorKind::InvalidData,
545                                format!("unknown message type {}", msg_type),
546                            ));
547                        }
548                    };
549                    src.reserve(5);
550                    self.decode_state = DecodeState::Head;
551                    return Ok(Some(msg));
552                }
553            }
554        }
555    }
556}
557
558fn decode_terminate(mut _buf: Cursor) -> Result<FrontendMessage, io::Error> {
559    // Nothing more to decode.
560    Ok(FrontendMessage::Terminate)
561}
562
563fn decode_auth(mut buf: Cursor) -> Result<FrontendMessage, io::Error> {
564    let mut value = Vec::new();
565    while let Ok(b) = buf.read_byte() {
566        value.push(b);
567    }
568    Ok(FrontendMessage::RawAuthentication(value))
569}
570
571fn expect(buf: &mut Cursor, expected: &[u8]) -> Result<(), io::Error> {
572    for i in 0..expected.len() {
573        if buf.read_byte()? != expected[i] {
574            return Err(input_err(format!(
575                "Invalid SASL initial response: expected '{}'",
576                std::str::from_utf8(expected).unwrap_or("invalid UTF-8")
577            )));
578        }
579    }
580    Ok(())
581}
582
583fn read_until_comma(buf: &mut Cursor) -> Result<Vec<u8>, io::Error> {
584    let mut v = Vec::new();
585    while let Ok(b) = buf.peek_byte() {
586        if b == b',' {
587            break;
588        }
589        v.push(buf.read_byte()?);
590    }
591    Ok(v)
592}
593
594// All SASL parsing is based on RFC 5802, [section 7](https://datatracker.ietf.org/doc/html/rfc5802#section-7)
595
596//   extensions = attr-val *("," attr-val)
597//                     ;; All extensions are optional,
598//                     ;; i.e., unrecognized attributes
599//                     ;; not defined in this document
600//                     ;; MUST be ignored.
601//   reserved-mext  = "m=" 1*(value-char)
602//                     ;; Reserved for signaling mandatory extensions.
603//                     ;; The exact syntax will be defined in
604//                     ;; the future.
605//   gs2-cbind-flag  = ("p=" cb-name) / "n" / "y"
606//                     ;; "n" -> client doesn't support channel binding.
607//                     ;; "y" -> client does support channel binding
608//                     ;;        but thinks the server does not.
609//                     ;; "p" -> client requires channel binding.
610//                     ;; The selected channel binding follows "p=".
611//
612//   gs2-header      = gs2-cbind-flag "," [ authzid ] ","
613//                     ;; GS2 header for SCRAM
614//                     ;; (the actual GS2 header includes an optional
615//                     ;; flag to indicate that the GSS mechanism is not
616//                     ;; "standard", but since SCRAM is "standard", we
617//                     ;; don't include that flag).
618//   client-first-message-bare =
619//                     [reserved-mext ","]
620//                     username "," nonce ["," extensions]
621//
622//   client-first-message =
623//                     gs2-header client-first-message-bare
624pub fn decode_sasl_client_first_message(mut buf: Cursor) -> Result<SASLInitialResponse, io::Error> {
625    // 1) GS2 cbind flag
626    let cbind_flag = match buf.read_byte()? {
627        b'n' => ChannelBinding::None,
628        b'y' => ChannelBinding::ClientSupported,
629        b'p' => {
630            // must be "p=" then cbname up to next comma
631            expect(&mut buf, b"=")?;
632            let cbname = String::from_utf8(read_until_comma(&mut buf)?)
633                .map_err(|_| input_err("invalid cbname utf8"))?;
634            ChannelBinding::Required(cbname)
635        }
636        other => {
637            return Err(input_err(format!(
638                "Invalid channel binding flag: {}",
639                other
640            )));
641        }
642    };
643    expect(&mut buf, b",")?;
644
645    // 2) Optional authzid: either empty, or "a=" up to next comma
646    let mut authzid = None;
647    if buf.peek_byte()? == b'a' {
648        expect(&mut buf, b"a=")?;
649        let a = String::from_utf8(read_until_comma(&mut buf)?)
650            .map_err(|_| input_err("invalid authzid utf8"))?;
651        authzid = Some(a);
652    }
653    expect(&mut buf, b",")?;
654
655    let mut client_first_message_bare_raw = String::new();
656
657    // 3) Optional reserved "m=" extension before n=
658    let mut reserved_mext = None;
659    if buf.peek_byte()? == b'm' {
660        expect(&mut buf, b"m=")?;
661        let mext_val = String::from_utf8(read_until_comma(&mut buf)?)
662            .map_err(|_| input_err("invalid m ext utf8"))?;
663        client_first_message_bare_raw.push_str(&format!("m={},", mext_val));
664        reserved_mext = Some(mext_val);
665        expect(&mut buf, b",")?;
666    }
667
668    // 4) Username: must be "n=" then saslname
669    expect(&mut buf, b"n=")?;
670    // Postgres doesn't use the username here, so we just consume
671    let username = String::from_utf8(read_until_comma(&mut buf)?)
672        .map_err(|_| input_err("invalid username utf8"))?;
673    expect(&mut buf, b",")?;
674    client_first_message_bare_raw.push_str(&format!("n={},", username));
675
676    // 5) Nonce: must be "r=" then value up to next comma or end
677    expect(&mut buf, b"r=")?;
678    let nonce = String::from_utf8(read_until_comma(&mut buf)?)
679        .map_err(|_| input_err("invalid nonce utf8"))?;
680    client_first_message_bare_raw.push_str(&format!("r={}", nonce));
681
682    // 6) Optional extensions: "," key=value chunks
683    let mut extensions = Vec::new();
684    while let Ok(b',') = buf.peek_byte().map(|b| b) {
685        expect(&mut buf, b",")?;
686        let ext = String::from_utf8(read_until_comma(&mut buf)?)
687            .map_err(|_| input_err("invalid ext utf8"))?;
688        if !ext.is_empty() {
689            client_first_message_bare_raw.push_str(&format!(",{}", ext));
690            extensions.push(ext);
691        }
692    }
693
694    Ok(SASLInitialResponse {
695        gs2_header: GS2Header {
696            cbind_flag,
697            authzid,
698        },
699        nonce,
700        extensions,
701        reserved_mext,
702        client_first_message_bare_raw,
703    })
704}
705
706pub fn decode_sasl_initial_response(mut buf: Cursor) -> Result<FrontendMessage, io::Error> {
707    let mechanism = buf.read_cstr()?;
708    let initial_resp_len = buf.read_i32()?;
709    if initial_resp_len < 0 {
710        // -1 means no response? We bail here
711        return Err(input_err("No initial response"));
712    }
713
714    let initial_response = decode_sasl_client_first_message(buf)?;
715    Ok(FrontendMessage::SASLInitialResponse {
716        gs2_header: initial_response.gs2_header.clone(),
717        mechanism: mechanism.to_owned(),
718        initial_response,
719    })
720}
721
722//   proof           = "p=" base64
723//
724//   channel-binding = "c=" base64
725//                     ;; base64 encoding of cbind-input.
726//   client-final-message-without-proof =
727//                     channel-binding "," nonce [","
728//                     extensions]
729//
730//   client-final-message =
731//                     client-final-message-without-proof "," proof
732pub fn decode_sasl_response(mut buf: Cursor) -> Result<FrontendMessage, io::Error> {
733    // --- client-final-message-without-proof ---
734    let mut client_final_message_bare_raw = String::new();
735    // channel-binding: "c=" <base64>, up to the next comma
736    expect(&mut buf, b"c=")?;
737    let channel_binding = String::from_utf8(read_until_comma(&mut buf)?)
738        .map_err(|_| input_err("invalid channel-binding utf8"))?;
739    expect(&mut buf, b",")?;
740    client_final_message_bare_raw.push_str(&format!("c={},", channel_binding));
741
742    // nonce: "r=" <printable>, up to the next comma
743    expect(&mut buf, b"r=")?;
744    let nonce = String::from_utf8(read_until_comma(&mut buf)?)
745        .map_err(|_| input_err("invalid nonce utf8"))?;
746    client_final_message_bare_raw.push_str(&format!("r={}", nonce));
747
748    // after reading channel-binding and nonce
749    let mut extensions = Vec::new();
750
751    // Keep reading ",<token>" until we see ",p="
752    while buf.peek_byte()? == b',' {
753        expect(&mut buf, b",")?;
754        if buf.peek_byte()? == b'p' {
755            break;
756        }
757        let ext = String::from_utf8(read_until_comma(&mut buf)?)
758            .map_err(|_| input_err("invalid extension utf8"))?;
759        if !ext.is_empty() {
760            client_final_message_bare_raw.push_str(&format!(",{}", ext));
761            extensions.push(ext);
762        }
763    }
764
765    // Proof is mandatory and last
766    expect(&mut buf, b"p=")?;
767    let proof = String::from_utf8(read_until_comma(&mut buf)?)
768        .map_err(|_| input_err("invalid proof utf8"))?;
769
770    Ok(FrontendMessage::SASLResponse(SASLClientFinalResponse {
771        channel_binding,
772        nonce,
773        extensions,
774        proof,
775        client_final_message_bare_raw,
776    }))
777}
778
779pub fn decode_password(mut buf: Cursor) -> Result<FrontendMessage, io::Error> {
780    Ok(FrontendMessage::Password {
781        password: buf.read_cstr()?.to_owned(),
782    })
783}
784
785fn decode_query(mut buf: Cursor) -> Result<FrontendMessage, io::Error> {
786    Ok(FrontendMessage::Query {
787        sql: buf.read_cstr()?.to_string(),
788    })
789}
790
791fn decode_parse(mut buf: Cursor) -> Result<FrontendMessage, io::Error> {
792    let name = buf.read_cstr()?;
793    let sql = buf.read_cstr()?;
794
795    let mut param_types = vec![];
796    for _ in 0..buf.read_i16()? {
797        param_types.push(buf.read_u32()?);
798    }
799
800    Ok(FrontendMessage::Parse {
801        name: name.into(),
802        sql: sql.into(),
803        param_types,
804    })
805}
806
807fn decode_close(mut buf: Cursor) -> Result<FrontendMessage, io::Error> {
808    match buf.read_byte()? {
809        b'S' => Ok(FrontendMessage::CloseStatement {
810            name: buf.read_cstr()?.to_owned(),
811        }),
812        b'P' => Ok(FrontendMessage::ClosePortal {
813            name: buf.read_cstr()?.to_owned(),
814        }),
815        b => Err(input_err(format!(
816            "invalid type byte in close message: {}",
817            b
818        ))),
819    }
820}
821
822fn decode_describe(mut buf: Cursor) -> Result<FrontendMessage, io::Error> {
823    let first_char = buf.read_byte()?;
824    let name = buf.read_cstr()?.to_string();
825    match first_char {
826        b'S' => Ok(FrontendMessage::DescribeStatement { name }),
827        b'P' => Ok(FrontendMessage::DescribePortal { name }),
828        other => Err(input_err(format!("Invalid describe type: {:#x?}", other))),
829    }
830}
831
832fn decode_bind(mut buf: Cursor) -> Result<FrontendMessage, io::Error> {
833    let portal_name = buf.read_cstr()?.to_string();
834    let statement_name = buf.read_cstr()?.to_string();
835
836    let mut param_formats = Vec::new();
837    for _ in 0..buf.read_i16()? {
838        param_formats.push(buf.read_format()?);
839    }
840
841    let mut raw_params = Vec::new();
842    for _ in 0..buf.read_i16()? {
843        let len = buf.read_i32()?;
844        if len == -1 {
845            raw_params.push(None); // NULL
846        } else {
847            // TODO(benesch): this should use bytes::Bytes to avoid the copy.
848            let mut value = Vec::new();
849            for _ in 0..len {
850                value.push(buf.read_byte()?);
851            }
852            raw_params.push(Some(value));
853        }
854    }
855
856    let mut result_formats = Vec::new();
857    for _ in 0..buf.read_i16()? {
858        result_formats.push(buf.read_format()?);
859    }
860
861    Ok(FrontendMessage::Bind {
862        portal_name,
863        statement_name,
864        param_formats,
865        raw_params,
866        result_formats,
867    })
868}
869
870fn decode_execute(mut buf: Cursor) -> Result<FrontendMessage, io::Error> {
871    let portal_name = buf.read_cstr()?.to_string();
872    let max_rows = buf.read_i32()?;
873    Ok(FrontendMessage::Execute {
874        portal_name,
875        max_rows,
876    })
877}
878
879fn decode_flush(mut _buf: Cursor) -> Result<FrontendMessage, io::Error> {
880    // Nothing more to decode.
881    Ok(FrontendMessage::Flush)
882}
883
884fn decode_sync(mut _buf: Cursor) -> Result<FrontendMessage, io::Error> {
885    // Nothing more to decode.
886    Ok(FrontendMessage::Sync)
887}
888
889fn decode_copy_data(mut buf: Cursor, frame_len: usize) -> Result<FrontendMessage, io::Error> {
890    let mut data = Vec::with_capacity(frame_len);
891    for _ in 0..frame_len {
892        data.push(buf.read_byte()?);
893    }
894    Ok(FrontendMessage::CopyData(data))
895}
896
897fn decode_copy_done(mut _buf: Cursor) -> Result<FrontendMessage, io::Error> {
898    // Nothing more to decode.
899    Ok(FrontendMessage::CopyDone)
900}
901
902fn decode_copy_fail(mut buf: Cursor) -> Result<FrontendMessage, io::Error> {
903    Ok(FrontendMessage::CopyFail(buf.read_cstr()?.to_string()))
904}
mz_pgwire/codec.rs

mz_pgwire/
codec.rs
