Skip to main content

powdb_server/
protocol.rs

1use tokio::io::{AsyncReadExt, AsyncWriteExt};
2use zeroize::Zeroizing;
3
4const MSG_CONNECT: u8 = 0x01;
5const MSG_CONNECT_OK: u8 = 0x02;
6const MSG_QUERY: u8 = 0x03;
7/// Query carrying positional `$N` parameters (Task 4). Pure protocol
8/// addition: old clients never send it, and old servers reject it with the
9/// existing "unknown message type" error — no existing frame changes shape.
10const MSG_QUERY_PARAMS: u8 = 0x04;
11const MSG_RESULT_ROWS: u8 = 0x07;
12const MSG_RESULT_SCALAR: u8 = 0x08;
13const MSG_RESULT_OK: u8 = 0x09;
14const MSG_ERROR: u8 = 0x0A;
15const MSG_RESULT_MSG: u8 = 0x0B;
16const MSG_DISCONNECT: u8 = 0x10;
17const MSG_PING: u8 = 0x11;
18const MSG_PONG: u8 = 0x12;
19
20/// Maximum payload size accepted from the wire (64 MB).
21const MAX_PAYLOAD_SIZE: usize = 64 * 1024 * 1024;
22
23/// Maximum payload size for pre-auth CONNECT messages (4 KB).
24/// Only a database name and password are needed before authentication.
25const MAX_CONNECT_PAYLOAD_SIZE: usize = 4096;
26
27/// Maximum number of columns allowed in a result set.
28const MAX_COLUMNS: usize = 4096;
29
30/// Maximum number of rows allowed in a single result message.
31const MAX_ROWS: usize = 10_000_000;
32
33/// Maximum number of bound parameters in a single QueryWithParams message.
34const MAX_PARAMS: usize = 4096;
35
36/// A positional parameter value carried by [`Message::QueryWithParams`].
37///
38/// Wire encoding per param: a 1-byte tag followed by the body —
39///   `0` null (no body), `1` int (8B LE i64), `2` float (8B LE f64),
40///   `3` bool (1B), `4` str (length-prefixed UTF-8).
41#[derive(Debug, Clone, PartialEq)]
42pub enum WireParam {
43    Null,
44    Int(i64),
45    Float(f64),
46    Bool(bool),
47    Str(String),
48}
49
50#[derive(Debug, Clone)]
51pub enum Message {
52    Connect {
53        db_name: String,
54        /// Client-supplied candidate password. Wrapped in `Zeroizing` so the
55        /// raw bytes from the wire are wiped from memory once the `Message`
56        /// (and thus this field) is dropped after the constant-time compare —
57        /// the candidate never lingers in a plain `String`.
58        password: Option<Zeroizing<String>>,
59        /// Optional user name for multi-user authentication. Appended after the
60        /// password on the wire so the format is a pure, backward-compatible
61        /// extension: old clients that omit it decode as `None`.
62        username: Option<String>,
63    },
64    ConnectOk {
65        version: String,
66    },
67    Query {
68        query: String,
69    },
70    /// A query string with positional `$N` parameters bound at the server.
71    QueryWithParams {
72        query: String,
73        params: Vec<WireParam>,
74    },
75    ResultRows {
76        columns: Vec<String>,
77        rows: Vec<Vec<String>>,
78    },
79    ResultScalar {
80        value: String,
81    },
82    ResultOk {
83        affected: u64,
84    },
85    /// A descriptive status message (e.g. "type User created", "index dropped").
86    ResultMessage {
87        message: String,
88    },
89    Error {
90        message: String,
91    },
92    Disconnect,
93    Ping,
94    Pong,
95}
96
97impl Message {
98    /// Encode message into wire format: [type(1)][flags(1)][len(4)][payload]
99    pub fn encode(&self) -> Vec<u8> {
100        let (msg_type, payload) = match self {
101            Message::Connect {
102                db_name,
103                password,
104                username,
105            } => {
106                let mut buf = encode_string(db_name);
107                // Password is encoded as a length-prefixed string. Empty (len=0) means None.
108                match password {
109                    Some(p) => buf.extend_from_slice(&encode_string(p)),
110                    None => buf.extend_from_slice(&0u32.to_le_bytes()),
111                }
112                // Username is appended after the password (append-only extension).
113                // Length-prefixed string; empty (len=0) means None.
114                match username {
115                    Some(u) => buf.extend_from_slice(&encode_string(u)),
116                    None => buf.extend_from_slice(&0u32.to_le_bytes()),
117                }
118                (MSG_CONNECT, buf)
119            }
120            Message::ConnectOk { version } => (MSG_CONNECT_OK, encode_string(version)),
121            Message::Query { query } => (MSG_QUERY, encode_string(query)),
122            Message::QueryWithParams { query, params } => {
123                let mut buf = encode_string(query);
124                buf.extend_from_slice(&(params.len() as u16).to_le_bytes());
125                for p in params {
126                    match p {
127                        WireParam::Null => buf.push(0),
128                        WireParam::Int(v) => {
129                            buf.push(1);
130                            buf.extend_from_slice(&v.to_le_bytes());
131                        }
132                        WireParam::Float(v) => {
133                            buf.push(2);
134                            buf.extend_from_slice(&v.to_le_bytes());
135                        }
136                        WireParam::Bool(v) => {
137                            buf.push(3);
138                            buf.push(if *v { 1 } else { 0 });
139                        }
140                        WireParam::Str(s) => {
141                            buf.push(4);
142                            buf.extend_from_slice(&encode_string(s));
143                        }
144                    }
145                }
146                (MSG_QUERY_PARAMS, buf)
147            }
148            Message::ResultRows { columns, rows } => {
149                let mut buf = Vec::new();
150                buf.extend_from_slice(&(columns.len() as u16).to_le_bytes());
151                for col in columns {
152                    buf.extend_from_slice(&encode_string(col));
153                }
154                buf.extend_from_slice(&(rows.len() as u32).to_le_bytes());
155                for row in rows {
156                    for val in row {
157                        buf.extend_from_slice(&encode_string(val));
158                    }
159                }
160                (MSG_RESULT_ROWS, buf)
161            }
162            Message::ResultScalar { value } => (MSG_RESULT_SCALAR, encode_string(value)),
163            Message::ResultOk { affected } => (MSG_RESULT_OK, affected.to_le_bytes().to_vec()),
164            Message::ResultMessage { message } => (MSG_RESULT_MSG, encode_string(message)),
165            Message::Error { message } => (MSG_ERROR, encode_string(message)),
166            Message::Disconnect => (MSG_DISCONNECT, Vec::new()),
167            Message::Ping => (MSG_PING, Vec::new()),
168            Message::Pong => (MSG_PONG, Vec::new()),
169        };
170
171        let mut frame = Vec::with_capacity(6 + payload.len());
172        frame.push(msg_type);
173        frame.push(0); // flags
174        frame.extend_from_slice(&(payload.len() as u32).to_le_bytes());
175        frame.extend_from_slice(&payload);
176        frame
177    }
178
179    /// Decode message from wire format.
180    pub fn decode(data: &[u8]) -> Result<Message, String> {
181        if data.len() < 6 {
182            return Err("frame too short".into());
183        }
184        let msg_type = data[0];
185        let _flags = data[1];
186        let len_bytes: [u8; 4] = data[2..6]
187            .try_into()
188            .map_err(|_| "invalid header length field".to_string())?;
189        let payload_len = u32::from_le_bytes(len_bytes) as usize;
190        if 6 + payload_len > data.len() {
191            return Err("payload length exceeds frame".into());
192        }
193        let payload = &data[6..6 + payload_len];
194
195        match msg_type {
196            MSG_CONNECT => {
197                let mut pos = 0;
198                let db_name = decode_string(payload, &mut pos)?;
199                // Password is optional. If there are no more bytes, treat as None
200                // (backwards compatible with old clients that don't send a password).
201                let password = if pos < payload.len() {
202                    // Wrap the candidate immediately so the only owned copy of
203                    // the secret lives inside `Zeroizing` and is wiped on drop.
204                    let p = Zeroizing::new(decode_string(payload, &mut pos)?);
205                    if p.is_empty() {
206                        None
207                    } else {
208                        Some(p)
209                    }
210                } else {
211                    None
212                };
213                // Username is optional and appended after the password. Old
214                // clients omit it entirely → no more bytes → None.
215                let username = if pos < payload.len() {
216                    let u = decode_string(payload, &mut pos)?;
217                    if u.is_empty() {
218                        None
219                    } else {
220                        Some(u)
221                    }
222                } else {
223                    None
224                };
225                Ok(Message::Connect {
226                    db_name,
227                    password,
228                    username,
229                })
230            }
231            MSG_CONNECT_OK => {
232                let version = decode_string(payload, &mut 0)?;
233                Ok(Message::ConnectOk { version })
234            }
235            MSG_QUERY => {
236                let query = decode_string(payload, &mut 0)?;
237                Ok(Message::Query { query })
238            }
239            MSG_QUERY_PARAMS => {
240                let mut pos = 0;
241                let query = decode_string(payload, &mut pos)?;
242                if pos + 2 > payload.len() {
243                    return Err("truncated param count".into());
244                }
245                let count_bytes: [u8; 2] = payload[pos..pos + 2]
246                    .try_into()
247                    .map_err(|_| "invalid param count bytes".to_string())?;
248                let count = u16::from_le_bytes(count_bytes) as usize;
249                pos += 2;
250                if count > MAX_PARAMS {
251                    return Err("too many parameters".into());
252                }
253                let mut params = Vec::with_capacity(count);
254                for _ in 0..count {
255                    if pos >= payload.len() {
256                        return Err("truncated param tag".into());
257                    }
258                    let tag = payload[pos];
259                    pos += 1;
260                    let p = match tag {
261                        0 => WireParam::Null,
262                        1 => {
263                            if pos + 8 > payload.len() {
264                                return Err("truncated int param".into());
265                            }
266                            let b: [u8; 8] = payload[pos..pos + 8]
267                                .try_into()
268                                .map_err(|_| "invalid int param bytes".to_string())?;
269                            pos += 8;
270                            WireParam::Int(i64::from_le_bytes(b))
271                        }
272                        2 => {
273                            if pos + 8 > payload.len() {
274                                return Err("truncated float param".into());
275                            }
276                            let b: [u8; 8] = payload[pos..pos + 8]
277                                .try_into()
278                                .map_err(|_| "invalid float param bytes".to_string())?;
279                            pos += 8;
280                            WireParam::Float(f64::from_le_bytes(b))
281                        }
282                        3 => {
283                            if pos + 1 > payload.len() {
284                                return Err("truncated bool param".into());
285                            }
286                            let v = payload[pos] != 0;
287                            pos += 1;
288                            WireParam::Bool(v)
289                        }
290                        4 => WireParam::Str(decode_string(payload, &mut pos)?),
291                        other => return Err(format!("unknown param tag: {other}")),
292                    };
293                    params.push(p);
294                }
295                Ok(Message::QueryWithParams { query, params })
296            }
297            MSG_RESULT_ROWS => {
298                let mut pos = 0;
299                if pos + 2 > payload.len() {
300                    return Err("truncated column count".into());
301                }
302                let col_bytes: [u8; 2] = payload[pos..pos + 2]
303                    .try_into()
304                    .map_err(|_| "invalid column count bytes".to_string())?;
305                let col_count = u16::from_le_bytes(col_bytes) as usize;
306                pos += 2;
307                if col_count > MAX_COLUMNS {
308                    return Err("too many columns".into());
309                }
310                let mut columns = Vec::with_capacity(col_count);
311                for _ in 0..col_count {
312                    columns.push(decode_string(payload, &mut pos)?);
313                }
314                if pos + 4 > payload.len() {
315                    return Err("truncated row count".into());
316                }
317                let row_bytes: [u8; 4] = payload[pos..pos + 4]
318                    .try_into()
319                    .map_err(|_| "invalid row count bytes".to_string())?;
320                let row_count = u32::from_le_bytes(row_bytes) as usize;
321                pos += 4;
322                if row_count > MAX_ROWS {
323                    return Err("too many rows".into());
324                }
325                let mut rows = Vec::with_capacity(row_count);
326                for _ in 0..row_count {
327                    let mut row = Vec::with_capacity(col_count);
328                    for _ in 0..col_count {
329                        row.push(decode_string(payload, &mut pos)?);
330                    }
331                    rows.push(row);
332                }
333                Ok(Message::ResultRows { columns, rows })
334            }
335            MSG_RESULT_SCALAR => {
336                let value = decode_string(payload, &mut 0)?;
337                Ok(Message::ResultScalar { value })
338            }
339            MSG_RESULT_OK => {
340                if payload.len() < 8 {
341                    return Err("truncated result ok payload".into());
342                }
343                let aff_bytes: [u8; 8] = payload[0..8]
344                    .try_into()
345                    .map_err(|_| "invalid affected count bytes".to_string())?;
346                let affected = u64::from_le_bytes(aff_bytes);
347                Ok(Message::ResultOk { affected })
348            }
349            MSG_RESULT_MSG => {
350                let message = decode_string(payload, &mut 0)?;
351                Ok(Message::ResultMessage { message })
352            }
353            MSG_ERROR => {
354                let message = decode_string(payload, &mut 0)?;
355                Ok(Message::Error { message })
356            }
357            MSG_DISCONNECT => Ok(Message::Disconnect),
358            MSG_PING => Ok(Message::Ping),
359            MSG_PONG => Ok(Message::Pong),
360            _ => Err(format!("unknown message type: {msg_type:#x}")),
361        }
362    }
363
364    /// Write this message to an async writer.
365    pub async fn write_to<W: AsyncWriteExt + Unpin>(&self, writer: &mut W) -> std::io::Result<()> {
366        let bytes = self.encode();
367        writer.write_all(&bytes).await
368    }
369
370    /// Read a message from an async reader.
371    pub async fn read_from<R: AsyncReadExt + Unpin>(
372        reader: &mut R,
373    ) -> std::io::Result<Option<Message>> {
374        Self::read_from_with_limit(reader, MAX_PAYLOAD_SIZE).await
375    }
376
377    /// Read a pre-auth message with a smaller payload limit (4 KB).
378    /// Use this before authentication is complete to prevent oversized
379    /// CONNECT payloads from consuming server memory.
380    pub async fn read_from_preauth<R: AsyncReadExt + Unpin>(
381        reader: &mut R,
382    ) -> std::io::Result<Option<Message>> {
383        Self::read_from_with_limit(reader, MAX_CONNECT_PAYLOAD_SIZE).await
384    }
385
386    /// Read a message from an async reader with a configurable payload limit.
387    async fn read_from_with_limit<R: AsyncReadExt + Unpin>(
388        reader: &mut R,
389        max_payload: usize,
390    ) -> std::io::Result<Option<Message>> {
391        let mut header = [0u8; 6];
392        match reader.read_exact(&mut header).await {
393            Ok(_) => {}
394            Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => return Ok(None),
395            Err(e) => return Err(e),
396        }
397        let len_bytes: [u8; 4] = header[2..6].try_into().map_err(|_| {
398            std::io::Error::new(
399                std::io::ErrorKind::InvalidData,
400                "invalid header length field",
401            )
402        })?;
403        let payload_len = u32::from_le_bytes(len_bytes) as usize;
404        if payload_len > max_payload {
405            return Err(std::io::Error::new(
406                std::io::ErrorKind::InvalidData,
407                format!("payload too large: {payload_len} bytes (max {max_payload})"),
408            ));
409        }
410        let mut payload = vec![0u8; payload_len];
411        if payload_len > 0 {
412            reader.read_exact(&mut payload).await?;
413        }
414
415        let mut full = Vec::with_capacity(6 + payload_len);
416        full.extend_from_slice(&header);
417        full.extend_from_slice(&payload);
418
419        Message::decode(&full)
420            .map(Some)
421            .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))
422    }
423}
424
425fn encode_string(s: &str) -> Vec<u8> {
426    let mut buf = Vec::with_capacity(4 + s.len());
427    buf.extend_from_slice(&(s.len() as u32).to_le_bytes());
428    buf.extend_from_slice(s.as_bytes());
429    buf
430}
431
432fn decode_string(data: &[u8], pos: &mut usize) -> Result<String, String> {
433    if *pos + 4 > data.len() {
434        return Err("truncated string length".into());
435    }
436    let len_bytes: [u8; 4] = data[*pos..*pos + 4]
437        .try_into()
438        .map_err(|_| "invalid string length bytes".to_string())?;
439    let len = u32::from_le_bytes(len_bytes) as usize;
440    *pos += 4;
441    if *pos + len > data.len() {
442        return Err("truncated string data".into());
443    }
444    let s = String::from_utf8_lossy(&data[*pos..*pos + len]).into_owned();
445    *pos += len;
446    Ok(s)
447}
448
449#[cfg(test)]
450mod tests {
451    use super::*;
452
453    #[test]
454    fn test_encode_decode_query() {
455        let msg = Message::Query {
456            query: "User filter .age > 30".into(),
457        };
458        let bytes = msg.encode();
459        let decoded = Message::decode(&bytes).unwrap();
460        match decoded {
461            Message::Query { query } => assert_eq!(query, "User filter .age > 30"),
462            _ => panic!("expected Query"),
463        }
464    }
465
466    #[test]
467    fn test_encode_decode_connect_with_username() {
468        let msg = Message::Connect {
469            db_name: "mydb".into(),
470            password: Some(Zeroizing::new("secret".into())),
471            username: Some("alice".into()),
472        };
473        let bytes = msg.encode();
474        match Message::decode(&bytes).unwrap() {
475            Message::Connect {
476                db_name,
477                password,
478                username,
479            } => {
480                assert_eq!(db_name, "mydb");
481                assert_eq!(password.as_deref().map(|s| s.as_str()), Some("secret"));
482                assert_eq!(username.as_deref(), Some("alice"));
483            }
484            other => panic!("expected Connect, got {other:?}"),
485        }
486    }
487
488    #[test]
489    fn test_encode_decode_connect_without_username() {
490        // New-format Connect that explicitly carries no username.
491        let msg = Message::Connect {
492            db_name: "mydb".into(),
493            password: Some(Zeroizing::new("secret".into())),
494            username: None,
495        };
496        let bytes = msg.encode();
497        match Message::decode(&bytes).unwrap() {
498            Message::Connect {
499                db_name,
500                password,
501                username,
502            } => {
503                assert_eq!(db_name, "mydb");
504                assert_eq!(password.as_deref().map(|s| s.as_str()), Some("secret"));
505                assert_eq!(username, None);
506            }
507            other => panic!("expected Connect, got {other:?}"),
508        }
509    }
510
511    #[test]
512    fn test_decode_old_client_connect_db_and_password_only() {
513        // Simulate an OLD client frame: db_name + password, with NO username
514        // bytes at all. Must decode with username: None (backward compat).
515        let mut payload = encode_string("mydb");
516        payload.extend_from_slice(&encode_string("pw"));
517        let mut frame = vec![MSG_CONNECT, 0];
518        frame.extend_from_slice(&(payload.len() as u32).to_le_bytes());
519        frame.extend_from_slice(&payload);
520        match Message::decode(&frame).unwrap() {
521            Message::Connect {
522                db_name,
523                password,
524                username,
525            } => {
526                assert_eq!(db_name, "mydb");
527                assert_eq!(password.as_deref().map(|s| s.as_str()), Some("pw"));
528                assert_eq!(username, None, "old-client frame must yield username=None");
529            }
530            other => panic!("expected Connect, got {other:?}"),
531        }
532    }
533
534    #[test]
535    fn test_decode_old_client_connect_db_only() {
536        // Oldest client: db_name only, no password and no username bytes.
537        let payload = encode_string("mydb");
538        let mut frame = vec![MSG_CONNECT, 0];
539        frame.extend_from_slice(&(payload.len() as u32).to_le_bytes());
540        frame.extend_from_slice(&payload);
541        match Message::decode(&frame).unwrap() {
542            Message::Connect {
543                db_name,
544                password,
545                username,
546            } => {
547                assert_eq!(db_name, "mydb");
548                assert_eq!(password, None);
549                assert_eq!(username, None);
550            }
551            other => panic!("expected Connect, got {other:?}"),
552        }
553    }
554
555    #[test]
556    fn test_encode_decode_result_rows() {
557        let msg = Message::ResultRows {
558            columns: vec!["name".into(), "age".into()],
559            rows: vec![
560                vec!["Alice".into(), "30".into()],
561                vec!["Bob".into(), "25".into()],
562            ],
563        };
564        let bytes = msg.encode();
565        let decoded = Message::decode(&bytes).unwrap();
566        match decoded {
567            Message::ResultRows { columns, rows } => {
568                assert_eq!(columns, vec!["name", "age"]);
569                assert_eq!(rows.len(), 2);
570            }
571            _ => panic!("expected ResultRows"),
572        }
573    }
574
575    #[test]
576    fn test_encode_decode_result_message() {
577        let msg = Message::ResultMessage {
578            message: "type User created".into(),
579        };
580        let bytes = msg.encode();
581        let decoded = Message::decode(&bytes).unwrap();
582        match decoded {
583            Message::ResultMessage { message } => assert_eq!(message, "type User created"),
584            _ => panic!("expected ResultMessage"),
585        }
586    }
587
588    #[test]
589    fn test_encode_decode_error() {
590        let msg = Message::Error {
591            message: "table not found".into(),
592        };
593        let bytes = msg.encode();
594        let decoded = Message::decode(&bytes).unwrap();
595        match decoded {
596            Message::Error { message } => assert_eq!(message, "table not found"),
597            _ => panic!("expected Error"),
598        }
599    }
600
601    #[test]
602    fn test_encode_decode_query_with_params() {
603        let msg = Message::QueryWithParams {
604            query: "insert User { name := $1, age := $2, ok := $3, note := $4 }".into(),
605            params: vec![
606                WireParam::Str(r#"a"b\c; drop User"#.into()),
607                WireParam::Int(-7),
608                WireParam::Bool(true),
609                WireParam::Null,
610            ],
611        };
612        let bytes = msg.encode();
613        // The new frame must use the dedicated 0x04 tag.
614        assert_eq!(bytes[0], 0x04);
615        match Message::decode(&bytes).unwrap() {
616            Message::QueryWithParams { query, params } => {
617                assert!(query.contains("$1"));
618                assert_eq!(params.len(), 4);
619                assert!(matches!(&params[0], WireParam::Str(s) if s == r#"a"b\c; drop User"#));
620                assert!(matches!(&params[1], WireParam::Int(-7)));
621                assert!(matches!(&params[2], WireParam::Bool(true)));
622                assert!(matches!(&params[3], WireParam::Null));
623            }
624            other => panic!("expected QueryWithParams, got {other:?}"),
625        }
626    }
627
628    #[test]
629    fn test_query_with_params_float_round_trip() {
630        let msg = Message::QueryWithParams {
631            query: "T filter .f = $1".into(),
632            params: vec![WireParam::Float(2.5)],
633        };
634        match Message::decode(&msg.encode()).unwrap() {
635            Message::QueryWithParams { params, .. } => {
636                assert!(matches!(&params[0], WireParam::Float(f) if (*f - 2.5).abs() < 1e-12));
637            }
638            other => panic!("expected QueryWithParams, got {other:?}"),
639        }
640    }
641
642    #[test]
643    fn test_decode_garbage_never_panics() {
644        // Feed a wide range of malformed/truncated byte sequences to the
645        // wire-protocol decode path. Every one must return Err, never panic:
646        // a malformed client message must not crash the server.
647        let cases: Vec<Vec<u8>> = vec![
648            vec![],                                   // empty
649            vec![0x03],                               // 1 byte, shorter than header
650            vec![0x03, 0x00, 0x00, 0x00, 0x00],       // 5 bytes, header truncated by one
651            vec![0xFF, 0x00, 0x00, 0x00, 0x00, 0x00], // unknown message type
652            // QUERY with payload_len far exceeding the frame.
653            vec![0x03, 0x00, 0xFF, 0xFF, 0xFF, 0xFF],
654            // CONNECT claiming a string len of 0xFFFFFFFF but no data.
655            vec![0x01, 0x00, 0x04, 0x00, 0x00, 0x00, 0xFF, 0xFF, 0xFF, 0xFF],
656            // RESULT_ROWS claiming a huge column count with no data.
657            vec![0x07, 0x00, 0x02, 0x00, 0x00, 0x00, 0xFF, 0xFF],
658            // RESULT_OK with a truncated 8-byte affected field.
659            vec![0x09, 0x00, 0x03, 0x00, 0x00, 0x00, 0x01, 0x02, 0x03],
660            // QUERY_PARAMS (0x04) claiming a query string len with no data.
661            vec![0x04, 0x00, 0x04, 0x00, 0x00, 0x00, 0xFF, 0xFF, 0xFF, 0xFF],
662            // QUERY_PARAMS: empty query string, claims 1 param, no param bytes.
663            vec![
664                0x04, 0x00, 0x06, 0x00, 0x00, 0x00, // header, payload_len=6
665                0x00, 0x00, 0x00, 0x00, // query string len = 0
666                0x01, 0x00, // param count = 1, then nothing
667            ],
668            // QUERY_PARAMS: 1 int param with a truncated i64 body.
669            vec![
670                0x04, 0x00, 0x0B, 0x00, 0x00, 0x00, // header, payload_len=11
671                0x00, 0x00, 0x00, 0x00, // query len = 0
672                0x01, 0x00, // param count = 1
673                0x01, // tag = int, then only 3 of 8 bytes
674                0x01, 0x02, 0x03,
675            ],
676            // QUERY_PARAMS: 1 str param with a truncated string body.
677            vec![
678                0x04, 0x00, 0x0F, 0x00, 0x00, 0x00, // header, payload_len=15
679                0x00, 0x00, 0x00, 0x00, // query len = 0
680                0x01, 0x00, // param count = 1
681                0x04, // tag = str
682                0xFF, 0xFF, 0xFF, 0xFF, // str len huge, no data
683            ],
684            // QUERY_PARAMS: unknown param tag byte.
685            vec![
686                0x04, 0x00, 0x0B, 0x00, 0x00, 0x00, // header, payload_len=11
687                0x00, 0x00, 0x00, 0x00, // query len = 0
688                0x01, 0x00, // param count = 1
689                0x63, // bogus tag
690            ],
691        ];
692        for bytes in cases {
693            let result = Message::decode(&bytes);
694            assert!(
695                result.is_err(),
696                "expected Err for malformed input {bytes:?}, got {result:?}"
697            );
698        }
699    }
700
701    #[test]
702    fn test_decode_arbitrary_prefixes_never_panic() {
703        // Exhaustively walk every truncation of a well-formed frame plus a
704        // few byte mutations. None may panic.
705        let valid = Message::ResultRows {
706            columns: vec!["a".into(), "b".into()],
707            rows: vec![vec!["1".into(), "2".into()]],
708        }
709        .encode();
710        for end in 0..valid.len() {
711            // Every prefix that is not the full frame must be rejected, not panic.
712            let _ = Message::decode(&valid[..end]);
713        }
714        // Mutate each byte to a high value and confirm no panic.
715        for i in 0..valid.len() {
716            let mut m = valid.clone();
717            m[i] = 0xFF;
718            let _ = Message::decode(&m);
719        }
720    }
721
722    #[test]
723    fn test_frame_length() {
724        let msg = Message::Query {
725            query: "User".into(),
726        };
727        let bytes = msg.encode();
728        assert!(bytes.len() >= 6);
729        let payload_len = u32::from_le_bytes(bytes[2..6].try_into().unwrap()) as usize;
730        assert_eq!(bytes.len(), 6 + payload_len);
731    }
732}