1use powdb_storage::pj1::pj1_validate;
2use powdb_storage::types::{TypeId, Value};
3use tokio::io::{AsyncReadExt, AsyncWriteExt};
4use zeroize::Zeroizing;
5
6const MSG_CONNECT: u8 = 0x01;
7const MSG_CONNECT_OK: u8 = 0x02;
8const MSG_QUERY: u8 = 0x03;
9const MSG_QUERY_PARAMS: u8 = 0x04;
13const MSG_QUERY_SQL: u8 = 0x05;
15const MSG_SYNC_STATUS: u8 = 0x20;
18const MSG_SYNC_PULL: u8 = 0x21;
19const MSG_SYNC_ACK: u8 = 0x22;
20const MSG_SYNC_STATUS_RESULT: u8 = 0x23;
21const MSG_SYNC_PULL_RESULT: u8 = 0x24;
22const MSG_SYNC_ACK_RESULT: u8 = 0x25;
23const MSG_RESULT_ROWS: u8 = 0x07;
24const MSG_RESULT_SCALAR: u8 = 0x08;
25const MSG_RESULT_OK: u8 = 0x09;
26const MSG_ERROR: u8 = 0x0A;
27const MSG_RESULT_MSG: u8 = 0x0B;
28const MSG_DISCONNECT: u8 = 0x10;
29const MSG_PING: u8 = 0x11;
30const MSG_PONG: u8 = 0x12;
31pub const MSG_QUERY_NATIVE: u8 = 0x13;
34pub const MSG_QUERY_PARAMS_NATIVE: u8 = 0x14;
36pub const MSG_QUERY_SQL_NATIVE: u8 = 0x15;
38pub const MSG_RESULT_ROWS_NATIVE: u8 = 0x16;
40pub const MSG_RESULT_SCALAR_NATIVE: u8 = 0x17;
42
43const MAX_PAYLOAD_SIZE: usize = 64 * 1024 * 1024;
45
46const MAX_CONNECT_PAYLOAD_SIZE: usize = 4096;
49
50const MAX_COLUMNS: usize = 4096;
52
53const MAX_ROWS: usize = 10_000_000;
55
56const MAX_PARAMS: usize = 4096;
58
59const MAX_SYNC_UNITS: usize = 4096;
61
62const STRING_LEN_PREFIX: usize = 4; #[derive(Debug, Clone, Copy, PartialEq, Eq)]
81#[repr(u8)]
82pub enum ErrorClass {
83 Internal = 0,
86 Parse = 1,
88 Execution = 2,
91 Timeout = 3,
94 LimitExceeded = 4,
97 ReadonlyRefused = 5,
100 AuthFailed = 6,
102 RateLimited = 7,
104 ConstraintViolation = 8,
106 Cancelled = 9,
108}
109
110impl ErrorClass {
111 pub fn as_u8(self) -> u8 {
113 self as u8
114 }
115
116 pub fn from_u8(raw: u8) -> Option<ErrorClass> {
119 Some(match raw {
120 0 => ErrorClass::Internal,
121 1 => ErrorClass::Parse,
122 2 => ErrorClass::Execution,
123 3 => ErrorClass::Timeout,
124 4 => ErrorClass::LimitExceeded,
125 5 => ErrorClass::ReadonlyRefused,
126 6 => ErrorClass::AuthFailed,
127 7 => ErrorClass::RateLimited,
128 8 => ErrorClass::ConstraintViolation,
129 9 => ErrorClass::Cancelled,
130 _ => return None,
131 })
132 }
133}
134
135pub fn decode_error_class(frame: &[u8]) -> Option<u8> {
139 if frame.len() < 6 || frame[0] != MSG_ERROR {
140 return None;
141 }
142 let payload_len = u32::from_le_bytes(frame[2..6].try_into().ok()?) as usize;
143 let payload = frame.get(6..6 + payload_len)?;
144 let string_len = u32::from_le_bytes(payload.get(..4)?.try_into().ok()?) as usize;
145 if payload.len() == 4 + string_len + 1 {
147 Some(payload[4 + string_len])
148 } else {
149 None
150 }
151}
152
153#[derive(Debug, Clone, PartialEq)]
159pub enum WireParam {
160 Null,
161 Int(i64),
162 Float(f64),
163 Bool(bool),
164 Str(String),
165}
166
167#[derive(Debug, Clone, Copy, PartialEq, Eq)]
168pub enum WireSyncRepairAction {
169 None,
170 Pull,
171 AwaitArchive,
172 Rebootstrap,
173}
174
175#[derive(Debug, Clone, PartialEq, Eq)]
176pub struct WireSyncStatus {
177 pub replica_id: String,
178 pub active: bool,
179 pub last_applied_lsn: Option<u64>,
180 pub remote_lsn: u64,
181 pub servable_lsn: Option<u64>,
182 pub unarchived_lsn: Option<u64>,
183 pub lag_lsn: Option<u64>,
184 pub lag_bytes: Option<u64>,
185 pub lag_ms: Option<u64>,
186 pub stale: bool,
187 pub repair_action: WireSyncRepairAction,
188 pub last_sync_error: Option<String>,
189}
190
191#[derive(Debug, Clone, PartialEq, Eq)]
192pub struct WireRetainedUnit {
193 pub tx_id: u64,
194 pub record_type: u8,
195 pub lsn: u64,
196 pub data: Vec<u8>,
197}
198
199impl WireRetainedUnit {
200 pub fn encoded_len(&self) -> Result<u64, String> {
204 let data_len = u64::try_from(self.data.len())
205 .map_err(|_| "sync retained unit payload too large".to_string())?;
206 8u64.checked_add(1)
207 .and_then(|n| n.checked_add(8))
208 .and_then(|n| n.checked_add(4))
209 .and_then(|n| n.checked_add(data_len))
210 .ok_or_else(|| "sync retained unit encoded length overflow".to_string())
211 }
212}
213
214#[derive(Debug, Clone)]
215pub enum Message {
216 Connect {
217 db_name: String,
218 password: Option<Zeroizing<String>>,
223 username: Option<String>,
227 },
228 ConnectOk {
229 version: String,
230 },
231 Query {
232 query: String,
233 },
234 QuerySql {
236 query: String,
237 },
238 QueryWithParams {
240 query: String,
241 params: Vec<WireParam>,
242 },
243 QueryNative {
245 query: String,
246 },
247 QueryWithParamsNative {
249 query: String,
250 params: Vec<WireParam>,
251 },
252 QuerySqlNative {
254 query: String,
255 },
256 SyncStatus {
258 replica_id: String,
259 },
260 SyncPull {
262 replica_id: String,
263 since_lsn: u64,
264 max_units: u32,
265 max_bytes: u64,
266 database_id: [u8; 16],
267 primary_generation: u64,
268 wal_format_version: u16,
269 catalog_version: u16,
270 segment_format_version: u16,
271 },
272 SyncAck {
274 replica_id: String,
275 applied_lsn: u64,
276 remote_lsn: u64,
277 },
278 SyncStatusResult {
279 status: WireSyncStatus,
280 },
281 SyncPullResult {
282 status: WireSyncStatus,
283 units: Vec<WireRetainedUnit>,
284 has_more: bool,
285 },
286 SyncAckResult {
287 previous_applied_lsn: u64,
288 applied_lsn: u64,
289 remote_lsn: u64,
290 advanced: bool,
291 status: WireSyncStatus,
292 },
293 ResultRows {
294 columns: Vec<String>,
295 rows: Vec<Vec<String>>,
296 },
297 ResultScalar {
298 value: String,
299 },
300 ResultRowsNative {
301 columns: Vec<String>,
302 rows: Vec<Vec<Value>>,
303 },
304 ResultScalarNative {
305 value: Value,
306 },
307 ResultOk {
308 affected: u64,
309 },
310 ResultMessage {
312 message: String,
313 },
314 Error {
315 message: String,
316 },
317 ErrorWithClass {
323 message: String,
324 class: ErrorClass,
325 },
326 Disconnect,
327 Ping,
328 Pong,
329}
330
331impl Message {
332 pub fn encode(&self) -> Vec<u8> {
334 let (msg_type, payload) = match self {
335 Message::Connect {
336 db_name,
337 password,
338 username,
339 } => {
340 let mut buf = encode_string(db_name);
341 match password {
343 Some(p) => buf.extend_from_slice(&encode_string(p)),
344 None => buf.extend_from_slice(&0u32.to_le_bytes()),
345 }
346 match username {
349 Some(u) => buf.extend_from_slice(&encode_string(u)),
350 None => buf.extend_from_slice(&0u32.to_le_bytes()),
351 }
352 (MSG_CONNECT, buf)
353 }
354 Message::ConnectOk { version } => (MSG_CONNECT_OK, encode_string(version)),
355 Message::Query { query } => (MSG_QUERY, encode_string(query)),
356 Message::QuerySql { query } => (MSG_QUERY_SQL, encode_string(query)),
357 Message::QueryWithParams { query, params } => {
358 (MSG_QUERY_PARAMS, encode_query_with_params(query, params))
359 }
360 Message::QueryNative { query } => (MSG_QUERY_NATIVE, encode_string(query)),
361 Message::QueryWithParamsNative { query, params } => (
362 MSG_QUERY_PARAMS_NATIVE,
363 encode_query_with_params(query, params),
364 ),
365 Message::QuerySqlNative { query } => (MSG_QUERY_SQL_NATIVE, encode_string(query)),
366 Message::SyncStatus { replica_id } => (MSG_SYNC_STATUS, encode_string(replica_id)),
367 Message::SyncPull {
368 replica_id,
369 since_lsn,
370 max_units,
371 max_bytes,
372 database_id,
373 primary_generation,
374 wal_format_version,
375 catalog_version,
376 segment_format_version,
377 } => {
378 let mut buf = encode_string(replica_id);
379 buf.extend_from_slice(&since_lsn.to_le_bytes());
380 buf.extend_from_slice(&max_units.to_le_bytes());
381 buf.extend_from_slice(&max_bytes.to_le_bytes());
382 buf.extend_from_slice(database_id);
383 buf.extend_from_slice(&primary_generation.to_le_bytes());
384 buf.extend_from_slice(&wal_format_version.to_le_bytes());
385 buf.extend_from_slice(&catalog_version.to_le_bytes());
386 buf.extend_from_slice(&segment_format_version.to_le_bytes());
387 (MSG_SYNC_PULL, buf)
388 }
389 Message::SyncAck {
390 replica_id,
391 applied_lsn,
392 remote_lsn,
393 } => {
394 let mut buf = encode_string(replica_id);
395 buf.extend_from_slice(&applied_lsn.to_le_bytes());
396 buf.extend_from_slice(&remote_lsn.to_le_bytes());
397 (MSG_SYNC_ACK, buf)
398 }
399 Message::SyncStatusResult { status } => {
400 (MSG_SYNC_STATUS_RESULT, encode_sync_status(status))
401 }
402 Message::SyncPullResult {
403 status,
404 units,
405 has_more,
406 } => {
407 let mut buf = encode_sync_status(status);
408 buf.extend_from_slice(&(units.len() as u32).to_le_bytes());
409 for unit in units {
410 encode_retained_unit(&mut buf, unit);
411 }
412 buf.push(u8::from(*has_more));
413 (MSG_SYNC_PULL_RESULT, buf)
414 }
415 Message::SyncAckResult {
416 previous_applied_lsn,
417 applied_lsn,
418 remote_lsn,
419 advanced,
420 status,
421 } => {
422 let mut buf = Vec::new();
423 buf.extend_from_slice(&previous_applied_lsn.to_le_bytes());
424 buf.extend_from_slice(&applied_lsn.to_le_bytes());
425 buf.extend_from_slice(&remote_lsn.to_le_bytes());
426 buf.push(u8::from(*advanced));
427 buf.extend_from_slice(&encode_sync_status(status));
428 (MSG_SYNC_ACK_RESULT, buf)
429 }
430 Message::ResultRows { columns, rows } => {
431 let mut buf = Vec::new();
432 buf.extend_from_slice(&(columns.len() as u16).to_le_bytes());
433 for col in columns {
434 buf.extend_from_slice(&encode_string(col));
435 }
436 buf.extend_from_slice(&(rows.len() as u32).to_le_bytes());
437 for row in rows {
438 for val in row {
439 buf.extend_from_slice(&encode_string(val));
440 }
441 }
442 (MSG_RESULT_ROWS, buf)
443 }
444 Message::ResultScalar { value } => (MSG_RESULT_SCALAR, encode_string(value)),
445 Message::ResultRowsNative { columns, rows } => {
446 let mut buf = Vec::new();
447 buf.extend_from_slice(&(columns.len() as u16).to_le_bytes());
448 for column in columns {
449 buf.extend_from_slice(&encode_string(column));
450 }
451 buf.extend_from_slice(&(rows.len() as u32).to_le_bytes());
452 for row in rows {
453 for value in row {
454 encode_typed_value(&mut buf, value);
455 }
456 }
457 (MSG_RESULT_ROWS_NATIVE, buf)
458 }
459 Message::ResultScalarNative { value } => {
460 let mut buf = Vec::new();
461 encode_typed_value(&mut buf, value);
462 (MSG_RESULT_SCALAR_NATIVE, buf)
463 }
464 Message::ResultOk { affected } => (MSG_RESULT_OK, affected.to_le_bytes().to_vec()),
465 Message::ResultMessage { message } => (MSG_RESULT_MSG, encode_string(message)),
466 Message::Error { message } => (MSG_ERROR, encode_string(message)),
467 Message::ErrorWithClass { message, class } => {
468 let mut buf = encode_string(message);
469 buf.push(class.as_u8());
470 (MSG_ERROR, buf)
471 }
472 Message::Disconnect => (MSG_DISCONNECT, Vec::new()),
473 Message::Ping => (MSG_PING, Vec::new()),
474 Message::Pong => (MSG_PONG, Vec::new()),
475 };
476
477 let mut frame = Vec::with_capacity(6 + payload.len());
478 frame.push(msg_type);
479 frame.push(0); frame.extend_from_slice(&(payload.len() as u32).to_le_bytes());
481 frame.extend_from_slice(&payload);
482 frame
483 }
484
485 pub fn decode(data: &[u8]) -> Result<Message, String> {
487 if data.len() < 6 {
488 return Err("frame too short".into());
489 }
490 let msg_type = data[0];
491 let _flags = data[1];
492 let len_bytes: [u8; 4] = data[2..6]
493 .try_into()
494 .map_err(|_| "invalid header length field".to_string())?;
495 let payload_len = u32::from_le_bytes(len_bytes) as usize;
496 if 6 + payload_len > data.len() {
497 return Err("payload length exceeds frame".into());
498 }
499 let payload = &data[6..6 + payload_len];
500
501 match msg_type {
502 MSG_CONNECT => {
503 let mut pos = 0;
504 let db_name = decode_string(payload, &mut pos)?;
505 let password = if pos < payload.len() {
508 let p = Zeroizing::new(decode_string(payload, &mut pos)?);
511 if p.is_empty() {
512 None
513 } else {
514 Some(p)
515 }
516 } else {
517 None
518 };
519 let username = if pos < payload.len() {
522 let u = decode_string(payload, &mut pos)?;
523 if u.is_empty() {
524 None
525 } else {
526 Some(u)
527 }
528 } else {
529 None
530 };
531 Ok(Message::Connect {
532 db_name,
533 password,
534 username,
535 })
536 }
537 MSG_CONNECT_OK => {
538 let version = decode_string(payload, &mut 0)?;
539 Ok(Message::ConnectOk { version })
540 }
541 MSG_QUERY => {
542 let query = decode_string(payload, &mut 0)?;
543 Ok(Message::Query { query })
544 }
545 MSG_QUERY_SQL => {
546 let query = decode_string(payload, &mut 0)?;
547 Ok(Message::QuerySql { query })
548 }
549 MSG_QUERY_PARAMS => {
550 let mut pos = 0;
551 let query = decode_string(payload, &mut pos)?;
552 if pos + 2 > payload.len() {
553 return Err("truncated param count".into());
554 }
555 let count_bytes: [u8; 2] = payload[pos..pos + 2]
556 .try_into()
557 .map_err(|_| "invalid param count bytes".to_string())?;
558 let count = u16::from_le_bytes(count_bytes) as usize;
559 pos += 2;
560 if count > MAX_PARAMS {
561 return Err("too many parameters".into());
562 }
563 let mut params = Vec::with_capacity(count.min(payload.len() - pos));
564 for _ in 0..count {
565 if pos >= payload.len() {
566 return Err("truncated param tag".into());
567 }
568 let tag = payload[pos];
569 pos += 1;
570 let p = match tag {
571 0 => WireParam::Null,
572 1 => {
573 if pos + 8 > payload.len() {
574 return Err("truncated int param".into());
575 }
576 let b: [u8; 8] = payload[pos..pos + 8]
577 .try_into()
578 .map_err(|_| "invalid int param bytes".to_string())?;
579 pos += 8;
580 WireParam::Int(i64::from_le_bytes(b))
581 }
582 2 => {
583 if pos + 8 > payload.len() {
584 return Err("truncated float param".into());
585 }
586 let b: [u8; 8] = payload[pos..pos + 8]
587 .try_into()
588 .map_err(|_| "invalid float param bytes".to_string())?;
589 pos += 8;
590 WireParam::Float(f64::from_le_bytes(b))
591 }
592 3 => {
593 if pos + 1 > payload.len() {
594 return Err("truncated bool param".into());
595 }
596 let v = payload[pos] != 0;
597 pos += 1;
598 WireParam::Bool(v)
599 }
600 4 => WireParam::Str(decode_string(payload, &mut pos)?),
601 other => return Err(format!("unknown param tag: {other}")),
602 };
603 params.push(p);
604 }
605 Ok(Message::QueryWithParams { query, params })
606 }
607 MSG_QUERY_NATIVE => {
608 let query = decode_exact_string(payload, "native PowQL query")?;
609 Ok(Message::QueryNative { query })
610 }
611 MSG_QUERY_PARAMS_NATIVE => {
612 let (query, params) = decode_query_with_params_exact(payload)?;
613 Ok(Message::QueryWithParamsNative { query, params })
614 }
615 MSG_QUERY_SQL_NATIVE => {
616 let query = decode_exact_string(payload, "native SQL query")?;
617 Ok(Message::QuerySqlNative { query })
618 }
619 MSG_SYNC_STATUS => {
620 let replica_id = decode_string(payload, &mut 0)?;
621 Ok(Message::SyncStatus { replica_id })
622 }
623 MSG_SYNC_PULL => {
624 let mut pos = 0;
625 let replica_id = decode_string(payload, &mut pos)?;
626 let since_lsn = decode_u64(payload, &mut pos, "sync pull since LSN")?;
627 let max_units = decode_u32(payload, &mut pos, "sync pull max units")?;
628 let max_bytes = decode_u64(payload, &mut pos, "sync pull max bytes")?;
629 let database_id = decode_16_bytes(payload, &mut pos, "sync database id")?;
630 let primary_generation = decode_u64(payload, &mut pos, "sync primary generation")?;
631 let wal_format_version = decode_u16(payload, &mut pos, "sync WAL format version")?;
632 let catalog_version = decode_u16(payload, &mut pos, "sync catalog version")?;
633 let segment_format_version =
634 decode_u16(payload, &mut pos, "sync segment format version")?;
635 Ok(Message::SyncPull {
636 replica_id,
637 since_lsn,
638 max_units,
639 max_bytes,
640 database_id,
641 primary_generation,
642 wal_format_version,
643 catalog_version,
644 segment_format_version,
645 })
646 }
647 MSG_SYNC_ACK => {
648 let mut pos = 0;
649 let replica_id = decode_string(payload, &mut pos)?;
650 let applied_lsn = decode_u64(payload, &mut pos, "sync ack applied LSN")?;
651 let remote_lsn = decode_u64(payload, &mut pos, "sync ack remote LSN")?;
652 Ok(Message::SyncAck {
653 replica_id,
654 applied_lsn,
655 remote_lsn,
656 })
657 }
658 MSG_SYNC_STATUS_RESULT => {
659 let mut pos = 0;
660 let status = decode_sync_status(payload, &mut pos)?;
661 Ok(Message::SyncStatusResult { status })
662 }
663 MSG_SYNC_PULL_RESULT => {
664 let mut pos = 0;
665 let status = decode_sync_status(payload, &mut pos)?;
666 let count = decode_u32(payload, &mut pos, "sync retained unit count")? as usize;
667 if count > MAX_SYNC_UNITS {
668 return Err("too many retained units".into());
669 }
670 let mut units = Vec::with_capacity(count.min(payload.len().saturating_sub(pos)));
671 for _ in 0..count {
672 units.push(decode_retained_unit(payload, &mut pos)?);
673 }
674 let has_more = decode_bool(payload, &mut pos, "sync has_more")?;
675 Ok(Message::SyncPullResult {
676 status,
677 units,
678 has_more,
679 })
680 }
681 MSG_SYNC_ACK_RESULT => {
682 let mut pos = 0;
683 let previous_applied_lsn = decode_u64(payload, &mut pos, "previous applied LSN")?;
684 let applied_lsn = decode_u64(payload, &mut pos, "applied LSN")?;
685 let remote_lsn = decode_u64(payload, &mut pos, "remote LSN")?;
686 let advanced = decode_bool(payload, &mut pos, "sync ack advanced")?;
687 let status = decode_sync_status(payload, &mut pos)?;
688 Ok(Message::SyncAckResult {
689 previous_applied_lsn,
690 applied_lsn,
691 remote_lsn,
692 advanced,
693 status,
694 })
695 }
696 MSG_RESULT_ROWS => {
697 let mut pos = 0;
698 if pos + 2 > payload.len() {
699 return Err("truncated column count".into());
700 }
701 let col_bytes: [u8; 2] = payload[pos..pos + 2]
702 .try_into()
703 .map_err(|_| "invalid column count bytes".to_string())?;
704 let col_count = u16::from_le_bytes(col_bytes) as usize;
705 pos += 2;
706 if col_count > MAX_COLUMNS {
707 return Err("too many columns".into());
708 }
709 let mut columns =
710 Vec::with_capacity(col_count.min((payload.len() - pos) / STRING_LEN_PREFIX));
711 for _ in 0..col_count {
712 columns.push(decode_string(payload, &mut pos)?);
713 }
714 if pos + 4 > payload.len() {
715 return Err("truncated row count".into());
716 }
717 let row_bytes: [u8; 4] = payload[pos..pos + 4]
718 .try_into()
719 .map_err(|_| "invalid row count bytes".to_string())?;
720 let row_count = u32::from_le_bytes(row_bytes) as usize;
721 pos += 4;
722 if row_count > MAX_ROWS {
723 return Err("too many rows".into());
724 }
725 let max_rows = match col_count.checked_mul(STRING_LEN_PREFIX) {
731 Some(0) | None => 0,
732 Some(per_row) => (payload.len() - pos) / per_row,
733 };
734 if row_count > max_rows {
735 return Err("row count exceeds payload size".into());
736 }
737 let mut rows = Vec::with_capacity(row_count);
738 for _ in 0..row_count {
739 let mut row = Vec::with_capacity(col_count);
740 for _ in 0..col_count {
741 row.push(decode_string(payload, &mut pos)?);
742 }
743 rows.push(row);
744 }
745 Ok(Message::ResultRows { columns, rows })
746 }
747 MSG_RESULT_SCALAR => {
748 let value = decode_string(payload, &mut 0)?;
749 Ok(Message::ResultScalar { value })
750 }
751 MSG_RESULT_ROWS_NATIVE => decode_native_rows(payload),
752 MSG_RESULT_SCALAR_NATIVE => {
753 let mut pos = 0;
754 let value = decode_typed_value(payload, &mut pos)?;
755 require_payload_end(payload, pos, "native scalar")?;
756 Ok(Message::ResultScalarNative { value })
757 }
758 MSG_RESULT_OK => {
759 if payload.len() < 8 {
760 return Err("truncated result ok payload".into());
761 }
762 let aff_bytes: [u8; 8] = payload[0..8]
763 .try_into()
764 .map_err(|_| "invalid affected count bytes".to_string())?;
765 let affected = u64::from_le_bytes(aff_bytes);
766 Ok(Message::ResultOk { affected })
767 }
768 MSG_RESULT_MSG => {
769 let message = decode_string(payload, &mut 0)?;
770 Ok(Message::ResultMessage { message })
771 }
772 MSG_ERROR => {
773 let message = decode_string(payload, &mut 0)?;
774 Ok(Message::Error { message })
775 }
776 MSG_DISCONNECT => Ok(Message::Disconnect),
777 MSG_PING => Ok(Message::Ping),
778 MSG_PONG => Ok(Message::Pong),
779 _ => Err(format!("unknown message type: {msg_type:#x}")),
780 }
781 }
782
783 pub async fn write_to<W: AsyncWriteExt + Unpin>(&self, writer: &mut W) -> std::io::Result<()> {
785 let bytes = self.encode();
786 writer.write_all(&bytes).await
787 }
788
789 pub async fn read_from<R: AsyncReadExt + Unpin>(
791 reader: &mut R,
792 ) -> std::io::Result<Option<Message>> {
793 Self::read_from_with_limit(reader, MAX_PAYLOAD_SIZE).await
794 }
795
796 pub async fn read_from_preauth<R: AsyncReadExt + Unpin>(
800 reader: &mut R,
801 ) -> std::io::Result<Option<Message>> {
802 Self::read_from_with_limit(reader, MAX_CONNECT_PAYLOAD_SIZE).await
803 }
804
805 async fn read_from_with_limit<R: AsyncReadExt + Unpin>(
807 reader: &mut R,
808 max_payload: usize,
809 ) -> std::io::Result<Option<Message>> {
810 let mut header = [0u8; 6];
811 match reader.read_exact(&mut header).await {
812 Ok(_) => {}
813 Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => return Ok(None),
814 Err(e) => return Err(e),
815 }
816 let len_bytes: [u8; 4] = header[2..6].try_into().map_err(|_| {
817 std::io::Error::new(
818 std::io::ErrorKind::InvalidData,
819 "invalid header length field",
820 )
821 })?;
822 let payload_len = u32::from_le_bytes(len_bytes) as usize;
823 if payload_len > max_payload {
824 return Err(std::io::Error::new(
825 std::io::ErrorKind::InvalidData,
826 format!("payload too large: {payload_len} bytes (max {max_payload})"),
827 ));
828 }
829 let mut payload = vec![0u8; payload_len];
830 if payload_len > 0 {
831 reader.read_exact(&mut payload).await?;
832 }
833
834 let mut full = Vec::with_capacity(6 + payload_len);
835 full.extend_from_slice(&header);
836 full.extend_from_slice(&payload);
837
838 Message::decode(&full)
839 .map(Some)
840 .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))
841 }
842}
843
844fn encode_query_with_params(query: &str, params: &[WireParam]) -> Vec<u8> {
845 let mut buf = encode_string(query);
846 buf.extend_from_slice(&(params.len() as u16).to_le_bytes());
847 for param in params {
848 match param {
849 WireParam::Null => buf.push(0),
850 WireParam::Int(value) => {
851 buf.push(1);
852 buf.extend_from_slice(&value.to_le_bytes());
853 }
854 WireParam::Float(value) => {
855 buf.push(2);
856 buf.extend_from_slice(&value.to_le_bytes());
857 }
858 WireParam::Bool(value) => {
859 buf.push(3);
860 buf.push(u8::from(*value));
861 }
862 WireParam::Str(value) => {
863 buf.push(4);
864 buf.extend_from_slice(&encode_string(value));
865 }
866 }
867 }
868 buf
869}
870
871fn decode_query_with_params_exact(payload: &[u8]) -> Result<(String, Vec<WireParam>), String> {
872 let mut pos = 0;
873 let query = decode_string_strict(payload, &mut pos, "native query")?;
874 let count = decode_u16(payload, &mut pos, "native param count")? as usize;
875 if count > MAX_PARAMS {
876 return Err("too many parameters".into());
877 }
878 if count > payload.len().saturating_sub(pos) {
880 return Err("parameter count exceeds payload size".into());
881 }
882 let mut params = Vec::with_capacity(count);
883 for _ in 0..count {
884 if pos >= payload.len() {
885 return Err("truncated param tag".into());
886 }
887 let tag = payload[pos];
888 pos += 1;
889 params.push(match tag {
890 0 => WireParam::Null,
891 1 => {
892 let bytes = take_exact(payload, &mut pos, 8, "int param")?;
893 WireParam::Int(i64::from_le_bytes(bytes.try_into().expect("8 bytes")))
894 }
895 2 => {
896 let bytes = take_exact(payload, &mut pos, 8, "float param")?;
897 WireParam::Float(f64::from_le_bytes(bytes.try_into().expect("8 bytes")))
898 }
899 3 => WireParam::Bool(decode_bool(payload, &mut pos, "bool param")?),
900 4 => WireParam::Str(decode_string_strict(payload, &mut pos, "string param")?),
901 other => return Err(format!("unknown param tag: {other}")),
902 });
903 }
904 require_payload_end(payload, pos, "native parameterized query")?;
905 Ok((query, params))
906}
907
908fn encode_typed_value(out: &mut Vec<u8>, value: &Value) {
909 out.push(value.type_id() as u8);
910 let body_len = match value {
911 Value::Empty => 0,
912 Value::Int(_) | Value::Float(_) | Value::DateTime(_) => 8,
913 Value::Bool(_) => 1,
914 Value::Str(value) => value.len(),
915 Value::Uuid(_) => 16,
916 Value::Bytes(value) => value.len(),
917 Value::Json(value) => value.len(),
918 };
919 out.extend_from_slice(&(body_len as u32).to_le_bytes());
920 match value {
921 Value::Empty => {}
922 Value::Int(value) | Value::DateTime(value) => out.extend_from_slice(&value.to_le_bytes()),
923 Value::Float(value) => out.extend_from_slice(&value.to_le_bytes()),
924 Value::Bool(value) => out.push(u8::from(*value)),
925 Value::Str(value) => out.extend_from_slice(value.as_bytes()),
926 Value::Uuid(value) => out.extend_from_slice(value),
927 Value::Bytes(value) => out.extend_from_slice(value),
928 Value::Json(value) => out.extend_from_slice(value),
929 }
930}
931
932fn decode_typed_value(data: &[u8], pos: &mut usize) -> Result<Value, String> {
933 if *pos >= data.len() {
934 return Err("truncated typed value tag".into());
935 }
936 let raw_type = data[*pos];
937 *pos += 1;
938 let type_id =
939 TypeId::from_u8(raw_type).ok_or_else(|| format!("unknown typed value tag: {raw_type}"))?;
940 let body_len = decode_u32(data, pos, "typed value body length")? as usize;
941 let body = take_exact(data, pos, body_len, "typed value body")?;
942
943 let require_len = |expected: usize| {
944 if body_len == expected {
945 Ok(())
946 } else {
947 Err(format!(
948 "invalid {type_id:?} typed value length: expected {expected}, got {body_len}"
949 ))
950 }
951 };
952 match type_id {
953 TypeId::Empty => {
954 require_len(0)?;
955 Ok(Value::Empty)
956 }
957 TypeId::Int => {
958 require_len(8)?;
959 Ok(Value::Int(i64::from_le_bytes(
960 body.try_into().expect("validated 8-byte int"),
961 )))
962 }
963 TypeId::Float => {
964 require_len(8)?;
965 Ok(Value::Float(f64::from_le_bytes(
966 body.try_into().expect("validated 8-byte float"),
967 )))
968 }
969 TypeId::Bool => {
970 require_len(1)?;
971 match body[0] {
972 0 => Ok(Value::Bool(false)),
973 1 => Ok(Value::Bool(true)),
974 other => Err(format!("invalid typed boolean: {other}")),
975 }
976 }
977 TypeId::Str => Ok(Value::Str(
978 std::str::from_utf8(body)
979 .map_err(|error| format!("invalid UTF-8 in typed string: {error}"))?
980 .to_owned(),
981 )),
982 TypeId::DateTime => {
983 require_len(8)?;
984 Ok(Value::DateTime(i64::from_le_bytes(
985 body.try_into().expect("validated 8-byte datetime"),
986 )))
987 }
988 TypeId::Uuid => {
989 require_len(16)?;
990 Ok(Value::Uuid(
991 body.try_into().expect("validated 16-byte UUID"),
992 ))
993 }
994 TypeId::Bytes => Ok(Value::Bytes(body.to_vec())),
995 TypeId::Json => {
996 pj1_validate(body).map_err(|error| format!("invalid typed PJ1 JSON: {error}"))?;
997 Ok(Value::Json(body.into()))
998 }
999 }
1000}
1001
1002fn decode_native_rows(payload: &[u8]) -> Result<Message, String> {
1003 let mut pos = 0;
1004 let col_count = decode_u16(payload, &mut pos, "native column count")? as usize;
1005 if col_count > MAX_COLUMNS {
1006 return Err("too many columns".into());
1007 }
1008 let mut columns = Vec::with_capacity(col_count);
1009 for _ in 0..col_count {
1010 columns.push(decode_string_strict(
1011 payload,
1012 &mut pos,
1013 "native column name",
1014 )?);
1015 }
1016 let row_count = decode_u32(payload, &mut pos, "native row count")? as usize;
1017 if row_count > MAX_ROWS {
1018 return Err("too many rows".into());
1019 }
1020 let minimum_row_len = col_count
1022 .checked_mul(5)
1023 .ok_or_else(|| "native row width overflow".to_string())?;
1024 if minimum_row_len == 0 {
1025 if row_count != 0 {
1026 return Err("nonzero native row count with zero columns".into());
1027 }
1028 } else if row_count > payload.len().saturating_sub(pos) / minimum_row_len {
1029 return Err("native row count exceeds payload size".into());
1030 }
1031
1032 let mut rows = Vec::with_capacity(row_count);
1033 for _ in 0..row_count {
1034 let mut row = Vec::with_capacity(col_count);
1035 for _ in 0..col_count {
1036 row.push(decode_typed_value(payload, &mut pos)?);
1037 }
1038 rows.push(row);
1039 }
1040 require_payload_end(payload, pos, "native rows")?;
1041 Ok(Message::ResultRowsNative { columns, rows })
1042}
1043
1044fn take_exact<'a>(
1045 data: &'a [u8],
1046 pos: &mut usize,
1047 len: usize,
1048 label: &str,
1049) -> Result<&'a [u8], String> {
1050 let end = pos
1051 .checked_add(len)
1052 .ok_or_else(|| format!("{label} length overflow"))?;
1053 if end > data.len() {
1054 return Err(format!("truncated {label}"));
1055 }
1056 let bytes = &data[*pos..end];
1057 *pos = end;
1058 Ok(bytes)
1059}
1060
1061fn require_payload_end(payload: &[u8], pos: usize, label: &str) -> Result<(), String> {
1062 if pos == payload.len() {
1063 Ok(())
1064 } else {
1065 Err(format!("trailing bytes in {label} payload"))
1066 }
1067}
1068
1069fn decode_exact_string(payload: &[u8], label: &str) -> Result<String, String> {
1070 let mut pos = 0;
1071 let value = decode_string_strict(payload, &mut pos, label)?;
1072 require_payload_end(payload, pos, label)?;
1073 Ok(value)
1074}
1075
1076fn decode_string_strict(data: &[u8], pos: &mut usize, label: &str) -> Result<String, String> {
1077 let len = decode_u32(data, pos, label)? as usize;
1078 let bytes = take_exact(data, pos, len, label)?;
1079 std::str::from_utf8(bytes)
1080 .map(str::to_owned)
1081 .map_err(|error| format!("invalid UTF-8 in {label}: {error}"))
1082}
1083
1084fn encode_string(s: &str) -> Vec<u8> {
1085 let mut buf = Vec::with_capacity(4 + s.len());
1086 buf.extend_from_slice(&(s.len() as u32).to_le_bytes());
1087 buf.extend_from_slice(s.as_bytes());
1088 buf
1089}
1090
1091fn decode_string(data: &[u8], pos: &mut usize) -> Result<String, String> {
1092 if *pos + 4 > data.len() {
1093 return Err("truncated string length".into());
1094 }
1095 let len_bytes: [u8; 4] = data[*pos..*pos + 4]
1096 .try_into()
1097 .map_err(|_| "invalid string length bytes".to_string())?;
1098 let len = u32::from_le_bytes(len_bytes) as usize;
1099 *pos += 4;
1100 if *pos + len > data.len() {
1101 return Err("truncated string data".into());
1102 }
1103 let s = String::from_utf8_lossy(&data[*pos..*pos + len]).into_owned();
1104 *pos += len;
1105 Ok(s)
1106}
1107
1108fn encode_option_u64(out: &mut Vec<u8>, value: Option<u64>) {
1109 match value {
1110 Some(value) => {
1111 out.push(1);
1112 out.extend_from_slice(&value.to_le_bytes());
1113 }
1114 None => out.push(0),
1115 }
1116}
1117
1118fn decode_option_u64(data: &[u8], pos: &mut usize, label: &str) -> Result<Option<u64>, String> {
1119 let present = decode_bool(data, pos, label)?;
1120 if present {
1121 Ok(Some(decode_u64(data, pos, label)?))
1122 } else {
1123 Ok(None)
1124 }
1125}
1126
1127fn encode_option_string(out: &mut Vec<u8>, value: Option<&String>) {
1128 match value {
1129 Some(value) => {
1130 out.push(1);
1131 out.extend_from_slice(&encode_string(value));
1132 }
1133 None => out.push(0),
1134 }
1135}
1136
1137fn decode_option_string(
1138 data: &[u8],
1139 pos: &mut usize,
1140 label: &str,
1141) -> Result<Option<String>, String> {
1142 let present = decode_bool(data, pos, label)?;
1143 if present {
1144 Ok(Some(decode_string(data, pos)?))
1145 } else {
1146 Ok(None)
1147 }
1148}
1149
1150fn encode_sync_status(status: &WireSyncStatus) -> Vec<u8> {
1151 let mut out = Vec::new();
1152 out.extend_from_slice(&encode_string(&status.replica_id));
1153 out.push(u8::from(status.active));
1154 encode_option_u64(&mut out, status.last_applied_lsn);
1155 out.extend_from_slice(&status.remote_lsn.to_le_bytes());
1156 encode_option_u64(&mut out, status.servable_lsn);
1157 encode_option_u64(&mut out, status.unarchived_lsn);
1158 encode_option_u64(&mut out, status.lag_lsn);
1159 encode_option_u64(&mut out, status.lag_bytes);
1160 encode_option_u64(&mut out, status.lag_ms);
1161 out.push(u8::from(status.stale));
1162 out.push(match status.repair_action {
1163 WireSyncRepairAction::None => 0,
1164 WireSyncRepairAction::Pull => 1,
1165 WireSyncRepairAction::AwaitArchive => 2,
1166 WireSyncRepairAction::Rebootstrap => 3,
1167 });
1168 encode_option_string(&mut out, status.last_sync_error.as_ref());
1169 out
1170}
1171
1172fn decode_sync_status(data: &[u8], pos: &mut usize) -> Result<WireSyncStatus, String> {
1173 let replica_id = decode_string(data, pos)?;
1174 let active = decode_bool(data, pos, "sync status active")?;
1175 let last_applied_lsn = decode_option_u64(data, pos, "sync status last applied LSN")?;
1176 let remote_lsn = decode_u64(data, pos, "sync status remote LSN")?;
1177 let servable_lsn = decode_option_u64(data, pos, "sync status servable LSN")?;
1178 let unarchived_lsn = decode_option_u64(data, pos, "sync status unarchived LSN")?;
1179 let lag_lsn = decode_option_u64(data, pos, "sync status lag LSN")?;
1180 let lag_bytes = decode_option_u64(data, pos, "sync status lag bytes")?;
1181 let lag_ms = decode_option_u64(data, pos, "sync status lag milliseconds")?;
1182 let stale = decode_bool(data, pos, "sync status stale")?;
1183 if *pos >= data.len() {
1184 return Err("truncated sync repair action".into());
1185 }
1186 let repair_action = match data[*pos] {
1187 0 => WireSyncRepairAction::None,
1188 1 => WireSyncRepairAction::Pull,
1189 2 => WireSyncRepairAction::AwaitArchive,
1190 3 => WireSyncRepairAction::Rebootstrap,
1191 other => return Err(format!("unknown sync repair action: {other}")),
1192 };
1193 *pos += 1;
1194 let last_sync_error = decode_option_string(data, pos, "sync status last error")?;
1195 Ok(WireSyncStatus {
1196 replica_id,
1197 active,
1198 last_applied_lsn,
1199 remote_lsn,
1200 servable_lsn,
1201 unarchived_lsn,
1202 lag_lsn,
1203 lag_bytes,
1204 lag_ms,
1205 stale,
1206 repair_action,
1207 last_sync_error,
1208 })
1209}
1210
1211fn encode_retained_unit(out: &mut Vec<u8>, unit: &WireRetainedUnit) {
1212 out.extend_from_slice(&unit.tx_id.to_le_bytes());
1213 out.push(unit.record_type);
1214 out.extend_from_slice(&unit.lsn.to_le_bytes());
1215 out.extend_from_slice(&(unit.data.len() as u32).to_le_bytes());
1216 out.extend_from_slice(&unit.data);
1217}
1218
1219fn decode_retained_unit(data: &[u8], pos: &mut usize) -> Result<WireRetainedUnit, String> {
1220 let tx_id = decode_u64(data, pos, "sync retained unit tx id")?;
1221 if *pos >= data.len() {
1222 return Err("truncated sync retained unit record type".into());
1223 }
1224 let record_type = data[*pos];
1225 *pos += 1;
1226 let lsn = decode_u64(data, pos, "sync retained unit LSN")?;
1227 let data = decode_bytes(data, pos, "sync retained unit payload")?;
1228 Ok(WireRetainedUnit {
1229 tx_id,
1230 record_type,
1231 lsn,
1232 data,
1233 })
1234}
1235
1236fn decode_bool(data: &[u8], pos: &mut usize, label: &str) -> Result<bool, String> {
1237 if *pos >= data.len() {
1238 return Err(format!("truncated {label}"));
1239 }
1240 let raw = data[*pos];
1241 *pos += 1;
1242 match raw {
1243 0 => Ok(false),
1244 1 => Ok(true),
1245 other => Err(format!("invalid boolean for {label}: {other}")),
1246 }
1247}
1248
1249fn decode_u16(data: &[u8], pos: &mut usize, label: &str) -> Result<u16, String> {
1250 if *pos + 2 > data.len() {
1251 return Err(format!("truncated {label}"));
1252 }
1253 let bytes: [u8; 2] = data[*pos..*pos + 2]
1254 .try_into()
1255 .map_err(|_| format!("invalid {label} bytes"))?;
1256 *pos += 2;
1257 Ok(u16::from_le_bytes(bytes))
1258}
1259
1260fn decode_u32(data: &[u8], pos: &mut usize, label: &str) -> Result<u32, String> {
1261 if *pos + 4 > data.len() {
1262 return Err(format!("truncated {label}"));
1263 }
1264 let bytes: [u8; 4] = data[*pos..*pos + 4]
1265 .try_into()
1266 .map_err(|_| format!("invalid {label} bytes"))?;
1267 *pos += 4;
1268 Ok(u32::from_le_bytes(bytes))
1269}
1270
1271fn decode_u64(data: &[u8], pos: &mut usize, label: &str) -> Result<u64, String> {
1272 if *pos + 8 > data.len() {
1273 return Err(format!("truncated {label}"));
1274 }
1275 let bytes: [u8; 8] = data[*pos..*pos + 8]
1276 .try_into()
1277 .map_err(|_| format!("invalid {label} bytes"))?;
1278 *pos += 8;
1279 Ok(u64::from_le_bytes(bytes))
1280}
1281
1282fn decode_16_bytes(data: &[u8], pos: &mut usize, label: &str) -> Result<[u8; 16], String> {
1283 if *pos + 16 > data.len() {
1284 return Err(format!("truncated {label}"));
1285 }
1286 let mut out = [0u8; 16];
1287 out.copy_from_slice(&data[*pos..*pos + 16]);
1288 *pos += 16;
1289 Ok(out)
1290}
1291
1292fn decode_bytes(data: &[u8], pos: &mut usize, label: &str) -> Result<Vec<u8>, String> {
1293 let len = decode_u32(data, pos, label)? as usize;
1294 if *pos + len > data.len() {
1295 return Err(format!("truncated {label}"));
1296 }
1297 let out = data[*pos..*pos + len].to_vec();
1298 *pos += len;
1299 Ok(out)
1300}
1301
1302#[cfg(test)]
1303mod tests {
1304 use super::*;
1305
1306 #[test]
1307 fn test_encode_decode_query() {
1308 let msg = Message::Query {
1309 query: "User filter .age > 30".into(),
1310 };
1311 let bytes = msg.encode();
1312 let decoded = Message::decode(&bytes).unwrap();
1313 match decoded {
1314 Message::Query { query } => assert_eq!(query, "User filter .age > 30"),
1315 _ => panic!("expected Query"),
1316 }
1317 }
1318
1319 #[test]
1320 fn test_encode_decode_connect_with_username() {
1321 let msg = Message::Connect {
1322 db_name: "mydb".into(),
1323 password: Some(Zeroizing::new("secret".into())),
1324 username: Some("alice".into()),
1325 };
1326 let bytes = msg.encode();
1327 match Message::decode(&bytes).unwrap() {
1328 Message::Connect {
1329 db_name,
1330 password,
1331 username,
1332 } => {
1333 assert_eq!(db_name, "mydb");
1334 assert_eq!(password.as_deref().map(|s| s.as_str()), Some("secret"));
1335 assert_eq!(username.as_deref(), Some("alice"));
1336 }
1337 other => panic!("expected Connect, got {other:?}"),
1338 }
1339 }
1340
1341 #[test]
1342 fn test_encode_decode_connect_without_username() {
1343 let msg = Message::Connect {
1345 db_name: "mydb".into(),
1346 password: Some(Zeroizing::new("secret".into())),
1347 username: None,
1348 };
1349 let bytes = msg.encode();
1350 match Message::decode(&bytes).unwrap() {
1351 Message::Connect {
1352 db_name,
1353 password,
1354 username,
1355 } => {
1356 assert_eq!(db_name, "mydb");
1357 assert_eq!(password.as_deref().map(|s| s.as_str()), Some("secret"));
1358 assert_eq!(username, None);
1359 }
1360 other => panic!("expected Connect, got {other:?}"),
1361 }
1362 }
1363
1364 #[test]
1365 fn test_decode_old_client_connect_db_and_password_only() {
1366 let mut payload = encode_string("mydb");
1369 payload.extend_from_slice(&encode_string("pw"));
1370 let mut frame = vec![MSG_CONNECT, 0];
1371 frame.extend_from_slice(&(payload.len() as u32).to_le_bytes());
1372 frame.extend_from_slice(&payload);
1373 match Message::decode(&frame).unwrap() {
1374 Message::Connect {
1375 db_name,
1376 password,
1377 username,
1378 } => {
1379 assert_eq!(db_name, "mydb");
1380 assert_eq!(password.as_deref().map(|s| s.as_str()), Some("pw"));
1381 assert_eq!(username, None, "old-client frame must yield username=None");
1382 }
1383 other => panic!("expected Connect, got {other:?}"),
1384 }
1385 }
1386
1387 #[test]
1388 fn test_decode_old_client_connect_db_only() {
1389 let payload = encode_string("mydb");
1391 let mut frame = vec![MSG_CONNECT, 0];
1392 frame.extend_from_slice(&(payload.len() as u32).to_le_bytes());
1393 frame.extend_from_slice(&payload);
1394 match Message::decode(&frame).unwrap() {
1395 Message::Connect {
1396 db_name,
1397 password,
1398 username,
1399 } => {
1400 assert_eq!(db_name, "mydb");
1401 assert_eq!(password, None);
1402 assert_eq!(username, None);
1403 }
1404 other => panic!("expected Connect, got {other:?}"),
1405 }
1406 }
1407
1408 #[test]
1409 fn test_encode_decode_result_rows() {
1410 let msg = Message::ResultRows {
1411 columns: vec!["name".into(), "age".into()],
1412 rows: vec![
1413 vec!["Alice".into(), "30".into()],
1414 vec!["Bob".into(), "25".into()],
1415 ],
1416 };
1417 let bytes = msg.encode();
1418 let decoded = Message::decode(&bytes).unwrap();
1419 match decoded {
1420 Message::ResultRows { columns, rows } => {
1421 assert_eq!(columns, vec!["name", "age"]);
1422 assert_eq!(rows.len(), 2);
1423 }
1424 _ => panic!("expected ResultRows"),
1425 }
1426 }
1427
1428 #[test]
1429 fn legacy_rows_frame_is_byte_identical() {
1430 let encoded = Message::ResultRows {
1431 columns: vec!["x".into()],
1432 rows: vec![vec!["y".into()]],
1433 }
1434 .encode();
1435 assert_eq!(
1436 encoded,
1437 vec![
1438 0x07, 0x00, 0x10, 0x00, 0x00, 0x00, 0x01, 0x00, 0x01, 0x00, 0x00, 0x00, b'x', 0x01, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, b'y', ]
1444 );
1445 }
1446
1447 #[test]
1448 fn native_request_tags_and_params_round_trip() {
1449 let cases = [
1450 (
1451 MSG_QUERY_NATIVE,
1452 Message::QueryNative {
1453 query: "T { .x }".into(),
1454 },
1455 ),
1456 (
1457 MSG_QUERY_SQL_NATIVE,
1458 Message::QuerySqlNative {
1459 query: "SELECT x FROM T".into(),
1460 },
1461 ),
1462 ];
1463 for (tag, message) in cases {
1464 let encoded = message.encode();
1465 assert_eq!(encoded[0], tag);
1466 match Message::decode(&encoded).expect("native request round trip") {
1467 Message::QueryNative { query } => assert_eq!(query, "T { .x }"),
1468 Message::QuerySqlNative { query } => assert_eq!(query, "SELECT x FROM T"),
1469 other => panic!("unexpected native request: {other:?}"),
1470 }
1471 }
1472
1473 let encoded = Message::QueryWithParamsNative {
1474 query: "T filter .x = $1".into(),
1475 params: vec![WireParam::Int(7), WireParam::Bool(false)],
1476 }
1477 .encode();
1478 assert_eq!(encoded[0], MSG_QUERY_PARAMS_NATIVE);
1479 match Message::decode(&encoded).expect("native params round trip") {
1480 Message::QueryWithParamsNative { query, params } => {
1481 assert_eq!(query, "T filter .x = $1");
1482 assert_eq!(params, vec![WireParam::Int(7), WireParam::Bool(false)]);
1483 }
1484 other => panic!("unexpected native parameterized request: {other:?}"),
1485 }
1486 }
1487
1488 fn every_native_value() -> Vec<Value> {
1489 vec![
1490 Value::Empty,
1491 Value::Int(-9_007_199_254_740_993),
1492 Value::Float(2.5),
1493 Value::Bool(true),
1494 Value::Str("héllo".into()),
1495 Value::DateTime(1_723_650_123_456_789),
1496 Value::Uuid([
1497 0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99, 0xaa, 0xbb, 0xcc, 0xdd,
1498 0xee, 0xff,
1499 ]),
1500 Value::Bytes(vec![0x00, 0x7f, 0x80, 0xff]),
1501 Value::Json(
1502 powdb_storage::pj1::parse_json_text("9007199254740993")
1503 .expect("PJ1 fixture")
1504 .into_boxed_slice(),
1505 ),
1506 ]
1507 }
1508
1509 #[test]
1510 fn native_rows_and_scalar_round_trip_every_type() {
1511 let values = every_native_value();
1512 let columns = (0..values.len()).map(|index| format!("c{index}")).collect();
1513 let encoded = Message::ResultRowsNative {
1514 columns,
1515 rows: vec![values.clone()],
1516 }
1517 .encode();
1518 assert_eq!(encoded[0], MSG_RESULT_ROWS_NATIVE);
1519 match Message::decode(&encoded).expect("native rows round trip") {
1520 Message::ResultRowsNative { columns, rows } => {
1521 assert_eq!(columns.len(), values.len());
1522 assert_eq!(rows, vec![values.clone()]);
1523 }
1524 other => panic!("unexpected native rows: {other:?}"),
1525 }
1526
1527 for value in values {
1528 let encoded = Message::ResultScalarNative {
1529 value: value.clone(),
1530 }
1531 .encode();
1532 assert_eq!(encoded[0], MSG_RESULT_SCALAR_NATIVE);
1533 match Message::decode(&encoded).expect("native scalar round trip") {
1534 Message::ResultScalarNative { value: decoded } => assert_eq!(decoded, value),
1535 other => panic!("unexpected native scalar: {other:?}"),
1536 }
1537 }
1538 }
1539
1540 #[test]
1541 fn native_mixed_row_matches_cross_client_golden() {
1542 let encoded = Message::ResultRowsNative {
1543 columns: ["e", "i", "f", "b", "s", "d", "u", "x", "j"]
1544 .into_iter()
1545 .map(str::to_owned)
1546 .collect(),
1547 rows: vec![every_native_value()],
1548 }
1549 .encode();
1550 let hex = encoded
1551 .iter()
1552 .map(|byte| format!("{byte:02x}"))
1553 .collect::<String>();
1554 assert_eq!(
1555 hex,
1556 "16009c000000090001000000650100000069010000006601000000620100000073010000006401000000750100000078010000006a0100000000000000000108000000ffffffffffffdfff02080000000000000000000440030100000001040600000068c3a96c6c6f050800000015615391a61f0600061000000000112233445566778899aabbccddeeff0704000000007f80ff0809000000030100000000002000"
1557 );
1558 }
1559
1560 fn frame(tag: u8, payload: &[u8]) -> Vec<u8> {
1561 let mut frame = vec![tag, 0];
1562 frame.extend_from_slice(&(payload.len() as u32).to_le_bytes());
1563 frame.extend_from_slice(payload);
1564 frame
1565 }
1566
1567 fn typed_cell(tag: u8, body: &[u8]) -> Vec<u8> {
1568 let mut payload = vec![tag];
1569 payload.extend_from_slice(&(body.len() as u32).to_le_bytes());
1570 payload.extend_from_slice(body);
1571 payload
1572 }
1573
1574 #[test]
1575 fn native_scalar_rejects_malformed_cells() {
1576 let malformed = [
1577 typed_cell(0xff, &[]),
1578 typed_cell(TypeId::Int as u8, &[0; 7]),
1579 typed_cell(TypeId::Bool as u8, &[2]),
1580 typed_cell(TypeId::Str as u8, &[0xff]),
1581 typed_cell(TypeId::Json as u8, &[0xff]),
1582 typed_cell(TypeId::Json as u8, &[0, 0]),
1583 ];
1584 for payload in malformed {
1585 assert!(Message::decode(&frame(MSG_RESULT_SCALAR_NATIVE, &payload)).is_err());
1586 }
1587
1588 let mut trailing = typed_cell(TypeId::Empty as u8, &[]);
1589 trailing.push(0);
1590 assert!(Message::decode(&frame(MSG_RESULT_SCALAR_NATIVE, &trailing)).is_err());
1591 }
1592
1593 #[test]
1594 fn native_rows_rejects_bad_counts_and_trailing_data() {
1595 let mut payload = Vec::new();
1596 payload.extend_from_slice(&1u16.to_le_bytes());
1597 payload.extend_from_slice(&encode_string("x"));
1598 payload.extend_from_slice(&2u32.to_le_bytes());
1599 payload.extend_from_slice(&typed_cell(TypeId::Empty as u8, &[]));
1600 assert!(Message::decode(&frame(MSG_RESULT_ROWS_NATIVE, &payload)).is_err());
1601
1602 payload[7..11].copy_from_slice(&1u32.to_le_bytes());
1603 payload.extend_from_slice(&[0xaa]);
1604 assert!(Message::decode(&frame(MSG_RESULT_ROWS_NATIVE, &payload)).is_err());
1605 }
1606
1607 #[test]
1608 fn test_encode_decode_result_message() {
1609 let msg = Message::ResultMessage {
1610 message: "type User created".into(),
1611 };
1612 let bytes = msg.encode();
1613 let decoded = Message::decode(&bytes).unwrap();
1614 match decoded {
1615 Message::ResultMessage { message } => assert_eq!(message, "type User created"),
1616 _ => panic!("expected ResultMessage"),
1617 }
1618 }
1619
1620 #[test]
1621 fn test_encode_decode_error() {
1622 let msg = Message::Error {
1623 message: "table not found".into(),
1624 };
1625 let bytes = msg.encode();
1626 let decoded = Message::decode(&bytes).unwrap();
1627 match decoded {
1628 Message::Error { message } => assert_eq!(message, "table not found"),
1629 _ => panic!("expected Error"),
1630 }
1631 }
1632
1633 #[test]
1634 fn error_with_class_appends_one_trailing_byte() {
1635 let plain = Message::Error {
1636 message: "table 'users' not found".into(),
1637 }
1638 .encode();
1639 let classed = Message::ErrorWithClass {
1640 message: "table 'users' not found".into(),
1641 class: ErrorClass::Execution,
1642 }
1643 .encode();
1644 assert_eq!(classed[0], plain[0]);
1646 assert_eq!(classed.len(), plain.len() + 1);
1647 assert_eq!(&classed[6..plain.len()], &plain[6..]);
1648 assert_eq!(*classed.last().unwrap(), ErrorClass::Execution.as_u8());
1649 }
1650
1651 #[test]
1652 fn old_style_decode_of_classed_error_frame_yields_message() {
1653 let frame = Message::ErrorWithClass {
1657 message: "query timeout after 75ms".into(),
1658 class: ErrorClass::Timeout,
1659 }
1660 .encode();
1661 match Message::decode(&frame).unwrap() {
1662 Message::Error { message } => assert_eq!(message, "query timeout after 75ms"),
1663 other => panic!("expected Error, got {other:?}"),
1664 }
1665 }
1666
1667 #[test]
1668 fn decode_error_class_reads_trailing_byte_when_present() {
1669 let classed = Message::ErrorWithClass {
1670 message: "authentication failed".into(),
1671 class: ErrorClass::AuthFailed,
1672 }
1673 .encode();
1674 assert_eq!(decode_error_class(&classed), Some(6));
1675
1676 let plain = Message::Error {
1678 message: "authentication failed".into(),
1679 }
1680 .encode();
1681 assert_eq!(decode_error_class(&plain), None);
1682
1683 let ok = Message::ResultOk { affected: 1 }.encode();
1685 assert_eq!(decode_error_class(&ok), None);
1686 }
1687
1688 #[test]
1689 fn error_class_bytes_are_stable() {
1690 let expected: [(ErrorClass, u8); 10] = [
1693 (ErrorClass::Internal, 0),
1694 (ErrorClass::Parse, 1),
1695 (ErrorClass::Execution, 2),
1696 (ErrorClass::Timeout, 3),
1697 (ErrorClass::LimitExceeded, 4),
1698 (ErrorClass::ReadonlyRefused, 5),
1699 (ErrorClass::AuthFailed, 6),
1700 (ErrorClass::RateLimited, 7),
1701 (ErrorClass::ConstraintViolation, 8),
1702 (ErrorClass::Cancelled, 9),
1703 ];
1704 for (class, byte) in expected {
1705 assert_eq!(class.as_u8(), byte, "{class:?}");
1706 assert_eq!(ErrorClass::from_u8(byte), Some(class));
1707 }
1708 assert_eq!(ErrorClass::from_u8(10), None, "future bytes must be None");
1709 assert_eq!(ErrorClass::from_u8(255), None);
1710 }
1711
1712 #[test]
1713 fn test_encode_decode_query_sql() {
1714 let msg = Message::QuerySql {
1715 query: "SELECT * FROM User".into(),
1716 };
1717 let decoded = Message::decode(&msg.encode()).unwrap();
1718 match decoded {
1719 Message::QuerySql { query } => assert_eq!(query, "SELECT * FROM User"),
1720 other => panic!("expected QuerySql, got {other:?}"),
1721 }
1722 }
1723
1724 #[test]
1725 fn test_encode_decode_query_with_params() {
1726 let msg = Message::QueryWithParams {
1727 query: "insert User { name := $1, age := $2, ok := $3, note := $4 }".into(),
1728 params: vec![
1729 WireParam::Str(r#"a"b\c; drop User"#.into()),
1730 WireParam::Int(-7),
1731 WireParam::Bool(true),
1732 WireParam::Null,
1733 ],
1734 };
1735 let bytes = msg.encode();
1736 assert_eq!(bytes[0], 0x04);
1738 match Message::decode(&bytes).unwrap() {
1739 Message::QueryWithParams { query, params } => {
1740 assert!(query.contains("$1"));
1741 assert_eq!(params.len(), 4);
1742 assert!(matches!(¶ms[0], WireParam::Str(s) if s == r#"a"b\c; drop User"#));
1743 assert!(matches!(¶ms[1], WireParam::Int(-7)));
1744 assert!(matches!(¶ms[2], WireParam::Bool(true)));
1745 assert!(matches!(¶ms[3], WireParam::Null));
1746 }
1747 other => panic!("expected QueryWithParams, got {other:?}"),
1748 }
1749 }
1750
1751 #[test]
1752 fn test_query_with_params_float_round_trip() {
1753 let msg = Message::QueryWithParams {
1754 query: "T filter .f = $1".into(),
1755 params: vec![WireParam::Float(2.5)],
1756 };
1757 match Message::decode(&msg.encode()).unwrap() {
1758 Message::QueryWithParams { params, .. } => {
1759 assert!(matches!(¶ms[0], WireParam::Float(f) if (*f - 2.5).abs() < 1e-12));
1760 }
1761 other => panic!("expected QueryWithParams, got {other:?}"),
1762 }
1763 }
1764
1765 fn sample_sync_status() -> WireSyncStatus {
1766 WireSyncStatus {
1767 replica_id: "replica-a".into(),
1768 active: true,
1769 last_applied_lsn: Some(7),
1770 remote_lsn: 10,
1771 servable_lsn: Some(10),
1772 unarchived_lsn: Some(0),
1773 lag_lsn: Some(3),
1774 lag_bytes: Some(2048),
1775 lag_ms: Some(5000),
1776 stale: true,
1777 repair_action: WireSyncRepairAction::Pull,
1778 last_sync_error: None,
1779 }
1780 }
1781
1782 #[test]
1783 fn test_encode_decode_sync_requests() {
1784 let database_id = *b"sync-protocol!!!";
1785 let pull = Message::SyncPull {
1786 replica_id: "replica-a".into(),
1787 since_lsn: 7,
1788 max_units: 128,
1789 max_bytes: 4096,
1790 database_id,
1791 primary_generation: 9,
1792 wal_format_version: 1,
1793 catalog_version: 2,
1794 segment_format_version: 1,
1795 };
1796 let bytes = pull.encode();
1797 assert_eq!(bytes[0], MSG_SYNC_PULL);
1798 match Message::decode(&bytes).unwrap() {
1799 Message::SyncPull {
1800 replica_id,
1801 since_lsn,
1802 max_units,
1803 max_bytes,
1804 database_id: decoded_database_id,
1805 primary_generation,
1806 wal_format_version,
1807 catalog_version,
1808 segment_format_version,
1809 } => {
1810 assert_eq!(replica_id, "replica-a");
1811 assert_eq!(since_lsn, 7);
1812 assert_eq!(max_units, 128);
1813 assert_eq!(max_bytes, 4096);
1814 assert_eq!(decoded_database_id, database_id);
1815 assert_eq!(primary_generation, 9);
1816 assert_eq!(wal_format_version, 1);
1817 assert_eq!(catalog_version, 2);
1818 assert_eq!(segment_format_version, 1);
1819 }
1820 other => panic!("expected SyncPull, got {other:?}"),
1821 }
1822
1823 let ack = Message::SyncAck {
1824 replica_id: "replica-a".into(),
1825 applied_lsn: 10,
1826 remote_lsn: 10,
1827 };
1828 match Message::decode(&ack.encode()).unwrap() {
1829 Message::SyncAck {
1830 replica_id,
1831 applied_lsn,
1832 remote_lsn,
1833 } => {
1834 assert_eq!(replica_id, "replica-a");
1835 assert_eq!(applied_lsn, 10);
1836 assert_eq!(remote_lsn, 10);
1837 }
1838 other => panic!("expected SyncAck, got {other:?}"),
1839 }
1840 }
1841
1842 #[test]
1843 fn test_encode_decode_sync_results() {
1844 let status = sample_sync_status();
1845 match Message::decode(
1846 &Message::SyncStatusResult {
1847 status: status.clone(),
1848 }
1849 .encode(),
1850 )
1851 .unwrap()
1852 {
1853 Message::SyncStatusResult { status: decoded } => assert_eq!(decoded, status),
1854 other => panic!("expected SyncStatusResult, got {other:?}"),
1855 }
1856
1857 let await_archive_status = WireSyncStatus {
1858 servable_lsn: Some(7),
1859 unarchived_lsn: Some(3),
1860 repair_action: WireSyncRepairAction::AwaitArchive,
1861 last_sync_error: Some("primary WAL is not yet archived".into()),
1862 ..status.clone()
1863 };
1864 match Message::decode(
1865 &Message::SyncStatusResult {
1866 status: await_archive_status.clone(),
1867 }
1868 .encode(),
1869 )
1870 .unwrap()
1871 {
1872 Message::SyncStatusResult { status: decoded } => {
1873 assert_eq!(decoded, await_archive_status)
1874 }
1875 other => panic!("expected AwaitArchive SyncStatusResult, got {other:?}"),
1876 }
1877
1878 let units = vec![
1879 WireRetainedUnit {
1880 tx_id: 1,
1881 record_type: 4,
1882 lsn: 8,
1883 data: vec![1, 2, 3],
1884 },
1885 WireRetainedUnit {
1886 tx_id: 1,
1887 record_type: 4,
1888 lsn: 9,
1889 data: vec![4, 5],
1890 },
1891 ];
1892 let empty_pull_len = Message::SyncPullResult {
1893 status: status.clone(),
1894 units: Vec::new(),
1895 has_more: true,
1896 }
1897 .encode()
1898 .len();
1899 let populated_pull_len = Message::SyncPullResult {
1900 status: status.clone(),
1901 units: units.clone(),
1902 has_more: true,
1903 }
1904 .encode()
1905 .len();
1906 let expected_unit_len: u64 = units
1907 .iter()
1908 .map(WireRetainedUnit::encoded_len)
1909 .collect::<Result<Vec<_>, _>>()
1910 .unwrap()
1911 .into_iter()
1912 .sum();
1913 assert_eq!(
1914 u64::try_from(populated_pull_len - empty_pull_len).unwrap(),
1915 expected_unit_len,
1916 "retained-unit encoded length must track SyncPullResult wire shape"
1917 );
1918
1919 match Message::decode(
1920 &Message::SyncPullResult {
1921 status: status.clone(),
1922 units: units.clone(),
1923 has_more: true,
1924 }
1925 .encode(),
1926 )
1927 .unwrap()
1928 {
1929 Message::SyncPullResult {
1930 status: decoded_status,
1931 units: decoded_units,
1932 has_more,
1933 } => {
1934 assert_eq!(decoded_status, status);
1935 assert_eq!(decoded_units, units);
1936 assert!(has_more);
1937 }
1938 other => panic!("expected SyncPullResult, got {other:?}"),
1939 }
1940
1941 match Message::decode(
1942 &Message::SyncAckResult {
1943 previous_applied_lsn: 7,
1944 applied_lsn: 10,
1945 remote_lsn: 10,
1946 advanced: true,
1947 status: WireSyncStatus {
1948 stale: false,
1949 repair_action: WireSyncRepairAction::None,
1950 lag_lsn: Some(0),
1951 lag_bytes: Some(0),
1952 lag_ms: Some(0),
1953 ..status
1954 },
1955 }
1956 .encode(),
1957 )
1958 .unwrap()
1959 {
1960 Message::SyncAckResult {
1961 previous_applied_lsn,
1962 applied_lsn,
1963 remote_lsn,
1964 advanced,
1965 status,
1966 } => {
1967 assert_eq!(previous_applied_lsn, 7);
1968 assert_eq!(applied_lsn, 10);
1969 assert_eq!(remote_lsn, 10);
1970 assert!(advanced);
1971 assert!(!status.stale);
1972 }
1973 other => panic!("expected SyncAckResult, got {other:?}"),
1974 }
1975 }
1976
1977 #[test]
1978 fn test_decode_garbage_never_panics() {
1979 let cases: Vec<Vec<u8>> = vec![
1983 vec![], vec![0x03], vec![0x03, 0x00, 0x00, 0x00, 0x00], vec![0xFF, 0x00, 0x00, 0x00, 0x00, 0x00], vec![0x03, 0x00, 0xFF, 0xFF, 0xFF, 0xFF],
1989 {
1991 let mut payload = encode_string("replica-a");
1992 let mut frame = vec![MSG_SYNC_PULL, 0];
1993 frame.extend_from_slice(&(payload.len() as u32).to_le_bytes());
1994 frame.append(&mut payload);
1995 frame
1996 },
1997 {
1999 let mut payload = encode_sync_status(&sample_sync_status());
2000 payload.extend_from_slice(&((MAX_SYNC_UNITS as u32) + 1).to_le_bytes());
2001 let mut frame = vec![MSG_SYNC_PULL_RESULT, 0];
2002 frame.extend_from_slice(&(payload.len() as u32).to_le_bytes());
2003 frame.extend_from_slice(&payload);
2004 frame
2005 },
2006 vec![0x01, 0x00, 0x04, 0x00, 0x00, 0x00, 0xFF, 0xFF, 0xFF, 0xFF],
2008 vec![0x07, 0x00, 0x02, 0x00, 0x00, 0x00, 0xFF, 0xFF],
2010 vec![0x09, 0x00, 0x03, 0x00, 0x00, 0x00, 0x01, 0x02, 0x03],
2012 vec![0x04, 0x00, 0x04, 0x00, 0x00, 0x00, 0xFF, 0xFF, 0xFF, 0xFF],
2014 vec![
2016 0x04, 0x00, 0x06, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, ],
2020 vec![
2022 0x04, 0x00, 0x0B, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x01, 0x01, 0x02, 0x03,
2027 ],
2028 vec![
2030 0x04, 0x00, 0x0F, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x04, 0xFF, 0xFF, 0xFF, 0xFF, ],
2036 vec![
2038 0x04, 0x00, 0x0B, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x63, ],
2043 ];
2044 for bytes in cases {
2045 let result = Message::decode(&bytes);
2046 assert!(
2047 result.is_err(),
2048 "expected Err for malformed input {bytes:?}, got {result:?}"
2049 );
2050 }
2051 }
2052
2053 #[test]
2054 fn test_decode_arbitrary_prefixes_never_panic() {
2055 let valid = Message::ResultRows {
2058 columns: vec!["a".into(), "b".into()],
2059 rows: vec![vec!["1".into(), "2".into()]],
2060 }
2061 .encode();
2062 for end in 0..valid.len() {
2063 let _ = Message::decode(&valid[..end]);
2065 }
2066 for i in 0..valid.len() {
2068 let mut m = valid.clone();
2069 m[i] = 0xFF;
2070 let _ = Message::decode(&m);
2071 }
2072 }
2073
2074 #[test]
2075 fn test_decode_result_rows_rejects_amplified_row_count() {
2076 let mut payload = Vec::new();
2081 payload.extend_from_slice(&0u16.to_le_bytes()); payload.extend_from_slice(&10_000_000u32.to_le_bytes()); let mut frame = vec![MSG_RESULT_ROWS, 0];
2084 frame.extend_from_slice(&(payload.len() as u32).to_le_bytes());
2085 frame.extend_from_slice(&payload);
2086 assert!(
2087 Message::decode(&frame).is_err(),
2088 "amplified row count must be rejected"
2089 );
2090
2091 let msg = Message::ResultRows {
2093 columns: vec!["a".into()],
2094 rows: vec![vec!["x".into()]],
2095 };
2096 match Message::decode(&msg.encode()).unwrap() {
2097 Message::ResultRows { columns, rows } => {
2098 assert_eq!(columns, vec!["a"]);
2099 assert_eq!(rows, vec![vec!["x".to_string()]]);
2100 }
2101 other => panic!("expected ResultRows, got {other:?}"),
2102 }
2103 }
2104
2105 #[test]
2106 fn test_frame_length() {
2107 let msg = Message::Query {
2108 query: "User".into(),
2109 };
2110 let bytes = msg.encode();
2111 assert!(bytes.len() >= 6);
2112 let payload_len = u32::from_le_bytes(bytes[2..6].try_into().unwrap()) as usize;
2113 assert_eq!(bytes.len(), 6 + payload_len);
2114 }
2115}