LCOV - code coverage report
Current view: top level - libs/pq_proto/src - lib.rs (source / functions) Coverage Total Hit
Test: c639aa5f7ab62b43d647b10f40d15a15686ce8a9.info Lines: 86.2 % 589 508
Test Date: 2024-02-12 20:26:03 Functions: 77.4 % 124 96

            Line data    Source code
       1              : //! Postgres protocol messages serialization-deserialization. See
       2              : //! <https://www.postgresql.org/docs/devel/protocol-message-formats.html>
       3              : //! on message formats.
       4              : #![deny(clippy::undocumented_unsafe_blocks)]
       5              : 
       6              : pub mod framed;
       7              : 
       8              : use byteorder::{BigEndian, ReadBytesExt};
       9              : use bytes::{Buf, BufMut, Bytes, BytesMut};
      10              : use std::{borrow::Cow, collections::HashMap, fmt, io, str};
      11              : 
      12              : // re-export for use in utils pageserver_feedback.rs
      13              : pub use postgres_protocol::PG_EPOCH;
      14              : 
      15              : pub type Oid = u32;
      16              : pub type SystemId = u64;
      17              : 
      18              : pub const INT8_OID: Oid = 20;
      19              : pub const INT4_OID: Oid = 23;
      20              : pub const TEXT_OID: Oid = 25;
      21              : 
      22           30 : #[derive(Debug)]
      23              : pub enum FeMessage {
      24              :     // Simple query.
      25              :     Query(Bytes),
      26              :     // Extended query protocol.
      27              :     Parse(FeParseMessage),
      28              :     Describe(FeDescribeMessage),
      29              :     Bind(FeBindMessage),
      30              :     Execute(FeExecuteMessage),
      31              :     Close(FeCloseMessage),
      32              :     Sync,
      33              :     Terminate,
      34              :     CopyData(Bytes),
      35              :     CopyDone,
      36              :     CopyFail,
      37              :     PasswordMessage(Bytes),
      38              : }
      39              : 
      40          230 : #[derive(Debug)]
      41              : pub enum FeStartupPacket {
      42              :     CancelRequest(CancelKeyData),
      43              :     SslRequest,
      44              :     GssEncRequest,
      45              :     StartupMessage {
      46              :         major_version: u32,
      47              :         minor_version: u32,
      48              :         params: StartupMessageParams,
      49              :     },
      50              : }
      51              : 
      52          126 : #[derive(Debug)]
      53              : pub struct StartupMessageParams {
      54              :     params: HashMap<String, String>,
      55              : }
      56              : 
      57              : impl StartupMessageParams {
      58              :     /// Get parameter's value by its name.
      59         7196 :     pub fn get(&self, name: &str) -> Option<&str> {
      60         7196 :         self.params.get(name).map(|s| s.as_str())
      61         7196 :     }
      62              : 
      63              :     /// Split command-line options according to PostgreSQL's logic,
      64              :     /// taking into account all escape sequences but leaving them as-is.
      65              :     /// [`None`] means that there's no `options` in [`Self`].
      66         3487 :     pub fn options_raw(&self) -> Option<impl Iterator<Item = &str>> {
      67         3487 :         self.get("options").map(Self::parse_options_raw)
      68         3487 :     }
      69              : 
      70              :     /// Split command-line options according to PostgreSQL's logic,
      71              :     /// applying all escape sequences (using owned strings as needed).
      72              :     /// [`None`] means that there's no `options` in [`Self`].
      73           10 :     pub fn options_escaped(&self) -> Option<impl Iterator<Item = Cow<'_, str>>> {
      74           10 :         self.get("options").map(Self::parse_options_escaped)
      75           10 :     }
      76              : 
      77              :     /// Split command-line options according to PostgreSQL's logic,
      78              :     /// taking into account all escape sequences but leaving them as-is.
      79         3487 :     pub fn parse_options_raw(input: &str) -> impl Iterator<Item = &str> {
      80         3487 :         // See `postgres: pg_split_opts`.
      81         3487 :         let mut last_was_escape = false;
      82         3487 :         input
      83       302246 :             .split(move |c: char| {
      84              :                 // We split by non-escaped whitespace symbols.
      85       302246 :                 let should_split = c.is_ascii_whitespace() && !last_was_escape;
      86       302246 :                 last_was_escape = c == '\\' && !last_was_escape;
      87       302246 :                 should_split
      88       302246 :             })
      89        10321 :             .filter(|s| !s.is_empty())
      90         3487 :     }
      91              : 
      92              :     /// Split command-line options according to PostgreSQL's logic,
      93              :     /// applying all escape sequences (using owned strings as needed).
      94            8 :     pub fn parse_options_escaped(input: &str) -> impl Iterator<Item = Cow<'_, str>> {
      95            8 :         // See `postgres: pg_split_opts`.
      96           14 :         Self::parse_options_raw(input).map(|s| {
      97           14 :             let mut preserve_next_escape = false;
      98           34 :             let escape = |c| {
      99              :                 // We should remove '\\' unless it's preceded by '\\'.
     100           34 :                 let should_remove = c == '\\' && !preserve_next_escape;
     101           34 :                 preserve_next_escape = should_remove;
     102           34 :                 should_remove
     103           34 :             };
     104              : 
     105           14 :             match s.contains('\\') {
     106            6 :                 true => Cow::Owned(s.replace(escape, "")),
     107            8 :                 false => Cow::Borrowed(s),
     108              :             }
     109           14 :         })
     110            8 :     }
     111              : 
     112              :     /// Iterate through key-value pairs in an arbitrary order.
     113           14 :     pub fn iter(&self) -> impl Iterator<Item = (&str, &str)> {
     114           42 :         self.params.iter().map(|(k, v)| (k.as_str(), v.as_str()))
     115           14 :     }
     116              : 
     117              :     // This function is mostly useful in tests.
     118              :     #[doc(hidden)]
     119           46 :     pub fn new<'a, const N: usize>(pairs: [(&'a str, &'a str); N]) -> Self {
     120           46 :         Self {
     121           62 :             params: pairs.map(|(k, v)| (k.to_owned(), v.to_owned())).into(),
     122           46 :         }
     123           46 :     }
     124              : }
     125              : 
     126          301 : #[derive(Debug, Hash, PartialEq, Eq, Clone, Copy)]
     127              : pub struct CancelKeyData {
     128              :     pub backend_pid: i32,
     129              :     pub cancel_key: i32,
     130              : }
     131              : 
     132              : impl fmt::Display for CancelKeyData {
     133          164 :     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
     134          164 :         let hi = (self.backend_pid as u64) << 32;
     135          164 :         let lo = self.cancel_key as u64;
     136          164 :         let id = hi | lo;
     137          164 : 
     138          164 :         // This format is more compact and might work better for logs.
     139          164 :         f.debug_tuple("CancelKeyData")
     140          164 :             .field(&format_args!("{:x}", id))
     141          164 :             .finish()
     142          164 :     }
     143              : }
     144              : 
     145              : use rand::distributions::{Distribution, Standard};
     146              : impl Distribution<CancelKeyData> for Standard {
     147           43 :     fn sample<R: rand::Rng + ?Sized>(&self, rng: &mut R) -> CancelKeyData {
     148           43 :         CancelKeyData {
     149           43 :             backend_pid: rng.gen(),
     150           43 :             cancel_key: rng.gen(),
     151           43 :         }
     152           43 :     }
     153              : }
     154              : 
     155              : // We only support the simple case of Parse on unnamed prepared statement and
     156              : // no params
     157            0 : #[derive(Debug)]
     158              : pub struct FeParseMessage {
     159              :     pub query_string: Bytes,
     160              : }
     161              : 
     162            0 : #[derive(Debug)]
     163              : pub struct FeDescribeMessage {
     164              :     pub kind: u8, // 'S' to describe a prepared statement; or 'P' to describe a portal.
     165              :                   // we only support unnamed prepared stmt or portal
     166              : }
     167              : 
     168              : // we only support unnamed prepared stmt and portal
     169            0 : #[derive(Debug)]
     170              : pub struct FeBindMessage;
     171              : 
     172              : // we only support unnamed prepared stmt or portal
     173            0 : #[derive(Debug)]
     174              : pub struct FeExecuteMessage {
     175              :     /// max # of rows
     176              :     pub maxrows: i32,
     177              : }
     178              : 
     179              : // we only support unnamed prepared stmt and portal
     180            0 : #[derive(Debug)]
     181              : pub struct FeCloseMessage;
     182              : 
     183              : /// An error occurred while parsing or serializing raw stream into Postgres
     184              : /// messages.
     185            6 : #[derive(thiserror::Error, Debug)]
     186              : pub enum ProtocolError {
     187              :     /// Invalid packet was received from the client (e.g. unexpected message
     188              :     /// type or broken len).
     189              :     #[error("Protocol error: {0}")]
     190              :     Protocol(String),
     191              :     /// Failed to parse or, (unlikely), serialize a protocol message.
     192              :     #[error("Message parse error: {0}")]
     193              :     BadMessage(String),
     194              : }
     195              : 
     196              : impl ProtocolError {
     197              :     /// Proxy stream.rs uses only io::Error; provide it.
     198            0 :     pub fn into_io_error(self) -> io::Error {
     199            0 :         io::Error::new(io::ErrorKind::Other, self.to_string())
     200            0 :     }
     201              : }
     202              : 
     203              : impl FeMessage {
     204              :     /// Read and parse one message from the `buf` input buffer. If there is at
     205              :     /// least one valid message, returns it, advancing `buf`; redundant copies
     206              :     /// are avoided, as thanks to `bytes` crate ptrs in parsed message point
     207              :     /// directly into the `buf` (processed data is garbage collected after
     208              :     /// parsed message is dropped).
     209              :     ///
     210              :     /// Returns None if `buf` doesn't contain enough data for a single message.
     211              :     /// For efficiency, tries to reserve large enough space in `buf` for the
     212              :     /// next message in this case to save the repeated calls.
     213              :     ///
     214              :     /// Returns Error if message is malformed, the only possible ErrorKind is
     215              :     /// InvalidInput.
     216              :     //
     217              :     // Inspired by rust-postgres Message::parse.
     218     15313487 :     pub fn parse(buf: &mut BytesMut) -> Result<Option<FeMessage>, ProtocolError> {
     219     15313487 :         // Every message contains message type byte and 4 bytes len; can't do
     220     15313487 :         // much without them.
     221     15313487 :         if buf.len() < 5 {
     222      6972091 :             let to_read = 5 - buf.len();
     223      6972091 :             buf.reserve(to_read);
     224      6972091 :             return Ok(None);
     225      8341396 :         }
     226      8341396 : 
     227      8341396 :         // We shouldn't advance `buf` as probably full message is not there yet,
     228      8341396 :         // so can't directly use Bytes::get_u32 etc.
     229      8341396 :         let tag = buf[0];
     230      8341396 :         let len = (&buf[1..5]).read_u32::<BigEndian>().unwrap();
     231      8341396 :         if len < 4 {
     232            0 :             return Err(ProtocolError::Protocol(format!(
     233            0 :                 "invalid message length {}",
     234            0 :                 len
     235            0 :             )));
     236      8341396 :         }
     237      8341396 : 
     238      8341396 :         // length field includes itself, but not message type.
     239      8341396 :         let total_len = len as usize + 1;
     240      8341396 :         if buf.len() < total_len {
     241              :             // Don't have full message yet.
     242       175935 :             let to_read = total_len - buf.len();
     243       175935 :             buf.reserve(to_read);
     244       175935 :             return Ok(None);
     245      8165461 :         }
     246      8165461 : 
     247      8165461 :         // got the message, advance buffer
     248      8165461 :         let mut msg = buf.split_to(total_len).freeze();
     249      8165461 :         msg.advance(5); // consume message type and len
     250      8165461 : 
     251      8165461 :         match tag {
     252        13996 :             b'Q' => Ok(Some(FeMessage::Query(msg))),
     253          630 :             b'P' => Ok(Some(FeParseMessage::parse(msg)?)),
     254          630 :             b'D' => Ok(Some(FeDescribeMessage::parse(msg)?)),
     255          630 :             b'E' => Ok(Some(FeExecuteMessage::parse(msg)?)),
     256          630 :             b'B' => Ok(Some(FeBindMessage::parse(msg)?)),
     257          629 :             b'C' => Ok(Some(FeCloseMessage::parse(msg)?)),
     258         1909 :             b'S' => Ok(Some(FeMessage::Sync)),
     259         2165 :             b'X' => Ok(Some(FeMessage::Terminate)),
     260      8143833 :             b'd' => Ok(Some(FeMessage::CopyData(msg))),
     261           13 :             b'c' => Ok(Some(FeMessage::CopyDone)),
     262           61 :             b'f' => Ok(Some(FeMessage::CopyFail)),
     263          335 :             b'p' => Ok(Some(FeMessage::PasswordMessage(msg))),
     264            0 :             tag => Err(ProtocolError::Protocol(format!(
     265            0 :                 "unknown message tag: {tag},'{msg:?}'"
     266            0 :             ))),
     267              :         }
     268     15313487 :     }
     269              : }
     270              : 
     271              : impl FeStartupPacket {
     272              :     /// Read and parse startup message from the `buf` input buffer. It is
     273              :     /// different from [`FeMessage::parse`] because startup messages don't have
     274              :     /// message type byte; otherwise, its comments apply.
     275        51848 :     pub fn parse(buf: &mut BytesMut) -> Result<Option<FeStartupPacket>, ProtocolError> {
     276        51848 :         const MAX_STARTUP_PACKET_LENGTH: usize = 10000;
     277        51848 :         const RESERVED_INVALID_MAJOR_VERSION: u32 = 1234;
     278        51848 :         const CANCEL_REQUEST_CODE: u32 = 5678;
     279        51848 :         const NEGOTIATE_SSL_CODE: u32 = 5679;
     280        51848 :         const NEGOTIATE_GSS_CODE: u32 = 5680;
     281        51848 : 
     282        51848 :         // need at least 4 bytes with packet len
     283        51848 :         if buf.len() < 4 {
     284        25924 :             let to_read = 4 - buf.len();
     285        25924 :             buf.reserve(to_read);
     286        25924 :             return Ok(None);
     287        25924 :         }
     288        25924 : 
     289        25924 :         // We shouldn't advance `buf` as probably full message is not there yet,
     290        25924 :         // so can't directly use Bytes::get_u32 etc.
     291        25924 :         let len = (&buf[0..4]).read_u32::<BigEndian>().unwrap() as usize;
     292        25924 :         // The proposed replacement is `!(8..=MAX_STARTUP_PACKET_LENGTH).contains(&len)`
     293        25924 :         // which is less readable
     294        25924 :         #[allow(clippy::manual_range_contains)]
     295        25924 :         if len < 8 || len > MAX_STARTUP_PACKET_LENGTH {
     296            2 :             return Err(ProtocolError::Protocol(format!(
     297            2 :                 "invalid startup packet message length {}",
     298            2 :                 len
     299            2 :             )));
     300        25922 :         }
     301        25922 : 
     302        25922 :         if buf.len() < len {
     303              :             // Don't have full message yet.
     304            0 :             let to_read = len - buf.len();
     305            0 :             buf.reserve(to_read);
     306            0 :             return Ok(None);
     307        25922 :         }
     308        25922 : 
     309        25922 :         // got the message, advance buffer
     310        25922 :         let mut msg = buf.split_to(len).freeze();
     311        25922 :         msg.advance(4); // consume len
     312        25922 : 
     313        25922 :         let request_code = msg.get_u32();
     314        25922 :         let req_hi = request_code >> 16;
     315        25922 :         let req_lo = request_code & ((1 << 16) - 1);
     316              :         // StartupMessage, CancelRequest, SSLRequest etc are differentiated by request code.
     317        25922 :         let message = match (req_hi, req_lo) {
     318              :             (RESERVED_INVALID_MAJOR_VERSION, CANCEL_REQUEST_CODE) => {
     319            0 :                 if msg.remaining() != 8 {
     320            0 :                     return Err(ProtocolError::BadMessage(
     321            0 :                         "CancelRequest message is malformed, backend PID / secret key missing"
     322            0 :                             .to_owned(),
     323            0 :                     ));
     324            0 :                 }
     325            0 :                 FeStartupPacket::CancelRequest(CancelKeyData {
     326            0 :                     backend_pid: msg.get_i32(),
     327            0 :                     cancel_key: msg.get_i32(),
     328            0 :                 })
     329              :             }
     330              :             (RESERVED_INVALID_MAJOR_VERSION, NEGOTIATE_SSL_CODE) => {
     331              :                 // Requested upgrade to SSL (aka TLS)
     332        11941 :                 FeStartupPacket::SslRequest
     333              :             }
     334              :             (RESERVED_INVALID_MAJOR_VERSION, NEGOTIATE_GSS_CODE) => {
     335              :                 // Requested upgrade to GSSAPI
     336            0 :                 FeStartupPacket::GssEncRequest
     337              :             }
     338            0 :             (RESERVED_INVALID_MAJOR_VERSION, unrecognized_code) => {
     339            0 :                 return Err(ProtocolError::Protocol(format!(
     340            0 :                     "Unrecognized request code {unrecognized_code}"
     341            0 :                 )));
     342              :             }
     343              :             // TODO bail if protocol major_version is not 3?
     344        13981 :             (major_version, minor_version) => {
     345              :                 // StartupMessage
     346              : 
     347              :                 // Parse pairs of null-terminated strings (key, value).
     348              :                 // See `postgres: ProcessStartupPacket, build_startup_packet`.
     349        13981 :                 let mut tokens = str::from_utf8(&msg)
     350        13981 :                     .map_err(|_e| {
     351            0 :                         ProtocolError::BadMessage("StartupMessage params: invalid utf-8".to_owned())
     352        13981 :                     })?
     353        13981 :                     .strip_suffix('\0') // drop packet's own null
     354        13981 :                     .ok_or_else(|| {
     355            0 :                         ProtocolError::Protocol(
     356            0 :                             "StartupMessage params: missing null terminator".to_string(),
     357            0 :                         )
     358        13981 :                     })?
     359        13981 :                     .split_terminator('\0');
     360        13981 : 
     361        13981 :                 let mut params = HashMap::new();
     362        46154 :                 while let Some(name) = tokens.next() {
     363        32173 :                     let value = tokens.next().ok_or_else(|| {
     364            0 :                         ProtocolError::Protocol(
     365            0 :                             "StartupMessage params: key without value".to_string(),
     366            0 :                         )
     367        32173 :                     })?;
     368              : 
     369        32173 :                     params.insert(name.to_owned(), value.to_owned());
     370              :                 }
     371              : 
     372        13981 :                 FeStartupPacket::StartupMessage {
     373        13981 :                     major_version,
     374        13981 :                     minor_version,
     375        13981 :                     params: StartupMessageParams { params },
     376        13981 :                 }
     377              :             }
     378              :         };
     379        25922 :         Ok(Some(message))
     380        51848 :     }
     381              : }
     382              : 
     383              : impl FeParseMessage {
     384          630 :     fn parse(mut buf: Bytes) -> Result<FeMessage, ProtocolError> {
     385              :         // FIXME: the rust-postgres driver uses a named prepared statement
     386              :         // for copy_out(). We're not prepared to handle that correctly. For
     387              :         // now, just ignore the statement name, assuming that the client never
     388              :         // uses more than one prepared statement at a time.
     389              : 
     390          630 :         let _pstmt_name = read_cstr(&mut buf)?;
     391          630 :         let query_string = read_cstr(&mut buf)?;
     392          630 :         if buf.remaining() < 2 {
     393            0 :             return Err(ProtocolError::BadMessage(
     394            0 :                 "Parse message is malformed, nparams missing".to_string(),
     395            0 :             ));
     396          630 :         }
     397          630 :         let nparams = buf.get_i16();
     398          630 : 
     399          630 :         if nparams != 0 {
     400            0 :             return Err(ProtocolError::BadMessage(
     401            0 :                 "query params not implemented".to_string(),
     402            0 :             ));
     403          630 :         }
     404          630 : 
     405          630 :         Ok(FeMessage::Parse(FeParseMessage { query_string }))
     406          630 :     }
     407              : }
     408              : 
     409              : impl FeDescribeMessage {
     410          630 :     fn parse(mut buf: Bytes) -> Result<FeMessage, ProtocolError> {
     411          630 :         let kind = buf.get_u8();
     412          630 :         let _pstmt_name = read_cstr(&mut buf)?;
     413              : 
     414              :         // FIXME: see FeParseMessage::parse
     415          630 :         if kind != b'S' {
     416            0 :             return Err(ProtocolError::BadMessage(
     417            0 :                 "only prepared statemement Describe is implemented".to_string(),
     418            0 :             ));
     419          630 :         }
     420          630 : 
     421          630 :         Ok(FeMessage::Describe(FeDescribeMessage { kind }))
     422          630 :     }
     423              : }
     424              : 
     425              : impl FeExecuteMessage {
     426          630 :     fn parse(mut buf: Bytes) -> Result<FeMessage, ProtocolError> {
     427          630 :         let portal_name = read_cstr(&mut buf)?;
     428          630 :         if buf.remaining() < 4 {
     429            0 :             return Err(ProtocolError::BadMessage(
     430            0 :                 "FeExecuteMessage message is malformed, maxrows missing".to_string(),
     431            0 :             ));
     432          630 :         }
     433          630 :         let maxrows = buf.get_i32();
     434          630 : 
     435          630 :         if !portal_name.is_empty() {
     436            0 :             return Err(ProtocolError::BadMessage(
     437            0 :                 "named portals not implemented".to_string(),
     438            0 :             ));
     439          630 :         }
     440          630 :         if maxrows != 0 {
     441            0 :             return Err(ProtocolError::BadMessage(
     442            0 :                 "row limit in Execute message not implemented".to_string(),
     443            0 :             ));
     444          630 :         }
     445          630 : 
     446          630 :         Ok(FeMessage::Execute(FeExecuteMessage { maxrows }))
     447          630 :     }
     448              : }
     449              : 
     450              : impl FeBindMessage {
     451          630 :     fn parse(mut buf: Bytes) -> Result<FeMessage, ProtocolError> {
     452          630 :         let portal_name = read_cstr(&mut buf)?;
     453          630 :         let _pstmt_name = read_cstr(&mut buf)?;
     454              : 
     455              :         // FIXME: see FeParseMessage::parse
     456          630 :         if !portal_name.is_empty() {
     457            0 :             return Err(ProtocolError::BadMessage(
     458            0 :                 "named portals not implemented".to_string(),
     459            0 :             ));
     460          630 :         }
     461          630 : 
     462          630 :         Ok(FeMessage::Bind(FeBindMessage))
     463          630 :     }
     464              : }
     465              : 
     466              : impl FeCloseMessage {
     467          629 :     fn parse(mut buf: Bytes) -> Result<FeMessage, ProtocolError> {
     468          629 :         let _kind = buf.get_u8();
     469          629 :         let _pstmt_or_portal_name = read_cstr(&mut buf)?;
     470              : 
     471              :         // FIXME: we do nothing with Close
     472          629 :         Ok(FeMessage::Close(FeCloseMessage))
     473          629 :     }
     474              : }
     475              : 
     476              : // Backend
     477              : 
     478            0 : #[derive(Debug)]
     479              : pub enum BeMessage<'a> {
     480              :     AuthenticationOk,
     481              :     AuthenticationMD5Password([u8; 4]),
     482              :     AuthenticationSasl(BeAuthenticationSaslMessage<'a>),
     483              :     AuthenticationCleartextPassword,
     484              :     BackendKeyData(CancelKeyData),
     485              :     BindComplete,
     486              :     CommandComplete(&'a [u8]),
     487              :     CopyData(&'a [u8]),
     488              :     CopyDone,
     489              :     CopyFail,
     490              :     CopyInResponse,
     491              :     CopyOutResponse,
     492              :     CopyBothResponse,
     493              :     CloseComplete,
     494              :     // None means column is NULL
     495              :     DataRow(&'a [Option<&'a [u8]>]),
     496              :     // None errcode means internal_error will be sent.
     497              :     ErrorResponse(&'a str, Option<&'a [u8; 5]>),
     498              :     /// Single byte - used in response to SSLRequest/GSSENCRequest.
     499              :     EncryptionResponse(bool),
     500              :     NoData,
     501              :     ParameterDescription,
     502              :     ParameterStatus {
     503              :         name: &'a [u8],
     504              :         value: &'a [u8],
     505              :     },
     506              :     ParseComplete,
     507              :     ReadyForQuery,
     508              :     RowDescription(&'a [RowDescriptor<'a>]),
     509              :     XLogData(XLogDataBody<'a>),
     510              :     NoticeResponse(&'a str),
     511              :     KeepAlive(WalSndKeepAlive),
     512              : }
     513              : 
     514              : /// Common shorthands.
     515              : impl<'a> BeMessage<'a> {
     516              :     /// A [`BeMessage::ParameterStatus`] holding the client encoding, i.e. UTF-8.
     517              :     /// This is a sensible default, given that:
     518              :     ///  * rust strings only support this encoding out of the box.
     519              :     ///  * tokio-postgres, postgres-jdbc (and probably more) mandate it.
     520              :     ///
     521              :     /// TODO: do we need to report `server_encoding` as well?
     522              :     pub const CLIENT_ENCODING: Self = Self::ParameterStatus {
     523              :         name: b"client_encoding",
     524              :         value: b"UTF8",
     525              :     };
     526              : 
     527              :     pub const INTEGER_DATETIMES: Self = Self::ParameterStatus {
     528              :         name: b"integer_datetimes",
     529              :         value: b"on",
     530              :     };
     531              : 
     532              :     /// Build a [`BeMessage::ParameterStatus`] holding the server version.
     533        13656 :     pub fn server_version(version: &'a str) -> Self {
     534        13656 :         Self::ParameterStatus {
     535        13656 :             name: b"server_version",
     536        13656 :             value: version.as_bytes(),
     537        13656 :         }
     538        13656 :     }
     539              : }
     540              : 
     541            0 : #[derive(Debug)]
     542              : pub enum BeAuthenticationSaslMessage<'a> {
     543              :     Methods(&'a [&'a str]),
     544              :     Continue(&'a [u8]),
     545              :     Final(&'a [u8]),
     546              : }
     547              : 
     548            0 : #[derive(Debug)]
     549              : pub enum BeParameterStatusMessage<'a> {
     550              :     Encoding(&'a str),
     551              :     ServerVersion(&'a str),
     552              : }
     553              : 
     554              : // One row description in RowDescription packet.
     555            0 : #[derive(Debug)]
     556              : pub struct RowDescriptor<'a> {
     557              :     pub name: &'a [u8],
     558              :     pub tableoid: Oid,
     559              :     pub attnum: i16,
     560              :     pub typoid: Oid,
     561              :     pub typlen: i16,
     562              :     pub typmod: i32,
     563              :     pub formatcode: i16,
     564              : }
     565              : 
     566              : impl Default for RowDescriptor<'_> {
     567         2975 :     fn default() -> RowDescriptor<'static> {
     568         2975 :         RowDescriptor {
     569         2975 :             name: b"",
     570         2975 :             tableoid: 0,
     571         2975 :             attnum: 0,
     572         2975 :             typoid: 0,
     573         2975 :             typlen: 0,
     574         2975 :             typmod: 0,
     575         2975 :             formatcode: 0,
     576         2975 :         }
     577         2975 :     }
     578              : }
     579              : 
     580              : impl RowDescriptor<'_> {
     581              :     /// Convenience function to create a RowDescriptor message for an int8 column
     582           45 :     pub const fn int8_col(name: &[u8]) -> RowDescriptor {
     583           45 :         RowDescriptor {
     584           45 :             name,
     585           45 :             tableoid: 0,
     586           45 :             attnum: 0,
     587           45 :             typoid: INT8_OID,
     588           45 :             typlen: 8,
     589           45 :             typmod: 0,
     590           45 :             formatcode: 0,
     591           45 :         }
     592           45 :     }
     593              : 
     594         1344 :     pub const fn text_col(name: &[u8]) -> RowDescriptor {
     595         1344 :         RowDescriptor {
     596         1344 :             name,
     597         1344 :             tableoid: 0,
     598         1344 :             attnum: 0,
     599         1344 :             typoid: TEXT_OID,
     600         1344 :             typlen: -1,
     601         1344 :             typmod: 0,
     602         1344 :             formatcode: 0,
     603         1344 :         }
     604         1344 :     }
     605              : }
     606              : 
     607            0 : #[derive(Debug)]
     608              : pub struct XLogDataBody<'a> {
     609              :     pub wal_start: u64,
     610              :     pub wal_end: u64, // current end of WAL on the server
     611              :     pub timestamp: i64,
     612              :     pub data: &'a [u8],
     613              : }
     614              : 
     615            0 : #[derive(Debug)]
     616              : pub struct WalSndKeepAlive {
     617              :     pub wal_end: u64, // current end of WAL on the server
     618              :     pub timestamp: i64,
     619              :     pub request_reply: bool,
     620              : }
     621              : 
     622              : pub static HELLO_WORLD_ROW: BeMessage = BeMessage::DataRow(&[Some(b"hello world")]);
     623              : 
     624              : // single text column
     625              : pub static SINGLE_COL_ROWDESC: BeMessage = BeMessage::RowDescription(&[RowDescriptor {
     626              :     name: b"data",
     627              :     tableoid: 0,
     628              :     attnum: 0,
     629              :     typoid: TEXT_OID,
     630              :     typlen: -1,
     631              :     typmod: 0,
     632              :     formatcode: 0,
     633              : }]);
     634              : 
     635              : /// Call f() to write body of the message and prepend it with 4-byte len as
     636              : /// prescribed by the protocol.
     637      8117563 : fn write_body<R>(buf: &mut BytesMut, f: impl FnOnce(&mut BytesMut) -> R) -> R {
     638      8117563 :     let base = buf.len();
     639      8117563 :     buf.extend_from_slice(&[0; 4]);
     640      8117563 : 
     641      8117563 :     let res = f(buf);
     642      8117563 : 
     643      8117563 :     let size = i32::try_from(buf.len() - base).expect("message too big to transmit");
     644      8117563 :     (&mut buf[base..]).put_slice(&size.to_be_bytes());
     645      8117563 : 
     646      8117563 :     res
     647      8117563 : }
     648              : 
     649              : /// Safe write of s into buf as cstring (String in the protocol).
     650        90318 : fn write_cstr(s: impl AsRef<[u8]>, buf: &mut BytesMut) -> Result<(), ProtocolError> {
     651        90318 :     let bytes = s.as_ref();
     652        90318 :     if bytes.contains(&0) {
     653            0 :         return Err(ProtocolError::BadMessage(
     654            0 :             "string contains embedded null".to_owned(),
     655            0 :         ));
     656        90318 :     }
     657        90318 :     buf.put_slice(bytes);
     658        90318 :     buf.put_u8(0);
     659        90318 :     Ok(())
     660        90318 : }
     661              : 
     662              : /// Read cstring from buf, advancing it.
     663      3990716 : pub fn read_cstr(buf: &mut Bytes) -> Result<Bytes, ProtocolError> {
     664      3990716 :     let pos = buf
     665      3990716 :         .iter()
     666     56666619 :         .position(|x| *x == 0)
     667      3990716 :         .ok_or_else(|| ProtocolError::BadMessage("missing cstring terminator".to_owned()))?;
     668      3990716 :     let result = buf.split_to(pos);
     669      3990716 :     buf.advance(1); // drop the null terminator
     670      3990716 :     Ok(result)
     671      3990716 : }
     672              : 
     673              : pub const SQLSTATE_INTERNAL_ERROR: &[u8; 5] = b"XX000";
     674              : pub const SQLSTATE_ADMIN_SHUTDOWN: &[u8; 5] = b"57P01";
     675              : pub const SQLSTATE_SUCCESSFUL_COMPLETION: &[u8; 5] = b"00000";
     676              : 
     677              : impl<'a> BeMessage<'a> {
     678              :     /// Serialize `message` to the given `buf`.
     679              :     /// Apart from smart memory managemet, BytesMut is good here as msg len
     680              :     /// precedes its body and it is handy to write it down first and then fill
     681              :     /// the length. With Write we would have to either calc it manually or have
     682              :     /// one more buffer.
     683      8129504 :     pub fn write(buf: &mut BytesMut, message: &BeMessage) -> Result<(), ProtocolError> {
     684      8129504 :         match message {
     685        13922 :             BeMessage::AuthenticationOk => {
     686        13922 :                 buf.put_u8(b'R');
     687        13922 :                 write_body(buf, |buf| {
     688        13922 :                     buf.put_i32(0); // Specifies that the authentication was successful.
     689        13922 :                 });
     690        13922 :             }
     691              : 
     692          221 :             BeMessage::AuthenticationCleartextPassword => {
     693          221 :                 buf.put_u8(b'R');
     694          221 :                 write_body(buf, |buf| {
     695          221 :                     buf.put_i32(3); // Specifies that clear text password is required.
     696          221 :                 });
     697          221 :             }
     698              : 
     699            0 :             BeMessage::AuthenticationMD5Password(salt) => {
     700            0 :                 buf.put_u8(b'R');
     701            0 :                 write_body(buf, |buf| {
     702            0 :                     buf.put_i32(5); // Specifies that an MD5-encrypted password is required.
     703            0 :                     buf.put_slice(&salt[..]);
     704            0 :                 });
     705            0 :             }
     706              : 
     707          168 :             BeMessage::AuthenticationSasl(msg) => {
     708          168 :                 buf.put_u8(b'R');
     709          168 :                 write_body(buf, |buf| {
     710          168 :                     use BeAuthenticationSaslMessage::*;
     711          168 :                     match msg {
     712           63 :                         Methods(methods) => {
     713           63 :                             buf.put_i32(10); // Specifies that SASL auth method is used.
     714          126 :                             for method in methods.iter() {
     715          126 :                                 write_cstr(method, buf)?;
     716              :                             }
     717           63 :                             buf.put_u8(0); // zero terminator for the list
     718              :                         }
     719           59 :                         Continue(extra) => {
     720           59 :                             buf.put_i32(11); // Continue SASL auth.
     721           59 :                             buf.put_slice(extra);
     722           59 :                         }
     723           46 :                         Final(extra) => {
     724           46 :                             buf.put_i32(12); // Send final SASL message.
     725           46 :                             buf.put_slice(extra);
     726           46 :                         }
     727              :                     }
     728          168 :                     Ok(())
     729          168 :                 })?;
     730              :             }
     731              : 
     732           41 :             BeMessage::BackendKeyData(key_data) => {
     733           41 :                 buf.put_u8(b'K');
     734           41 :                 write_body(buf, |buf| {
     735           41 :                     buf.put_i32(key_data.backend_pid);
     736           41 :                     buf.put_i32(key_data.cancel_key);
     737           41 :                 });
     738           41 :             }
     739              : 
     740          630 :             BeMessage::BindComplete => {
     741          630 :                 buf.put_u8(b'2');
     742          630 :                 write_body(buf, |_| {});
     743          630 :             }
     744              : 
     745          629 :             BeMessage::CloseComplete => {
     746          629 :                 buf.put_u8(b'3');
     747          629 :                 write_body(buf, |_| {});
     748          629 :             }
     749              : 
     750         2066 :             BeMessage::CommandComplete(cmd) => {
     751         2066 :                 buf.put_u8(b'C');
     752         2066 :                 write_body(buf, |buf| write_cstr(cmd, buf))?;
     753              :             }
     754              : 
     755      7207959 :             BeMessage::CopyData(data) => {
     756      7207959 :                 buf.put_u8(b'd');
     757      7207959 :                 write_body(buf, |buf| {
     758      7207959 :                     buf.put_slice(data);
     759      7207959 :                 });
     760      7207959 :             }
     761              : 
     762          596 :             BeMessage::CopyDone => {
     763          596 :                 buf.put_u8(b'c');
     764          596 :                 write_body(buf, |_| {});
     765          596 :             }
     766              : 
     767            0 :             BeMessage::CopyFail => {
     768            0 :                 buf.put_u8(b'f');
     769            0 :                 write_body(buf, |_| {});
     770            0 :             }
     771              : 
     772           14 :             BeMessage::CopyInResponse => {
     773           14 :                 buf.put_u8(b'G');
     774           14 :                 write_body(buf, |buf| {
     775           14 :                     buf.put_u8(1); // copy_is_binary
     776           14 :                     buf.put_i16(0); // numAttributes
     777           14 :                 });
     778           14 :             }
     779              : 
     780          608 :             BeMessage::CopyOutResponse => {
     781          608 :                 buf.put_u8(b'H');
     782          608 :                 write_body(buf, |buf| {
     783          608 :                     buf.put_u8(0); // copy_is_binary
     784          608 :                     buf.put_i16(0); // numAttributes
     785          608 :                 });
     786          608 :             }
     787              : 
     788        12464 :             BeMessage::CopyBothResponse => {
     789        12464 :                 buf.put_u8(b'W');
     790        12464 :                 write_body(buf, |buf| {
     791        12464 :                     // doesn't matter, used only for replication
     792        12464 :                     buf.put_u8(0); // copy_is_binary
     793        12464 :                     buf.put_i16(0); // numAttributes
     794        12464 :                 });
     795        12464 :             }
     796              : 
     797          957 :             BeMessage::DataRow(vals) => {
     798          957 :                 buf.put_u8(b'D');
     799          957 :                 write_body(buf, |buf| {
     800          957 :                     buf.put_u16(vals.len() as u16); // num of cols
     801         3425 :                     for val_opt in vals.iter() {
     802         3425 :                         if let Some(val) = val_opt {
     803         2682 :                             buf.put_u32(val.len() as u32);
     804         2682 :                             buf.put_slice(val);
     805         2682 :                         } else {
     806          743 :                             buf.put_i32(-1);
     807          743 :                         }
     808              :                     }
     809          957 :                 });
     810          957 :             }
     811              : 
     812              :             // ErrorResponse is a zero-terminated array of zero-terminated fields.
     813              :             // First byte of each field represents type of this field. Set just enough fields
     814              :             // to satisfy rust-postgres client: 'S' -- severity, 'C' -- error, 'M' -- error
     815              :             // message text.
     816          295 :             BeMessage::ErrorResponse(error_msg, pg_error_code) => {
     817          295 :                 // 'E' signalizes ErrorResponse messages
     818          295 :                 buf.put_u8(b'E');
     819          295 :                 write_body(buf, |buf| {
     820          295 :                     buf.put_u8(b'S'); // severity
     821          295 :                     buf.put_slice(b"ERROR\0");
     822          295 : 
     823          295 :                     buf.put_u8(b'C'); // SQLSTATE error code
     824          295 :                     buf.put_slice(&terminate_code(
     825          295 :                         pg_error_code.unwrap_or(SQLSTATE_INTERNAL_ERROR),
     826          295 :                     ));
     827          295 : 
     828          295 :                     buf.put_u8(b'M'); // the message
     829          295 :                     write_cstr(error_msg, buf)?;
     830              : 
     831          295 :                     buf.put_u8(0); // terminator
     832          295 :                     Ok(())
     833          295 :                 })?;
     834              :             }
     835              : 
     836              :             // NoticeResponse has the same format as ErrorResponse. From doc: "The frontend should display the
     837              :             // message but continue listening for ReadyForQuery or ErrorResponse"
     838            6 :             BeMessage::NoticeResponse(error_msg) => {
     839            6 :                 // For all the errors set Severity to Error and error code to
     840            6 :                 // 'internal error'.
     841            6 : 
     842            6 :                 // 'N' signalizes NoticeResponse messages
     843            6 :                 buf.put_u8(b'N');
     844            6 :                 write_body(buf, |buf| {
     845            6 :                     buf.put_u8(b'S'); // severity
     846            6 :                     buf.put_slice(b"NOTICE\0");
     847            6 : 
     848            6 :                     buf.put_u8(b'C'); // SQLSTATE error code
     849            6 :                     buf.put_slice(&terminate_code(SQLSTATE_INTERNAL_ERROR));
     850            6 : 
     851            6 :                     buf.put_u8(b'M'); // the message
     852            6 :                     write_cstr(error_msg.as_bytes(), buf)?;
     853              : 
     854            6 :                     buf.put_u8(0); // terminator
     855            6 :                     Ok(())
     856            6 :                 })?;
     857              :             }
     858              : 
     859          630 :             BeMessage::NoData => {
     860          630 :                 buf.put_u8(b'n');
     861          630 :                 write_body(buf, |_| {});
     862          630 :             }
     863              : 
     864        11941 :             BeMessage::EncryptionResponse(should_negotiate) => {
     865        11941 :                 let response = if *should_negotiate { b'S' } else { b'N' };
     866        11941 :                 buf.put_u8(response);
     867              :             }
     868              : 
     869        41729 :             BeMessage::ParameterStatus { name, value } => {
     870        41729 :                 buf.put_u8(b'S');
     871        41729 :                 write_body(buf, |buf| {
     872        41729 :                     write_cstr(name, buf)?;
     873        41729 :                     write_cstr(value, buf)
     874        41729 :                 })?;
     875              :             }
     876              : 
     877          630 :             BeMessage::ParameterDescription => {
     878          630 :                 buf.put_u8(b't');
     879          630 :                 write_body(buf, |buf| {
     880          630 :                     // we don't support params, so always 0
     881          630 :                     buf.put_i16(0);
     882          630 :                 });
     883          630 :             }
     884              : 
     885          630 :             BeMessage::ParseComplete => {
     886          630 :                 buf.put_u8(b'1');
     887          630 :                 write_body(buf, |_| {});
     888          630 :             }
     889              : 
     890        29161 :             BeMessage::ReadyForQuery => {
     891        29161 :                 buf.put_u8(b'Z');
     892        29161 :                 write_body(buf, |buf| {
     893        29161 :                     buf.put_u8(b'I');
     894        29161 :                 });
     895        29161 :             }
     896              : 
     897         1428 :             BeMessage::RowDescription(rows) => {
     898         1428 :                 buf.put_u8(b'T');
     899         1428 :                 write_body(buf, |buf| {
     900         1428 :                     buf.put_i16(rows.len() as i16); // # of fields
     901         4367 :                     for row in rows.iter() {
     902         4367 :                         write_cstr(row.name, buf)?;
     903         4367 :                         buf.put_i32(0); /* table oid */
     904         4367 :                         buf.put_i16(0); /* attnum */
     905         4367 :                         buf.put_u32(row.typoid);
     906         4367 :                         buf.put_i16(row.typlen);
     907         4367 :                         buf.put_i32(-1); /* typmod */
     908         4367 :                         buf.put_i16(0); /* format code */
     909              :                     }
     910         1428 :                     Ok(())
     911         1428 :                 })?;
     912              :             }
     913              : 
     914       799473 :             BeMessage::XLogData(body) => {
     915       799473 :                 buf.put_u8(b'd');
     916       799473 :                 write_body(buf, |buf| {
     917       799473 :                     buf.put_u8(b'w');
     918       799473 :                     buf.put_u64(body.wal_start);
     919       799473 :                     buf.put_u64(body.wal_end);
     920       799473 :                     buf.put_i64(body.timestamp);
     921       799473 :                     buf.put_slice(body.data);
     922       799473 :                 });
     923       799473 :             }
     924              : 
     925         3306 :             BeMessage::KeepAlive(req) => {
     926         3306 :                 buf.put_u8(b'd');
     927         3306 :                 write_body(buf, |buf| {
     928         3306 :                     buf.put_u8(b'k');
     929         3306 :                     buf.put_u64(req.wal_end);
     930         3306 :                     buf.put_i64(req.timestamp);
     931         3306 :                     buf.put_u8(u8::from(req.request_reply));
     932         3306 :                 });
     933         3306 :             }
     934              :         }
     935      8129504 :         Ok(())
     936      8129504 :     }
     937              : }
     938              : 
     939          301 : fn terminate_code(code: &[u8; 5]) -> [u8; 6] {
     940          301 :     let mut terminated = [0; 6];
     941         1505 :     for (i, &elem) in code.iter().enumerate() {
     942         1505 :         terminated[i] = elem;
     943         1505 :     }
     944              : 
     945          301 :     terminated
     946          301 : }
     947              : 
     948              : #[cfg(test)]
     949              : mod tests {
     950              :     use super::*;
     951              : 
     952            2 :     #[test]
     953            2 :     fn test_startup_message_params_options_escaped() {
     954            8 :         fn split_options(params: &StartupMessageParams) -> Vec<Cow<'_, str>> {
     955            8 :             params
     956            8 :                 .options_escaped()
     957            8 :                 .expect("options are None")
     958            8 :                 .collect()
     959            8 :         }
     960            2 : 
     961            8 :         let make_params = |options| StartupMessageParams::new([("options", options)]);
     962              : 
     963            2 :         let params = StartupMessageParams::new([]);
     964            2 :         assert!(params.options_escaped().is_none());
     965              : 
     966            2 :         let params = make_params("");
     967            2 :         assert!(split_options(&params).is_empty());
     968              : 
     969            2 :         let params = make_params("foo");
     970            2 :         assert_eq!(split_options(&params), ["foo"]);
     971              : 
     972            2 :         let params = make_params(" foo  bar ");
     973            2 :         assert_eq!(split_options(&params), ["foo", "bar"]);
     974              : 
     975            2 :         let params = make_params("foo\\ bar \\ \\\\ baz\\  lol");
     976            2 :         assert_eq!(split_options(&params), ["foo bar", " \\", "baz ", "lol"]);
     977            2 :     }
     978              : 
     979            2 :     #[test]
     980            2 :     fn parse_fe_startup_packet_regression() {
     981            2 :         let data = [0, 0, 0, 7, 0, 0, 0, 0];
     982            2 :         FeStartupPacket::parse(&mut BytesMut::from_iter(data)).unwrap_err();
     983            2 :     }
     984              : }
        

Generated by: LCOV version 2.1-beta