Do not panic when PrepareOk fails to decode (#2572)

launchbadge · Jun 30, 2023 · 3fdb79d · 3fdb79d
1 parent 8c7f541
commit 3fdb79d
Showing 1 changed file with 13 additions and 7 deletions.
diff --git a/sqlx-mysql/src/protocol/statement/prepare_ok.rs b/sqlx-mysql/src/protocol/statement/prepare_ok.rs
@@ -16,22 +16,28 @@ pub(crate) struct PrepareOk {
 }
 
 impl Decode<'_, Capabilities> for PrepareOk {
-    fn decode_with(mut buf: Bytes, _: Capabilities) -> Result<Self, Error> {
-        let status = buf.get_u8();
+    fn decode_with(buf: Bytes, _: Capabilities) -> Result<Self, Error> {
+        const SIZE: usize = 12;
+
+        let mut slice = buf.get(..SIZE).ok_or_else(|| {
+            err_protocol!("PrepareOk expected 12 bytes but got {} bytes", buf.len())
+        })?;
+
+        let status = slice.get_u8();
         if status != 0x00 {
             return Err(err_protocol!(
                 "expected 0x00 (COM_STMT_PREPARE_OK) but found 0x{:02x}",
                 status
             ));
         }
 
-        let statement_id = buf.get_u32_le();
-        let columns = buf.get_u16_le();
-        let params = buf.get_u16_le();
+        let statement_id = slice.get_u32_le();
+        let columns = slice.get_u16_le();
+        let params = slice.get_u16_le();
 
-        buf.advance(1); // reserved: string<1>
+        slice.advance(1); // reserved: string<1>
 
-        let warnings = buf.get_u16_le();
+        let warnings = slice.get_u16_le();
 
         Ok(Self {
             statement_id,