Reject malformed element counts in message decoders

src/error.rs

@@ -17,6 +17,8 @@ pub enum PgWireError {
     InvalidMessageType(u8),
     #[error("Invalid message length, expected max {0}, actual: {1}")]
     MessageTooLarge(usize, usize),
+    #[error("Invalid element count {0}: {1} bytes required, {2} remaining")]
+    InvalidElementCount(usize, usize, usize),
     #[error("Invalid target type, received {0}")]
     InvalidTargetType(u8),
     #[error("Invalid transaction status, received {0}")]
@@ -327,6 +329,9 @@ impl From<PgWireError> for ErrorInfo {
             PgWireError::MessageTooLarge(..) => {
                 ErrorInfo::new("FATAL".to_owned(), "08P01".to_owned(), error.to_string())
             }
+            PgWireError::InvalidElementCount(..) => {
+                ErrorInfo::new("FATAL".to_owned(), "08P01".to_owned(), error.to_string())
+            }
             PgWireError::InvalidTransactionStatus(_) => {
                 ErrorInfo::new("FATAL".to_owned(), "08P01".to_owned(), error.to_string())
             }

src/messages/codec.rs

@@ -58,6 +58,19 @@ pub(crate) fn get_length(buf: &BytesMut, offset: usize) -> Option<usize> {
     }
 }
 
+/// Validate an on-wire element count before it is used to pre-allocate a
+/// collection. Counts are unsigned; `count` elements of `elem_size` bytes each
+/// must fit in the remaining buffer, or the message cannot be real.
+pub(crate) fn ensure_count(count: usize, elem_size: usize, buf: &BytesMut) -> PgWireResult<()> {
+    let remaining = buf.remaining();
+    let needed = count.saturating_mul(elem_size);
+    if needed > remaining {
+        Err(PgWireError::InvalidElementCount(count, needed, remaining))
+    } else {
+        Ok(())
+    }
+}
+
 /// Check if message_length matches and move the cursor to right position then
 /// call the `decode_fn` for the body
 pub(crate) fn decode_packet<T, F>(

src/messages/copy.rs

@@ -136,6 +136,7 @@ impl Message for CopyInResponse {
     fn decode_body(buf: &mut BytesMut, _len: usize, _ctx: &DecodeContext) -> PgWireResult<Self> {
         let format = buf.get_i8();
         let columns = buf.get_i16();
+        codec::ensure_count(columns as usize, 2, buf)?;
         let mut column_formats = Vec::with_capacity(columns as usize);
         for _ in 0..columns {
             column_formats.push(buf.get_i16());
@@ -183,13 +184,14 @@ impl Message for CopyOutResponse {
 
     fn decode_body(buf: &mut BytesMut, _len: usize, _ctx: &DecodeContext) -> PgWireResult<Self> {
         let format = buf.get_i8();
-        let columns = buf.get_i16();
+        let columns = buf.get_u16();
+        codec::ensure_count(columns as usize, 2, buf)?;
         let mut column_formats = Vec::with_capacity(columns as usize);
         for _ in 0..columns {
             column_formats.push(buf.get_i16());
         }
 
-        Ok(Self::new(format, columns, column_formats))
+        Ok(Self::new(format, columns as i16, column_formats))
     }
 }
 
@@ -231,12 +233,13 @@ impl Message for CopyBothResponse {
 
     fn decode_body(buf: &mut BytesMut, _len: usize, _ctx: &DecodeContext) -> PgWireResult<Self> {
         let format = buf.get_i8();
-        let columns = buf.get_i16();
+        let columns = buf.get_u16();
+        codec::ensure_count(columns as usize, 2, buf)?;
         let mut column_formats = Vec::with_capacity(columns as usize);
         for _ in 0..columns {
             column_formats.push(buf.get_i16());
         }
 
-        Ok(Self::new(format, columns, column_formats))
+        Ok(Self::new(format, columns as i16, column_formats))
     }
 }

src/messages/data.rs

@@ -75,7 +75,9 @@ impl Message for RowDescription {
     }
 
     fn decode_body(buf: &mut BytesMut, _: usize, _ctx: &DecodeContext) -> PgWireResult<Self> {
-        let fields_len = buf.get_i16();
+        let fields_len = buf.get_u16();
+        // Each field: C-string name (>= 1 byte) + 18 fixed bytes.
+        codec::ensure_count(fields_len as usize, 19, buf)?;
         let mut fields = Vec::with_capacity(fields_len as usize);
 
         for _ in 0..fields_len {

src/messages/extendedquery.rs

@@ -280,7 +280,8 @@ impl Message for Bind {
             }
         }
 
-        let result_column_format_code_len = buf.get_i16();
+        let result_column_format_code_len = buf.get_u16();
+        codec::ensure_count(result_column_format_code_len as usize, 2, buf)?;
         let mut result_column_format_codes =
             Vec::with_capacity(result_column_format_code_len as usize);
         for _ in 0..result_column_format_code_len {

src/messages/mod.rs

@@ -1029,4 +1029,55 @@ mod test {
         assert_eq!(196608i32, i32::from(ProtocolVersion::PROTOCOL3_0));
         assert_eq!(196610i32, i32::from(ProtocolVersion::PROTOCOL3_2));
     }
+
+    // A count read as signed (0xffff -> -1 -> usize::MAX) must become a decode
+    // error, not a `capacity overflow` panic. Completes #192 for the decoders it
+    // did not reach.
+    #[test]
+    fn test_element_count_capacity_overflow_rejected() {
+        let ctx = DecodeContext::new(ProtocolVersion::PROTOCOL3_0);
+
+        // Backend messages decoded by the client/proxy from an untrusted server.
+        for bytes in [
+            &b"T\0\0\0\x06\xff\xff"[..],                 // RowDescription
+            &b"G\0\0\0\x07\0\xff\xff"[..],               // CopyInResponse
+            &b"H\0\0\0\x07\0\xff\xff"[..],               // CopyOutResponse
+            &b"W\0\0\0\x07\0\xff\xff"[..],               // CopyBothResponse
+            &b"v\0\0\0\x0c\0\0\0\0\xff\xff\xff\xff"[..], // NegotiateProtocolVersion
+        ] {
+            let mut buf = BytesMut::from(bytes);
+            assert!(super::PgWireBackendMessage::decode(&mut buf, &ctx).is_err());
+        }
+
+        // Bind decoded by the server from an untrusted client. #192 widened the
+        // first two of Bind's three counts; this covers the third.
+        let mut fctx = DecodeContext::new(ProtocolVersion::PROTOCOL3_0);
+        fctx.awaiting_frontend_ssl = false;
+        fctx.awaiting_frontend_startup = false;
+        let mut buf = BytesMut::from(&b"B\0\0\0\x0c\0\0\0\0\0\0\xff\xff"[..]);
+        assert!(super::PgWireFrontendMessage::decode(&mut buf, &fctx).is_err());
+
+        // Bind with 5 result format codes but only 6 bytes left: the codes are
+        // Int16s, so the element-aware bound must reject (5 * 2 > 6) instead of
+        // passing a byte-count check and panicking on the truncated read.
+        let mut buf = BytesMut::from(&b"B\0\0\0\x0e\0\0\0\0\0\0\0\x05\0\0\0\0\0\0"[..]);
+        assert!(super::PgWireFrontendMessage::decode(&mut buf, &fctx).is_err());
+    }
+
+    #[test]
+    fn test_valid_element_count_still_decodes() {
+        let ctx = DecodeContext::new(ProtocolVersion::PROTOCOL3_0);
+
+        // Zero count decodes to an empty collection.
+        let mut buf = BytesMut::from(&b"T\0\0\0\x06\0\0"[..]);
+        let row = RowDescription::decode(&mut buf, &ctx).unwrap().unwrap();
+        assert!(row.fields.is_empty());
+
+        // A well-formed positive count round-trips.
+        let row = RowDescription::new(vec![
+            FieldDescription::new("id".to_owned(), 0, 0, 23, 4, -1, 0),
+            FieldDescription::new("name".to_owned(), 0, 0, 25, -1, -1, 0),
+        ]);
+        roundtrip!(row, RowDescription, &ctx);
+    }
 }

src/messages/startup.rs

@@ -820,7 +820,8 @@ impl Message for NegotiateProtocolVersion {
         _ctx: &DecodeContext,
     ) -> PgWireResult<Self> {
         let version = buf.get_i32();
-        let option_count = buf.get_i32();
+        let option_count = buf.get_u32();
+        codec::ensure_count(option_count as usize, 1, buf)?;
         let mut options = Vec::with_capacity(option_count as usize);
 
         for _ in 0..option_count {