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