Skip to main content

powdb_server/
handler.rs

1use crate::metrics::{Metrics, QueryOutcome, SyncOperation, SyncOutcome, SyncRepairLabel};
2use crate::protocol::{Message, WireParam, WireRetainedUnit, WireSyncRepairAction, WireSyncStatus};
3use powdb_auth::{Permission, Role, UserStore};
4use powdb_query::executor::{is_read_only_statement, Engine, WalDurabilityTicket};
5use powdb_query::parser;
6use powdb_query::result::{QueryError, QueryResult};
7use powdb_query::sql;
8use powdb_storage::types::Value;
9use powdb_sync::{
10    acknowledge_replica_apply, read_identity, read_units_through, replica_sync_status,
11    retained_segments_dir, validate_retained_tail_available, validate_v1_retained_units_applyable,
12    ReplicaSyncStatus, RetainedUnit, SyncRepairAction, RETAINED_SEGMENT_FORMAT_VERSION,
13};
14use std::collections::HashMap;
15use std::net::IpAddr;
16use std::path::{Path, PathBuf};
17use std::sync::{Arc, Mutex, RwLock};
18use std::time::{Duration, Instant};
19use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt, BufReader, BufWriter};
20use tokio::sync::{watch, OwnedSemaphorePermit, Semaphore};
21use tracing::{debug, error, info, warn};
22use zeroize::Zeroizing;
23
24/// Tracks per-IP authentication failure counts for rate limiting.
25pub type AuthRateLimiter = Arc<Mutex<HashMap<IpAddr, (u32, Instant)>>>;
26
27/// Gate that serializes wire-protocol statements while an explicit
28/// transaction is open on any connection. The connection that runs `begin`
29/// keeps an owned permit until `commit`, `rollback`, disconnect, or timeout,
30/// preventing other connections from observing or joining uncommitted state.
31pub type TxGate = Arc<Semaphore>;
32
33/// Create a transaction gate for a shared engine.
34pub fn new_tx_gate() -> TxGate {
35    Arc::new(Semaphore::new(1))
36}
37
38/// Maximum query text length accepted from the wire (1 MB).
39const MAX_QUERY_LENGTH: usize = 1024 * 1024;
40
41/// Server-side cap for private sync pull batches.
42const MAX_SYNC_PULL_UNITS: u32 = 4096;
43
44/// Server-side cap for retained-unit payload bytes in a private sync pull.
45const MAX_SYNC_PULL_BYTES: u64 = 16 * 1024 * 1024;
46
47/// Maximum encoded response payload size (64 MB). The wire format is still a
48/// single frame today, so oversized result sets must fail cleanly instead of
49/// building an unbounded `Vec<Vec<String>>` and frame in memory.
50#[cfg(not(test))]
51const MAX_RESPONSE_PAYLOAD_SIZE: usize = 64 * 1024 * 1024;
52#[cfg(test)]
53const MAX_RESPONSE_PAYLOAD_SIZE: usize = 1024;
54
55/// Timeout for writing a response to the client. Prevents slow-drain
56/// clients from blocking the handler indefinitely.
57const WRITE_TIMEOUT: Duration = Duration::from_secs(30);
58
59/// Maximum number of auth failures per IP within the rate-limit window.
60const MAX_AUTH_FAILURES: u32 = 5;
61
62/// Window during which auth failures are counted (60 seconds).
63const AUTH_FAILURE_WINDOW: Duration = Duration::from_secs(60);
64
65/// Create a new shared rate limiter.
66pub fn new_rate_limiter() -> AuthRateLimiter {
67    Arc::new(Mutex::new(HashMap::new()))
68}
69
70/// Check whether an IP is rate-limited and record a failure if requested.
71/// Returns `true` if the IP should be rejected.
72fn is_rate_limited(limiter: &AuthRateLimiter, ip: IpAddr) -> bool {
73    let mut map = limiter.lock().unwrap_or_else(|e| e.into_inner());
74    // Clean up stale entries while we have the lock.
75    let now = Instant::now();
76    map.retain(|_, (_, ts)| now.duration_since(*ts) < AUTH_FAILURE_WINDOW);
77
78    if let Some((count, _)) = map.get(&ip) {
79        *count >= MAX_AUTH_FAILURES
80    } else {
81        false
82    }
83}
84
85/// Record an auth failure for the given IP.
86fn record_auth_failure(limiter: &AuthRateLimiter, ip: IpAddr) {
87    let mut map = limiter.lock().unwrap_or_else(|e| e.into_inner());
88    let now = Instant::now();
89    let entry = map.entry(ip).or_insert((0, now));
90    // Reset counter if the window has elapsed.
91    if now.duration_since(entry.1) >= AUTH_FAILURE_WINDOW {
92        *entry = (1, now);
93    } else {
94        entry.0 += 1;
95    }
96}
97
98/// Clear the failure counter on successful auth.
99fn clear_auth_failures(limiter: &AuthRateLimiter, ip: IpAddr) {
100    let mut map = limiter.lock().unwrap_or_else(|e| e.into_inner());
101    map.remove(&ip);
102}
103
104/// Constant-time password comparison. Hashes both inputs to fixed-size
105/// SHA-256 digests so neither length nor content leaks through timing.
106fn constant_time_eq(a: &[u8], b: &[u8]) -> bool {
107    use sha2::{Digest, Sha256};
108    let ha = Sha256::digest(a);
109    let hb = Sha256::digest(b);
110    let mut diff = 0u8;
111    for (x, y) in ha.iter().zip(hb.iter()) {
112        diff |= x ^ y;
113    }
114    diff == 0
115}
116
117/// An authenticated connection's identity. Bound at connect time and consulted
118/// on every query by `dispatch_query` to enforce the user's role: a
119/// `readonly` principal may only execute read statements.
120#[derive(Debug, Clone, PartialEq, Eq)]
121pub struct Principal {
122    pub name: String,
123    pub role: String,
124}
125
126/// Whether a parsed statement is data-definition (schema) work: creating,
127/// altering, or dropping a type or view. `explain <ddl>` is classified by its
128/// inner statement so `explain drop User` needs the same permission as
129/// `drop User`. Mutations that change *rows* (insert/update/delete/upsert/
130/// refresh) and transaction control are NOT DDL — they fall under `Write`.
131fn is_ddl_statement(stmt: &powdb_query::ast::Statement) -> bool {
132    use powdb_query::ast::Statement;
133    let inner = match stmt {
134        Statement::Explain(inner) => inner.as_ref(),
135        other => other,
136    };
137    matches!(
138        inner,
139        Statement::CreateType(_)
140            | Statement::AlterTable(_)
141            | Statement::DropTable(_)
142            | Statement::CreateView(_)
143            | Statement::DropView(_)
144    )
145}
146
147/// The capability a parsed statement requires under the RBAC lattice
148/// (`crates/auth/src/role.rs`). Reads need [`Permission::Read`]; schema
149/// definition needs [`Permission::Ddl`]; every other mutation needs
150/// [`Permission::Write`]. [`Permission::Admin`] is reserved for user/role
151/// management, which is CLI-only today and never reaches this wire path.
152fn required_permission(stmt: &powdb_query::ast::Statement) -> Permission {
153    if is_read_only_statement(stmt) {
154        Permission::Read
155    } else if is_ddl_statement(stmt) {
156        Permission::Ddl
157    } else {
158        Permission::Write
159    }
160}
161
162/// Enforce the principal's role against a parsed statement using the full
163/// permission lattice. Reads are always permitted (any authenticated role can
164/// read — unknown role names still read but fail closed on any mutation).
165/// Mutations require the specific capability the statement maps to: row
166/// mutations need `Write`, schema changes need `Ddl`. Unknown role names
167/// resolve to no builtin and therefore grant nothing beyond reads.
168///
169/// Classification uses the parsed AST via
170/// [`powdb_query::executor::is_read_only_statement`] — the exact same
171/// classifier the RwLock read/write split relies on — so the permission
172/// boundary and the concurrency boundary can never disagree.
173fn check_statement_permitted(
174    principal: Option<&Principal>,
175    stmt: &powdb_query::ast::Statement,
176) -> Result<(), QueryError> {
177    let Some(p) = principal else {
178        // No per-user identity (shared-password or open mode): full access,
179        // byte-identical to the pre-RBAC behavior.
180        return Ok(());
181    };
182    // Reads are permitted for every authenticated principal (preserves the
183    // pre-lattice contract that any connected role may run read-only queries).
184    if is_read_only_statement(stmt) {
185        return Ok(());
186    }
187    let needed = required_permission(stmt);
188    if Role::builtin(&p.role).is_some_and(|r| r.allows(needed)) {
189        return Ok(());
190    }
191    let kind = if needed == Permission::Ddl {
192        "schema-definition"
193    } else {
194        "write"
195    };
196    Err(QueryError::Execution(format!(
197        "permission denied: role '{}' cannot execute {kind} statements",
198        p.role
199    )))
200}
201
202/// Result of the connect-time authentication decision.
203#[derive(Debug, Clone, PartialEq, Eq)]
204pub enum AuthOutcome {
205    /// Authenticated. `principal` is `Some` when a named user authenticated via
206    /// the UserStore, and `None` for the legacy shared-password / open paths
207    /// where there is no per-user identity.
208    Authenticated { principal: Option<Principal> },
209    /// Rejected. The caller sends a generic "authentication failed" error and
210    /// records a rate-limit failure — it must not reveal which check failed.
211    Rejected,
212}
213
214/// Pure, exhaustively-testable authentication decision for a CONNECT handshake.
215///
216/// Policy:
217/// - If `users` has at least one user, multi-user auth is in force: a
218///   `username` is required and `users.authenticate(username, password)` must
219///   succeed. Unknown user, wrong password, or a missing username all reject
220///   with an indistinguishable `Rejected` (no user-vs-password leak).
221/// - If `users` is empty, fall back verbatim to the legacy behavior: when
222///   `expected_password` is `Some`, the candidate must match it (constant time);
223///   when `None`, no auth is required (open). The `username` is ignored here so
224///   that a new client talking to a shared-password server still connects.
225pub fn authenticate_connect(
226    users: &UserStore,
227    expected_password: Option<&str>,
228    username: Option<&str>,
229    password: Option<&str>,
230) -> AuthOutcome {
231    if !users.is_empty() {
232        // Multi-user mode: a username is mandatory.
233        let Some(name) = username else {
234            return AuthOutcome::Rejected;
235        };
236        let Some(candidate) = password else {
237            return AuthOutcome::Rejected;
238        };
239        match users.authenticate(name, candidate) {
240            Some(user) => AuthOutcome::Authenticated {
241                principal: Some(Principal {
242                    name: user.name.clone(),
243                    role: user.role.clone(),
244                }),
245            },
246            None => AuthOutcome::Rejected,
247        }
248    } else {
249        // Legacy shared-password fallback (byte-identical to prior behavior).
250        match expected_password {
251            Some(expected) => {
252                if password.is_some_and(|p| constant_time_eq(p.as_bytes(), expected.as_bytes())) {
253                    AuthOutcome::Authenticated { principal: None }
254                } else {
255                    AuthOutcome::Rejected
256                }
257            }
258            None => AuthOutcome::Authenticated { principal: None },
259        }
260    }
261}
262
263/// Error messages that are safe to forward to the client verbatim.
264const SAFE_ERROR_PREFIXES: &[&str] = &[
265    "table not found",
266    "column not found",
267    "parse error",
268    "type mismatch",
269    "unknown table",
270    "unknown column",
271    "unknown function",
272    "syntax error",
273    "expected",
274    "unexpected",
275    "missing",
276    "duplicate",
277    "invalid",
278    "cannot",
279    "no such",
280    "already exists",
281    "permission denied",
282    "row too large",
283    "unique constraint violation",
284    // Resource-limit errors carry actionable guidance (e.g. "add a LIMIT
285    // clause") and leak no internal state, so surface them verbatim instead
286    // of masking them to the generic message. See QueryError::{SortLimit,
287    // JoinLimit,MemoryLimit}Exceeded in crates/query/src/result.rs.
288    "sort input exceeds",
289    "join result exceeds",
290    "query exceeded memory budget",
291    "result too large",
292    // A failed covering fsync means the statement executed in memory but was
293    // never made durable — the client MUST be able to distinguish this from
294    // an ordinary failed query (the statement may still be visible until the
295    // server restarts). The io::Error detail leaks no internal state.
296    "wal durability sync failed",
297];
298
299/// Sanitize an error message before sending it to the client.
300/// Known safe errors are passed through; everything else is replaced
301/// with a generic message to avoid leaking internal details.
302fn sanitize_error(e: &str) -> String {
303    let lower = e.to_lowercase();
304    for prefix in SAFE_ERROR_PREFIXES {
305        if lower.starts_with(prefix) {
306            return e.to_string();
307        }
308    }
309    "query execution error".into()
310}
311
312/// Write a message to the client with a timeout. Returns false if the
313/// write failed or timed out (caller should close the connection).
314async fn write_msg<W: AsyncWrite + Unpin>(writer: &mut BufWriter<W>, msg: &Message) -> bool {
315    let write_fut = async {
316        if msg.write_to(writer).await.is_err() {
317            return false;
318        }
319        writer.flush().await.is_ok()
320    };
321    tokio::time::timeout(WRITE_TIMEOUT, write_fut)
322        .await
323        .unwrap_or_default()
324}
325
326/// Options for a single connection, bundled to keep `handle_connection`'s
327/// argument list short.
328pub struct ConnOpts<'a> {
329    pub engine: Arc<RwLock<Engine>>,
330    pub tx_gate: TxGate,
331    /// Expected client password. Wrapped in `Zeroizing` so the secret is wiped
332    /// from memory on drop (defends against leaking via a core dump).
333    pub expected_password: Option<Zeroizing<String>>,
334    /// Multi-user store loaded from the data dir at startup. When it has users,
335    /// the handshake authenticates `(username, password)` against it; when empty
336    /// the server falls back to `expected_password`. Shared across connections.
337    pub users: Arc<UserStore>,
338    pub shutdown_rx: &'a mut watch::Receiver<bool>,
339    pub idle_timeout: Duration,
340    pub query_timeout: Duration,
341    pub rate_limiter: Option<&'a AuthRateLimiter>,
342    pub peer_addr: Option<std::net::SocketAddr>,
343    /// Shared server metrics. Always present; tests pass `Arc::new(Metrics::new())`.
344    pub metrics: Arc<Metrics>,
345}
346
347/// Execute a query against the engine under the RwLock. Read-only
348/// statements acquire `.read()` so concurrent SELECTs can scan in
349/// parallel; mutations acquire `.write()`.
350///
351/// When `principal` is `Some`, the user's role is enforced first: a role
352/// without the `Write` permission (i.e. `readonly`) gets a clean
353/// "permission denied" error for any non-read statement, before any lock
354/// is taken or any engine state is touched.
355/// A statement's execution result plus its (not-yet-waited) WAL durability
356/// obligation. The ticket is settled by [`finalize_durability`] AFTER the
357/// TxGate permit is dropped, so overlapping committers on other connections
358/// can share a single fsync (group commit).
359type DispatchOutcome = (Result<QueryResult, QueryError>, Option<WalDurabilityTicket>);
360
361/// Execute a mutating statement under the engine write lock with WAL group
362/// commit: the commit's fsync obligation is registered (not performed) while
363/// the lock is held, then the lock is dropped and the un-waited ticket is
364/// returned. The caller must settle the ticket (via [`finalize_durability`])
365/// before acknowledging the statement — and must do so with the TxGate
366/// permit already released, or committers can never overlap.
367fn execute_write_deferred(
368    engine: &Arc<RwLock<Engine>>,
369    f: impl FnOnce(&mut Engine) -> Result<QueryResult, QueryError>,
370) -> DispatchOutcome {
371    let mut eng = match engine.write() {
372        Ok(eng) => eng,
373        Err(e) => {
374            return (
375                Err(QueryError::Execution(format!("lock poisoned: {e}"))),
376                None,
377            )
378        }
379    };
380    eng.run_with_deferred_durability(f)
381}
382
383fn dispatch_query(
384    engine: &Arc<RwLock<Engine>>,
385    query: &str,
386    principal: Option<&Principal>,
387) -> DispatchOutcome {
388    let stmt_result = parser::parse(query).map_err(|e| e.to_string());
389
390    // Role enforcement happens on the parsed AST. Statements that fail to
391    // parse fall through — the engine returns the parse error itself and
392    // can never execute anything for them.
393    if let Ok(stmt) = &stmt_result {
394        if let Err(e) = check_statement_permitted(principal, stmt) {
395            return (Err(e), None);
396        }
397    }
398
399    let can_try_read = matches!(&stmt_result, Ok(s) if is_read_only_statement(s));
400    if can_try_read {
401        let res = {
402            let eng = match engine.read() {
403                Ok(eng) => eng,
404                Err(e) => {
405                    return (
406                        Err(QueryError::Execution(format!("lock poisoned: {e}"))),
407                        None,
408                    )
409                }
410            };
411            eng.execute_powql_readonly(query)
412        };
413        match res {
414            Ok(r) => return (Ok(r), None),
415            Err(QueryError::ReadonlyNeedsWrite) => {
416                // Escalate: fall through to the write path below.
417            }
418            Err(e) => return (Err(e), None),
419        }
420    }
421
422    if matches!(
423        parsed_transaction_control(&stmt_result),
424        Some(TransactionControl::Rollback)
425    ) {
426        let mut eng = match engine.write() {
427            Ok(eng) => eng,
428            Err(e) => {
429                return (
430                    Err(QueryError::Execution(format!("lock poisoned: {e}"))),
431                    None,
432                )
433            }
434        };
435        return (execute_rollback_preserving_sync_if_needed(&mut eng), None);
436    }
437    execute_write_deferred(engine, |eng| eng.execute_powql(query))
438}
439
440fn dispatch_sql_query(
441    engine: &Arc<RwLock<Engine>>,
442    query: &str,
443    principal: Option<&Principal>,
444) -> DispatchOutcome {
445    let stmt_result = sql::parse_sql(query).map_err(|e| e.to_string());
446
447    if let Ok(stmt) = &stmt_result {
448        if let Err(e) = check_statement_permitted(principal, stmt) {
449            return (Err(e), None);
450        }
451    }
452
453    let can_try_read = matches!(&stmt_result, Ok(s) if is_read_only_statement(s));
454    if can_try_read {
455        let res = {
456            let eng = match engine.read() {
457                Ok(eng) => eng,
458                Err(e) => {
459                    return (
460                        Err(QueryError::Execution(format!("lock poisoned: {e}"))),
461                        None,
462                    )
463                }
464            };
465            eng.execute_sql_readonly(query)
466        };
467        match res {
468            Ok(r) => return (Ok(r), None),
469            Err(QueryError::ReadonlyNeedsWrite) => {}
470            Err(e) => return (Err(e), None),
471        }
472    }
473
474    if matches!(
475        parsed_transaction_control(&stmt_result),
476        Some(TransactionControl::Rollback)
477    ) {
478        let mut eng = match engine.write() {
479            Ok(eng) => eng,
480            Err(e) => {
481                return (
482                    Err(QueryError::Execution(format!("lock poisoned: {e}"))),
483                    None,
484                )
485            }
486        };
487        return (execute_rollback_preserving_sync_if_needed(&mut eng), None);
488    }
489    execute_write_deferred(engine, |eng| eng.execute_sql(query))
490}
491
492/// Convert a wire parameter into the query-crate [`ParamValue`] used for
493/// token-level binding.
494fn wire_param_to_value(p: &WireParam) -> powdb_query::ast::ParamValue {
495    use powdb_query::ast::ParamValue;
496    match p {
497        WireParam::Null => ParamValue::Null,
498        WireParam::Int(v) => ParamValue::Int(*v),
499        WireParam::Float(v) => ParamValue::Float(*v),
500        WireParam::Bool(v) => ParamValue::Bool(*v),
501        WireParam::Str(s) => ParamValue::Str(s.clone()),
502    }
503}
504
505/// Parameterized counterpart of [`dispatch_query`]. Routes through the exact
506/// same role-enforcement and read/write escalation logic, but binds the
507/// `$N` placeholders at the token level via the query crate's
508/// `parse_with_params` path. A string parameter can never change the query's
509/// shape — it is substituted as a literal token, not interpolated text.
510#[derive(Debug, Clone, Copy, PartialEq, Eq)]
511enum TransactionControl {
512    Begin,
513    Commit,
514    Rollback,
515}
516
517fn transaction_control(stmt: &powdb_query::ast::Statement) -> Option<TransactionControl> {
518    use powdb_query::ast::Statement;
519    match stmt {
520        Statement::Begin => Some(TransactionControl::Begin),
521        Statement::Commit => Some(TransactionControl::Commit),
522        Statement::Rollback => Some(TransactionControl::Rollback),
523        _ => None,
524    }
525}
526
527fn classify_query_transaction_control(query: &str) -> Option<TransactionControl> {
528    parser::parse(query)
529        .ok()
530        .and_then(|stmt| transaction_control(&stmt))
531}
532
533fn classify_sql_transaction_control(query: &str) -> Option<TransactionControl> {
534    sql::parse_sql(query)
535        .ok()
536        .and_then(|stmt| transaction_control(&stmt))
537}
538
539fn classify_params_transaction_control(
540    query: &str,
541    params: &[WireParam],
542) -> Option<TransactionControl> {
543    let bound: Vec<powdb_query::ast::ParamValue> = params.iter().map(wire_param_to_value).collect();
544    parser::parse_with_params(query, &bound)
545        .ok()
546        .and_then(|stmt| transaction_control(&stmt))
547}
548
549fn parsed_transaction_control(
550    stmt_result: &Result<powdb_query::ast::Statement, String>,
551) -> Option<TransactionControl> {
552    stmt_result.as_ref().ok().and_then(transaction_control)
553}
554
555fn execute_rollback_preserving_sync_if_needed(
556    engine: &mut Engine,
557) -> Result<QueryResult, QueryError> {
558    engine.rollback_transaction_preserving_wal_archive()
559}
560
561fn dispatch_query_with_params(
562    engine: &Arc<RwLock<Engine>>,
563    query: &str,
564    params: &[WireParam],
565    principal: Option<&Principal>,
566) -> DispatchOutcome {
567    let bound: Vec<powdb_query::ast::ParamValue> = params.iter().map(wire_param_to_value).collect();
568
569    // Parse once (with params bound) so role enforcement and read/write
570    // classification see exactly the statement that will execute.
571    let stmt_result = parser::parse_with_params(query, &bound).map_err(|e| e.to_string());
572
573    if let Ok(stmt) = &stmt_result {
574        if let Err(e) = check_statement_permitted(principal, stmt) {
575            return (Err(e), None);
576        }
577    }
578
579    let can_try_read = matches!(&stmt_result, Ok(s) if is_read_only_statement(s));
580    if can_try_read {
581        let res = {
582            let eng = match engine.read() {
583                Ok(eng) => eng,
584                Err(e) => {
585                    return (
586                        Err(QueryError::Execution(format!("lock poisoned: {e}"))),
587                        None,
588                    )
589                }
590            };
591            eng.execute_powql_readonly_with_params(query, &bound)
592        };
593        match res {
594            Ok(r) => return (Ok(r), None),
595            Err(QueryError::ReadonlyNeedsWrite) => {
596                // Escalate to the write path below.
597            }
598            Err(e) => return (Err(e), None),
599        }
600    }
601
602    if matches!(
603        parsed_transaction_control(&stmt_result),
604        Some(TransactionControl::Rollback)
605    ) {
606        let mut eng = match engine.write() {
607            Ok(eng) => eng,
608            Err(e) => {
609                return (
610                    Err(QueryError::Execution(format!("lock poisoned: {e}"))),
611                    None,
612                )
613            }
614        };
615        return (execute_rollback_preserving_sync_if_needed(&mut eng), None);
616    }
617    execute_write_deferred(engine, |eng| eng.execute_powql_with_params(query, &bound))
618}
619
620#[derive(Debug, Clone)]
621struct SyncPullRequest {
622    replica_id: String,
623    since_lsn: u64,
624    max_units: u32,
625    max_bytes: u64,
626    database_id: [u8; 16],
627    primary_generation: u64,
628    wal_format_version: u16,
629    catalog_version: u16,
630    segment_format_version: u16,
631}
632
633#[derive(Debug, Clone, Copy, PartialEq, Eq)]
634enum SyncErrorClass {
635    AuthRequired,
636    PermissionDenied,
637    InvalidReplicaId,
638    ActiveTransaction,
639    QueryExecution,
640    InvalidMaxUnits,
641    InvalidMaxBytes,
642    SyncContext,
643    StatusRead,
644    CursorLsnMismatch,
645    IdentityRead,
646    IdentityOrFormatMismatch,
647    RetainedRead,
648    RetainedUnitEncoding,
649    RetainedChunkNotApplyable,
650    LsnAheadOfRemote,
651    AckValidation,
652    AckUpdate,
653    Internal,
654}
655
656impl SyncErrorClass {
657    const fn as_label(self) -> &'static str {
658        match self {
659            Self::AuthRequired => "auth_required",
660            Self::PermissionDenied => "permission_denied",
661            Self::InvalidReplicaId => "invalid_replica_id",
662            Self::ActiveTransaction => "active_transaction",
663            Self::QueryExecution => "query_execution",
664            Self::InvalidMaxUnits => "invalid_max_units",
665            Self::InvalidMaxBytes => "invalid_max_bytes",
666            Self::SyncContext => "sync_context",
667            Self::StatusRead => "status_read",
668            Self::CursorLsnMismatch => "cursor_lsn_mismatch",
669            Self::IdentityRead => "identity_read",
670            Self::IdentityOrFormatMismatch => "identity_or_format_mismatch",
671            Self::RetainedRead => "retained_read",
672            Self::RetainedUnitEncoding => "retained_unit_encoding",
673            Self::RetainedChunkNotApplyable => "retained_chunk_not_applyable",
674            Self::LsnAheadOfRemote => "lsn_ahead_of_remote",
675            Self::AckValidation => "ack_validation",
676            Self::AckUpdate => "ack_update",
677            Self::Internal => "internal",
678        }
679    }
680}
681
682#[derive(Debug, Clone)]
683struct SyncDecision {
684    message: Message,
685    error_class: Option<SyncErrorClass>,
686}
687
688impl SyncDecision {
689    fn ok(message: Message) -> Self {
690        Self {
691            message,
692            error_class: None,
693        }
694    }
695
696    fn error(class: SyncErrorClass, message: impl Into<String>) -> Self {
697        Self {
698            message: Message::Error {
699                message: message.into(),
700            },
701            error_class: Some(class),
702        }
703    }
704}
705
706fn check_sync_protocol_permitted(
707    credential_authenticated: bool,
708    principal: Option<&Principal>,
709) -> Result<(), (SyncErrorClass, String)> {
710    if !credential_authenticated {
711        return Err((
712            SyncErrorClass::AuthRequired,
713            "sync protocol requires authentication".to_string(),
714        ));
715    }
716    if let Some(principal) = principal {
717        let allowed =
718            Role::builtin(&principal.role).is_some_and(|role| role.allows(Permission::Write));
719        if !allowed {
720            return Err((
721                SyncErrorClass::PermissionDenied,
722                format!(
723                    "permission denied: role '{}' cannot use sync protocol",
724                    principal.role
725                ),
726            ));
727        }
728    }
729    Ok(())
730}
731
732fn validate_wire_replica_id(replica_id: &str) -> Result<(), String> {
733    if replica_id.is_empty() {
734        return Err("replica id must be non-empty".to_string());
735    }
736    if replica_id.len() > 128 {
737        return Err("replica id must be at most 128 bytes".to_string());
738    }
739    if !replica_id
740        .bytes()
741        .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'-' | b'_' | b'.' | b':'))
742    {
743        return Err("replica id contains unsupported characters".to_string());
744    }
745    Ok(())
746}
747
748fn sync_context(engine: &Arc<RwLock<Engine>>) -> Result<(PathBuf, u64), String> {
749    let engine = engine
750        .read()
751        .map_err(|e| format!("lock poisoned while reading sync context: {e}"))?;
752    let catalog = engine.catalog();
753    Ok((catalog.data_dir().to_path_buf(), catalog.max_lsn()))
754}
755
756fn wire_repair_action(action: SyncRepairAction) -> WireSyncRepairAction {
757    match action {
758        SyncRepairAction::None => WireSyncRepairAction::None,
759        SyncRepairAction::Pull => WireSyncRepairAction::Pull,
760        SyncRepairAction::AwaitArchive => WireSyncRepairAction::AwaitArchive,
761        SyncRepairAction::Rebootstrap => WireSyncRepairAction::Rebootstrap,
762    }
763}
764
765fn wire_sync_status(status: ReplicaSyncStatus) -> WireSyncStatus {
766    WireSyncStatus {
767        replica_id: status.replica_id,
768        active: status.active,
769        last_applied_lsn: status.last_applied_lsn,
770        remote_lsn: status.remote_lsn,
771        servable_lsn: status.servable_lsn,
772        unarchived_lsn: status.unarchived_lsn,
773        lag_lsn: status.lag_lsn,
774        lag_bytes: status.lag_bytes,
775        lag_ms: status.lag_ms,
776        stale: status.stale,
777        repair_action: wire_repair_action(status.repair_action),
778        last_sync_error: status.last_sync_error,
779    }
780}
781
782fn wire_retained_unit(unit: RetainedUnit) -> WireRetainedUnit {
783    WireRetainedUnit {
784        tx_id: unit.tx_id,
785        record_type: unit.record_type,
786        lsn: unit.lsn,
787        data: unit.data,
788    }
789}
790
791fn sync_operation_outcome(message: &Message) -> SyncOutcome {
792    match message {
793        Message::SyncStatusResult { .. }
794        | Message::SyncPullResult { .. }
795        | Message::SyncAckResult { .. } => SyncOutcome::Ok,
796        _ => SyncOutcome::Error,
797    }
798}
799
800fn sync_operation_label(operation: SyncOperation) -> &'static str {
801    match operation {
802        SyncOperation::Status => "status",
803        SyncOperation::Pull => "pull",
804        SyncOperation::Ack => "ack",
805    }
806}
807
808fn sync_pull_payload_bytes(units: &[WireRetainedUnit]) -> u64 {
809    units.iter().fold(0, |total, unit| {
810        total.saturating_add(unit.encoded_len().unwrap_or(0))
811    })
812}
813
814fn sync_repair_label(action: WireSyncRepairAction) -> SyncRepairLabel {
815    match action {
816        WireSyncRepairAction::None => SyncRepairLabel::None,
817        WireSyncRepairAction::Pull => SyncRepairLabel::Pull,
818        WireSyncRepairAction::AwaitArchive => SyncRepairLabel::AwaitArchive,
819        WireSyncRepairAction::Rebootstrap => SyncRepairLabel::Rebootstrap,
820    }
821}
822
823fn wire_repair_action_label(action: WireSyncRepairAction) -> &'static str {
824    match action {
825        WireSyncRepairAction::None => "none",
826        WireSyncRepairAction::Pull => "pull",
827        WireSyncRepairAction::AwaitArchive => "await_archive",
828        WireSyncRepairAction::Rebootstrap => "rebootstrap",
829    }
830}
831
832const FNV1A64_OFFSET: u64 = 0xcbf29ce484222325;
833const FNV1A64_PRIME: u64 = 0x100000001b3;
834const INVALID_REPLICA_FINGERPRINT: &str = "invalid";
835
836fn replica_fingerprint(replica_id: &str) -> String {
837    let mut hash = FNV1A64_OFFSET;
838    for byte in replica_id.as_bytes() {
839        hash ^= u64::from(*byte);
840        hash = hash.wrapping_mul(FNV1A64_PRIME);
841    }
842    format!("{hash:016x}")
843}
844
845fn log_replica_fingerprint(replica_id: &str) -> String {
846    if validate_wire_replica_id(replica_id).is_ok() {
847        replica_fingerprint(replica_id)
848    } else {
849        INVALID_REPLICA_FINGERPRINT.to_string()
850    }
851}
852
853#[derive(Debug, Clone)]
854struct SyncLogContext {
855    replica_fingerprint: String,
856    since_lsn: Option<u64>,
857    applied_lsn: Option<u64>,
858    observed_remote_lsn: Option<u64>,
859    max_units: Option<u32>,
860    max_bytes: Option<u64>,
861}
862
863impl SyncLogContext {
864    fn status(replica_id: &str) -> Self {
865        Self::base(replica_id)
866    }
867
868    fn pull(request: &SyncPullRequest) -> Self {
869        Self {
870            replica_fingerprint: log_replica_fingerprint(&request.replica_id),
871            since_lsn: Some(request.since_lsn),
872            applied_lsn: None,
873            observed_remote_lsn: None,
874            max_units: Some(request.max_units),
875            max_bytes: Some(request.max_bytes),
876        }
877    }
878
879    fn ack(replica_id: &str, applied_lsn: u64, observed_remote_lsn: u64) -> Self {
880        Self {
881            replica_fingerprint: log_replica_fingerprint(replica_id),
882            since_lsn: None,
883            applied_lsn: Some(applied_lsn),
884            observed_remote_lsn: Some(observed_remote_lsn),
885            max_units: None,
886            max_bytes: None,
887        }
888    }
889
890    fn base(replica_id: &str) -> Self {
891        Self {
892            replica_fingerprint: log_replica_fingerprint(replica_id),
893            since_lsn: None,
894            applied_lsn: None,
895            observed_remote_lsn: None,
896            max_units: None,
897            max_bytes: None,
898        }
899    }
900}
901
902struct SyncExecutionContext<'a> {
903    tx_gate: TxGate,
904    connection_has_transaction: bool,
905    operation: SyncOperation,
906    log_context: SyncLogContext,
907    metrics: &'a Arc<Metrics>,
908    query_timeout: Duration,
909}
910
911fn log_sync_decision(
912    operation: SyncOperation,
913    context: &SyncLogContext,
914    elapsed: Duration,
915    decision: &SyncDecision,
916) {
917    let operation = sync_operation_label(operation);
918    let elapsed_ms = elapsed.as_secs_f64() * 1000.0;
919    match &decision.message {
920        Message::SyncStatusResult { status } => {
921            let repair_action = wire_repair_action_label(status.repair_action);
922            if status.stale || status.repair_action != WireSyncRepairAction::None {
923                info!(
924                    operation = operation,
925                    replica_fingerprint = %context.replica_fingerprint,
926                    remote_lsn = status.remote_lsn,
927                    last_applied_lsn = ?status.last_applied_lsn,
928                    servable_lsn = ?status.servable_lsn,
929                    unarchived_lsn = ?status.unarchived_lsn,
930                    lag_lsn = ?status.lag_lsn,
931                    lag_bytes = ?status.lag_bytes,
932                    lag_ms = ?status.lag_ms,
933                    stale = status.stale,
934                    repair_action,
935                    elapsed_ms,
936                    "sync decision"
937                );
938            } else {
939                debug!(
940                    operation = operation,
941                    replica_fingerprint = %context.replica_fingerprint,
942                    remote_lsn = status.remote_lsn,
943                    last_applied_lsn = ?status.last_applied_lsn,
944                    repair_action,
945                    elapsed_ms,
946                    "sync decision"
947                );
948            }
949        }
950        Message::SyncPullResult {
951            status,
952            units,
953            has_more,
954        } => {
955            let repair_action = wire_repair_action_label(status.repair_action);
956            let units_len = units.len();
957            let payload_bytes = sync_pull_payload_bytes(units);
958            if *has_more || status.stale || status.repair_action != WireSyncRepairAction::None {
959                info!(
960                    operation = operation,
961                    replica_fingerprint = %context.replica_fingerprint,
962                    since_lsn = ?context.since_lsn,
963                    max_units = ?context.max_units,
964                    max_bytes = ?context.max_bytes,
965                    units = units_len,
966                    payload_bytes,
967                    has_more = *has_more,
968                    remote_lsn = status.remote_lsn,
969                    last_applied_lsn = ?status.last_applied_lsn,
970                    servable_lsn = ?status.servable_lsn,
971                    unarchived_lsn = ?status.unarchived_lsn,
972                    lag_lsn = ?status.lag_lsn,
973                    stale = status.stale,
974                    repair_action,
975                    elapsed_ms,
976                    "sync decision"
977                );
978            } else {
979                debug!(
980                    operation = operation,
981                    replica_fingerprint = %context.replica_fingerprint,
982                    since_lsn = ?context.since_lsn,
983                    units = units_len,
984                    payload_bytes,
985                    has_more = *has_more,
986                    remote_lsn = status.remote_lsn,
987                    repair_action,
988                    elapsed_ms,
989                    "sync decision"
990                );
991            }
992        }
993        Message::SyncAckResult {
994            previous_applied_lsn,
995            applied_lsn,
996            remote_lsn,
997            advanced,
998            status,
999        } => {
1000            let repair_action = wire_repair_action_label(status.repair_action);
1001            if *advanced || status.stale || status.repair_action != WireSyncRepairAction::None {
1002                info!(
1003                    operation = operation,
1004                    replica_fingerprint = %context.replica_fingerprint,
1005                    requested_applied_lsn = ?context.applied_lsn,
1006                    observed_remote_lsn = ?context.observed_remote_lsn,
1007                    previous_applied_lsn = *previous_applied_lsn,
1008                    applied_lsn = *applied_lsn,
1009                    remote_lsn = *remote_lsn,
1010                    advanced = *advanced,
1011                    stale = status.stale,
1012                    repair_action,
1013                    elapsed_ms,
1014                    "sync decision"
1015                );
1016            } else {
1017                debug!(
1018                    operation = operation,
1019                    replica_fingerprint = %context.replica_fingerprint,
1020                    requested_applied_lsn = ?context.applied_lsn,
1021                    previous_applied_lsn = *previous_applied_lsn,
1022                    applied_lsn = *applied_lsn,
1023                    remote_lsn = *remote_lsn,
1024                    advanced = *advanced,
1025                    repair_action,
1026                    elapsed_ms,
1027                    "sync decision"
1028                );
1029            }
1030        }
1031        Message::Error { .. } => {
1032            warn!(
1033                operation = operation,
1034                replica_fingerprint = %context.replica_fingerprint,
1035                since_lsn = ?context.since_lsn,
1036                applied_lsn = ?context.applied_lsn,
1037                observed_remote_lsn = ?context.observed_remote_lsn,
1038                max_units = ?context.max_units,
1039                max_bytes = ?context.max_bytes,
1040                error_class = decision
1041                    .error_class
1042                    .unwrap_or(SyncErrorClass::Internal)
1043                    .as_label(),
1044                elapsed_ms,
1045                "sync decision rejected"
1046            );
1047        }
1048        _ => {
1049            debug!(
1050                operation = operation,
1051                replica_fingerprint = %context.replica_fingerprint,
1052                elapsed_ms,
1053                "unexpected sync decision response"
1054            );
1055        }
1056    }
1057}
1058
1059fn trim_to_applyable_v1_prefix(
1060    raw_units: &mut Vec<RetainedUnit>,
1061    wire_units: &mut Vec<WireRetainedUnit>,
1062) -> Result<(), String> {
1063    let mut last_error = None;
1064    while !raw_units.is_empty() {
1065        match validate_v1_retained_units_applyable(raw_units) {
1066            Ok(()) => return Ok(()),
1067            Err(err) => {
1068                last_error = Some(err.to_string());
1069                raw_units.pop();
1070                wire_units.pop();
1071            }
1072        }
1073    }
1074    if let Some(error) = last_error {
1075        return Err(format!(
1076            "sync pull cannot serve an applyable V1 retained chunk with current limits: {error}"
1077        ));
1078    }
1079    Ok(())
1080}
1081
1082fn validate_sync_ack_applyable_boundary(
1083    data_dir: &Path,
1084    replica_id: &str,
1085    applied_lsn: u64,
1086    remote_lsn: u64,
1087) -> Result<(), String> {
1088    let status =
1089        replica_sync_status(data_dir, replica_id, remote_lsn).map_err(|err| err.to_string())?;
1090    let Some(previous_lsn) = status.last_applied_lsn else {
1091        return Ok(());
1092    };
1093    if applied_lsn <= previous_lsn {
1094        return Ok(());
1095    }
1096    let range_len = applied_lsn - previous_lsn;
1097    if range_len > u64::from(MAX_SYNC_PULL_UNITS) {
1098        return Err(format!(
1099            "sync ack range contains {range_len} units; acknowledge ranges no larger than {MAX_SYNC_PULL_UNITS}"
1100        ));
1101    }
1102    let max_units =
1103        usize::try_from(range_len).map_err(|_| "sync ack range is too large to validate")?;
1104    let identity = read_identity(data_dir).map_err(|err| err.to_string())?;
1105    let segment_dir = retained_segments_dir(data_dir);
1106    let units = read_units_through(
1107        &segment_dir,
1108        identity.segment_identity(),
1109        previous_lsn,
1110        applied_lsn,
1111        max_units,
1112    )
1113    .map_err(|err| err.to_string())?;
1114    if units.len() != max_units || units.last().map(|unit| unit.lsn) != Some(applied_lsn) {
1115        return Err(
1116            "sync ack does not cover a complete retained-unit range; rebootstrap required".into(),
1117        );
1118    }
1119    validate_v1_retained_units_applyable(&units).map_err(|err| err.to_string())
1120}
1121
1122#[cfg(test)]
1123fn dispatch_sync_status(
1124    engine: &Arc<RwLock<Engine>>,
1125    replica_id: String,
1126    credential_authenticated: bool,
1127    principal: Option<&Principal>,
1128) -> Message {
1129    dispatch_sync_status_decision(engine, replica_id, credential_authenticated, principal).message
1130}
1131
1132fn dispatch_sync_status_decision(
1133    engine: &Arc<RwLock<Engine>>,
1134    replica_id: String,
1135    credential_authenticated: bool,
1136    principal: Option<&Principal>,
1137) -> SyncDecision {
1138    if let Err((class, message)) =
1139        check_sync_protocol_permitted(credential_authenticated, principal)
1140    {
1141        return SyncDecision::error(class, message);
1142    }
1143    if let Err(message) = validate_wire_replica_id(&replica_id) {
1144        return SyncDecision::error(SyncErrorClass::InvalidReplicaId, message);
1145    }
1146    let (data_dir, remote_lsn) = match sync_context(engine) {
1147        Ok(context) => context,
1148        Err(message) => return SyncDecision::error(SyncErrorClass::SyncContext, message),
1149    };
1150    match replica_sync_status(&data_dir, &replica_id, remote_lsn) {
1151        Ok(status) => SyncDecision::ok(Message::SyncStatusResult {
1152            status: wire_sync_status(status),
1153        }),
1154        Err(err) => SyncDecision::error(SyncErrorClass::StatusRead, err.to_string()),
1155    }
1156}
1157
1158#[cfg(test)]
1159fn dispatch_sync_pull(
1160    engine: &Arc<RwLock<Engine>>,
1161    request: SyncPullRequest,
1162    credential_authenticated: bool,
1163    principal: Option<&Principal>,
1164) -> Message {
1165    dispatch_sync_pull_decision(engine, request, credential_authenticated, principal).message
1166}
1167
1168fn dispatch_sync_pull_decision(
1169    engine: &Arc<RwLock<Engine>>,
1170    request: SyncPullRequest,
1171    credential_authenticated: bool,
1172    principal: Option<&Principal>,
1173) -> SyncDecision {
1174    if let Err((class, message)) =
1175        check_sync_protocol_permitted(credential_authenticated, principal)
1176    {
1177        return SyncDecision::error(class, message);
1178    }
1179    if let Err(message) = validate_wire_replica_id(&request.replica_id) {
1180        return SyncDecision::error(SyncErrorClass::InvalidReplicaId, message);
1181    }
1182    if request.max_units == 0 || request.max_units > MAX_SYNC_PULL_UNITS {
1183        return SyncDecision::error(
1184            SyncErrorClass::InvalidMaxUnits,
1185            format!("sync pull maxUnits must be between 1 and {MAX_SYNC_PULL_UNITS}"),
1186        );
1187    }
1188    if request.max_bytes == 0 || request.max_bytes > MAX_SYNC_PULL_BYTES {
1189        return SyncDecision::error(
1190            SyncErrorClass::InvalidMaxBytes,
1191            format!("sync pull maxBytes must be between 1 and {MAX_SYNC_PULL_BYTES}"),
1192        );
1193    }
1194
1195    let (data_dir, remote_lsn) = match sync_context(engine) {
1196        Ok(context) => context,
1197        Err(message) => return SyncDecision::error(SyncErrorClass::SyncContext, message),
1198    };
1199    let status = match replica_sync_status(&data_dir, &request.replica_id, remote_lsn) {
1200        Ok(status) => status,
1201        Err(err) => {
1202            return SyncDecision::error(SyncErrorClass::StatusRead, err.to_string());
1203        }
1204    };
1205    let Some(cursor_lsn) = status.last_applied_lsn else {
1206        return SyncDecision::ok(Message::SyncPullResult {
1207            status: wire_sync_status(status),
1208            units: Vec::new(),
1209            has_more: false,
1210        });
1211    };
1212    if status.repair_action != SyncRepairAction::Pull {
1213        return SyncDecision::ok(Message::SyncPullResult {
1214            status: wire_sync_status(status),
1215            units: Vec::new(),
1216            has_more: false,
1217        });
1218    }
1219    if request.since_lsn != cursor_lsn {
1220        return SyncDecision::error(
1221            SyncErrorClass::CursorLsnMismatch,
1222            format!(
1223                "sync pull sinceLsn {} does not match primary cursor LSN {cursor_lsn}",
1224                request.since_lsn
1225            ),
1226        );
1227    }
1228
1229    let identity = match read_identity(&data_dir) {
1230        Ok(identity) => identity,
1231        Err(err) => {
1232            return SyncDecision::error(SyncErrorClass::IdentityRead, err.to_string());
1233        }
1234    };
1235    let expected = identity.segment_identity();
1236    if request.database_id != expected.database_id
1237        || request.primary_generation != expected.primary_generation
1238        || request.wal_format_version != expected.wal_format_version
1239        || request.catalog_version != expected.catalog_version
1240        || request.segment_format_version != RETAINED_SEGMENT_FORMAT_VERSION
1241    {
1242        return SyncDecision::error(
1243            SyncErrorClass::IdentityOrFormatMismatch,
1244            "sync pull identity or format version mismatch; rebootstrap required",
1245        );
1246    }
1247
1248    let effective_max_units = request.max_units.min(MAX_SYNC_PULL_UNITS) as usize;
1249    let requested_through_lsn = request
1250        .since_lsn
1251        .saturating_add(request.max_units as u64)
1252        .min(remote_lsn)
1253        .min(status.servable_lsn.unwrap_or(request.since_lsn));
1254    let segment_dir = retained_segments_dir(&data_dir);
1255    if requested_through_lsn > request.since_lsn {
1256        if let Err(err) = validate_retained_tail_available(
1257            &segment_dir,
1258            expected,
1259            request.since_lsn,
1260            requested_through_lsn,
1261        ) {
1262            let mut rebootstrap_status = status;
1263            rebootstrap_status.stale = true;
1264            rebootstrap_status.repair_action = SyncRepairAction::Rebootstrap;
1265            rebootstrap_status.last_sync_error = Some(format!(
1266                "retained history is unavailable; rebootstrap required: {err}"
1267            ));
1268            return SyncDecision::ok(Message::SyncPullResult {
1269                status: wire_sync_status(rebootstrap_status),
1270                units: Vec::new(),
1271                has_more: false,
1272            });
1273        }
1274    }
1275
1276    let raw_units = match read_units_through(
1277        &segment_dir,
1278        expected,
1279        request.since_lsn,
1280        requested_through_lsn,
1281        effective_max_units,
1282    ) {
1283        Ok(units) => units,
1284        Err(err) => {
1285            return SyncDecision::error(SyncErrorClass::RetainedRead, err.to_string());
1286        }
1287    };
1288
1289    let mut selected_raw = Vec::new();
1290    let mut selected = Vec::new();
1291    let mut selected_bytes = 0u64;
1292    for unit in raw_units {
1293        let wire_unit = wire_retained_unit(unit.clone());
1294        let unit_bytes = match wire_unit.encoded_len() {
1295            Ok(bytes) => bytes,
1296            Err(message) => {
1297                return SyncDecision::error(SyncErrorClass::RetainedUnitEncoding, message);
1298            }
1299        };
1300        if selected_bytes.saturating_add(unit_bytes) > request.max_bytes {
1301            if selected.is_empty() {
1302                return SyncDecision::error(
1303                    SyncErrorClass::InvalidMaxBytes,
1304                    "sync pull maxBytes is too small for the next retained unit",
1305                );
1306            }
1307            break;
1308        }
1309        selected_bytes += unit_bytes;
1310        selected_raw.push(unit);
1311        selected.push(wire_unit);
1312    }
1313    if let Err(message) = trim_to_applyable_v1_prefix(&mut selected_raw, &mut selected) {
1314        return SyncDecision::error(SyncErrorClass::RetainedChunkNotApplyable, message);
1315    }
1316
1317    let fetchable_through_lsn = status.servable_lsn.unwrap_or(remote_lsn).min(remote_lsn);
1318    let has_more = selected
1319        .last()
1320        .is_some_and(|unit| unit.lsn < fetchable_through_lsn);
1321    SyncDecision::ok(Message::SyncPullResult {
1322        status: wire_sync_status(status),
1323        units: selected,
1324        has_more,
1325    })
1326}
1327
1328#[cfg(test)]
1329fn dispatch_sync_ack(
1330    engine: &Arc<RwLock<Engine>>,
1331    replica_id: String,
1332    applied_lsn: u64,
1333    observed_remote_lsn: u64,
1334    credential_authenticated: bool,
1335    principal: Option<&Principal>,
1336) -> Message {
1337    dispatch_sync_ack_decision(
1338        engine,
1339        replica_id,
1340        applied_lsn,
1341        observed_remote_lsn,
1342        credential_authenticated,
1343        principal,
1344    )
1345    .message
1346}
1347
1348fn dispatch_sync_ack_decision(
1349    engine: &Arc<RwLock<Engine>>,
1350    replica_id: String,
1351    applied_lsn: u64,
1352    observed_remote_lsn: u64,
1353    credential_authenticated: bool,
1354    principal: Option<&Principal>,
1355) -> SyncDecision {
1356    if let Err((class, message)) =
1357        check_sync_protocol_permitted(credential_authenticated, principal)
1358    {
1359        return SyncDecision::error(class, message);
1360    }
1361    if let Err(message) = validate_wire_replica_id(&replica_id) {
1362        return SyncDecision::error(SyncErrorClass::InvalidReplicaId, message);
1363    }
1364    if applied_lsn > observed_remote_lsn {
1365        return SyncDecision::error(
1366            SyncErrorClass::LsnAheadOfRemote,
1367            format!(
1368                "sync ack appliedLsn {applied_lsn} is ahead of observed remoteLsn {observed_remote_lsn}"
1369            ),
1370        );
1371    }
1372    let (data_dir, remote_lsn) = match sync_context(engine) {
1373        Ok(context) => context,
1374        Err(message) => return SyncDecision::error(SyncErrorClass::SyncContext, message),
1375    };
1376    if observed_remote_lsn > remote_lsn {
1377        return SyncDecision::error(
1378            SyncErrorClass::LsnAheadOfRemote,
1379            format!(
1380                "sync ack remoteLsn {observed_remote_lsn} is ahead of primary LSN {remote_lsn}"
1381            ),
1382        );
1383    }
1384    if let Err(message) =
1385        validate_sync_ack_applyable_boundary(&data_dir, &replica_id, applied_lsn, remote_lsn)
1386    {
1387        return SyncDecision::error(SyncErrorClass::AckValidation, message);
1388    }
1389    match acknowledge_replica_apply(&data_dir, &replica_id, applied_lsn, remote_lsn) {
1390        Ok(summary) => SyncDecision::ok(Message::SyncAckResult {
1391            previous_applied_lsn: summary.previous_applied_lsn,
1392            applied_lsn: summary.applied_lsn,
1393            remote_lsn: summary.remote_lsn,
1394            advanced: summary.advanced,
1395            status: wire_sync_status(summary.status),
1396        }),
1397        Err(err) => SyncDecision::error(SyncErrorClass::AckUpdate, err.to_string()),
1398    }
1399}
1400
1401async fn run_blocking_sync<T, F>(input: T, query_timeout: Duration, f: F) -> SyncDecision
1402where
1403    T: Send + 'static,
1404    F: FnOnce(T) -> SyncDecision + Send + 'static,
1405{
1406    let mut handle = tokio::task::spawn_blocking(move || f(input));
1407    tokio::select! {
1408        result = &mut handle => match result {
1409            Ok(decision) => decision,
1410            Err(err) => SyncDecision::error(SyncErrorClass::Internal, format!("internal error: {err}")),
1411        },
1412        _ = tokio::time::sleep(query_timeout) => match handle.await {
1413            Ok(decision) => decision,
1414            Err(err) => SyncDecision::error(SyncErrorClass::Internal, format!("internal error: {err}")),
1415        },
1416    }
1417}
1418
1419async fn execute_gated_sync<T, F>(context: SyncExecutionContext<'_>, input: T, f: F) -> Message
1420where
1421    T: Send + 'static,
1422    F: FnOnce(T) -> SyncDecision + Send + 'static,
1423{
1424    let SyncExecutionContext {
1425        tx_gate,
1426        connection_has_transaction,
1427        operation,
1428        log_context,
1429        metrics,
1430        query_timeout,
1431    } = context;
1432    let start = Instant::now();
1433    if connection_has_transaction {
1434        let decision = SyncDecision::error(
1435            SyncErrorClass::ActiveTransaction,
1436            "sync protocol is unavailable inside an active transaction",
1437        );
1438        let elapsed = start.elapsed();
1439        log_sync_decision(operation, &log_context, elapsed, &decision);
1440        metrics.record_sync_operation(operation, elapsed, SyncOutcome::Error);
1441        return decision.message;
1442    }
1443
1444    let permit = match tx_gate.acquire_owned().await {
1445        Ok(permit) => permit,
1446        Err(_) => {
1447            let decision =
1448                SyncDecision::error(SyncErrorClass::QueryExecution, "query execution error");
1449            let elapsed = start.elapsed();
1450            log_sync_decision(operation, &log_context, elapsed, &decision);
1451            metrics.record_sync_operation(operation, elapsed, SyncOutcome::Error);
1452            return decision.message;
1453        }
1454    };
1455    let decision = run_blocking_sync(input, query_timeout, f).await;
1456    drop(permit);
1457    match &decision.message {
1458        Message::SyncStatusResult { status } => {
1459            metrics.record_sync_repair_action(operation, sync_repair_label(status.repair_action));
1460        }
1461        Message::SyncPullResult { status, units, .. } => {
1462            metrics.record_sync_repair_action(operation, sync_repair_label(status.repair_action));
1463            metrics.record_sync_pull_payload(units.len() as u64, sync_pull_payload_bytes(units));
1464        }
1465        Message::SyncAckResult {
1466            advanced, status, ..
1467        } => {
1468            metrics.record_sync_repair_action(operation, sync_repair_label(status.repair_action));
1469            if *advanced {
1470                metrics.inc_sync_ack_advanced();
1471            }
1472        }
1473        _ => {}
1474    }
1475    let elapsed = start.elapsed();
1476    log_sync_decision(operation, &log_context, elapsed, &decision);
1477    metrics.record_sync_operation(
1478        operation,
1479        elapsed,
1480        sync_operation_outcome(&decision.message),
1481    );
1482    decision.message
1483}
1484
1485/// Execute one wire query frame and return the response plus its un-waited
1486/// WAL durability ticket. The TxGate permit is managed here and — crucially —
1487/// is already released (bare statements, commit/rollback) by the time this
1488/// returns, so the caller's `finalize_durability` wait happens OUTSIDE the
1489/// gate and overlapping committers can share an fsync.
1490async fn execute_wire_query(
1491    engine: Arc<RwLock<Engine>>,
1492    tx_gate: TxGate,
1493    tx_permit: &mut Option<OwnedSemaphorePermit>,
1494    query: String,
1495    principal: Option<Principal>,
1496    query_timeout: Duration,
1497    metrics: &Arc<Metrics>,
1498) -> (Message, Option<PendingDurability>) {
1499    match classify_query_transaction_control(&query) {
1500        Some(TransactionControl::Begin) => {
1501            if tx_permit.is_some() {
1502                return (
1503                    Message::Error {
1504                        message: sanitize_error("transaction already active"),
1505                    },
1506                    None,
1507                );
1508            }
1509            let permit = match tx_gate.acquire_owned().await {
1510                Ok(permit) => permit,
1511                Err(_) => {
1512                    return (
1513                        Message::Error {
1514                            message: "query execution error".into(),
1515                        },
1516                        None,
1517                    )
1518                }
1519            };
1520            let (response, ticket) = run_blocking_query(
1521                engine,
1522                query,
1523                principal,
1524                query_timeout,
1525                metrics,
1526                |engine, query, principal| dispatch_query(&engine, &query, principal.as_ref()),
1527            )
1528            .await;
1529            if is_success_response(&response) {
1530                *tx_permit = Some(permit);
1531            }
1532            (response, ticket)
1533        }
1534        Some(TransactionControl::Commit | TransactionControl::Rollback) => {
1535            let (response, ticket) = run_blocking_query(
1536                engine,
1537                query,
1538                principal,
1539                query_timeout,
1540                metrics,
1541                |engine, query, principal| dispatch_query(&engine, &query, principal.as_ref()),
1542            )
1543            .await;
1544            if is_success_response(&response) {
1545                // Release the gate BEFORE the caller waits on the commit's
1546                // ticket: the engine work is done and WAL order is fixed, so
1547                // another connection's commit can start (and share the fsync)
1548                // while this one waits.
1549                tx_permit.take();
1550            }
1551            (response, ticket)
1552        }
1553        None if tx_permit.is_some() => {
1554            run_blocking_query(
1555                engine,
1556                query,
1557                principal,
1558                query_timeout,
1559                metrics,
1560                |engine, query, principal| dispatch_query(&engine, &query, principal.as_ref()),
1561            )
1562            .await
1563        }
1564        None => {
1565            let permit = match tx_gate.acquire_owned().await {
1566                Ok(permit) => permit,
1567                Err(_) => {
1568                    return (
1569                        Message::Error {
1570                            message: "query execution error".into(),
1571                        },
1572                        None,
1573                    )
1574                }
1575            };
1576            let out = run_blocking_query(
1577                engine,
1578                query,
1579                principal,
1580                query_timeout,
1581                metrics,
1582                |engine, query, principal| dispatch_query(&engine, &query, principal.as_ref()),
1583            )
1584            .await;
1585            drop(permit);
1586            out
1587        }
1588    }
1589}
1590
1591async fn execute_wire_query_sql(
1592    engine: Arc<RwLock<Engine>>,
1593    tx_gate: TxGate,
1594    tx_permit: &mut Option<OwnedSemaphorePermit>,
1595    query: String,
1596    principal: Option<Principal>,
1597    query_timeout: Duration,
1598    metrics: &Arc<Metrics>,
1599) -> (Message, Option<PendingDurability>) {
1600    match classify_sql_transaction_control(&query) {
1601        Some(TransactionControl::Begin) => {
1602            if tx_permit.is_some() {
1603                return (
1604                    Message::Error {
1605                        message: sanitize_error("transaction already active"),
1606                    },
1607                    None,
1608                );
1609            }
1610            let permit = match tx_gate.acquire_owned().await {
1611                Ok(permit) => permit,
1612                Err(_) => {
1613                    return (
1614                        Message::Error {
1615                            message: "query execution error".into(),
1616                        },
1617                        None,
1618                    )
1619                }
1620            };
1621            let (response, ticket) = run_blocking_query(
1622                engine,
1623                query,
1624                principal,
1625                query_timeout,
1626                metrics,
1627                |engine, query, principal| dispatch_sql_query(&engine, &query, principal.as_ref()),
1628            )
1629            .await;
1630            if is_success_response(&response) {
1631                *tx_permit = Some(permit);
1632            }
1633            (response, ticket)
1634        }
1635        Some(TransactionControl::Commit | TransactionControl::Rollback) => {
1636            let (response, ticket) = run_blocking_query(
1637                engine,
1638                query,
1639                principal,
1640                query_timeout,
1641                metrics,
1642                |engine, query, principal| dispatch_sql_query(&engine, &query, principal.as_ref()),
1643            )
1644            .await;
1645            if is_success_response(&response) {
1646                // See execute_wire_query: release the gate before the
1647                // caller's durability wait so commits can coalesce.
1648                tx_permit.take();
1649            }
1650            (response, ticket)
1651        }
1652        None if tx_permit.is_some() => {
1653            run_blocking_query(
1654                engine,
1655                query,
1656                principal,
1657                query_timeout,
1658                metrics,
1659                |engine, query, principal| dispatch_sql_query(&engine, &query, principal.as_ref()),
1660            )
1661            .await
1662        }
1663        None => {
1664            let permit = match tx_gate.acquire_owned().await {
1665                Ok(permit) => permit,
1666                Err(_) => {
1667                    return (
1668                        Message::Error {
1669                            message: "query execution error".into(),
1670                        },
1671                        None,
1672                    )
1673                }
1674            };
1675            let out = run_blocking_query(
1676                engine,
1677                query,
1678                principal,
1679                query_timeout,
1680                metrics,
1681                |engine, query, principal| dispatch_sql_query(&engine, &query, principal.as_ref()),
1682            )
1683            .await;
1684            drop(permit);
1685            out
1686        }
1687    }
1688}
1689
1690// One over clippy's default arg limit: the metrics handle was threaded through
1691// to instrument the typed query result. Bundling these into a struct would add
1692// more noise than it removes for an internal dispatcher.
1693#[allow(clippy::too_many_arguments)]
1694async fn execute_wire_query_with_params(
1695    engine: Arc<RwLock<Engine>>,
1696    tx_gate: TxGate,
1697    tx_permit: &mut Option<OwnedSemaphorePermit>,
1698    query: String,
1699    params: Vec<WireParam>,
1700    principal: Option<Principal>,
1701    query_timeout: Duration,
1702    metrics: &Arc<Metrics>,
1703) -> (Message, Option<PendingDurability>) {
1704    match classify_params_transaction_control(&query, &params) {
1705        Some(TransactionControl::Begin) => {
1706            if tx_permit.is_some() {
1707                return (
1708                    Message::Error {
1709                        message: sanitize_error("transaction already active"),
1710                    },
1711                    None,
1712                );
1713            }
1714            let permit = match tx_gate.acquire_owned().await {
1715                Ok(permit) => permit,
1716                Err(_) => {
1717                    return (
1718                        Message::Error {
1719                            message: "query execution error".into(),
1720                        },
1721                        None,
1722                    )
1723                }
1724            };
1725            let (response, ticket) = run_blocking_query(
1726                engine,
1727                (query, params),
1728                principal,
1729                query_timeout,
1730                metrics,
1731                |engine, (query, params), principal| {
1732                    dispatch_query_with_params(&engine, &query, &params, principal.as_ref())
1733                },
1734            )
1735            .await;
1736            if is_success_response(&response) {
1737                *tx_permit = Some(permit);
1738            }
1739            (response, ticket)
1740        }
1741        Some(TransactionControl::Commit | TransactionControl::Rollback) => {
1742            let (response, ticket) = run_blocking_query(
1743                engine,
1744                (query, params),
1745                principal,
1746                query_timeout,
1747                metrics,
1748                |engine, (query, params), principal| {
1749                    dispatch_query_with_params(&engine, &query, &params, principal.as_ref())
1750                },
1751            )
1752            .await;
1753            if is_success_response(&response) {
1754                // See execute_wire_query: release the gate before the
1755                // caller's durability wait so commits can coalesce.
1756                tx_permit.take();
1757            }
1758            (response, ticket)
1759        }
1760        None if tx_permit.is_some() => {
1761            run_blocking_query(
1762                engine,
1763                (query, params),
1764                principal,
1765                query_timeout,
1766                metrics,
1767                |engine, (query, params), principal| {
1768                    dispatch_query_with_params(&engine, &query, &params, principal.as_ref())
1769                },
1770            )
1771            .await
1772        }
1773        None => {
1774            let permit = match tx_gate.acquire_owned().await {
1775                Ok(permit) => permit,
1776                Err(_) => {
1777                    return (
1778                        Message::Error {
1779                            message: "query execution error".into(),
1780                        },
1781                        None,
1782                    )
1783                }
1784            };
1785            let out = run_blocking_query(
1786                engine,
1787                (query, params),
1788                principal,
1789                query_timeout,
1790                metrics,
1791                |engine, (query, params), principal| {
1792                    dispatch_query_with_params(&engine, &query, &params, principal.as_ref())
1793                },
1794            )
1795            .await;
1796            drop(permit);
1797            out
1798        }
1799    }
1800}
1801
1802/// A statement's metric sample whose recording is deferred until its WAL
1803/// durability obligation settles: a Full-mode fsync failure downgrades the
1804/// client's success reply to an error, and the metrics must tell the same
1805/// story (and the latency must include the wait the client observed).
1806struct DeferredQueryMetric {
1807    start: Instant,
1808    outcome: QueryOutcome,
1809    exceeded_timeout: bool,
1810}
1811
1812/// Durability ticket + the deferred metric of the statement that produced it.
1813type PendingDurability = (WalDurabilityTicket, DeferredQueryMetric);
1814
1815async fn run_blocking_query<T, F>(
1816    engine: Arc<RwLock<Engine>>,
1817    input: T,
1818    principal: Option<Principal>,
1819    query_timeout: Duration,
1820    metrics: &Arc<Metrics>,
1821    f: F,
1822) -> (Message, Option<PendingDurability>)
1823where
1824    T: Send + 'static,
1825    F: FnOnce(Arc<RwLock<Engine>>, T, Option<Principal>) -> DispatchOutcome + Send + 'static,
1826{
1827    let _in_flight = metrics.in_flight_guard();
1828    let start = Instant::now();
1829    let mut handle = tokio::task::spawn_blocking(move || f(engine, input, principal));
1830    let mut exceeded_timeout = false;
1831    let join_result = tokio::select! {
1832        result = &mut handle => result,
1833        _ = tokio::time::sleep(query_timeout) => {
1834            exceeded_timeout = true;
1835            // `spawn_blocking` tasks that have started cannot be aborted safely.
1836            // Wait for completion before replying so a client never receives a
1837            // timeout while the same query keeps running and possibly mutating
1838            // state in the background.
1839            handle.await
1840        }
1841    };
1842
1843    let (message, ticket, outcome) = match join_result {
1844        Ok((Ok(result), ticket)) => match query_result_to_message(result) {
1845            Ok(message) => (message, ticket, QueryOutcome::Ok),
1846            Err(e) => (
1847                Message::Error {
1848                    message: sanitize_error(&e.to_string()),
1849                },
1850                ticket,
1851                QueryOutcome::Error,
1852            ),
1853        },
1854        Ok((Err(e), ticket)) => {
1855            let outcome = if matches!(e, QueryError::MemoryLimitExceeded { .. }) {
1856                QueryOutcome::MemoryLimit
1857            } else {
1858                QueryOutcome::Error
1859            };
1860            (
1861                Message::Error {
1862                    message: sanitize_error(&e.to_string()),
1863                },
1864                ticket,
1865                outcome,
1866            )
1867        }
1868        Err(e) => (
1869            Message::Error {
1870                message: format!("internal error: {e}"),
1871            },
1872            None,
1873            QueryOutcome::Error,
1874        ),
1875    };
1876    match ticket {
1877        // The statement's durability (and thus its true outcome and the
1878        // latency the client observes) settles at batch end — defer the
1879        // metric to the settlement site instead of recording a success that
1880        // a failed fsync would falsify.
1881        Some(ticket) => (
1882            message,
1883            Some((
1884                ticket,
1885                DeferredQueryMetric {
1886                    start,
1887                    outcome,
1888                    exceeded_timeout,
1889                },
1890            )),
1891        ),
1892        None => {
1893            if exceeded_timeout {
1894                metrics.record_query(start.elapsed(), QueryOutcome::Timeout);
1895            } else {
1896                metrics.record_query(start.elapsed(), outcome);
1897            }
1898            (message, None)
1899        }
1900    }
1901}
1902
1903/// Settle a WAL durability ticket off the async path, AFTER the TxGate
1904/// permit has been dropped — that ordering is what lets committers on other
1905/// connections append (and share the fsync) while this one waits.
1906///
1907/// Returns `None` when the covering fsync succeeded, or `Some(client-facing
1908/// error message)` when it failed — in which case no statement the ticket
1909/// covers may be acknowledged as durable (it executed in memory only).
1910async fn settle_durability_ticket(ticket: WalDurabilityTicket) -> Option<String> {
1911    match tokio::task::spawn_blocking(move || ticket.wait()).await {
1912        Ok(Ok(())) => None,
1913        Ok(Err(e)) => Some(sanitize_error(&format!("WAL durability sync failed: {e}"))),
1914        Err(e) => Some(format!("internal error: {e}")),
1915    }
1916}
1917
1918fn is_success_response(msg: &Message) -> bool {
1919    matches!(
1920        msg,
1921        Message::ResultRows { .. }
1922            | Message::ResultScalar { .. }
1923            | Message::ResultOk { .. }
1924            | Message::ResultMessage { .. }
1925    )
1926}
1927
1928fn rollback_open_transaction(engine: Arc<RwLock<Engine>>, principal: Option<Principal>) {
1929    let (res, ticket) = dispatch_query(&engine, "rollback", principal.as_ref());
1930    let _ = res;
1931    // Rollback takes the sync-preserving path (no ticket), but settle one
1932    // defensively if it ever appears so the durability watermark stays honest.
1933    if let Some(ticket) = ticket {
1934        let _ = ticket.wait();
1935    }
1936}
1937
1938pub async fn handle_connection<S>(stream: S, opts: ConnOpts<'_>)
1939where
1940    S: AsyncRead + AsyncWrite + Unpin,
1941{
1942    let ConnOpts {
1943        engine,
1944        tx_gate,
1945        expected_password,
1946        users,
1947        shutdown_rx,
1948        idle_timeout,
1949        query_timeout,
1950        rate_limiter,
1951        peer_addr,
1952        metrics,
1953    } = opts;
1954
1955    let peer = peer_addr
1956        .map(|a| a.to_string())
1957        .unwrap_or_else(|| "unknown".into());
1958    let peer_ip = peer_addr.map(|a| a.ip());
1959
1960    let (reader, writer) = tokio::io::split(stream);
1961    let mut reader = BufReader::new(reader);
1962    let mut writer = BufWriter::new(writer);
1963
1964    // Wait for Connect message (with idle timeout).
1965    // Accept Ping messages before authentication so load balancers can
1966    // health-check without completing a full CONNECT handshake.
1967    // Uses the smaller pre-auth payload limit (4 KB) to prevent memory abuse.
1968    let connect_msg = loop {
1969        match tokio::time::timeout(idle_timeout, Message::read_from_preauth(&mut reader)).await {
1970            Ok(Ok(Some(Message::Ping))) => {
1971                debug!(peer = %peer, "pre-auth ping");
1972                if !write_msg(&mut writer, &Message::Pong).await {
1973                    return;
1974                }
1975                continue;
1976            }
1977            Ok(Ok(Some(msg))) => break msg,
1978            Ok(Ok(None)) => {
1979                debug!(peer = %peer, "client closed before CONNECT");
1980                return;
1981            }
1982            Ok(Err(e)) => {
1983                error!(peer = %peer, error = %e, "error reading CONNECT");
1984                return;
1985            }
1986            Err(_) => {
1987                warn!(peer = %peer, "idle timeout waiting for CONNECT");
1988                return;
1989            }
1990        }
1991    };
1992
1993    // The authenticated identity for this connection. Bound at connect time
1994    // and enforced on every query by `dispatch_query`.
1995    let principal: Option<Principal>;
1996    let credential_auth_configured = !users.is_empty() || expected_password.is_some();
1997    match connect_msg {
1998        Message::Connect {
1999            db_name,
2000            password,
2001            username,
2002        } => {
2003            // Check rate limiting before verifying credentials.
2004            if let (Some(limiter), Some(ip)) = (rate_limiter, peer_ip) {
2005                if is_rate_limited(limiter, ip) {
2006                    warn!(peer = %peer, "rate limited: too many auth failures");
2007                    let err = Message::Error {
2008                        message: "too many auth failures, try again later".into(),
2009                    };
2010                    write_msg(&mut writer, &err).await;
2011                    return;
2012                }
2013            }
2014
2015            let outcome = authenticate_connect(
2016                &users,
2017                expected_password.as_ref().map(|p| p.as_str()),
2018                username.as_deref(),
2019                password.as_ref().map(|p| p.as_str()),
2020            );
2021
2022            match outcome {
2023                AuthOutcome::Rejected => {
2024                    warn!(peer = %peer, db = %db_name, "auth rejected");
2025                    metrics.inc_auth_failure();
2026                    // Record the failure for rate limiting.
2027                    if let (Some(limiter), Some(ip)) = (rate_limiter, peer_ip) {
2028                        record_auth_failure(limiter, ip);
2029                    }
2030                    let err = Message::Error {
2031                        message: "authentication failed".into(),
2032                    };
2033                    write_msg(&mut writer, &err).await;
2034                    return;
2035                }
2036                AuthOutcome::Authenticated {
2037                    principal: auth_principal,
2038                } => {
2039                    // Auth succeeded — clear any prior failure count.
2040                    if let (Some(limiter), Some(ip)) = (rate_limiter, peer_ip) {
2041                        clear_auth_failures(limiter, ip);
2042                    }
2043                    match &auth_principal {
2044                        Some(p) => {
2045                            info!(peer = %peer, db = %db_name, user = %p.name, role = %p.role, "authenticated");
2046                        }
2047                        None => {
2048                            info!(peer = %peer, db = %db_name, "client connected");
2049                        }
2050                    }
2051                    principal = auth_principal;
2052                }
2053            }
2054
2055            let ok = Message::ConnectOk {
2056                version: env!("CARGO_PKG_VERSION").into(),
2057            };
2058            if !write_msg(&mut writer, &ok).await {
2059                return;
2060            }
2061        }
2062        _ => {
2063            warn!(peer = %peer, "first message was not CONNECT");
2064            let err = Message::Error {
2065                message: "expected CONNECT".into(),
2066            };
2067            write_msg(&mut writer, &err).await;
2068            return;
2069        }
2070    }
2071
2072    let mut tx_permit: Option<OwnedSemaphorePermit> = None;
2073    // A non-query frame decoded during read-ahead batching, carried over to
2074    // the next iteration of the main loop.
2075    let mut carry: Option<Message> = None;
2076
2077    // Main query loop with idle timeout and shutdown awareness.
2078    'conn: loop {
2079        let msg = if let Some(m) = carry.take() {
2080            m
2081        } else {
2082            tokio::select! {
2083                // Read next message with idle timeout.
2084                result = tokio::time::timeout(idle_timeout, Message::read_from(&mut reader)) => {
2085                    match result {
2086                        Ok(Ok(Some(msg))) => msg,
2087                        Ok(Ok(None)) => break,
2088                        Ok(Err(e)) => {
2089                            error!(peer = %peer, error = %e, "read error");
2090                            break;
2091                        }
2092                        Err(_) => {
2093                            info!(peer = %peer, "idle timeout, closing connection");
2094                            let err = Message::Error { message: "idle timeout".into() };
2095                            write_msg(&mut writer, &err).await;
2096                            break;
2097                        }
2098                    }
2099                }
2100                // If server is shutting down, notify client and close.
2101                _ = shutdown_rx.changed() => {
2102                    if *shutdown_rx.borrow() {
2103                        info!(peer = %peer, "server shutting down, closing connection");
2104                        let err = Message::Error { message: "server shutting down".into() };
2105                        write_msg(&mut writer, &err).await;
2106                        break;
2107                    }
2108                    continue;
2109                }
2110            }
2111        };
2112
2113        // Plain query frames take the batched path: a pipelining client's
2114        // whole burst is executed with ONE durability wait at the end
2115        // (durability generations are cumulative, so the newest statement's
2116        // ticket covers every earlier one). Everything else is handled one
2117        // frame at a time exactly as before.
2118        if matches!(
2119            msg,
2120            Message::Query { .. } | Message::QuerySql { .. } | Message::QueryWithParams { .. }
2121        ) {
2122            /// Read-ahead cap per batch: bounds unflushed responses and keeps
2123            /// the reply latency of the first statement bounded.
2124            const MAX_PIPELINE_BATCH: usize = 128;
2125            /// Byte cap on retained (unflushed) response payloads per batch:
2126            /// large row results stop read-ahead, so one connection can never
2127            /// hold gigabytes of replies hostage to the batch's durability
2128            /// wait.
2129            const MAX_PIPELINE_BATCH_BYTES: usize = 4 << 20;
2130
2131            /// How the read-ahead loop stopped, when it stopped the whole
2132            /// connection rather than just the batch.
2133            enum BatchFatal {
2134                Closed,
2135                ReadError,
2136            }
2137
2138            /// Approximate encoded size of a response, for the batch byte
2139            /// cap. Counts the dominant string payloads; exact per-frame
2140            /// overhead is irrelevant at the 4 MiB cap.
2141            fn approx_response_bytes(msg: &Message) -> usize {
2142                match msg {
2143                    Message::ResultRows { columns, rows } => {
2144                        columns.iter().map(|c| c.len() + 4).sum::<usize>()
2145                            + rows
2146                                .iter()
2147                                .map(|r| r.iter().map(|v| v.len() + 4).sum::<usize>())
2148                                .sum::<usize>()
2149                    }
2150                    Message::ResultScalar { value } => value.len(),
2151                    Message::ResultMessage { message } | Message::Error { message } => {
2152                        message.len()
2153                    }
2154                    _ => 16,
2155                }
2156            }
2157
2158            /// Whether the reader's buffered bytes hold at least one COMPLETE
2159            /// frame (6-byte header + payload). Read-ahead must never await
2160            /// the socket: blocking on a partial frame would hold the batch's
2161            /// durability settlement and unflushed replies hostage to a slow
2162            /// (or malicious) client, up to the idle timeout.
2163            fn complete_frame_buffered(buf: &[u8]) -> bool {
2164                buf.len() >= 6 && {
2165                    let payload_len =
2166                        u32::from_le_bytes(buf[2..6].try_into().expect("4-byte slice")) as usize;
2167                    buf.len() - 6 >= payload_len
2168                }
2169            }
2170
2171            let mut responses: Vec<Message> = Vec::new();
2172            let mut response_bytes: usize = 0;
2173            let mut last_ticket: Option<WalDurabilityTicket> = None;
2174            let mut deferred_metrics: Vec<DeferredQueryMetric> = Vec::new();
2175            let mut fatal: Option<BatchFatal> = None;
2176            let mut current = msg;
2177            loop {
2178                let (response, ticket) = match current {
2179                    Message::Query { query } => {
2180                        if query.len() > MAX_QUERY_LENGTH {
2181                            (
2182                                Message::Error {
2183                                    message: format!(
2184                                        "query too large: {} bytes (max {})",
2185                                        query.len(),
2186                                        MAX_QUERY_LENGTH
2187                                    ),
2188                                },
2189                                None,
2190                            )
2191                        } else {
2192                            debug!(peer = %peer, query = %query, "received query");
2193                            execute_wire_query(
2194                                engine.clone(),
2195                                tx_gate.clone(),
2196                                &mut tx_permit,
2197                                query,
2198                                principal.clone(),
2199                                query_timeout,
2200                                &metrics,
2201                            )
2202                            .await
2203                        }
2204                    }
2205                    Message::QuerySql { query } => {
2206                        if query.len() > MAX_QUERY_LENGTH {
2207                            (
2208                                Message::Error {
2209                                    message: format!(
2210                                        "query too large: {} bytes (max {})",
2211                                        query.len(),
2212                                        MAX_QUERY_LENGTH
2213                                    ),
2214                                },
2215                                None,
2216                            )
2217                        } else {
2218                            debug!(peer = %peer, query = %query, "received SQL query");
2219                            execute_wire_query_sql(
2220                                engine.clone(),
2221                                tx_gate.clone(),
2222                                &mut tx_permit,
2223                                query,
2224                                principal.clone(),
2225                                query_timeout,
2226                                &metrics,
2227                            )
2228                            .await
2229                        }
2230                    }
2231                    Message::QueryWithParams { query, params } => {
2232                        if query.len() > MAX_QUERY_LENGTH {
2233                            (
2234                                Message::Error {
2235                                    message: format!(
2236                                        "query too large: {} bytes (max {})",
2237                                        query.len(),
2238                                        MAX_QUERY_LENGTH
2239                                    ),
2240                                },
2241                                None,
2242                            )
2243                        } else {
2244                            debug!(peer = %peer, query = %query, n_params = params.len(), "received parameterized query");
2245                            execute_wire_query_with_params(
2246                                engine.clone(),
2247                                tx_gate.clone(),
2248                                &mut tx_permit,
2249                                query,
2250                                params,
2251                                principal.clone(),
2252                                query_timeout,
2253                                &metrics,
2254                            )
2255                            .await
2256                        }
2257                    }
2258                    _ => unreachable!("batch loop only receives plain query frames"),
2259                };
2260                if let Some((t, m)) = ticket {
2261                    // Later tickets cover earlier generations — keep only the
2262                    // newest; the batch-end wait settles them all. Every
2263                    // deferred metric is kept: each records after settlement.
2264                    last_ticket = Some(t);
2265                    deferred_metrics.push(m);
2266                }
2267                response_bytes += approx_response_bytes(&response);
2268                responses.push(response);
2269
2270                // Read ahead only when a COMPLETE next frame is already
2271                // buffered (never await the socket mid-batch) and the
2272                // retained replies stay small. While an explicit transaction
2273                // is open the connection holds the TxGate, so batching would
2274                // only extend the exclusive window — flush instead.
2275                if tx_permit.is_some()
2276                    || responses.len() >= MAX_PIPELINE_BATCH
2277                    || response_bytes >= MAX_PIPELINE_BATCH_BYTES
2278                    || !complete_frame_buffered(reader.buffer())
2279                {
2280                    break;
2281                }
2282                // The full frame is buffered, so this returns without socket
2283                // I/O; the timeout is a defensive backstop only.
2284                match tokio::time::timeout(idle_timeout, Message::read_from(&mut reader)).await {
2285                    Ok(Ok(Some(
2286                        next @ (Message::Query { .. }
2287                        | Message::QuerySql { .. }
2288                        | Message::QueryWithParams { .. }),
2289                    ))) => {
2290                        // If another connection currently holds the TxGate,
2291                        // the next statement would block on the gate with
2292                        // this batch's replies still unflushed (pre-batching,
2293                        // they'd already have been written). Flush first and
2294                        // handle the frame on the next main-loop iteration.
2295                        // Benign TOCTOU: worst case is one early flush or one
2296                        // gate wait with an empty reply queue.
2297                        if tx_gate.available_permits() == 0 {
2298                            carry = Some(next);
2299                            break;
2300                        }
2301                        current = next;
2302                    }
2303                    Ok(Ok(Some(other))) => {
2304                        // Not a plain query — flush this batch, then handle
2305                        // the frame on the next main-loop iteration.
2306                        carry = Some(other);
2307                        break;
2308                    }
2309                    Ok(Ok(None)) => {
2310                        fatal = Some(BatchFatal::Closed);
2311                        break;
2312                    }
2313                    Ok(Err(e)) => {
2314                        error!(peer = %peer, error = %e, "read error");
2315                        fatal = Some(BatchFatal::ReadError);
2316                        break;
2317                    }
2318                    Err(_) => {
2319                        // Unreachable in practice: the frame was fully
2320                        // buffered, so read_from needs no socket I/O.
2321                        error!(peer = %peer, "timeout decoding fully-buffered frame");
2322                        fatal = Some(BatchFatal::ReadError);
2323                        break;
2324                    }
2325                }
2326            }
2327
2328            // ONE durability wait for the whole batch, then the deferred
2329            // metrics: a durability failure records Ok statements as errors,
2330            // and latency includes the settlement wait the client observed.
2331            let mut durability_failed = false;
2332            if let Some(ticket) = last_ticket {
2333                if let Some(message) = settle_durability_ticket(ticket).await {
2334                    // The covering fsync failed: nothing in this batch may be
2335                    // acknowledged as durable.
2336                    durability_failed = true;
2337                    for r in responses.iter_mut() {
2338                        if is_success_response(r) {
2339                            *r = Message::Error {
2340                                message: message.clone(),
2341                            };
2342                        }
2343                    }
2344                }
2345            }
2346            for m in deferred_metrics.drain(..) {
2347                let outcome = if m.exceeded_timeout {
2348                    QueryOutcome::Timeout
2349                } else if durability_failed && matches!(m.outcome, QueryOutcome::Ok) {
2350                    QueryOutcome::Error
2351                } else {
2352                    m.outcome
2353                };
2354                metrics.record_query(m.start.elapsed(), outcome);
2355            }
2356
2357            for r in &responses {
2358                if !write_msg(&mut writer, r).await {
2359                    break 'conn;
2360                }
2361            }
2362            match fatal {
2363                None => continue,
2364                Some(BatchFatal::Closed | BatchFatal::ReadError) => break,
2365            }
2366        }
2367
2368        let response = match msg {
2369            Message::Ping => {
2370                debug!(peer = %peer, "ping");
2371                Message::Pong
2372            }
2373            Message::SyncStatus { replica_id } => {
2374                let engine = engine.clone();
2375                let principal = principal.clone();
2376                let log_context = SyncLogContext::status(&replica_id);
2377                execute_gated_sync(
2378                    SyncExecutionContext {
2379                        tx_gate: tx_gate.clone(),
2380                        connection_has_transaction: tx_permit.is_some(),
2381                        operation: SyncOperation::Status,
2382                        log_context,
2383                        metrics: &metrics,
2384                        query_timeout,
2385                    },
2386                    (engine, replica_id, credential_auth_configured, principal),
2387                    |(engine, replica_id, credential_authenticated, principal)| {
2388                        dispatch_sync_status_decision(
2389                            &engine,
2390                            replica_id,
2391                            credential_authenticated,
2392                            principal.as_ref(),
2393                        )
2394                    },
2395                )
2396                .await
2397            }
2398            Message::SyncPull {
2399                replica_id,
2400                since_lsn,
2401                max_units,
2402                max_bytes,
2403                database_id,
2404                primary_generation,
2405                wal_format_version,
2406                catalog_version,
2407                segment_format_version,
2408            } => {
2409                let engine = engine.clone();
2410                let principal = principal.clone();
2411                let request = SyncPullRequest {
2412                    replica_id,
2413                    since_lsn,
2414                    max_units,
2415                    max_bytes,
2416                    database_id,
2417                    primary_generation,
2418                    wal_format_version,
2419                    catalog_version,
2420                    segment_format_version,
2421                };
2422                let log_context = SyncLogContext::pull(&request);
2423                execute_gated_sync(
2424                    SyncExecutionContext {
2425                        tx_gate: tx_gate.clone(),
2426                        connection_has_transaction: tx_permit.is_some(),
2427                        operation: SyncOperation::Pull,
2428                        log_context,
2429                        metrics: &metrics,
2430                        query_timeout,
2431                    },
2432                    (engine, request, credential_auth_configured, principal),
2433                    |(engine, request, credential_authenticated, principal)| {
2434                        dispatch_sync_pull_decision(
2435                            &engine,
2436                            request,
2437                            credential_authenticated,
2438                            principal.as_ref(),
2439                        )
2440                    },
2441                )
2442                .await
2443            }
2444            Message::SyncAck {
2445                replica_id,
2446                applied_lsn,
2447                remote_lsn,
2448            } => {
2449                let engine = engine.clone();
2450                let principal = principal.clone();
2451                let log_context = SyncLogContext::ack(&replica_id, applied_lsn, remote_lsn);
2452                execute_gated_sync(
2453                    SyncExecutionContext {
2454                        tx_gate: tx_gate.clone(),
2455                        connection_has_transaction: tx_permit.is_some(),
2456                        operation: SyncOperation::Ack,
2457                        log_context,
2458                        metrics: &metrics,
2459                        query_timeout,
2460                    },
2461                    (
2462                        engine,
2463                        replica_id,
2464                        applied_lsn,
2465                        remote_lsn,
2466                        credential_auth_configured,
2467                        principal,
2468                    ),
2469                    |(
2470                        engine,
2471                        replica_id,
2472                        applied_lsn,
2473                        observed_remote_lsn,
2474                        credential_authenticated,
2475                        principal,
2476                    )| {
2477                        dispatch_sync_ack_decision(
2478                            &engine,
2479                            replica_id,
2480                            applied_lsn,
2481                            observed_remote_lsn,
2482                            credential_authenticated,
2483                            principal.as_ref(),
2484                        )
2485                    },
2486                )
2487                .await
2488            }
2489            Message::Disconnect => {
2490                debug!(peer = %peer, "received DISCONNECT");
2491                break;
2492            }
2493            _ => Message::Error {
2494                message: "unexpected message type".into(),
2495            },
2496        };
2497
2498        if !write_msg(&mut writer, &response).await {
2499            break;
2500        }
2501    }
2502
2503    // Roll back any open transaction the client left behind on disconnect.
2504    // The permit must stay alive in `tx_permit` for the duration of the awaited
2505    // rollback and be released only afterwards — mirroring the query-timeout
2506    // path above. Using `tx_permit.take().is_some()` here would drop the permit
2507    // (freeing the TxGate) *before* the rollback runs, letting another
2508    // connection BEGIN a transaction that this stale rollback would then clobber.
2509    if tx_permit.is_some() {
2510        let engine = engine.clone();
2511        let principal = principal.clone();
2512        let _ =
2513            tokio::task::spawn_blocking(move || rollback_open_transaction(engine, principal)).await;
2514    }
2515    tx_permit.take();
2516
2517    info!(peer = %peer, "client disconnected");
2518}
2519
2520fn charge_response_bytes(total: &mut usize, bytes: usize) -> Result<(), QueryError> {
2521    *total = total.saturating_add(bytes);
2522    if *total > MAX_RESPONSE_PAYLOAD_SIZE {
2523        return Err(QueryError::Execution(format!(
2524            "result too large: encoded response exceeds {} bytes; add a limit or narrower projection",
2525            MAX_RESPONSE_PAYLOAD_SIZE
2526        )));
2527    }
2528    Ok(())
2529}
2530
2531fn query_result_to_message(result: QueryResult) -> Result<Message, QueryError> {
2532    match result {
2533        QueryResult::Rows { columns, rows } => {
2534            let mut encoded_bytes = 2usize; // column count
2535            let mut out_columns = Vec::with_capacity(columns.len());
2536            for col in columns {
2537                charge_response_bytes(&mut encoded_bytes, 4 + col.len())?;
2538                out_columns.push(col);
2539            }
2540            charge_response_bytes(&mut encoded_bytes, 4)?; // row count
2541
2542            let mut str_rows = Vec::with_capacity(rows.len());
2543            for row in rows {
2544                let mut str_row = Vec::with_capacity(row.len());
2545                for value in row {
2546                    let display = value_to_display(&value);
2547                    charge_response_bytes(&mut encoded_bytes, 4 + display.len())?;
2548                    str_row.push(display);
2549                }
2550                str_rows.push(str_row);
2551            }
2552            Ok(Message::ResultRows {
2553                columns: out_columns,
2554                rows: str_rows,
2555            })
2556        }
2557        QueryResult::Scalar(val) => Ok(Message::ResultScalar {
2558            value: value_to_display(&val),
2559        }),
2560        QueryResult::Modified(n) => Ok(Message::ResultOk { affected: n }),
2561        QueryResult::Created(name) => Ok(Message::ResultMessage {
2562            message: format!("type {name} created"),
2563        }),
2564        QueryResult::Executed { message } => Ok(Message::ResultMessage { message }),
2565    }
2566}
2567
2568// Canonical wire rendering lives on `Value` (`powdb_storage`) so the server,
2569// CLI, and embedded bindings render results identically. Kept as a thin alias
2570// to minimize churn at the call sites in this module.
2571fn value_to_display(v: &Value) -> String {
2572    v.to_wire_string()
2573}
2574
2575#[cfg(test)]
2576mod tests {
2577    use super::*;
2578    use powdb_storage::wal::WalRecordType;
2579    use powdb_sync::{
2580        write_identity_snapshot, write_segment_atomic, DatabaseIdentity, IdentitySnapshot,
2581        ReplicaCursor, RetainedSegment, RetainedUnit,
2582    };
2583
2584    // ---- Wire NULL rendering (Fix: remote protocol rendered NULL as `{}`) ----
2585
2586    #[test]
2587    fn null_serializes_as_null_bareword_on_wire() {
2588        assert_eq!(value_to_display(&Value::Empty), "null");
2589    }
2590
2591    // ---- Error sanitization allowlist ----
2592
2593    #[test]
2594    fn unique_violation_error_surfaces_to_remote_clients() {
2595        // The storage layer reports the actionable message; the server must
2596        // not replace it with the generic "query execution error".
2597        assert_eq!(
2598            sanitize_error("unique constraint violation on User.email"),
2599            "unique constraint violation on User.email"
2600        );
2601    }
2602
2603    #[test]
2604    fn internal_errors_stay_generic() {
2605        assert_eq!(
2606            sanitize_error("some internal io panic detail"),
2607            "query execution error"
2608        );
2609    }
2610
2611    #[test]
2612    fn resource_limit_errors_surface_actionable_hints() {
2613        // These carry user-actionable guidance and leak no internal state, so
2614        // they must reach the client verbatim — not be masked to the generic
2615        // message. The exact strings come from QueryError's Display impl
2616        // (crates/query/src/result.rs).
2617        for msg in [
2618            "sort input exceeds row limit — add a LIMIT clause",
2619            "join result exceeds row limit",
2620            "query exceeded memory budget: requested 100 bytes, limit 50 bytes",
2621            "result too large: encoded response exceeds 1024 bytes; add a limit or narrower projection",
2622        ] {
2623            assert_eq!(sanitize_error(msg), msg, "should pass through verbatim");
2624        }
2625    }
2626
2627    #[test]
2628    fn oversized_result_is_rejected_before_wire_encoding() {
2629        let long = "x".repeat(MAX_RESPONSE_PAYLOAD_SIZE);
2630        let result = QueryResult::Rows {
2631            columns: vec!["payload".into()],
2632            rows: vec![vec![Value::Str(long)]],
2633        };
2634        let err = query_result_to_message(result).unwrap_err();
2635        assert!(
2636            err.to_string().starts_with("result too large"),
2637            "unexpected error: {err}"
2638        );
2639    }
2640
2641    // ---- Role enforcement (Fix: readonly role was not enforced) ----
2642
2643    fn parsed(q: &str) -> powdb_query::ast::Statement {
2644        parser::parse(q).unwrap()
2645    }
2646
2647    fn principal(role: &str) -> Option<Principal> {
2648        Some(Principal {
2649            name: "u".into(),
2650            role: role.into(),
2651        })
2652    }
2653
2654    #[test]
2655    fn readonly_can_read_but_not_write() {
2656        let p = principal("readonly");
2657        // Reads pass.
2658        assert!(check_statement_permitted(p.as_ref(), &parsed("User")).is_ok());
2659        assert!(check_statement_permitted(p.as_ref(), &parsed("count(User)")).is_ok());
2660        assert!(check_statement_permitted(p.as_ref(), &parsed("explain User")).is_ok());
2661        // Writes, DDL, and transaction control are denied.
2662        for q in [
2663            r#"insert User { name := "x" }"#,
2664            "User filter .id = 1 update { age := 2 }",
2665            "User filter .id = 1 delete",
2666            "drop User",
2667            "alter User add column c: str",
2668            "type T { required id: int }",
2669            "begin",
2670            "commit",
2671            "rollback",
2672        ] {
2673            let err = check_statement_permitted(p.as_ref(), &parsed(q))
2674                .expect_err(&format!("must deny: {q}"));
2675            assert!(
2676                err.to_string().contains("permission denied"),
2677                "unexpected error for {q}: {err}"
2678            );
2679        }
2680    }
2681
2682    #[test]
2683    fn readwrite_and_admin_have_full_query_access() {
2684        for role in ["readwrite", "admin"] {
2685            let p = principal(role);
2686            assert!(check_statement_permitted(p.as_ref(), &parsed("User")).is_ok());
2687            assert!(check_statement_permitted(
2688                p.as_ref(),
2689                &parsed(r#"insert User { name := "x" }"#)
2690            )
2691            .is_ok());
2692            assert!(check_statement_permitted(p.as_ref(), &parsed("drop User")).is_ok());
2693        }
2694    }
2695
2696    #[test]
2697    fn unknown_role_fails_closed_for_writes() {
2698        let p = principal("mystery");
2699        assert!(check_statement_permitted(p.as_ref(), &parsed("User")).is_ok());
2700        assert!(
2701            check_statement_permitted(p.as_ref(), &parsed(r#"insert User { name := "x" }"#))
2702                .is_err()
2703        );
2704    }
2705
2706    #[test]
2707    fn no_principal_means_full_access() {
2708        // Shared-password / open mode: no per-user identity, no restriction.
2709        assert!(check_statement_permitted(None, &parsed("drop User")).is_ok());
2710        assert!(check_statement_permitted(None, &parsed(r#"insert User { name := "x" }"#)).is_ok());
2711    }
2712
2713    fn store_with_alice() -> UserStore {
2714        let mut s = UserStore::new();
2715        s.create_user("alice", "pw", "readwrite").unwrap();
2716        s
2717    }
2718
2719    // ---- Empty store: legacy shared-password fallback ----
2720
2721    #[test]
2722    fn empty_store_no_password_is_open() {
2723        let s = UserStore::new();
2724        assert_eq!(
2725            authenticate_connect(&s, None, None, None),
2726            AuthOutcome::Authenticated { principal: None }
2727        );
2728        // Even a stray username/password is accepted (legacy open behavior).
2729        assert_eq!(
2730            authenticate_connect(&s, None, Some("x"), Some("y")),
2731            AuthOutcome::Authenticated { principal: None }
2732        );
2733    }
2734
2735    #[test]
2736    fn empty_store_correct_shared_password_succeeds() {
2737        let s = UserStore::new();
2738        assert_eq!(
2739            authenticate_connect(&s, Some("pw"), None, Some("pw")),
2740            AuthOutcome::Authenticated { principal: None }
2741        );
2742    }
2743
2744    #[test]
2745    fn empty_store_wrong_shared_password_rejected() {
2746        let s = UserStore::new();
2747        assert_eq!(
2748            authenticate_connect(&s, Some("pw"), None, Some("bad")),
2749            AuthOutcome::Rejected
2750        );
2751    }
2752
2753    #[test]
2754    fn empty_store_missing_password_rejected_when_expected() {
2755        let s = UserStore::new();
2756        assert_eq!(
2757            authenticate_connect(&s, Some("pw"), None, None),
2758            AuthOutcome::Rejected
2759        );
2760    }
2761
2762    #[test]
2763    fn empty_store_ignores_username_for_shared_password() {
2764        // A new client may send a username even against a shared-password
2765        // server; the username is ignored and the password still governs.
2766        let s = UserStore::new();
2767        assert_eq!(
2768            authenticate_connect(&s, Some("pw"), Some("whoever"), Some("pw")),
2769            AuthOutcome::Authenticated { principal: None }
2770        );
2771    }
2772
2773    // ---- Populated store: multi-user auth ----
2774
2775    #[test]
2776    fn user_auth_success_binds_principal() {
2777        let s = store_with_alice();
2778        assert_eq!(
2779            authenticate_connect(&s, None, Some("alice"), Some("pw")),
2780            AuthOutcome::Authenticated {
2781                principal: Some(Principal {
2782                    name: "alice".into(),
2783                    role: "readwrite".into(),
2784                })
2785            }
2786        );
2787    }
2788
2789    #[test]
2790    fn user_auth_wrong_password_rejected() {
2791        let s = store_with_alice();
2792        assert_eq!(
2793            authenticate_connect(&s, None, Some("alice"), Some("bad")),
2794            AuthOutcome::Rejected
2795        );
2796    }
2797
2798    #[test]
2799    fn user_auth_unknown_user_rejected() {
2800        let s = store_with_alice();
2801        assert_eq!(
2802            authenticate_connect(&s, None, Some("mallory"), Some("pw")),
2803            AuthOutcome::Rejected
2804        );
2805    }
2806
2807    #[test]
2808    fn user_auth_missing_username_rejected() {
2809        let s = store_with_alice();
2810        assert_eq!(
2811            authenticate_connect(&s, None, None, Some("pw")),
2812            AuthOutcome::Rejected
2813        );
2814    }
2815
2816    #[test]
2817    fn user_auth_missing_password_rejected() {
2818        let s = store_with_alice();
2819        assert_eq!(
2820            authenticate_connect(&s, Some("pw"), Some("alice"), None),
2821            AuthOutcome::Rejected
2822        );
2823    }
2824
2825    #[test]
2826    fn user_auth_ignores_shared_password_when_users_present() {
2827        // With users present, the shared password is irrelevant: supplying it as
2828        // the password without a valid user must NOT authenticate.
2829        let s = store_with_alice();
2830        assert_eq!(
2831            authenticate_connect(&s, Some("shared"), None, Some("shared")),
2832            AuthOutcome::Rejected
2833        );
2834    }
2835
2836    #[test]
2837    fn replica_fingerprint_is_stable_and_redacted() {
2838        let replica_id = "customer-prod-replica-a";
2839        let fingerprint = replica_fingerprint(replica_id);
2840        assert_eq!(fingerprint, replica_fingerprint(replica_id));
2841        assert_eq!(fingerprint, log_replica_fingerprint(replica_id));
2842        assert_ne!(fingerprint, replica_fingerprint("customer-prod-replica-b"));
2843        assert_eq!(fingerprint.len(), 16);
2844        assert!(fingerprint.chars().all(|c| c.is_ascii_hexdigit()));
2845        assert!(!fingerprint.contains("customer"));
2846        assert!(!fingerprint.contains("replica"));
2847        assert!(!fingerprint.contains(replica_id));
2848    }
2849
2850    #[test]
2851    fn invalid_replica_ids_use_fixed_log_fingerprint() {
2852        assert_eq!(log_replica_fingerprint(""), INVALID_REPLICA_FINGERPRINT);
2853        assert_eq!(
2854            log_replica_fingerprint("customer/prod/replica"),
2855            INVALID_REPLICA_FINGERPRINT
2856        );
2857        assert_eq!(
2858            log_replica_fingerprint(&"a".repeat(4096)),
2859            INVALID_REPLICA_FINGERPRINT
2860        );
2861    }
2862
2863    #[test]
2864    fn sync_error_classes_use_bounded_labels() {
2865        assert_eq!(SyncErrorClass::AuthRequired.as_label(), "auth_required");
2866        assert_eq!(
2867            SyncErrorClass::PermissionDenied.as_label(),
2868            "permission_denied"
2869        );
2870        assert_eq!(
2871            SyncErrorClass::IdentityOrFormatMismatch.as_label(),
2872            "identity_or_format_mismatch"
2873        );
2874        assert_eq!(SyncErrorClass::AckValidation.as_label(), "ack_validation");
2875        assert_eq!(SyncErrorClass::Internal.as_label(), "internal");
2876    }
2877
2878    fn sync_identity() -> DatabaseIdentity {
2879        DatabaseIdentity {
2880            database_id: *b"server-sync-test",
2881            primary_generation: 1,
2882        }
2883    }
2884
2885    fn retained_unit(lsn: u64) -> RetainedUnit {
2886        RetainedUnit {
2887            tx_id: 1,
2888            record_type: 4,
2889            lsn,
2890            data: lsn.to_le_bytes().to_vec(),
2891        }
2892    }
2893
2894    fn retained_unit_with(tx_id: u64, record_type: WalRecordType, lsn: u64) -> RetainedUnit {
2895        RetainedUnit {
2896            tx_id,
2897            record_type: record_type as u8,
2898            lsn,
2899            data: lsn.to_le_bytes().to_vec(),
2900        }
2901    }
2902
2903    fn write_sync_identity_and_tail(data_dir: &std::path::Path, through_lsn: u64) {
2904        let identity = sync_identity();
2905        write_identity_snapshot(data_dir, &IdentitySnapshot::from_identity(identity, 1)).unwrap();
2906        let units = (1..=through_lsn).map(retained_unit).collect();
2907        let segment = RetainedSegment::new(identity.segment_identity(), units).unwrap();
2908        write_segment_atomic(&retained_segments_dir(data_dir), &segment).unwrap();
2909    }
2910
2911    fn write_sync_identity_and_units(data_dir: &std::path::Path, units: Vec<RetainedUnit>) {
2912        let identity = sync_identity();
2913        write_identity_snapshot(data_dir, &IdentitySnapshot::from_identity(identity, 1)).unwrap();
2914        let segment = RetainedSegment::new(identity.segment_identity(), units).unwrap();
2915        write_segment_atomic(&retained_segments_dir(data_dir), &segment).unwrap();
2916    }
2917
2918    fn write_sync_identity_only(data_dir: &std::path::Path) {
2919        let identity = sync_identity();
2920        write_identity_snapshot(data_dir, &IdentitySnapshot::from_identity(identity, 1)).unwrap();
2921    }
2922
2923    fn admin_principal() -> Principal {
2924        Principal {
2925            name: "admin".into(),
2926            role: "admin".into(),
2927        }
2928    }
2929
2930    #[test]
2931    fn sync_protocol_requires_credential_auth_and_rejects_readonly() {
2932        let dir = tempfile::tempdir().unwrap();
2933        let engine = Arc::new(RwLock::new(Engine::new(dir.path()).unwrap()));
2934
2935        match dispatch_sync_status(&engine, "replica-a".into(), false, None) {
2936            Message::Error { message } => {
2937                assert!(message.contains("requires authentication"));
2938            }
2939            other => panic!("expected auth error, got {other:?}"),
2940        }
2941
2942        let readonly = Principal {
2943            name: "reader".into(),
2944            role: "readonly".into(),
2945        };
2946        match dispatch_sync_status(&engine, "replica-a".into(), true, Some(&readonly)) {
2947            Message::Error { message } => {
2948                assert!(message.contains("permission denied"));
2949            }
2950            other => panic!("expected permission error, got {other:?}"),
2951        }
2952    }
2953
2954    #[test]
2955    fn sync_status_pull_and_ack_use_server_remote_lsn() {
2956        let dir = tempfile::tempdir().unwrap();
2957        let mut engine = Engine::new(dir.path()).unwrap();
2958        engine
2959            .execute_powql("type SyncT { required id: int, v: str }")
2960            .unwrap();
2961        engine
2962            .execute_powql(r#"insert SyncT { id := 1, v := "one" }"#)
2963            .unwrap();
2964        let remote_lsn = engine.catalog().max_lsn();
2965        assert!(remote_lsn > 0);
2966        write_sync_identity_and_tail(dir.path(), remote_lsn);
2967        powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
2968            .unwrap();
2969
2970        let engine = Arc::new(RwLock::new(engine));
2971        let principal = admin_principal();
2972        let status = match dispatch_sync_status(&engine, "replica-a".into(), true, Some(&principal))
2973        {
2974            Message::SyncStatusResult { status } => status,
2975            other => panic!("expected sync status, got {other:?}"),
2976        };
2977        assert_eq!(status.remote_lsn, remote_lsn);
2978        assert_eq!(status.servable_lsn, Some(remote_lsn));
2979        assert_eq!(status.unarchived_lsn, Some(0));
2980        assert_eq!(status.last_applied_lsn, Some(0));
2981        assert_eq!(status.repair_action, WireSyncRepairAction::Pull);
2982        assert!(status.stale);
2983
2984        let identity = sync_identity().segment_identity();
2985        let pull = SyncPullRequest {
2986            replica_id: "replica-a".into(),
2987            since_lsn: 0,
2988            max_units: MAX_SYNC_PULL_UNITS,
2989            max_bytes: MAX_SYNC_PULL_BYTES,
2990            database_id: identity.database_id,
2991            primary_generation: identity.primary_generation,
2992            wal_format_version: identity.wal_format_version,
2993            catalog_version: identity.catalog_version,
2994            segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
2995        };
2996        let units = match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
2997            Message::SyncPullResult {
2998                status,
2999                units,
3000                has_more,
3001            } => {
3002                assert_eq!(status.repair_action, WireSyncRepairAction::Pull);
3003                assert!(!has_more);
3004                units
3005            }
3006            other => panic!("expected sync pull result, got {other:?}"),
3007        };
3008        assert_eq!(units.len() as u64, remote_lsn);
3009        assert_eq!(units.last().unwrap().lsn, remote_lsn);
3010
3011        let ack = match dispatch_sync_ack(
3012            &engine,
3013            "replica-a".into(),
3014            remote_lsn,
3015            remote_lsn,
3016            true,
3017            Some(&principal),
3018        ) {
3019            Message::SyncAckResult {
3020                previous_applied_lsn,
3021                applied_lsn,
3022                remote_lsn: ack_remote_lsn,
3023                advanced,
3024                status,
3025            } => {
3026                assert_eq!(previous_applied_lsn, 0);
3027                assert_eq!(applied_lsn, remote_lsn);
3028                assert_eq!(ack_remote_lsn, remote_lsn);
3029                assert!(advanced);
3030                status
3031            }
3032            other => panic!("expected sync ack result, got {other:?}"),
3033        };
3034        assert_eq!(ack.repair_action, WireSyncRepairAction::None);
3035        assert!(!ack.stale);
3036        assert_eq!(ack.lag_lsn, Some(0));
3037    }
3038
3039    #[test]
3040    fn sync_pull_and_ack_reject_transaction_cut_boundaries() {
3041        let dir = tempfile::tempdir().unwrap();
3042        let mut engine = Engine::new(dir.path()).unwrap();
3043        engine
3044            .execute_powql("type SyncT { required id: int }")
3045            .unwrap();
3046        for id in 1..=3 {
3047            engine
3048                .execute_powql(&format!("insert SyncT {{ id := {id} }}"))
3049                .unwrap();
3050        }
3051        let remote_lsn = engine.catalog().max_lsn();
3052        assert!(remote_lsn >= 3);
3053        write_sync_identity_and_units(
3054            dir.path(),
3055            vec![
3056                retained_unit_with(77, WalRecordType::Begin, 1),
3057                retained_unit_with(77, WalRecordType::Insert, 2),
3058                retained_unit_with(77, WalRecordType::Commit, 3),
3059            ],
3060        );
3061        powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
3062            .unwrap();
3063
3064        let engine = Arc::new(RwLock::new(engine));
3065        let principal = admin_principal();
3066        let identity = sync_identity().segment_identity();
3067        let cut_pull = SyncPullRequest {
3068            replica_id: "replica-a".into(),
3069            since_lsn: 0,
3070            max_units: 2,
3071            max_bytes: MAX_SYNC_PULL_BYTES,
3072            database_id: identity.database_id,
3073            primary_generation: identity.primary_generation,
3074            wal_format_version: identity.wal_format_version,
3075            catalog_version: identity.catalog_version,
3076            segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
3077        };
3078        match dispatch_sync_pull(&engine, cut_pull, true, Some(&principal)) {
3079            Message::Error { message } => assert!(message.contains("cuts through transaction")),
3080            other => panic!("expected transaction-cut pull error, got {other:?}"),
3081        }
3082
3083        let cut_bytes_pull = SyncPullRequest {
3084            replica_id: "replica-a".into(),
3085            since_lsn: 0,
3086            max_units: 3,
3087            max_bytes: 58,
3088            database_id: identity.database_id,
3089            primary_generation: identity.primary_generation,
3090            wal_format_version: identity.wal_format_version,
3091            catalog_version: identity.catalog_version,
3092            segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
3093        };
3094        match dispatch_sync_pull(&engine, cut_bytes_pull, true, Some(&principal)) {
3095            Message::Error { message } => assert!(message.contains("cuts through transaction")),
3096            other => panic!("expected byte-capped transaction-cut pull error, got {other:?}"),
3097        }
3098
3099        let full_pull = SyncPullRequest {
3100            replica_id: "replica-a".into(),
3101            since_lsn: 0,
3102            max_units: 3,
3103            max_bytes: MAX_SYNC_PULL_BYTES,
3104            database_id: identity.database_id,
3105            primary_generation: identity.primary_generation,
3106            wal_format_version: identity.wal_format_version,
3107            catalog_version: identity.catalog_version,
3108            segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
3109        };
3110        match dispatch_sync_pull(&engine, full_pull, true, Some(&principal)) {
3111            Message::SyncPullResult { units, .. } => {
3112                assert_eq!(units.len(), 3);
3113                assert_eq!(units.last().unwrap().lsn, 3);
3114            }
3115            other => panic!("expected complete transaction pull, got {other:?}"),
3116        }
3117
3118        match dispatch_sync_ack(
3119            &engine,
3120            "replica-a".into(),
3121            2,
3122            remote_lsn,
3123            true,
3124            Some(&principal),
3125        ) {
3126            Message::Error { message } => assert!(message.contains("cuts through transaction")),
3127            other => panic!("expected transaction-cut ack error, got {other:?}"),
3128        }
3129        let cursor = powdb_sync::read_replica_cursors(dir.path()).unwrap();
3130        assert_eq!(cursor[0].applied_lsn, 0);
3131
3132        match dispatch_sync_ack(
3133            &engine,
3134            "replica-a".into(),
3135            3,
3136            remote_lsn,
3137            true,
3138            Some(&principal),
3139        ) {
3140            Message::SyncAckResult {
3141                previous_applied_lsn,
3142                applied_lsn,
3143                advanced,
3144                ..
3145            } => {
3146                assert_eq!(previous_applied_lsn, 0);
3147                assert_eq!(applied_lsn, 3);
3148                assert!(advanced);
3149            }
3150            other => panic!("expected complete transaction ack, got {other:?}"),
3151        }
3152    }
3153
3154    #[test]
3155    fn sync_pull_byte_cap_returns_applyable_prefix_with_reused_tx_id() {
3156        let dir = tempfile::tempdir().unwrap();
3157        let mut engine = Engine::new(dir.path()).unwrap();
3158        engine
3159            .execute_powql("type SyncT { required id: int }")
3160            .unwrap();
3161        for id in 1..=6 {
3162            engine
3163                .execute_powql(&format!("insert SyncT {{ id := {id} }}"))
3164                .unwrap();
3165        }
3166        let remote_lsn = engine.catalog().max_lsn();
3167        assert!(remote_lsn >= 6);
3168        write_sync_identity_and_units(
3169            dir.path(),
3170            vec![
3171                retained_unit_with(1, WalRecordType::Begin, 1),
3172                retained_unit_with(1, WalRecordType::Insert, 2),
3173                retained_unit_with(1, WalRecordType::Commit, 3),
3174                retained_unit_with(1, WalRecordType::Begin, 4),
3175                retained_unit_with(1, WalRecordType::Insert, 5),
3176                retained_unit_with(1, WalRecordType::Commit, 6),
3177            ],
3178        );
3179        powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
3180            .unwrap();
3181
3182        let engine = Arc::new(RwLock::new(engine));
3183        let principal = admin_principal();
3184        let identity = sync_identity().segment_identity();
3185        let pull = SyncPullRequest {
3186            replica_id: "replica-a".into(),
3187            since_lsn: 0,
3188            max_units: 6,
3189            max_bytes: 100,
3190            database_id: identity.database_id,
3191            primary_generation: identity.primary_generation,
3192            wal_format_version: identity.wal_format_version,
3193            catalog_version: identity.catalog_version,
3194            segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
3195        };
3196
3197        match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
3198            Message::SyncPullResult {
3199                status,
3200                units,
3201                has_more,
3202            } => {
3203                assert_eq!(status.repair_action, WireSyncRepairAction::Pull);
3204                assert_eq!(units.len(), 3);
3205                assert_eq!(units.last().unwrap().lsn, 3);
3206                assert!(has_more);
3207            }
3208            other => panic!("expected byte-capped applyable prefix, got {other:?}"),
3209        }
3210    }
3211
3212    #[test]
3213    fn sync_pull_never_serves_units_beyond_server_remote_lsn() {
3214        let dir = tempfile::tempdir().unwrap();
3215        let mut engine = Engine::new(dir.path()).unwrap();
3216        engine
3217            .execute_powql("type SyncT { required id: int }")
3218            .unwrap();
3219        engine.execute_powql("insert SyncT { id := 1 }").unwrap();
3220        let remote_lsn = engine.catalog().max_lsn();
3221        assert!(remote_lsn > 0);
3222        write_sync_identity_and_tail(dir.path(), remote_lsn + 2);
3223        powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
3224            .unwrap();
3225
3226        let engine = Arc::new(RwLock::new(engine));
3227        let principal = admin_principal();
3228        let identity = sync_identity().segment_identity();
3229        let pull = SyncPullRequest {
3230            replica_id: "replica-a".into(),
3231            since_lsn: 0,
3232            max_units: MAX_SYNC_PULL_UNITS,
3233            max_bytes: MAX_SYNC_PULL_BYTES,
3234            database_id: identity.database_id,
3235            primary_generation: identity.primary_generation,
3236            wal_format_version: identity.wal_format_version,
3237            catalog_version: identity.catalog_version,
3238            segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
3239        };
3240
3241        match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
3242            Message::SyncPullResult {
3243                status,
3244                units,
3245                has_more,
3246            } => {
3247                assert_eq!(status.remote_lsn, remote_lsn);
3248                assert_eq!(status.servable_lsn, Some(remote_lsn));
3249                assert_eq!(status.unarchived_lsn, Some(0));
3250                assert_eq!(status.repair_action, WireSyncRepairAction::Pull);
3251                assert!(!has_more);
3252                assert_eq!(units.len() as u64, remote_lsn);
3253                assert_eq!(units.last().unwrap().lsn, remote_lsn);
3254                assert!(units.iter().all(|unit| unit.lsn <= remote_lsn));
3255            }
3256            other => panic!("expected capped sync pull result, got {other:?}"),
3257        }
3258    }
3259
3260    #[test]
3261    fn sync_status_reports_await_archive_when_primary_outruns_retained_tail() {
3262        let dir = tempfile::tempdir().unwrap();
3263        let mut engine = Engine::new(dir.path()).unwrap();
3264        engine
3265            .execute_powql("type SyncT { required id: int }")
3266            .unwrap();
3267        engine.execute_powql("insert SyncT { id := 1 }").unwrap();
3268        let remote_lsn = engine.catalog().max_lsn();
3269        assert!(remote_lsn > 0);
3270        write_sync_identity_only(dir.path());
3271        powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
3272            .unwrap();
3273
3274        let engine = Arc::new(RwLock::new(engine));
3275        let principal = admin_principal();
3276        let identity = sync_identity().segment_identity();
3277        let status = match dispatch_sync_status(&engine, "replica-a".into(), true, Some(&principal))
3278        {
3279            Message::SyncStatusResult { status } => status,
3280            other => panic!("expected sync status, got {other:?}"),
3281        };
3282        assert_eq!(status.remote_lsn, remote_lsn);
3283        assert_eq!(status.servable_lsn, Some(0));
3284        assert_eq!(status.unarchived_lsn, Some(remote_lsn));
3285        assert_eq!(status.repair_action, WireSyncRepairAction::AwaitArchive);
3286        assert!(status
3287            .last_sync_error
3288            .as_deref()
3289            .unwrap()
3290            .contains("not yet archived"));
3291
3292        let pull = SyncPullRequest {
3293            replica_id: "replica-a".into(),
3294            since_lsn: 0,
3295            max_units: MAX_SYNC_PULL_UNITS,
3296            max_bytes: MAX_SYNC_PULL_BYTES,
3297            database_id: identity.database_id,
3298            primary_generation: identity.primary_generation,
3299            wal_format_version: identity.wal_format_version,
3300            catalog_version: identity.catalog_version,
3301            segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
3302        };
3303        match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
3304            Message::SyncPullResult {
3305                status,
3306                units,
3307                has_more,
3308            } => {
3309                assert_eq!(status.repair_action, WireSyncRepairAction::AwaitArchive);
3310                assert!(units.is_empty());
3311                assert!(!has_more);
3312            }
3313            other => panic!("expected await-archive sync pull result, got {other:?}"),
3314        }
3315    }
3316
3317    #[test]
3318    fn sync_pull_serves_partial_retained_prefix_when_archive_lags_remote_lsn() {
3319        let dir = tempfile::tempdir().unwrap();
3320        let mut engine = Engine::new(dir.path()).unwrap();
3321        engine
3322            .execute_powql("type SyncT { required id: int }")
3323            .unwrap();
3324        engine.execute_powql("insert SyncT { id := 1 }").unwrap();
3325        engine.execute_powql("insert SyncT { id := 2 }").unwrap();
3326        let remote_lsn = engine.catalog().max_lsn();
3327        assert!(remote_lsn > 1);
3328        let servable_lsn = remote_lsn - 1;
3329        write_sync_identity_and_tail(dir.path(), servable_lsn);
3330        powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
3331            .unwrap();
3332
3333        let engine = Arc::new(RwLock::new(engine));
3334        let principal = admin_principal();
3335        let identity = sync_identity().segment_identity();
3336        let pull = SyncPullRequest {
3337            replica_id: "replica-a".into(),
3338            since_lsn: 0,
3339            max_units: MAX_SYNC_PULL_UNITS,
3340            max_bytes: MAX_SYNC_PULL_BYTES,
3341            database_id: identity.database_id,
3342            primary_generation: identity.primary_generation,
3343            wal_format_version: identity.wal_format_version,
3344            catalog_version: identity.catalog_version,
3345            segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
3346        };
3347
3348        match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
3349            Message::SyncPullResult {
3350                status,
3351                units,
3352                has_more,
3353            } => {
3354                assert_eq!(status.remote_lsn, remote_lsn);
3355                assert_eq!(status.servable_lsn, Some(servable_lsn));
3356                assert_eq!(status.unarchived_lsn, Some(1));
3357                assert_eq!(status.repair_action, WireSyncRepairAction::Pull);
3358                assert!(!has_more);
3359                assert_eq!(units.len() as u64, servable_lsn);
3360                assert_eq!(units.last().unwrap().lsn, servable_lsn);
3361                assert!(units.iter().all(|unit| unit.lsn <= servable_lsn));
3362            }
3363            other => panic!("expected partial sync pull result, got {other:?}"),
3364        }
3365    }
3366
3367    #[test]
3368    fn sync_pull_rejects_cursor_or_format_mismatch() {
3369        let dir = tempfile::tempdir().unwrap();
3370        let mut engine = Engine::new(dir.path()).unwrap();
3371        engine
3372            .execute_powql("type SyncT { required id: int }")
3373            .unwrap();
3374        engine.execute_powql("insert SyncT { id := 1 }").unwrap();
3375        let remote_lsn = engine.catalog().max_lsn();
3376        write_sync_identity_and_tail(dir.path(), remote_lsn);
3377        powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
3378            .unwrap();
3379        let engine = Arc::new(RwLock::new(engine));
3380        let principal = admin_principal();
3381        let identity = sync_identity().segment_identity();
3382
3383        let wrong_cursor = SyncPullRequest {
3384            replica_id: "replica-a".into(),
3385            since_lsn: 1,
3386            max_units: 10,
3387            max_bytes: MAX_SYNC_PULL_BYTES,
3388            database_id: identity.database_id,
3389            primary_generation: identity.primary_generation,
3390            wal_format_version: identity.wal_format_version,
3391            catalog_version: identity.catalog_version,
3392            segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
3393        };
3394        match dispatch_sync_pull(&engine, wrong_cursor, true, Some(&principal)) {
3395            Message::Error { message } => assert!(message.contains("does not match")),
3396            other => panic!("expected cursor mismatch error, got {other:?}"),
3397        }
3398
3399        let wrong_format = SyncPullRequest {
3400            replica_id: "replica-a".into(),
3401            since_lsn: 0,
3402            max_units: 10,
3403            max_bytes: MAX_SYNC_PULL_BYTES,
3404            database_id: identity.database_id,
3405            primary_generation: identity.primary_generation,
3406            wal_format_version: identity.wal_format_version,
3407            catalog_version: identity.catalog_version,
3408            segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION + 1,
3409        };
3410        match dispatch_sync_pull(&engine, wrong_format, true, Some(&principal)) {
3411            Message::Error { message } => assert!(message.contains("rebootstrap required")),
3412            other => panic!("expected format mismatch error, got {other:?}"),
3413        }
3414    }
3415}