diff --git a/src/error.rs b/src/error.rs index 1e2a649..77b7c8c 100644 --- a/src/error.rs +++ b/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 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()) } diff --git a/src/messages/codec.rs b/src/messages/codec.rs index b58f42b..72558ea 100644 --- a/src/messages/codec.rs +++ b/src/messages/codec.rs @@ -58,6 +58,19 @@ pub(crate) fn get_length(buf: &BytesMut, offset: usize) -> Option { } } +/// 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( diff --git a/src/messages/copy.rs b/src/messages/copy.rs index 8473af9..d31921d 100644 --- a/src/messages/copy.rs +++ b/src/messages/copy.rs @@ -136,6 +136,7 @@ impl Message for CopyInResponse { fn decode_body(buf: &mut BytesMut, _len: usize, _ctx: &DecodeContext) -> PgWireResult { 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 { 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 { 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)) } } diff --git a/src/messages/data.rs b/src/messages/data.rs index 241e0f9..5832897 100644 --- a/src/messages/data.rs +++ b/src/messages/data.rs @@ -75,7 +75,9 @@ impl Message for RowDescription { } fn decode_body(buf: &mut BytesMut, _: usize, _ctx: &DecodeContext) -> PgWireResult { - 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 { diff --git a/src/messages/extendedquery.rs b/src/messages/extendedquery.rs index ff8129d..be0fe8e 100644 --- a/src/messages/extendedquery.rs +++ b/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 { diff --git a/src/messages/mod.rs b/src/messages/mod.rs index a334aba..c23a0f4 100644 --- a/src/messages/mod.rs +++ b/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); + } } diff --git a/src/messages/startup.rs b/src/messages/startup.rs index 9784412..1330134 100644 --- a/src/messages/startup.rs +++ b/src/messages/startup.rs @@ -820,7 +820,8 @@ impl Message for NegotiateProtocolVersion { _ctx: &DecodeContext, ) -> PgWireResult { 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 {