Skip to main content

powdb_server/
protocol.rs

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;
9/// Query carrying positional `$N` parameters (Task 4). Pure protocol
10/// addition: old clients never send it, and old servers reject it with the
11/// existing "unknown message type" error — no existing frame changes shape.
12const MSG_QUERY_PARAMS: u8 = 0x04;
13/// SQL query frame. Plain Query remains PowQL for backward compatibility.
14const MSG_QUERY_SQL: u8 = 0x05;
15/// Private sync control frames. These are append-only protocol extensions:
16/// legacy query/result frames keep their original tags and shape.
17const 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;
31/// Native typed PowQL request. The response is `MSG_RESULT_ROWS_NATIVE` or
32/// `MSG_RESULT_SCALAR_NATIVE`; legacy result frames remain byte-identical.
33pub const MSG_QUERY_NATIVE: u8 = 0x13;
34/// Native typed PowQL request with positional `$N` parameters.
35pub const MSG_QUERY_PARAMS_NATIVE: u8 = 0x14;
36/// Native typed SQL request.
37pub const MSG_QUERY_SQL_NATIVE: u8 = 0x15;
38/// Native typed row result.
39pub const MSG_RESULT_ROWS_NATIVE: u8 = 0x16;
40/// Native typed scalar result.
41pub const MSG_RESULT_SCALAR_NATIVE: u8 = 0x17;
42
43/// Maximum payload size accepted from the wire (64 MB).
44const MAX_PAYLOAD_SIZE: usize = 64 * 1024 * 1024;
45
46/// Maximum payload size for pre-auth CONNECT messages (4 KB).
47/// Only a database name and password are needed before authentication.
48const MAX_CONNECT_PAYLOAD_SIZE: usize = 4096;
49
50/// Maximum number of columns allowed in a result set.
51const MAX_COLUMNS: usize = 4096;
52
53/// Maximum number of rows allowed in a single result message.
54const MAX_ROWS: usize = 10_000_000;
55
56/// Maximum number of bound parameters in a single QueryWithParams message.
57const MAX_PARAMS: usize = 4096;
58
59/// Maximum retained units accepted in one sync pull frame.
60const MAX_SYNC_UNITS: usize = 4096;
61
62const STRING_LEN_PREFIX: usize = 4; // decode_string reads a 4-byte length prefix
63
64/// Stable 1-byte error classification carried at the tail of a `MSG_ERROR`
65/// payload.
66///
67/// Wire contract: the `MSG_ERROR` payload is a length-prefixed message string,
68/// optionally followed by exactly one trailing class byte. Every first-party
69/// decoder (this module, the TS client, the CLI via this module) reads the
70/// string by its length prefix and ignores trailing payload bytes, so:
71///   - old client + new server: the class byte is silently skipped,
72///   - new client + old server: the byte is absent and decodes as "no class".
73///
74/// The class is orthogonal to message sanitization: it is safe metadata even
75/// when the message text is masked to a generic string.
76///
77/// STABILITY: these numeric values are wire-stable and documented in
78/// docs/errors.md. Never renumber or reuse a value; only append new classes.
79/// Clients must treat unknown values as [`ErrorClass::Internal`].
80#[derive(Debug, Clone, Copy, PartialEq, Eq)]
81#[repr(u8)]
82pub enum ErrorClass {
83    /// Unclassified or internal server error (also the fallback for values a
84    /// client does not recognize).
85    Internal = 0,
86    /// The query text failed to lex or parse.
87    Parse = 1,
88    /// Planning or execution failed (unknown table or column, type
89    /// mismatch, unsupported statement, ...).
90    Execution = 2,
91    /// A time budget elapsed: per-query timeout, transaction-gate wait,
92    /// cooperative deadline cancellation, or idle-connection timeout.
93    Timeout = 3,
94    /// A memory or size limit was exceeded (sort/join row caps, per-query
95    /// memory budget, oversized query text, oversized result).
96    LimitExceeded = 4,
97    /// The server serves a read-only snapshot and the statement requires a
98    /// writer.
99    ReadonlyRefused = 5,
100    /// Authentication or database selection failed at CONNECT time.
101    AuthFailed = 6,
102    /// Too many failed authentication attempts from this address.
103    RateLimited = 7,
104    /// A constraint (e.g. a unique index) rejected the write.
105    ConstraintViolation = 8,
106    /// Execution was cancelled cooperatively (client disconnect).
107    Cancelled = 9,
108}
109
110impl ErrorClass {
111    /// The stable wire byte for this class.
112    pub fn as_u8(self) -> u8 {
113        self as u8
114    }
115
116    /// Parse a wire byte back into a class. Unknown (future) values return
117    /// `None`; callers should treat them as [`ErrorClass::Internal`].
118    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
135/// Extract the trailing error-class byte from a full `MSG_ERROR` frame, if
136/// present. Returns `None` for non-error frames, malformed frames, and
137/// legacy error frames that carry only the message string.
138pub 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    // Exactly one byte after the length-prefixed message string.
146    if payload.len() == 4 + string_len + 1 {
147        Some(payload[4 + string_len])
148    } else {
149        None
150    }
151}
152
153/// A positional parameter value carried by [`Message::QueryWithParams`].
154///
155/// Wire encoding per param: a 1-byte tag followed by the body —
156///   `0` null (no body), `1` int (8B LE i64), `2` float (8B LE f64),
157///   `3` bool (1B), `4` str (length-prefixed UTF-8).
158#[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    /// Encoded byte length of this retained unit inside a sync pull result
201    /// payload. Keep this next to `encode_retained_unit` so metrics and
202    /// max-byte enforcement evolve with the wire shape.
203    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        /// Client-supplied candidate password. Wrapped in `Zeroizing` so the
219        /// raw bytes from the wire are wiped from memory once the `Message`
220        /// (and thus this field) is dropped after the constant-time compare —
221        /// the candidate never lingers in a plain `String`.
222        password: Option<Zeroizing<String>>,
223        /// Optional user name for multi-user authentication. Appended after the
224        /// password on the wire so the format is a pure, backward-compatible
225        /// extension: old clients that omit it decode as `None`.
226        username: Option<String>,
227    },
228    ConnectOk {
229        version: String,
230    },
231    Query {
232        query: String,
233    },
234    /// A SQL query string.
235    QuerySql {
236        query: String,
237    },
238    /// A query string with positional `$N` parameters bound at the server.
239    QueryWithParams {
240        query: String,
241        params: Vec<WireParam>,
242    },
243    /// PowQL request whose result preserves storage value types.
244    QueryNative {
245        query: String,
246    },
247    /// Parameterized PowQL request whose result preserves storage value types.
248    QueryWithParamsNative {
249        query: String,
250        params: Vec<WireParam>,
251    },
252    /// SQL request whose result preserves storage value types.
253    QuerySqlNative {
254        query: String,
255    },
256    /// Request primary-side status for one embedded replica cursor.
257    SyncStatus {
258        replica_id: String,
259    },
260    /// Pull a bounded retained-unit chunk after the server-side replica cursor.
261    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    /// Acknowledge that the replica applied retained history through `applied_lsn`.
273    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    /// A descriptive status message (e.g. "type User created", "index dropped").
311    ResultMessage {
312        message: String,
313    },
314    Error {
315        message: String,
316    },
317    /// An error carrying its stable [`ErrorClass`] byte after the message
318    /// string. Encodes to the same `MSG_ERROR` tag as [`Message::Error`];
319    /// decoding a `MSG_ERROR` frame deliberately yields [`Message::Error`]
320    /// regardless (the class byte is a trailing extension that legacy
321    /// decoders skip; clients that want it call [`decode_error_class`]).
322    ErrorWithClass {
323        message: String,
324        class: ErrorClass,
325    },
326    Disconnect,
327    Ping,
328    Pong,
329}
330
331impl Message {
332    /// Encode message into wire format: `[type(1)][flags(1)][len(4)][payload]`
333    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                // Password is encoded as a length-prefixed string. Empty (len=0) means None.
342                match password {
343                    Some(p) => buf.extend_from_slice(&encode_string(p)),
344                    None => buf.extend_from_slice(&0u32.to_le_bytes()),
345                }
346                // Username is appended after the password (append-only extension).
347                // Length-prefixed string; empty (len=0) means None.
348                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); // flags
480        frame.extend_from_slice(&(payload.len() as u32).to_le_bytes());
481        frame.extend_from_slice(&payload);
482        frame
483    }
484
485    /// Decode message from wire format.
486    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                // Password is optional. If there are no more bytes, treat as None
506                // (backwards compatible with old clients that don't send a password).
507                let password = if pos < payload.len() {
508                    // Wrap the candidate immediately so the only owned copy of
509                    // the secret lives inside `Zeroizing` and is wiped on drop.
510                    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                // Username is optional and appended after the password. Old
520                // clients omit it entirely → no more bytes → None.
521                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                // Never preallocate (or iterate) proportional to an untrusted count: each row
726                // carries `col_count` length-prefixed strings of >= STRING_LEN_PREFIX bytes, so
727                // the remaining payload bounds how many rows can follow. A zero-column row
728                // consumes no bytes (vacuous bound), so a tiny frame could otherwise declare
729                // millions of rows and force a huge allocation (reachable pre-auth). Reject it.
730                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    /// Write this message to an async writer.
784    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    /// Read a message from an async reader.
790    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    /// Read a pre-auth message with a smaller payload limit (4 KB).
797    /// Use this before authentication is complete to prevent oversized
798    /// CONNECT payloads from consuming server memory.
799    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    /// Read a message from an async reader with a configurable payload limit.
806    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    // Every parameter consumes at least its one-byte tag.
879    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    // Every cell has at least a one-byte type and four-byte body length.
1021    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        // New-format Connect that explicitly carries no username.
1344        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        // Simulate an OLD client frame: db_name + password, with NO username
1367        // bytes at all. Must decode with username: None (backward compat).
1368        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        // Oldest client: db_name only, no password and no username bytes.
1390        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, // frame header
1439                0x01, 0x00, // one column
1440                0x01, 0x00, 0x00, 0x00, b'x', // column name
1441                0x01, 0x00, 0x00, 0x00, // one row
1442                0x01, 0x00, 0x00, 0x00, b'y', // one legacy string cell
1443            ]
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        // Same tag, payload one byte longer, identical prefix.
1645        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        // An old client's decoder (this decode path, unchanged) must read the
1654        // message string by its length prefix and skip the trailing class
1655        // byte a new server appends.
1656        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        // Legacy frame from an old server: no trailing byte, no class.
1677        let plain = Message::Error {
1678            message: "authentication failed".into(),
1679        }
1680        .encode();
1681        assert_eq!(decode_error_class(&plain), None);
1682
1683        // Non-error frames never yield a class.
1684        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        // These values are the documented wire contract (docs/errors.md).
1691        // Appending new classes is fine; renumbering is a protocol break.
1692        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        // The new frame must use the dedicated 0x04 tag.
1737        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!(&params[0], WireParam::Str(s) if s == r#"a"b\c; drop User"#));
1743                assert!(matches!(&params[1], WireParam::Int(-7)));
1744                assert!(matches!(&params[2], WireParam::Bool(true)));
1745                assert!(matches!(&params[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!(&params[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        // Feed a wide range of malformed/truncated byte sequences to the
1980        // wire-protocol decode path. Every one must return Err, never panic:
1981        // a malformed client message must not crash the server.
1982        let cases: Vec<Vec<u8>> = vec![
1983            vec![],                                   // empty
1984            vec![0x03],                               // 1 byte, shorter than header
1985            vec![0x03, 0x00, 0x00, 0x00, 0x00],       // 5 bytes, header truncated by one
1986            vec![0xFF, 0x00, 0x00, 0x00, 0x00, 0x00], // unknown message type
1987            // QUERY with payload_len far exceeding the frame.
1988            vec![0x03, 0x00, 0xFF, 0xFF, 0xFF, 0xFF],
1989            // SYNC_PULL with only a replica id and no fixed fields.
1990            {
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            // SYNC_PULL_RESULT with an amplified retained-unit count.
1998            {
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            // CONNECT claiming a string len of 0xFFFFFFFF but no data.
2007            vec![0x01, 0x00, 0x04, 0x00, 0x00, 0x00, 0xFF, 0xFF, 0xFF, 0xFF],
2008            // RESULT_ROWS claiming a huge column count with no data.
2009            vec![0x07, 0x00, 0x02, 0x00, 0x00, 0x00, 0xFF, 0xFF],
2010            // RESULT_OK with a truncated 8-byte affected field.
2011            vec![0x09, 0x00, 0x03, 0x00, 0x00, 0x00, 0x01, 0x02, 0x03],
2012            // QUERY_PARAMS (0x04) claiming a query string len with no data.
2013            vec![0x04, 0x00, 0x04, 0x00, 0x00, 0x00, 0xFF, 0xFF, 0xFF, 0xFF],
2014            // QUERY_PARAMS: empty query string, claims 1 param, no param bytes.
2015            vec![
2016                0x04, 0x00, 0x06, 0x00, 0x00, 0x00, // header, payload_len=6
2017                0x00, 0x00, 0x00, 0x00, // query string len = 0
2018                0x01, 0x00, // param count = 1, then nothing
2019            ],
2020            // QUERY_PARAMS: 1 int param with a truncated i64 body.
2021            vec![
2022                0x04, 0x00, 0x0B, 0x00, 0x00, 0x00, // header, payload_len=11
2023                0x00, 0x00, 0x00, 0x00, // query len = 0
2024                0x01, 0x00, // param count = 1
2025                0x01, // tag = int, then only 3 of 8 bytes
2026                0x01, 0x02, 0x03,
2027            ],
2028            // QUERY_PARAMS: 1 str param with a truncated string body.
2029            vec![
2030                0x04, 0x00, 0x0F, 0x00, 0x00, 0x00, // header, payload_len=15
2031                0x00, 0x00, 0x00, 0x00, // query len = 0
2032                0x01, 0x00, // param count = 1
2033                0x04, // tag = str
2034                0xFF, 0xFF, 0xFF, 0xFF, // str len huge, no data
2035            ],
2036            // QUERY_PARAMS: unknown param tag byte.
2037            vec![
2038                0x04, 0x00, 0x0B, 0x00, 0x00, 0x00, // header, payload_len=11
2039                0x00, 0x00, 0x00, 0x00, // query len = 0
2040                0x01, 0x00, // param count = 1
2041                0x63, // bogus tag
2042            ],
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        // Exhaustively walk every truncation of a well-formed frame plus a
2056        // few byte mutations. None may panic.
2057        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            // Every prefix that is not the full frame must be rejected, not panic.
2064            let _ = Message::decode(&valid[..end]);
2065        }
2066        // Mutate each byte to a high value and confirm no panic.
2067        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        // A tiny frame that declares col_count=0 and row_count=10_000_000.
2077        // With zero columns each row consumes no bytes, so the old decoder
2078        // would allocate/iterate 10M empty rows from ~12 bytes (reachable
2079        // pre-auth). The amplification guard must reject it.
2080        let mut payload = Vec::new();
2081        payload.extend_from_slice(&0u16.to_le_bytes()); // col_count = 0
2082        payload.extend_from_slice(&10_000_000u32.to_le_bytes()); // row_count = 10M
2083        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        // A normal small ResultRows must still round-trip unchanged.
2092        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}