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;
7const 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
20const MAX_PAYLOAD_SIZE: usize = 64 * 1024 * 1024;
22
23const MAX_CONNECT_PAYLOAD_SIZE: usize = 4096;
26
27const MAX_COLUMNS: usize = 4096;
29
30const MAX_ROWS: usize = 10_000_000;
32
33const MAX_PARAMS: usize = 4096;
35
36#[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 password: Option<Zeroizing<String>>,
59 username: Option<String>,
63 },
64 ConnectOk {
65 version: String,
66 },
67 Query {
68 query: String,
69 },
70 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 ResultMessage {
87 message: String,
88 },
89 Error {
90 message: String,
91 },
92 Disconnect,
93 Ping,
94 Pong,
95}
96
97impl Message {
98 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 match password {
109 Some(p) => buf.extend_from_slice(&encode_string(p)),
110 None => buf.extend_from_slice(&0u32.to_le_bytes()),
111 }
112 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); frame.extend_from_slice(&(payload.len() as u32).to_le_bytes());
175 frame.extend_from_slice(&payload);
176 frame
177 }
178
179 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 let password = if pos < payload.len() {
202 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 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 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 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 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 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 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 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 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 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!(¶ms[0], WireParam::Str(s) if s == r#"a"b\c; drop User"#));
620 assert!(matches!(¶ms[1], WireParam::Int(-7)));
621 assert!(matches!(¶ms[2], WireParam::Bool(true)));
622 assert!(matches!(¶ms[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!(¶ms[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 let cases: Vec<Vec<u8>> = vec![
648 vec![], vec![0x03], vec![0x03, 0x00, 0x00, 0x00, 0x00], vec![0xFF, 0x00, 0x00, 0x00, 0x00, 0x00], vec![0x03, 0x00, 0xFF, 0xFF, 0xFF, 0xFF],
654 vec![0x01, 0x00, 0x04, 0x00, 0x00, 0x00, 0xFF, 0xFF, 0xFF, 0xFF],
656 vec![0x07, 0x00, 0x02, 0x00, 0x00, 0x00, 0xFF, 0xFF],
658 vec![0x09, 0x00, 0x03, 0x00, 0x00, 0x00, 0x01, 0x02, 0x03],
660 vec![0x04, 0x00, 0x04, 0x00, 0x00, 0x00, 0xFF, 0xFF, 0xFF, 0xFF],
662 vec![
664 0x04, 0x00, 0x06, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, ],
668 vec![
670 0x04, 0x00, 0x0B, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x01, 0x01, 0x02, 0x03,
675 ],
676 vec![
678 0x04, 0x00, 0x0F, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x04, 0xFF, 0xFF, 0xFF, 0xFF, ],
684 vec![
686 0x04, 0x00, 0x0B, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x63, ],
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 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 let _ = Message::decode(&valid[..end]);
713 }
714 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}