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, SegmentIdentity, SyncRepairAction,
13    RETAINED_SEGMENT_FORMAT_VERSION,
14};
15use std::collections::{HashMap, VecDeque};
16use std::net::IpAddr;
17use std::path::{Path, PathBuf};
18use std::sync::{Arc, Mutex, RwLock};
19use std::time::{Duration, Instant};
20use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufReader, BufWriter};
21use tokio::sync::{watch, OwnedSemaphorePermit, Semaphore};
22use tracing::{debug, error, info, warn};
23use zeroize::Zeroizing;
24
25/// Tracks per-IP authentication failure counts for rate limiting.
26pub type AuthRateLimiter = Arc<Mutex<HashMap<IpAddr, (u32, Instant)>>>;
27
28/// Fixed reader-permit pool. Read-only autocommit statements take one permit;
29/// writers, sync operations, and explicit transactions take the entire pool.
30/// Tokio's fair semaphore queue prevents a waiting writer from being starved by
31/// later readers.
32#[derive(Clone)]
33pub struct TxGate {
34    semaphore: Arc<Semaphore>,
35    permit_count: u32,
36}
37
38pub const DEFAULT_TX_GATE_READER_PERMITS: u32 = 1024;
39
40/// Create a transaction gate for a shared engine.
41pub fn new_tx_gate() -> TxGate {
42    new_tx_gate_with_permits(DEFAULT_TX_GATE_READER_PERMITS)
43}
44
45/// Create a transaction gate with an explicit reader capacity.
46///
47/// The configurable constructor exists so benchmark and compatibility tests can
48/// reproduce the former single-permit admission policy using the exact same
49/// handler code. Production uses [`new_tx_gate`].
50pub fn new_tx_gate_with_permits(permit_count: u32) -> TxGate {
51    assert!(permit_count > 0, "transaction gate requires a permit");
52    TxGate {
53        semaphore: Arc::new(Semaphore::new(permit_count as usize)),
54        permit_count,
55    }
56}
57
58impl TxGate {
59    pub fn permit_count(&self) -> u32 {
60        self.permit_count
61    }
62
63    fn available_permits(&self) -> usize {
64        self.semaphore.available_permits()
65    }
66
67    async fn acquire_many_owned(
68        self,
69        permits: u32,
70    ) -> Result<OwnedSemaphorePermit, tokio::sync::AcquireError> {
71        self.semaphore.acquire_many_owned(permits).await
72    }
73
74    fn try_acquire_many_owned(
75        self,
76        permits: u32,
77    ) -> Result<OwnedSemaphorePermit, tokio::sync::TryAcquireError> {
78        self.semaphore.try_acquire_many_owned(permits)
79    }
80}
81
82/// Maximum query text length accepted from the wire (1 MB).
83const MAX_QUERY_LENGTH: usize = 1024 * 1024;
84
85/// Maximum payload accepted by the post-auth cancellation-safe frame reader.
86/// Keep this equal to the protocol reader's public wire limit.
87const MAX_WIRE_PAYLOAD_SIZE: usize = 64 * 1024 * 1024;
88
89/// Frames received while a query is executing are retained for normal
90/// pipelined processing. Both limits are deliberately much smaller than the
91/// ordinary 64 MiB per-frame protocol limit: in-flight read-ahead is merely a
92/// liveness aid, not a second request buffer. Reaching either cap cancels the
93/// query and closes the connection; socket monitoring is never disabled.
94const MAX_IN_FLIGHT_READ_AHEAD_FRAMES: usize = 128;
95const MAX_IN_FLIGHT_READ_AHEAD_BYTES: usize = 1024 * 1024;
96
97/// Server-side cap for private sync pull batches.
98const MAX_SYNC_PULL_UNITS: u32 = 4096;
99
100/// Server-side cap for retained-unit payload bytes in a private sync pull.
101const MAX_SYNC_PULL_BYTES: u64 = 16 * 1024 * 1024;
102
103/// Maximum encoded response payload size (64 MB). The wire format is still a
104/// single frame today, so oversized result sets must fail cleanly instead of
105/// building an unbounded `Vec<Vec<String>>` and frame in memory.
106#[cfg(not(test))]
107const MAX_RESPONSE_PAYLOAD_SIZE: usize = 64 * 1024 * 1024;
108#[cfg(test)]
109const MAX_RESPONSE_PAYLOAD_SIZE: usize = 1024;
110
111/// Timeout for writing a response to the client. Prevents slow-drain
112/// clients from blocking the handler indefinitely.
113const WRITE_TIMEOUT: Duration = Duration::from_secs(30);
114
115/// Maximum number of auth failures per IP within the rate-limit window.
116const MAX_AUTH_FAILURES: u32 = 5;
117
118/// Window during which auth failures are counted (60 seconds).
119const AUTH_FAILURE_WINDOW: Duration = Duration::from_secs(60);
120
121/// Create a new shared rate limiter.
122pub fn new_rate_limiter() -> AuthRateLimiter {
123    Arc::new(Mutex::new(HashMap::new()))
124}
125
126/// Check whether an IP is rate-limited and record a failure if requested.
127/// Returns `true` if the IP should be rejected.
128fn is_rate_limited(limiter: &AuthRateLimiter, ip: IpAddr) -> bool {
129    let mut map = limiter.lock().unwrap_or_else(|e| e.into_inner());
130    // Clean up stale entries while we have the lock.
131    let now = Instant::now();
132    map.retain(|_, (_, ts)| now.duration_since(*ts) < AUTH_FAILURE_WINDOW);
133
134    if let Some((count, _)) = map.get(&ip) {
135        *count >= MAX_AUTH_FAILURES
136    } else {
137        false
138    }
139}
140
141/// Record an auth failure for the given IP.
142fn record_auth_failure(limiter: &AuthRateLimiter, ip: IpAddr) {
143    let mut map = limiter.lock().unwrap_or_else(|e| e.into_inner());
144    let now = Instant::now();
145    let entry = map.entry(ip).or_insert((0, now));
146    // Reset counter if the window has elapsed.
147    if now.duration_since(entry.1) >= AUTH_FAILURE_WINDOW {
148        *entry = (1, now);
149    } else {
150        entry.0 += 1;
151    }
152}
153
154/// Clear the failure counter on successful auth.
155fn clear_auth_failures(limiter: &AuthRateLimiter, ip: IpAddr) {
156    let mut map = limiter.lock().unwrap_or_else(|e| e.into_inner());
157    map.remove(&ip);
158}
159
160/// Constant-time password comparison. Hashes both inputs to fixed-size
161/// SHA-256 digests so neither length nor content leaks through timing.
162fn constant_time_eq(a: &[u8], b: &[u8]) -> bool {
163    use sha2::{Digest, Sha256};
164    let ha = Sha256::digest(a);
165    let hb = Sha256::digest(b);
166    let mut diff = 0u8;
167    for (x, y) in ha.iter().zip(hb.iter()) {
168        diff |= x ^ y;
169    }
170    diff == 0
171}
172
173/// An authenticated connection's identity. Bound at connect time and consulted
174/// on every query by `dispatch_query` to enforce the user's role: a
175/// `readonly` principal may only execute read statements.
176#[derive(Debug, Clone, PartialEq, Eq)]
177pub struct Principal {
178    pub name: String,
179    pub role: String,
180}
181
182/// Whether a parsed statement is data-definition (schema) work: creating,
183/// altering, or dropping a type or view. `explain <ddl>` is classified by its
184/// inner statement so `explain drop User` needs the same permission as
185/// `drop User`. Mutations that change *rows* (insert/update/delete/upsert/
186/// refresh) and transaction control are NOT DDL — they fall under `Write`.
187fn is_ddl_statement(stmt: &powdb_query::ast::Statement) -> bool {
188    use powdb_query::ast::Statement;
189    let inner = match stmt {
190        Statement::Explain(inner) => inner.as_ref(),
191        other => other,
192    };
193    matches!(
194        inner,
195        Statement::CreateType(_)
196            | Statement::AlterTable(_)
197            | Statement::DropTable(_)
198            | Statement::CreateView(_)
199            | Statement::DropView(_)
200    )
201}
202
203/// The capability a parsed statement requires under the RBAC lattice
204/// (`crates/auth/src/role.rs`). Reads need [`Permission::Read`]; schema
205/// definition needs [`Permission::Ddl`]; every other mutation needs
206/// [`Permission::Write`]. [`Permission::Admin`] is reserved for user/role
207/// management, which is CLI-only today and never reaches this wire path.
208fn required_permission(stmt: &powdb_query::ast::Statement) -> Permission {
209    if is_read_only_statement(stmt) {
210        Permission::Read
211    } else if is_ddl_statement(stmt) {
212        Permission::Ddl
213    } else {
214        Permission::Write
215    }
216}
217
218/// Enforce the principal's role against a parsed statement using the full
219/// permission lattice. Reads are always permitted (any authenticated role can
220/// read — unknown role names still read but fail closed on any mutation).
221/// Mutations require the specific capability the statement maps to: row
222/// mutations need `Write`, schema changes need `Ddl`. Unknown role names
223/// resolve to no builtin and therefore grant nothing beyond reads.
224///
225/// Classification uses the parsed AST via
226/// [`powdb_query::executor::is_read_only_statement`] — the exact same
227/// classifier the RwLock read/write split relies on — so the permission
228/// boundary and the concurrency boundary can never disagree.
229fn check_statement_permitted(
230    principal: Option<&Principal>,
231    stmt: &powdb_query::ast::Statement,
232) -> Result<(), QueryError> {
233    let Some(p) = principal else {
234        // No per-user identity (shared-password or open mode): full access,
235        // byte-identical to the pre-RBAC behavior.
236        return Ok(());
237    };
238    // Reads are permitted for every authenticated principal (preserves the
239    // pre-lattice contract that any connected role may run read-only queries).
240    if is_read_only_statement(stmt) {
241        return Ok(());
242    }
243    let needed = required_permission(stmt);
244    if Role::builtin(&p.role).is_some_and(|r| r.allows(needed)) {
245        return Ok(());
246    }
247    let kind = if needed == Permission::Ddl {
248        "schema-definition"
249    } else {
250        "write"
251    };
252    Err(QueryError::Execution(format!(
253        "permission denied: role '{}' cannot execute {kind} statements",
254        p.role
255    )))
256}
257
258/// Result of the connect-time authentication decision.
259#[derive(Debug, Clone, PartialEq, Eq)]
260pub enum AuthOutcome {
261    /// Authenticated. `principal` is `Some` when a named user authenticated via
262    /// the UserStore, and `None` for the legacy shared-password / open paths
263    /// where there is no per-user identity.
264    Authenticated { principal: Option<Principal> },
265    /// Rejected. The caller sends a generic "authentication failed" error and
266    /// records a rate-limit failure — it must not reveal which check failed.
267    Rejected,
268}
269
270/// Pure, exhaustively-testable authentication decision for a CONNECT handshake.
271///
272/// Policy:
273/// - If `users` has at least one user, multi-user auth is in force: a
274///   `username` is required and `users.authenticate(username, password)` must
275///   succeed. Unknown user, wrong password, or a missing username all reject
276///   with an indistinguishable `Rejected` (no user-vs-password leak).
277/// - If `users` is empty, fall back verbatim to the legacy behavior: when
278///   `expected_password` is `Some`, the candidate must match it (constant time);
279///   when `None`, no auth is required (open). The `username` is ignored here so
280///   that a new client talking to a shared-password server still connects.
281pub fn authenticate_connect(
282    users: &UserStore,
283    expected_password: Option<&str>,
284    username: Option<&str>,
285    password: Option<&str>,
286) -> AuthOutcome {
287    if !users.is_empty() {
288        // Multi-user mode: a username is mandatory.
289        let Some(name) = username else {
290            return AuthOutcome::Rejected;
291        };
292        let Some(candidate) = password else {
293            return AuthOutcome::Rejected;
294        };
295        match users.authenticate(name, candidate) {
296            Some(user) => AuthOutcome::Authenticated {
297                principal: Some(Principal {
298                    name: user.name.clone(),
299                    role: user.role.clone(),
300                }),
301            },
302            None => AuthOutcome::Rejected,
303        }
304    } else {
305        // Legacy shared-password fallback (byte-identical to prior behavior).
306        match expected_password {
307            Some(expected) => {
308                if password.is_some_and(|p| constant_time_eq(p.as_bytes(), expected.as_bytes())) {
309                    AuthOutcome::Authenticated { principal: None }
310                } else {
311                    AuthOutcome::Rejected
312                }
313            }
314            None => AuthOutcome::Authenticated { principal: None },
315        }
316    }
317}
318
319/// The sentinel database name clients send when the user selected none. Both
320/// the CLI and the TS client default to this, so it means "no specific
321/// database" and is always accepted — even when the server is pinned to a name.
322const DEFAULT_DB_NAME: &str = "default";
323
324/// Decide whether a CONNECT's requested `db_name` is served by this process.
325///
326/// One server process serves exactly one global database. When it is pinned to
327/// a name (`configured = Some`), a request that *explicitly* names a different
328/// database is rejected so a client can never silently read/write the wrong
329/// store. An empty name or the client default sentinel (`"default"`) means "no
330/// specific database selected" and is always accepted. When unpinned (`None`)
331/// every name is accepted (0.9.x back-compat); the caller warns on a non-default
332/// name so the silent-mismatch footgun is at least visible in the logs.
333fn check_db_name(configured: Option<&str>, requested: &str) -> Result<(), String> {
334    if requested.is_empty() || requested == DEFAULT_DB_NAME {
335        return Ok(());
336    }
337    match configured {
338        None => Ok(()),
339        Some(name) if requested == name => Ok(()),
340        Some(name) => Err(format!(
341            "unknown database '{requested}'; this server serves '{name}'"
342        )),
343    }
344}
345
346/// Error messages that are safe to forward to the client verbatim.
347const SAFE_ERROR_PREFIXES: &[&str] = &[
348    "table not found",
349    // The executor's actual phrasing is `table 'X' not found`, which the
350    // bare prefix above never matches — keep both so the real message
351    // reaches clients.
352    "table '",
353    "type '",
354    "column not found",
355    // Lexer diagnostics (`at position N: unterminated quoted identifier`)
356    // are derived purely from the client's own query text.
357    "at position",
358    "parse error",
359    "type mismatch",
360    "unknown table",
361    "unknown column",
362    "unknown function",
363    "syntax error",
364    "expected",
365    "unexpected",
366    "missing",
367    "duplicate",
368    "invalid",
369    "cannot",
370    "no such",
371    "already exists",
372    "permission denied",
373    "row too large",
374    "unique constraint violation",
375    // Resource-limit errors carry actionable guidance (e.g. "add a LIMIT
376    // clause") and leak no internal state, so surface them verbatim instead
377    // of masking them to the generic message. See QueryError::{SortLimit,
378    // JoinLimit,MemoryLimit}Exceeded in crates/query/src/result.rs.
379    "sort input exceeds",
380    "join result exceeds",
381    "query exceeded memory budget",
382    "result too large",
383    // A failed covering fsync means the statement executed in memory but was
384    // never made durable — the client MUST be able to distinguish this from
385    // an ordinary failed query (the statement may still be visible until the
386    // server restarts). The io::Error detail leaks no internal state.
387    "wal durability sync failed",
388    // Cooperative query cancellation. Both messages are derived purely from
389    // the configured timeout / a client disconnect and leak no internal state.
390    // See QueryError::{Timeout,Cancelled} in crates/query/src/result.rs.
391    "query timeout after",
392    "query cancelled",
393];
394
395/// Sanitize an error message before sending it to the client.
396/// Known safe errors are passed through; everything else is replaced
397/// with a generic message to avoid leaking internal details.
398fn sanitize_error(e: &str) -> String {
399    let lower = e.to_lowercase();
400    for prefix in SAFE_ERROR_PREFIXES {
401        if lower.starts_with(prefix) {
402            return e.to_string();
403        }
404    }
405    "query execution error".into()
406}
407
408/// Write a message to the client with a timeout. Returns false if the
409/// write failed or timed out (caller should close the connection).
410async fn write_msg<W: AsyncWrite + Unpin>(writer: &mut BufWriter<W>, msg: &Message) -> bool {
411    let write_fut = async {
412        if msg.write_to(writer).await.is_err() {
413            return false;
414        }
415        writer.flush().await.is_ok()
416    };
417    tokio::time::timeout(WRITE_TIMEOUT, write_fut)
418        .await
419        .unwrap_or_default()
420}
421
422/// Options for a single connection, bundled to keep `handle_connection`'s
423/// argument list short.
424pub struct ConnOpts<'a> {
425    pub engine: Arc<RwLock<Engine>>,
426    pub tx_gate: TxGate,
427    /// Expected client password. Wrapped in `Zeroizing` so the secret is wiped
428    /// from memory on drop (defends against leaking via a core dump).
429    pub expected_password: Option<Zeroizing<String>>,
430    /// Multi-user store loaded from the data dir at startup. When it has users,
431    /// the handshake authenticates `(username, password)` against it; when empty
432    /// the server falls back to `expected_password`. Shared across connections.
433    pub users: Arc<UserStore>,
434    pub shutdown_rx: &'a mut watch::Receiver<bool>,
435    pub idle_timeout: Duration,
436    pub query_timeout: Duration,
437    pub rate_limiter: Option<&'a AuthRateLimiter>,
438    pub peer_addr: Option<std::net::SocketAddr>,
439    /// Shared server metrics. Always present; tests pass `Arc::new(Metrics::new())`.
440    pub metrics: Arc<Metrics>,
441    /// How long an explicit `begin` waits to acquire the transaction gate while
442    /// another connection holds an open explicit transaction, before giving up
443    /// with a clear timeout error instead of queueing indefinitely.
444    pub tx_wait_timeout: Duration,
445    /// When `Some`, the single database name this server serves. A CONNECT that
446    /// explicitly names a *different* database is rejected at connect time.
447    /// `None` accepts any name (0.9.x behavior) and only warns.
448    pub db_name: Option<String>,
449}
450
451/// Execute a query against the engine under the RwLock. Read-only
452/// statements acquire `.read()` so concurrent SELECTs can scan in
453/// parallel; mutations acquire `.write()`.
454///
455/// When `principal` is `Some`, the user's role is enforced first: a role
456/// without the `Write` permission (i.e. `readonly`) gets a clean
457/// "permission denied" error for any non-read statement, before any lock
458/// is taken or any engine state is touched.
459/// A statement's execution result plus its (not-yet-waited) WAL durability
460/// obligation. The ticket is settled by [`finalize_durability`] AFTER the
461/// TxGate permit is dropped, so overlapping committers on other connections
462/// can share a single fsync (group commit).
463type DispatchOutcome = (Result<QueryResult, QueryError>, Option<WalDurabilityTicket>);
464
465/// Execute a mutating statement under the engine write lock with WAL group
466/// commit: the commit's fsync obligation is registered (not performed) while
467/// the lock is held, then the lock is dropped and the un-waited ticket is
468/// returned. The caller must settle the ticket (via [`finalize_durability`])
469/// before acknowledging the statement — and must do so with the TxGate
470/// permit already released, or committers can never overlap.
471fn execute_write_deferred(
472    engine: &Arc<RwLock<Engine>>,
473    f: impl FnOnce(&mut Engine) -> Result<QueryResult, QueryError>,
474) -> DispatchOutcome {
475    let mut eng = match engine.write() {
476        Ok(eng) => eng,
477        Err(e) => {
478            return (
479                Err(QueryError::Execution(format!("lock poisoned: {e}"))),
480                None,
481            )
482        }
483    };
484    eng.run_with_deferred_durability(f)
485}
486
487fn dispatch_query(
488    engine: &Arc<RwLock<Engine>>,
489    query: &str,
490    principal: Option<&Principal>,
491    allow_readonly_escalation: bool,
492) -> DispatchOutcome {
493    let stmt_result = parser::parse(query).map_err(|e| e.to_string());
494    dispatch_query_parsed(
495        engine,
496        query,
497        &stmt_result,
498        principal,
499        allow_readonly_escalation,
500    )
501}
502
503fn dispatch_query_parsed(
504    engine: &Arc<RwLock<Engine>>,
505    query: &str,
506    stmt_result: &Result<powdb_query::ast::Statement, String>,
507    principal: Option<&Principal>,
508    allow_readonly_escalation: bool,
509) -> DispatchOutcome {
510    // Role enforcement happens on the parsed AST. Statements that fail to
511    // parse fall through — the engine returns the parse error itself and
512    // can never execute anything for them.
513    if let Ok(stmt) = &stmt_result {
514        if let Err(e) = check_statement_permitted(principal, stmt) {
515            return (Err(e), None);
516        }
517    }
518
519    let can_try_read = matches!(&stmt_result, Ok(s) if is_read_only_statement(s));
520    if can_try_read {
521        let res = {
522            let eng = match engine.read() {
523                Ok(eng) => eng,
524                Err(e) => {
525                    return (
526                        Err(QueryError::Execution(format!("lock poisoned: {e}"))),
527                        None,
528                    )
529                }
530            };
531            eng.execute_powql_readonly(query)
532        };
533        match res {
534            Ok(r) => return (Ok(r), None),
535            Err(QueryError::ReadonlyNeedsWrite) => {
536                if !allow_readonly_escalation {
537                    return (Err(QueryError::ReadonlyNeedsWrite), None);
538                }
539                // The caller already owns writer admission, so it is safe to
540                // fall through and retry under the engine write lock.
541            }
542            Err(e) => return (Err(e), None),
543        }
544    }
545
546    if matches!(
547        parsed_transaction_control(stmt_result),
548        Some(TransactionControl::Rollback)
549    ) {
550        let mut eng = match engine.write() {
551            Ok(eng) => eng,
552            Err(e) => {
553                return (
554                    Err(QueryError::Execution(format!("lock poisoned: {e}"))),
555                    None,
556                )
557            }
558        };
559        return (execute_rollback_preserving_sync_if_needed(&mut eng), None);
560    }
561    execute_write_deferred(engine, |eng| eng.execute_powql(query))
562}
563
564#[cfg(test)]
565fn dispatch_sql_query(
566    engine: &Arc<RwLock<Engine>>,
567    query: &str,
568    principal: Option<&Principal>,
569    allow_readonly_escalation: bool,
570) -> DispatchOutcome {
571    let stmt_result = sql::parse_sql(query).map_err(|e| e.to_string());
572    dispatch_sql_query_parsed(
573        engine,
574        query,
575        &stmt_result,
576        principal,
577        allow_readonly_escalation,
578    )
579}
580
581fn dispatch_sql_query_parsed(
582    engine: &Arc<RwLock<Engine>>,
583    query: &str,
584    stmt_result: &Result<powdb_query::ast::Statement, String>,
585    principal: Option<&Principal>,
586    allow_readonly_escalation: bool,
587) -> DispatchOutcome {
588    if let Ok(stmt) = &stmt_result {
589        if let Err(e) = check_statement_permitted(principal, stmt) {
590            return (Err(e), None);
591        }
592    }
593
594    let can_try_read = matches!(&stmt_result, Ok(s) if is_read_only_statement(s));
595    if can_try_read {
596        let res = {
597            let eng = match engine.read() {
598                Ok(eng) => eng,
599                Err(e) => {
600                    return (
601                        Err(QueryError::Execution(format!("lock poisoned: {e}"))),
602                        None,
603                    )
604                }
605            };
606            eng.execute_sql_readonly(query)
607        };
608        match res {
609            Ok(r) => return (Ok(r), None),
610            Err(QueryError::ReadonlyNeedsWrite) => {
611                if !allow_readonly_escalation {
612                    return (Err(QueryError::ReadonlyNeedsWrite), None);
613                }
614            }
615            Err(e) => return (Err(e), None),
616        }
617    }
618
619    if matches!(
620        parsed_transaction_control(stmt_result),
621        Some(TransactionControl::Rollback)
622    ) {
623        let mut eng = match engine.write() {
624            Ok(eng) => eng,
625            Err(e) => {
626                return (
627                    Err(QueryError::Execution(format!("lock poisoned: {e}"))),
628                    None,
629                )
630            }
631        };
632        return (execute_rollback_preserving_sync_if_needed(&mut eng), None);
633    }
634    execute_write_deferred(engine, |eng| eng.execute_sql(query))
635}
636
637/// Convert a wire parameter into the query-crate [`ParamValue`] used for
638/// token-level binding.
639fn wire_param_to_value(p: &WireParam) -> powdb_query::ast::ParamValue {
640    use powdb_query::ast::ParamValue;
641    match p {
642        WireParam::Null => ParamValue::Null,
643        WireParam::Int(v) => ParamValue::Int(*v),
644        WireParam::Float(v) => ParamValue::Float(*v),
645        WireParam::Bool(v) => ParamValue::Bool(*v),
646        WireParam::Str(s) => ParamValue::Str(s.clone()),
647    }
648}
649
650/// Parameterized counterpart of [`dispatch_query`]. Routes through the exact
651/// same role-enforcement and read/write escalation logic, but binds the
652/// `$N` placeholders at the token level via the query crate's
653/// `parse_with_params` path. A string parameter can never change the query's
654/// shape — it is substituted as a literal token, not interpolated text.
655#[derive(Debug, Clone, Copy, PartialEq, Eq)]
656enum TransactionControl {
657    Begin,
658    Commit,
659    Rollback,
660}
661
662#[derive(Debug, Clone, Copy, PartialEq, Eq)]
663enum AdmissionMode {
664    Reader,
665    Writer,
666}
667
668#[derive(Clone, Copy, Debug, PartialEq, Eq)]
669enum WireResultMode {
670    LegacyText,
671    Native,
672}
673
674fn statement_admission(stmt: &powdb_query::ast::Statement) -> AdmissionMode {
675    if is_read_only_statement(stmt) {
676        AdmissionMode::Reader
677    } else {
678        AdmissionMode::Writer
679    }
680}
681
682#[cfg(test)]
683fn classify_query_admission(query: &str) -> AdmissionMode {
684    parser::parse(query)
685        .map(|stmt| statement_admission(&stmt))
686        .unwrap_or(AdmissionMode::Writer)
687}
688
689#[cfg(test)]
690fn classify_sql_admission(query: &str) -> AdmissionMode {
691    sql::parse_sql(query)
692        .map(|stmt| statement_admission(&stmt))
693        .unwrap_or(AdmissionMode::Writer)
694}
695
696#[cfg(test)]
697fn classify_params_admission(query: &str, params: &[WireParam]) -> AdmissionMode {
698    let bound: Vec<powdb_query::ast::ParamValue> = params.iter().map(wire_param_to_value).collect();
699    parser::parse_with_params(query, &bound)
700        .map(|stmt| statement_admission(&stmt))
701        .unwrap_or(AdmissionMode::Writer)
702}
703
704fn transaction_control(stmt: &powdb_query::ast::Statement) -> Option<TransactionControl> {
705    use powdb_query::ast::Statement;
706    match stmt {
707        Statement::Begin => Some(TransactionControl::Begin),
708        Statement::Commit => Some(TransactionControl::Commit),
709        Statement::Rollback => Some(TransactionControl::Rollback),
710        _ => None,
711    }
712}
713
714fn parsed_transaction_control(
715    stmt_result: &Result<powdb_query::ast::Statement, String>,
716) -> Option<TransactionControl> {
717    stmt_result.as_ref().ok().and_then(transaction_control)
718}
719
720fn execute_rollback_preserving_sync_if_needed(
721    engine: &mut Engine,
722) -> Result<QueryResult, QueryError> {
723    engine.rollback_transaction_preserving_wal_archive()
724}
725
726#[cfg(test)]
727fn dispatch_query_with_params(
728    engine: &Arc<RwLock<Engine>>,
729    query: &str,
730    params: &[WireParam],
731    principal: Option<&Principal>,
732    allow_readonly_escalation: bool,
733) -> DispatchOutcome {
734    let bound: Vec<powdb_query::ast::ParamValue> = params.iter().map(wire_param_to_value).collect();
735
736    // Parse once (with params bound) so role enforcement and read/write
737    // classification see exactly the statement that will execute.
738    let stmt_result = parser::parse_with_params(query, &bound).map_err(|e| e.to_string());
739    dispatch_query_with_bound_params_parsed(
740        engine,
741        query,
742        &bound,
743        &stmt_result,
744        principal,
745        allow_readonly_escalation,
746    )
747}
748
749fn dispatch_query_with_bound_params_parsed(
750    engine: &Arc<RwLock<Engine>>,
751    query: &str,
752    bound: &[powdb_query::ast::ParamValue],
753    stmt_result: &Result<powdb_query::ast::Statement, String>,
754    principal: Option<&Principal>,
755    allow_readonly_escalation: bool,
756) -> DispatchOutcome {
757    if let Ok(stmt) = &stmt_result {
758        if let Err(e) = check_statement_permitted(principal, stmt) {
759            return (Err(e), None);
760        }
761    }
762
763    let can_try_read = matches!(&stmt_result, Ok(s) if is_read_only_statement(s));
764    if can_try_read {
765        let res = {
766            let eng = match engine.read() {
767                Ok(eng) => eng,
768                Err(e) => {
769                    return (
770                        Err(QueryError::Execution(format!("lock poisoned: {e}"))),
771                        None,
772                    )
773                }
774            };
775            eng.execute_powql_readonly_with_params(query, bound)
776        };
777        match res {
778            Ok(r) => return (Ok(r), None),
779            Err(QueryError::ReadonlyNeedsWrite) => {
780                if !allow_readonly_escalation {
781                    return (Err(QueryError::ReadonlyNeedsWrite), None);
782                }
783            }
784            Err(e) => return (Err(e), None),
785        }
786    }
787
788    if matches!(
789        parsed_transaction_control(stmt_result),
790        Some(TransactionControl::Rollback)
791    ) {
792        let mut eng = match engine.write() {
793            Ok(eng) => eng,
794            Err(e) => {
795                return (
796                    Err(QueryError::Execution(format!("lock poisoned: {e}"))),
797                    None,
798                )
799            }
800        };
801        return (execute_rollback_preserving_sync_if_needed(&mut eng), None);
802    }
803    execute_write_deferred(engine, |eng| eng.execute_powql_with_params(query, bound))
804}
805
806#[derive(Debug, Clone)]
807struct SyncPullRequest {
808    replica_id: String,
809    since_lsn: u64,
810    max_units: u32,
811    max_bytes: u64,
812    database_id: [u8; 16],
813    primary_generation: u64,
814    wal_format_version: u16,
815    catalog_version: u16,
816    segment_format_version: u16,
817}
818
819#[derive(Debug, Clone, Copy, PartialEq, Eq)]
820enum SyncErrorClass {
821    AuthRequired,
822    PermissionDenied,
823    InvalidReplicaId,
824    ActiveTransaction,
825    QueryExecution,
826    InvalidMaxUnits,
827    InvalidMaxBytes,
828    SyncContext,
829    StatusRead,
830    CursorLsnMismatch,
831    IdentityRead,
832    IdentityOrFormatMismatch,
833    RetainedRead,
834    RetainedUnitEncoding,
835    RetainedChunkNotApplyable,
836    LsnAheadOfRemote,
837    AckValidation,
838    AckUpdate,
839    Internal,
840}
841
842impl SyncErrorClass {
843    const fn as_label(self) -> &'static str {
844        match self {
845            Self::AuthRequired => "auth_required",
846            Self::PermissionDenied => "permission_denied",
847            Self::InvalidReplicaId => "invalid_replica_id",
848            Self::ActiveTransaction => "active_transaction",
849            Self::QueryExecution => "query_execution",
850            Self::InvalidMaxUnits => "invalid_max_units",
851            Self::InvalidMaxBytes => "invalid_max_bytes",
852            Self::SyncContext => "sync_context",
853            Self::StatusRead => "status_read",
854            Self::CursorLsnMismatch => "cursor_lsn_mismatch",
855            Self::IdentityRead => "identity_read",
856            Self::IdentityOrFormatMismatch => "identity_or_format_mismatch",
857            Self::RetainedRead => "retained_read",
858            Self::RetainedUnitEncoding => "retained_unit_encoding",
859            Self::RetainedChunkNotApplyable => "retained_chunk_not_applyable",
860            Self::LsnAheadOfRemote => "lsn_ahead_of_remote",
861            Self::AckValidation => "ack_validation",
862            Self::AckUpdate => "ack_update",
863            Self::Internal => "internal",
864        }
865    }
866}
867
868#[derive(Debug, Clone)]
869struct SyncDecision {
870    message: Message,
871    error_class: Option<SyncErrorClass>,
872}
873
874impl SyncDecision {
875    fn ok(message: Message) -> Self {
876        Self {
877            message,
878            error_class: None,
879        }
880    }
881
882    fn error(class: SyncErrorClass, message: impl Into<String>) -> Self {
883        Self {
884            message: Message::Error {
885                message: message.into(),
886            },
887            error_class: Some(class),
888        }
889    }
890}
891
892fn check_sync_protocol_permitted(
893    credential_authenticated: bool,
894    principal: Option<&Principal>,
895) -> Result<(), (SyncErrorClass, String)> {
896    if !credential_authenticated {
897        return Err((
898            SyncErrorClass::AuthRequired,
899            "sync protocol requires authentication".to_string(),
900        ));
901    }
902    if let Some(principal) = principal {
903        let allowed =
904            Role::builtin(&principal.role).is_some_and(|role| role.allows(Permission::Write));
905        if !allowed {
906            return Err((
907                SyncErrorClass::PermissionDenied,
908                format!(
909                    "permission denied: role '{}' cannot use sync protocol",
910                    principal.role
911                ),
912            ));
913        }
914    }
915    Ok(())
916}
917
918fn validate_wire_replica_id(replica_id: &str) -> Result<(), String> {
919    if replica_id.is_empty() {
920        return Err("replica id must be non-empty".to_string());
921    }
922    if replica_id.len() > 128 {
923        return Err("replica id must be at most 128 bytes".to_string());
924    }
925    if !replica_id
926        .bytes()
927        .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'-' | b'_' | b'.' | b':'))
928    {
929        return Err("replica id contains unsupported characters".to_string());
930    }
931    Ok(())
932}
933
934fn sync_context(engine: &Arc<RwLock<Engine>>) -> Result<SyncContext, String> {
935    let engine = engine
936        .read()
937        .map_err(|e| format!("lock poisoned while reading sync context: {e}"))?;
938    let catalog = engine.catalog();
939    Ok(SyncContext {
940        data_dir: catalog.data_dir().to_path_buf(),
941        remote_lsn: catalog.max_lsn(),
942        active_catalog_version: catalog.active_catalog_version(),
943    })
944}
945
946/// Snapshot of the primary's sync-relevant state captured under the engine lock.
947/// `active_catalog_version` is the database's *active* on-disk catalog format,
948/// used to stamp the expected segment identity a pulling replica is checked
949/// against (a database that never activated v6 keeps expecting v5).
950struct SyncContext {
951    data_dir: PathBuf,
952    remote_lsn: u64,
953    active_catalog_version: u16,
954}
955
956fn wire_repair_action(action: SyncRepairAction) -> WireSyncRepairAction {
957    match action {
958        SyncRepairAction::None => WireSyncRepairAction::None,
959        SyncRepairAction::Pull => WireSyncRepairAction::Pull,
960        SyncRepairAction::AwaitArchive => WireSyncRepairAction::AwaitArchive,
961        SyncRepairAction::Rebootstrap => WireSyncRepairAction::Rebootstrap,
962    }
963}
964
965fn wire_sync_status(status: ReplicaSyncStatus) -> WireSyncStatus {
966    WireSyncStatus {
967        replica_id: status.replica_id,
968        active: status.active,
969        last_applied_lsn: status.last_applied_lsn,
970        remote_lsn: status.remote_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: wire_repair_action(status.repair_action),
978        last_sync_error: status.last_sync_error,
979    }
980}
981
982fn wire_retained_unit(unit: RetainedUnit) -> WireRetainedUnit {
983    WireRetainedUnit {
984        tx_id: unit.tx_id,
985        record_type: unit.record_type,
986        lsn: unit.lsn,
987        data: unit.data,
988    }
989}
990
991fn sync_operation_outcome(message: &Message) -> SyncOutcome {
992    match message {
993        Message::SyncStatusResult { .. }
994        | Message::SyncPullResult { .. }
995        | Message::SyncAckResult { .. } => SyncOutcome::Ok,
996        _ => SyncOutcome::Error,
997    }
998}
999
1000fn sync_operation_label(operation: SyncOperation) -> &'static str {
1001    match operation {
1002        SyncOperation::Status => "status",
1003        SyncOperation::Pull => "pull",
1004        SyncOperation::Ack => "ack",
1005    }
1006}
1007
1008fn sync_pull_payload_bytes(units: &[WireRetainedUnit]) -> u64 {
1009    units.iter().fold(0, |total, unit| {
1010        total.saturating_add(unit.encoded_len().unwrap_or(0))
1011    })
1012}
1013
1014fn sync_repair_label(action: WireSyncRepairAction) -> SyncRepairLabel {
1015    match action {
1016        WireSyncRepairAction::None => SyncRepairLabel::None,
1017        WireSyncRepairAction::Pull => SyncRepairLabel::Pull,
1018        WireSyncRepairAction::AwaitArchive => SyncRepairLabel::AwaitArchive,
1019        WireSyncRepairAction::Rebootstrap => SyncRepairLabel::Rebootstrap,
1020    }
1021}
1022
1023fn wire_repair_action_label(action: WireSyncRepairAction) -> &'static str {
1024    match action {
1025        WireSyncRepairAction::None => "none",
1026        WireSyncRepairAction::Pull => "pull",
1027        WireSyncRepairAction::AwaitArchive => "await_archive",
1028        WireSyncRepairAction::Rebootstrap => "rebootstrap",
1029    }
1030}
1031
1032const FNV1A64_OFFSET: u64 = 0xcbf29ce484222325;
1033const FNV1A64_PRIME: u64 = 0x100000001b3;
1034const INVALID_REPLICA_FINGERPRINT: &str = "invalid";
1035
1036fn replica_fingerprint(replica_id: &str) -> String {
1037    let mut hash = FNV1A64_OFFSET;
1038    for byte in replica_id.as_bytes() {
1039        hash ^= u64::from(*byte);
1040        hash = hash.wrapping_mul(FNV1A64_PRIME);
1041    }
1042    format!("{hash:016x}")
1043}
1044
1045fn log_replica_fingerprint(replica_id: &str) -> String {
1046    if validate_wire_replica_id(replica_id).is_ok() {
1047        replica_fingerprint(replica_id)
1048    } else {
1049        INVALID_REPLICA_FINGERPRINT.to_string()
1050    }
1051}
1052
1053#[derive(Debug, Clone)]
1054struct SyncLogContext {
1055    replica_fingerprint: String,
1056    since_lsn: Option<u64>,
1057    applied_lsn: Option<u64>,
1058    observed_remote_lsn: Option<u64>,
1059    max_units: Option<u32>,
1060    max_bytes: Option<u64>,
1061}
1062
1063impl SyncLogContext {
1064    fn status(replica_id: &str) -> Self {
1065        Self::base(replica_id)
1066    }
1067
1068    fn pull(request: &SyncPullRequest) -> Self {
1069        Self {
1070            replica_fingerprint: log_replica_fingerprint(&request.replica_id),
1071            since_lsn: Some(request.since_lsn),
1072            applied_lsn: None,
1073            observed_remote_lsn: None,
1074            max_units: Some(request.max_units),
1075            max_bytes: Some(request.max_bytes),
1076        }
1077    }
1078
1079    fn ack(replica_id: &str, applied_lsn: u64, observed_remote_lsn: u64) -> Self {
1080        Self {
1081            replica_fingerprint: log_replica_fingerprint(replica_id),
1082            since_lsn: None,
1083            applied_lsn: Some(applied_lsn),
1084            observed_remote_lsn: Some(observed_remote_lsn),
1085            max_units: None,
1086            max_bytes: None,
1087        }
1088    }
1089
1090    fn base(replica_id: &str) -> Self {
1091        Self {
1092            replica_fingerprint: log_replica_fingerprint(replica_id),
1093            since_lsn: None,
1094            applied_lsn: None,
1095            observed_remote_lsn: None,
1096            max_units: None,
1097            max_bytes: None,
1098        }
1099    }
1100}
1101
1102struct SyncExecutionContext<'a> {
1103    tx_gate: TxGate,
1104    connection_has_transaction: bool,
1105    operation: SyncOperation,
1106    log_context: SyncLogContext,
1107    metrics: &'a Arc<Metrics>,
1108    query_timeout: Duration,
1109}
1110
1111fn log_sync_decision(
1112    operation: SyncOperation,
1113    context: &SyncLogContext,
1114    elapsed: Duration,
1115    decision: &SyncDecision,
1116) {
1117    let operation = sync_operation_label(operation);
1118    let elapsed_ms = elapsed.as_secs_f64() * 1000.0;
1119    match &decision.message {
1120        Message::SyncStatusResult { status } => {
1121            let repair_action = wire_repair_action_label(status.repair_action);
1122            if status.stale || status.repair_action != WireSyncRepairAction::None {
1123                info!(
1124                    operation = operation,
1125                    replica_fingerprint = %context.replica_fingerprint,
1126                    remote_lsn = status.remote_lsn,
1127                    last_applied_lsn = ?status.last_applied_lsn,
1128                    servable_lsn = ?status.servable_lsn,
1129                    unarchived_lsn = ?status.unarchived_lsn,
1130                    lag_lsn = ?status.lag_lsn,
1131                    lag_bytes = ?status.lag_bytes,
1132                    lag_ms = ?status.lag_ms,
1133                    stale = status.stale,
1134                    repair_action,
1135                    elapsed_ms,
1136                    "sync decision"
1137                );
1138            } else {
1139                debug!(
1140                    operation = operation,
1141                    replica_fingerprint = %context.replica_fingerprint,
1142                    remote_lsn = status.remote_lsn,
1143                    last_applied_lsn = ?status.last_applied_lsn,
1144                    repair_action,
1145                    elapsed_ms,
1146                    "sync decision"
1147                );
1148            }
1149        }
1150        Message::SyncPullResult {
1151            status,
1152            units,
1153            has_more,
1154        } => {
1155            let repair_action = wire_repair_action_label(status.repair_action);
1156            let units_len = units.len();
1157            let payload_bytes = sync_pull_payload_bytes(units);
1158            if *has_more || status.stale || status.repair_action != WireSyncRepairAction::None {
1159                info!(
1160                    operation = operation,
1161                    replica_fingerprint = %context.replica_fingerprint,
1162                    since_lsn = ?context.since_lsn,
1163                    max_units = ?context.max_units,
1164                    max_bytes = ?context.max_bytes,
1165                    units = units_len,
1166                    payload_bytes,
1167                    has_more = *has_more,
1168                    remote_lsn = status.remote_lsn,
1169                    last_applied_lsn = ?status.last_applied_lsn,
1170                    servable_lsn = ?status.servable_lsn,
1171                    unarchived_lsn = ?status.unarchived_lsn,
1172                    lag_lsn = ?status.lag_lsn,
1173                    stale = status.stale,
1174                    repair_action,
1175                    elapsed_ms,
1176                    "sync decision"
1177                );
1178            } else {
1179                debug!(
1180                    operation = operation,
1181                    replica_fingerprint = %context.replica_fingerprint,
1182                    since_lsn = ?context.since_lsn,
1183                    units = units_len,
1184                    payload_bytes,
1185                    has_more = *has_more,
1186                    remote_lsn = status.remote_lsn,
1187                    repair_action,
1188                    elapsed_ms,
1189                    "sync decision"
1190                );
1191            }
1192        }
1193        Message::SyncAckResult {
1194            previous_applied_lsn,
1195            applied_lsn,
1196            remote_lsn,
1197            advanced,
1198            status,
1199        } => {
1200            let repair_action = wire_repair_action_label(status.repair_action);
1201            if *advanced || status.stale || status.repair_action != WireSyncRepairAction::None {
1202                info!(
1203                    operation = operation,
1204                    replica_fingerprint = %context.replica_fingerprint,
1205                    requested_applied_lsn = ?context.applied_lsn,
1206                    observed_remote_lsn = ?context.observed_remote_lsn,
1207                    previous_applied_lsn = *previous_applied_lsn,
1208                    applied_lsn = *applied_lsn,
1209                    remote_lsn = *remote_lsn,
1210                    advanced = *advanced,
1211                    stale = status.stale,
1212                    repair_action,
1213                    elapsed_ms,
1214                    "sync decision"
1215                );
1216            } else {
1217                debug!(
1218                    operation = operation,
1219                    replica_fingerprint = %context.replica_fingerprint,
1220                    requested_applied_lsn = ?context.applied_lsn,
1221                    previous_applied_lsn = *previous_applied_lsn,
1222                    applied_lsn = *applied_lsn,
1223                    remote_lsn = *remote_lsn,
1224                    advanced = *advanced,
1225                    repair_action,
1226                    elapsed_ms,
1227                    "sync decision"
1228                );
1229            }
1230        }
1231        Message::Error { .. } => {
1232            warn!(
1233                operation = operation,
1234                replica_fingerprint = %context.replica_fingerprint,
1235                since_lsn = ?context.since_lsn,
1236                applied_lsn = ?context.applied_lsn,
1237                observed_remote_lsn = ?context.observed_remote_lsn,
1238                max_units = ?context.max_units,
1239                max_bytes = ?context.max_bytes,
1240                error_class = decision
1241                    .error_class
1242                    .unwrap_or(SyncErrorClass::Internal)
1243                    .as_label(),
1244                elapsed_ms,
1245                "sync decision rejected"
1246            );
1247        }
1248        _ => {
1249            debug!(
1250                operation = operation,
1251                replica_fingerprint = %context.replica_fingerprint,
1252                elapsed_ms,
1253                "unexpected sync decision response"
1254            );
1255        }
1256    }
1257}
1258
1259fn trim_to_applyable_v1_prefix(
1260    raw_units: &mut Vec<RetainedUnit>,
1261    wire_units: &mut Vec<WireRetainedUnit>,
1262) -> Result<(), String> {
1263    let mut last_error = None;
1264    while !raw_units.is_empty() {
1265        match validate_v1_retained_units_applyable(raw_units) {
1266            Ok(()) => return Ok(()),
1267            Err(err) => {
1268                last_error = Some(err.to_string());
1269                raw_units.pop();
1270                wire_units.pop();
1271            }
1272        }
1273    }
1274    if let Some(error) = last_error {
1275        return Err(format!(
1276            "sync pull cannot serve an applyable V1 retained chunk with current limits: {error}"
1277        ));
1278    }
1279    Ok(())
1280}
1281
1282fn validate_sync_ack_applyable_boundary(
1283    data_dir: &Path,
1284    replica_id: &str,
1285    applied_lsn: u64,
1286    remote_lsn: u64,
1287) -> Result<(), String> {
1288    let status =
1289        replica_sync_status(data_dir, replica_id, remote_lsn).map_err(|err| err.to_string())?;
1290    let Some(previous_lsn) = status.last_applied_lsn else {
1291        return Ok(());
1292    };
1293    if applied_lsn <= previous_lsn {
1294        return Ok(());
1295    }
1296    let range_len = applied_lsn - previous_lsn;
1297    if range_len > u64::from(MAX_SYNC_PULL_UNITS) {
1298        return Err(format!(
1299            "sync ack range contains {range_len} units; acknowledge ranges no larger than {MAX_SYNC_PULL_UNITS}"
1300        ));
1301    }
1302    let max_units =
1303        usize::try_from(range_len).map_err(|_| "sync ack range is too large to validate")?;
1304    let identity = read_identity(data_dir).map_err(|err| err.to_string())?;
1305    let segment_dir = retained_segments_dir(data_dir);
1306    let units = read_units_through(
1307        &segment_dir,
1308        identity.segment_identity(),
1309        previous_lsn,
1310        applied_lsn,
1311        max_units,
1312    )
1313    .map_err(|err| err.to_string())?;
1314    if units.len() != max_units || units.last().map(|unit| unit.lsn) != Some(applied_lsn) {
1315        return Err(
1316            "sync ack does not cover a complete retained-unit range; rebootstrap required".into(),
1317        );
1318    }
1319    validate_v1_retained_units_applyable(&units).map_err(|err| err.to_string())
1320}
1321
1322#[cfg(test)]
1323fn dispatch_sync_status(
1324    engine: &Arc<RwLock<Engine>>,
1325    replica_id: String,
1326    credential_authenticated: bool,
1327    principal: Option<&Principal>,
1328) -> Message {
1329    dispatch_sync_status_decision(engine, replica_id, credential_authenticated, principal).message
1330}
1331
1332fn dispatch_sync_status_decision(
1333    engine: &Arc<RwLock<Engine>>,
1334    replica_id: String,
1335    credential_authenticated: bool,
1336    principal: Option<&Principal>,
1337) -> SyncDecision {
1338    if let Err((class, message)) =
1339        check_sync_protocol_permitted(credential_authenticated, principal)
1340    {
1341        return SyncDecision::error(class, message);
1342    }
1343    if let Err(message) = validate_wire_replica_id(&replica_id) {
1344        return SyncDecision::error(SyncErrorClass::InvalidReplicaId, message);
1345    }
1346    let SyncContext {
1347        data_dir,
1348        remote_lsn,
1349        ..
1350    } = match sync_context(engine) {
1351        Ok(context) => context,
1352        Err(message) => return SyncDecision::error(SyncErrorClass::SyncContext, message),
1353    };
1354    match replica_sync_status(&data_dir, &replica_id, remote_lsn) {
1355        Ok(status) => SyncDecision::ok(Message::SyncStatusResult {
1356            status: wire_sync_status(status),
1357        }),
1358        Err(err) => SyncDecision::error(SyncErrorClass::StatusRead, err.to_string()),
1359    }
1360}
1361
1362#[cfg(test)]
1363fn dispatch_sync_pull(
1364    engine: &Arc<RwLock<Engine>>,
1365    request: SyncPullRequest,
1366    credential_authenticated: bool,
1367    principal: Option<&Principal>,
1368) -> Message {
1369    dispatch_sync_pull_decision(engine, request, credential_authenticated, principal).message
1370}
1371
1372fn dispatch_sync_pull_decision(
1373    engine: &Arc<RwLock<Engine>>,
1374    request: SyncPullRequest,
1375    credential_authenticated: bool,
1376    principal: Option<&Principal>,
1377) -> SyncDecision {
1378    if let Err((class, message)) =
1379        check_sync_protocol_permitted(credential_authenticated, principal)
1380    {
1381        return SyncDecision::error(class, message);
1382    }
1383    if let Err(message) = validate_wire_replica_id(&request.replica_id) {
1384        return SyncDecision::error(SyncErrorClass::InvalidReplicaId, message);
1385    }
1386    if request.max_units == 0 || request.max_units > MAX_SYNC_PULL_UNITS {
1387        return SyncDecision::error(
1388            SyncErrorClass::InvalidMaxUnits,
1389            format!("sync pull maxUnits must be between 1 and {MAX_SYNC_PULL_UNITS}"),
1390        );
1391    }
1392    if request.max_bytes == 0 || request.max_bytes > MAX_SYNC_PULL_BYTES {
1393        return SyncDecision::error(
1394            SyncErrorClass::InvalidMaxBytes,
1395            format!("sync pull maxBytes must be between 1 and {MAX_SYNC_PULL_BYTES}"),
1396        );
1397    }
1398
1399    let SyncContext {
1400        data_dir,
1401        remote_lsn,
1402        active_catalog_version,
1403    } = match sync_context(engine) {
1404        Ok(context) => context,
1405        Err(message) => return SyncDecision::error(SyncErrorClass::SyncContext, message),
1406    };
1407    let status = match replica_sync_status(&data_dir, &request.replica_id, remote_lsn) {
1408        Ok(status) => status,
1409        Err(err) => {
1410            return SyncDecision::error(SyncErrorClass::StatusRead, err.to_string());
1411        }
1412    };
1413    let Some(cursor_lsn) = status.last_applied_lsn else {
1414        return SyncDecision::ok(Message::SyncPullResult {
1415            status: wire_sync_status(status),
1416            units: Vec::new(),
1417            has_more: false,
1418        });
1419    };
1420    if status.repair_action != SyncRepairAction::Pull {
1421        return SyncDecision::ok(Message::SyncPullResult {
1422            status: wire_sync_status(status),
1423            units: Vec::new(),
1424            has_more: false,
1425        });
1426    }
1427    if request.since_lsn != cursor_lsn {
1428        return SyncDecision::error(
1429            SyncErrorClass::CursorLsnMismatch,
1430            format!(
1431                "sync pull sinceLsn {} does not match primary cursor LSN {cursor_lsn}",
1432                request.since_lsn
1433            ),
1434        );
1435    }
1436
1437    let identity = match read_identity(&data_dir) {
1438        Ok(identity) => identity,
1439        Err(err) => {
1440            return SyncDecision::error(SyncErrorClass::IdentityRead, err.to_string());
1441        }
1442    };
1443    // Stamp the expected identity with the database's *active* catalog version,
1444    // not this binary's compile-time maximum, so a database that never activated
1445    // v6 still expects v5 and accepts a v0.12 replica.
1446    let expected = SegmentIdentity::with_catalog_version(
1447        identity.database_id,
1448        identity.primary_generation,
1449        active_catalog_version,
1450    );
1451    if request.database_id != expected.database_id
1452        || request.primary_generation != expected.primary_generation
1453        || request.wal_format_version != expected.wal_format_version
1454        || request.segment_format_version != RETAINED_SEGMENT_FORMAT_VERSION
1455    {
1456        return SyncDecision::error(
1457            SyncErrorClass::IdentityOrFormatMismatch,
1458            "sync pull identity or format version mismatch; rebootstrap required",
1459        );
1460    }
1461    // The request states the maximum catalog format the replica can read. Accept
1462    // any replica that can read the active format (>= active); reject a replica
1463    // whose maximum is older than the data it would receive.
1464    if request.catalog_version < expected.catalog_version {
1465        return SyncDecision::error(
1466            SyncErrorClass::IdentityOrFormatMismatch,
1467            format!(
1468                "sync pull replica catalog format v{} cannot read this database's active catalog format v{}; rebootstrap with an upgraded replica required",
1469                request.catalog_version, expected.catalog_version
1470            ),
1471        );
1472    }
1473
1474    let effective_max_units = request.max_units.min(MAX_SYNC_PULL_UNITS) as usize;
1475    let requested_through_lsn = request
1476        .since_lsn
1477        .saturating_add(request.max_units as u64)
1478        .min(remote_lsn)
1479        .min(status.servable_lsn.unwrap_or(request.since_lsn));
1480    let segment_dir = retained_segments_dir(&data_dir);
1481    if requested_through_lsn > request.since_lsn {
1482        if let Err(err) = validate_retained_tail_available(
1483            &segment_dir,
1484            expected,
1485            request.since_lsn,
1486            requested_through_lsn,
1487        ) {
1488            let mut rebootstrap_status = status;
1489            rebootstrap_status.stale = true;
1490            rebootstrap_status.repair_action = SyncRepairAction::Rebootstrap;
1491            rebootstrap_status.last_sync_error = Some(format!(
1492                "retained history is unavailable; rebootstrap required: {err}"
1493            ));
1494            return SyncDecision::ok(Message::SyncPullResult {
1495                status: wire_sync_status(rebootstrap_status),
1496                units: Vec::new(),
1497                has_more: false,
1498            });
1499        }
1500    }
1501
1502    let raw_units = match read_units_through(
1503        &segment_dir,
1504        expected,
1505        request.since_lsn,
1506        requested_through_lsn,
1507        effective_max_units,
1508    ) {
1509        Ok(units) => units,
1510        Err(err) => {
1511            return SyncDecision::error(SyncErrorClass::RetainedRead, err.to_string());
1512        }
1513    };
1514
1515    let mut selected_raw = Vec::new();
1516    let mut selected = Vec::new();
1517    let mut selected_bytes = 0u64;
1518    for unit in raw_units {
1519        let wire_unit = wire_retained_unit(unit.clone());
1520        let unit_bytes = match wire_unit.encoded_len() {
1521            Ok(bytes) => bytes,
1522            Err(message) => {
1523                return SyncDecision::error(SyncErrorClass::RetainedUnitEncoding, message);
1524            }
1525        };
1526        if selected_bytes.saturating_add(unit_bytes) > request.max_bytes {
1527            if selected.is_empty() {
1528                return SyncDecision::error(
1529                    SyncErrorClass::InvalidMaxBytes,
1530                    "sync pull maxBytes is too small for the next retained unit",
1531                );
1532            }
1533            break;
1534        }
1535        selected_bytes += unit_bytes;
1536        selected_raw.push(unit);
1537        selected.push(wire_unit);
1538    }
1539    if let Err(message) = trim_to_applyable_v1_prefix(&mut selected_raw, &mut selected) {
1540        return SyncDecision::error(SyncErrorClass::RetainedChunkNotApplyable, message);
1541    }
1542
1543    let fetchable_through_lsn = status.servable_lsn.unwrap_or(remote_lsn).min(remote_lsn);
1544    let has_more = selected
1545        .last()
1546        .is_some_and(|unit| unit.lsn < fetchable_through_lsn);
1547    SyncDecision::ok(Message::SyncPullResult {
1548        status: wire_sync_status(status),
1549        units: selected,
1550        has_more,
1551    })
1552}
1553
1554#[cfg(test)]
1555fn dispatch_sync_ack(
1556    engine: &Arc<RwLock<Engine>>,
1557    replica_id: String,
1558    applied_lsn: u64,
1559    observed_remote_lsn: u64,
1560    credential_authenticated: bool,
1561    principal: Option<&Principal>,
1562) -> Message {
1563    dispatch_sync_ack_decision(
1564        engine,
1565        replica_id,
1566        applied_lsn,
1567        observed_remote_lsn,
1568        credential_authenticated,
1569        principal,
1570    )
1571    .message
1572}
1573
1574fn dispatch_sync_ack_decision(
1575    engine: &Arc<RwLock<Engine>>,
1576    replica_id: String,
1577    applied_lsn: u64,
1578    observed_remote_lsn: u64,
1579    credential_authenticated: bool,
1580    principal: Option<&Principal>,
1581) -> SyncDecision {
1582    if let Err((class, message)) =
1583        check_sync_protocol_permitted(credential_authenticated, principal)
1584    {
1585        return SyncDecision::error(class, message);
1586    }
1587    if let Err(message) = validate_wire_replica_id(&replica_id) {
1588        return SyncDecision::error(SyncErrorClass::InvalidReplicaId, message);
1589    }
1590    if applied_lsn > observed_remote_lsn {
1591        return SyncDecision::error(
1592            SyncErrorClass::LsnAheadOfRemote,
1593            format!(
1594                "sync ack appliedLsn {applied_lsn} is ahead of observed remoteLsn {observed_remote_lsn}"
1595            ),
1596        );
1597    }
1598    let SyncContext {
1599        data_dir,
1600        remote_lsn,
1601        ..
1602    } = match sync_context(engine) {
1603        Ok(context) => context,
1604        Err(message) => return SyncDecision::error(SyncErrorClass::SyncContext, message),
1605    };
1606    if observed_remote_lsn > remote_lsn {
1607        return SyncDecision::error(
1608            SyncErrorClass::LsnAheadOfRemote,
1609            format!(
1610                "sync ack remoteLsn {observed_remote_lsn} is ahead of primary LSN {remote_lsn}"
1611            ),
1612        );
1613    }
1614    if let Err(message) =
1615        validate_sync_ack_applyable_boundary(&data_dir, &replica_id, applied_lsn, remote_lsn)
1616    {
1617        return SyncDecision::error(SyncErrorClass::AckValidation, message);
1618    }
1619    match acknowledge_replica_apply(&data_dir, &replica_id, applied_lsn, remote_lsn) {
1620        Ok(summary) => SyncDecision::ok(Message::SyncAckResult {
1621            previous_applied_lsn: summary.previous_applied_lsn,
1622            applied_lsn: summary.applied_lsn,
1623            remote_lsn: summary.remote_lsn,
1624            advanced: summary.advanced,
1625            status: wire_sync_status(summary.status),
1626        }),
1627        Err(err) => SyncDecision::error(SyncErrorClass::AckUpdate, err.to_string()),
1628    }
1629}
1630
1631async fn run_blocking_sync<T, F>(input: T, query_timeout: Duration, f: F) -> SyncDecision
1632where
1633    T: Send + 'static,
1634    F: FnOnce(T) -> SyncDecision + Send + 'static,
1635{
1636    let mut handle = tokio::task::spawn_blocking(move || f(input));
1637    tokio::select! {
1638        result = &mut handle => match result {
1639            Ok(decision) => decision,
1640            Err(err) => SyncDecision::error(SyncErrorClass::Internal, format!("internal error: {err}")),
1641        },
1642        _ = tokio::time::sleep(query_timeout) => match handle.await {
1643            Ok(decision) => decision,
1644            Err(err) => SyncDecision::error(SyncErrorClass::Internal, format!("internal error: {err}")),
1645        },
1646    }
1647}
1648
1649async fn execute_gated_sync<T, F>(context: SyncExecutionContext<'_>, input: T, f: F) -> Message
1650where
1651    T: Send + 'static,
1652    F: FnOnce(T) -> SyncDecision + Send + 'static,
1653{
1654    let SyncExecutionContext {
1655        tx_gate,
1656        connection_has_transaction,
1657        operation,
1658        log_context,
1659        metrics,
1660        query_timeout,
1661    } = context;
1662    let start = Instant::now();
1663    if connection_has_transaction {
1664        let decision = SyncDecision::error(
1665            SyncErrorClass::ActiveTransaction,
1666            "sync protocol is unavailable inside an active transaction",
1667        );
1668        let elapsed = start.elapsed();
1669        log_sync_decision(operation, &log_context, elapsed, &decision);
1670        metrics.record_sync_operation(operation, elapsed, SyncOutcome::Error);
1671        return decision.message;
1672    }
1673
1674    let permit_count = tx_gate.permit_count();
1675    let permit = match tx_gate.acquire_many_owned(permit_count).await {
1676        Ok(permit) => permit,
1677        Err(_) => {
1678            let decision =
1679                SyncDecision::error(SyncErrorClass::QueryExecution, "query execution error");
1680            let elapsed = start.elapsed();
1681            log_sync_decision(operation, &log_context, elapsed, &decision);
1682            metrics.record_sync_operation(operation, elapsed, SyncOutcome::Error);
1683            return decision.message;
1684        }
1685    };
1686    let decision = run_blocking_sync(input, query_timeout, f).await;
1687    drop(permit);
1688    match &decision.message {
1689        Message::SyncStatusResult { status } => {
1690            metrics.record_sync_repair_action(operation, sync_repair_label(status.repair_action));
1691        }
1692        Message::SyncPullResult { status, units, .. } => {
1693            metrics.record_sync_repair_action(operation, sync_repair_label(status.repair_action));
1694            metrics.record_sync_pull_payload(units.len() as u64, sync_pull_payload_bytes(units));
1695        }
1696        Message::SyncAckResult {
1697            advanced, status, ..
1698        } => {
1699            metrics.record_sync_repair_action(operation, sync_repair_label(status.repair_action));
1700            if *advanced {
1701                metrics.inc_sync_ack_advanced();
1702            }
1703        }
1704        _ => {}
1705    }
1706    let elapsed = start.elapsed();
1707    log_sync_decision(operation, &log_context, elapsed, &decision);
1708    metrics.record_sync_operation(
1709        operation,
1710        elapsed,
1711        sync_operation_outcome(&decision.message),
1712    );
1713    decision.message
1714}
1715
1716/// Acquire the TxGate for an explicit `begin`, bounded by `tx_wait_timeout`.
1717/// Overlapping explicit transactions queue behind the permit rather than being
1718/// rejected, but a connection gives up with a clear, client-facing error once
1719/// the wait elapses — so a transaction stalled (or held open) on another
1720/// connection can never block this one indefinitely. A timeout is recorded so
1721/// `powdb_tx_gate_timeouts_total` (and the error total) stay truthful.
1722async fn acquire_begin_permit(
1723    tx_gate: &TxGate,
1724    tx_wait_timeout: Duration,
1725    metrics: &Arc<Metrics>,
1726) -> Result<OwnedSemaphorePermit, Message> {
1727    match tokio::time::timeout(
1728        tx_wait_timeout,
1729        tx_gate.clone().acquire_many_owned(tx_gate.permit_count()),
1730    )
1731    .await
1732    {
1733        Ok(Ok(permit)) => Ok(permit),
1734        Ok(Err(_)) => Err(Message::Error {
1735            message: "query execution error".into(),
1736        }),
1737        Err(_) => {
1738            metrics.inc_tx_gate_timeout();
1739            Err(Message::Error {
1740                message: format!(
1741                    "transaction gate timeout after {}ms waiting for concurrent transaction to complete",
1742                    tx_wait_timeout.as_millis()
1743                ),
1744            })
1745        }
1746    }
1747}
1748
1749/// Acquire the TxGate for a BARE autocommit statement, bounded by
1750/// `tx_wait_timeout` exactly like [`acquire_begin_permit`]. Autocommit writes
1751/// serialize through the same gate as explicit transactions, so a stalled (or
1752/// held-open) transaction on another connection would otherwise block this
1753/// write indefinitely. Bounding the acquire turns that indefinite wait into a
1754/// clear, client-facing timeout error and records the timeout so
1755/// `powdb_tx_gate_timeouts_total` (and the error total) stay truthful. This
1756/// only bounds the ACQUIRE; the permit is still dropped BEFORE the caller's
1757/// durability wait so overlapping committers can share an fsync.
1758async fn acquire_autocommit_permit(
1759    tx_gate: &TxGate,
1760    admission: AdmissionMode,
1761    tx_wait_timeout: Duration,
1762    metrics: &Arc<Metrics>,
1763) -> Result<OwnedSemaphorePermit, Message> {
1764    let permits = match admission {
1765        AdmissionMode::Reader => 1,
1766        AdmissionMode::Writer => tx_gate.permit_count(),
1767    };
1768
1769    // The uncontended path is overwhelmingly common for autocommit work.
1770    // Avoid constructing and polling a timeout-wrapped semaphore future when
1771    // the permits are already available. If a reader or writer is queued, the
1772    // try-acquire fails and we fall back to Tokio's fair semaphore queue, so a
1773    // waiting writer still cannot be bypassed by later readers.
1774    if let Ok(permit) = tx_gate.clone().try_acquire_many_owned(permits) {
1775        return Ok(permit);
1776    }
1777
1778    let acquire = tx_gate.clone().acquire_many_owned(permits);
1779    match tokio::time::timeout(tx_wait_timeout, acquire).await {
1780        Ok(Ok(permit)) => Ok(permit),
1781        Ok(Err(_)) => Err(Message::Error {
1782            message: "query execution error".into(),
1783        }),
1784        Err(_) => {
1785            metrics.inc_tx_gate_timeout();
1786            Err(Message::Error {
1787                message: format!(
1788                    "transaction gate timeout after {}ms waiting for concurrent transaction to complete",
1789                    tx_wait_timeout.as_millis()
1790                ),
1791            })
1792        }
1793    }
1794}
1795
1796/// Run the shared four-arm transaction-routing state machine for one wire
1797/// query frame, returning the response plus its un-waited WAL durability
1798/// ticket. The TxGate permit is managed here and, crucially, is already
1799/// released (bare statements, commit/rollback) by the time this returns, so
1800/// the caller's `finalize_durability` wait happens OUTSIDE the gate and
1801/// overlapping committers can share an fsync.
1802///
1803/// The three wire dialects (PowQL, SQL, parameterized PowQL) differ only in
1804/// how a frame is parsed and dispatched; they share this routing, so a
1805/// behavior fix (e.g. cancellation rollback parity) lands here exactly once.
1806/// `parsed_query` is the dialect's already-parsed frame, `tx_control` its
1807/// transaction-control classification, `autocommit_admission` the admission
1808/// mode for a bare (non-transaction) statement, and `dispatch` the closure
1809/// that executes the parsed frame (its `bool` argument is
1810/// `allow_readonly_escalation`).
1811#[allow(clippy::too_many_arguments)]
1812async fn run_wire_query_state_machine<Inner, D, R>(
1813    engine: Arc<RwLock<Engine>>,
1814    tx_gate: TxGate,
1815    tx_permit: &mut Option<OwnedSemaphorePermit>,
1816    parsed_query: Arc<Inner>,
1817    tx_control: Option<TransactionControl>,
1818    autocommit_admission: AdmissionMode,
1819    result_mode: WireResultMode,
1820    principal: Option<Principal>,
1821    query_timeout: Duration,
1822    query_deadline: Instant,
1823    tx_wait_timeout: Duration,
1824    metrics: &Arc<Metrics>,
1825    reader: &mut BufReader<R>,
1826    wire_read_buffer: &mut Vec<u8>,
1827    pending_messages: &mut InFlightReadAhead,
1828    dispatch: D,
1829) -> (
1830    Message,
1831    Option<PendingDurability>,
1832    Option<ConnectionTermination>,
1833)
1834where
1835    Inner: Send + Sync + 'static,
1836    D: Fn(Arc<RwLock<Engine>>, Arc<Inner>, Option<Principal>, bool) -> DispatchOutcome
1837        + Clone
1838        + Send
1839        + Sync
1840        + 'static,
1841    R: AsyncRead + Unpin,
1842{
1843    match tx_control {
1844        Some(TransactionControl::Begin) => {
1845            if tx_permit.is_some() {
1846                return (
1847                    Message::Error {
1848                        message: sanitize_error(
1849                            "cannot begin: a transaction is already active on this connection",
1850                        ),
1851                    },
1852                    None,
1853                    None,
1854                );
1855            }
1856            let permit = match acquire_begin_permit(&tx_gate, tx_wait_timeout, metrics).await {
1857                Ok(permit) => permit,
1858                Err(response) => return (response, None, None),
1859            };
1860            let dispatch_begin = dispatch.clone();
1861            let (response, ticket, mut termination, _) = run_blocking_query(
1862                engine.clone(),
1863                Arc::clone(&parsed_query),
1864                principal.clone(),
1865                result_mode,
1866                query_timeout,
1867                query_deadline,
1868                metrics,
1869                reader,
1870                wire_read_buffer,
1871                pending_messages,
1872                move |engine, parsed_query, principal| {
1873                    dispatch_begin(engine, parsed_query, principal, true)
1874                },
1875            )
1876            .await;
1877            if is_success_response(&response) {
1878                *tx_permit = Some(permit);
1879            } else if is_query_cancellation_response(&response) {
1880                // Parity with the commit/rollback and in-transaction arms: a
1881                // cancelled begin must not leave the engine's transaction
1882                // state ownerless. Install the just-acquired permit so the
1883                // shared rollback helper can undo any transaction the begin
1884                // opened, then release the permit and close the connection.
1885                *tx_permit = Some(permit);
1886                rollback_connection_transaction(engine, principal, tx_permit).await;
1887                termination = Some(ConnectionTermination::Closed);
1888            }
1889            (response, ticket, termination)
1890        }
1891        Some(TransactionControl::Commit | TransactionControl::Rollback) => {
1892            let standalone_permit = if tx_permit.is_none() {
1893                match acquire_autocommit_permit(
1894                    &tx_gate,
1895                    AdmissionMode::Writer,
1896                    tx_wait_timeout,
1897                    metrics,
1898                )
1899                .await
1900                {
1901                    Ok(permit) => Some(permit),
1902                    Err(response) => return (response, None, None),
1903                }
1904            } else {
1905                None
1906            };
1907            let dispatch_commit = dispatch.clone();
1908            let (response, ticket, mut termination, _) = run_blocking_query(
1909                engine.clone(),
1910                Arc::clone(&parsed_query),
1911                principal.clone(),
1912                result_mode,
1913                query_timeout,
1914                query_deadline,
1915                metrics,
1916                reader,
1917                wire_read_buffer,
1918                pending_messages,
1919                move |engine, parsed_query, principal| {
1920                    dispatch_commit(engine, parsed_query, principal, true)
1921                },
1922            )
1923            .await;
1924            if is_success_response(&response) {
1925                // Release the gate BEFORE the caller waits on the commit's
1926                // ticket: the engine work is done and WAL order is fixed, so
1927                // another connection's commit can start (and share the fsync)
1928                // while this one waits.
1929                tx_permit.take();
1930            } else if is_query_cancellation_response(&response) {
1931                rollback_connection_transaction(engine, principal, tx_permit).await;
1932                termination = Some(ConnectionTermination::Closed);
1933            }
1934            drop(standalone_permit);
1935            (response, ticket, termination)
1936        }
1937        None if tx_permit.is_some() => {
1938            let dispatch_in_tx = dispatch.clone();
1939            let mut out = run_blocking_query(
1940                engine.clone(),
1941                Arc::clone(&parsed_query),
1942                principal.clone(),
1943                result_mode,
1944                query_timeout,
1945                query_deadline,
1946                metrics,
1947                reader,
1948                wire_read_buffer,
1949                pending_messages,
1950                move |engine, parsed_query, principal| {
1951                    dispatch_in_tx(engine, parsed_query, principal, true)
1952                },
1953            )
1954            .await;
1955            if is_query_cancellation_response(&out.0) {
1956                rollback_connection_transaction(engine, principal, tx_permit).await;
1957                out.2 = Some(ConnectionTermination::Closed);
1958            }
1959            (out.0, out.1, out.2)
1960        }
1961        None => {
1962            let admission = autocommit_admission;
1963            let permit = match acquire_autocommit_permit(
1964                &tx_gate,
1965                admission,
1966                tx_wait_timeout,
1967                metrics,
1968            )
1969            .await
1970            {
1971                Ok(permit) => permit,
1972                Err(response) => return (response, None, None),
1973            };
1974            // The parsed AST can be large. Share the first parse with the
1975            // rare dirty-view retry instead of cloning it on every successful
1976            // autocommit read.
1977            let retry_engine = Arc::clone(&engine);
1978            let retry_parsed_query = Arc::clone(&parsed_query);
1979            let retry_principal = principal.clone();
1980            let allow_readonly_escalation = admission == AdmissionMode::Writer;
1981            let dispatch_first = dispatch.clone();
1982            let mut out = run_blocking_query(
1983                engine,
1984                parsed_query,
1985                principal,
1986                result_mode,
1987                query_timeout,
1988                query_deadline,
1989                metrics,
1990                reader,
1991                wire_read_buffer,
1992                pending_messages,
1993                move |engine, parsed_query, principal| {
1994                    dispatch_first(engine, parsed_query, principal, allow_readonly_escalation)
1995                },
1996            )
1997            .await;
1998            drop(permit);
1999            if out.3 {
2000                let writer_permit = match acquire_autocommit_permit(
2001                    &tx_gate,
2002                    AdmissionMode::Writer,
2003                    tx_wait_timeout,
2004                    metrics,
2005                )
2006                .await
2007                {
2008                    Ok(permit) => permit,
2009                    Err(response) => return (response, None, None),
2010                };
2011                let dispatch_retry = dispatch.clone();
2012                out = run_blocking_query(
2013                    retry_engine,
2014                    retry_parsed_query,
2015                    retry_principal,
2016                    result_mode,
2017                    query_timeout,
2018                    query_deadline,
2019                    metrics,
2020                    reader,
2021                    wire_read_buffer,
2022                    pending_messages,
2023                    move |engine, parsed_query, principal| {
2024                        dispatch_retry(engine, parsed_query, principal, true)
2025                    },
2026                )
2027                .await;
2028                drop(writer_permit);
2029            }
2030            (out.0, out.1, out.2)
2031        }
2032    }
2033}
2034
2035/// Execute one PowQL wire query frame. Thin dialect wrapper over
2036/// [`run_wire_query_state_machine`].
2037#[allow(clippy::too_many_arguments)]
2038async fn execute_wire_query<R>(
2039    engine: Arc<RwLock<Engine>>,
2040    tx_gate: TxGate,
2041    tx_permit: &mut Option<OwnedSemaphorePermit>,
2042    query: String,
2043    result_mode: WireResultMode,
2044    principal: Option<Principal>,
2045    query_timeout: Duration,
2046    tx_wait_timeout: Duration,
2047    metrics: &Arc<Metrics>,
2048    reader: &mut BufReader<R>,
2049    wire_read_buffer: &mut Vec<u8>,
2050    pending_messages: &mut InFlightReadAhead,
2051) -> (
2052    Message,
2053    Option<PendingDurability>,
2054    Option<ConnectionTermination>,
2055)
2056where
2057    R: AsyncRead + Unpin,
2058{
2059    let query_deadline = Instant::now() + query_timeout;
2060    // Parse each frame once for transaction routing, admission, and role
2061    // enforcement. The engine still canonicalizes/parses as needed for plan
2062    // cache execution, but the server no longer repeats the same parse in
2063    // three separate routing helpers before reaching it.
2064    let stmt_result = parser::parse(&query).map_err(|e| e.to_string());
2065    let parsed_query = Arc::new((query, stmt_result));
2066    let tx_control = parsed_transaction_control(&parsed_query.1);
2067    let autocommit_admission = parsed_query
2068        .1
2069        .as_ref()
2070        .map(statement_admission)
2071        .unwrap_or(AdmissionMode::Writer);
2072    run_wire_query_state_machine(
2073        engine,
2074        tx_gate,
2075        tx_permit,
2076        parsed_query,
2077        tx_control,
2078        autocommit_admission,
2079        result_mode,
2080        principal,
2081        query_timeout,
2082        query_deadline,
2083        tx_wait_timeout,
2084        metrics,
2085        reader,
2086        wire_read_buffer,
2087        pending_messages,
2088        |engine,
2089         parsed_query: Arc<(String, Result<powdb_query::ast::Statement, String>)>,
2090         principal: Option<Principal>,
2091         allow| {
2092            dispatch_query_parsed(
2093                &engine,
2094                &parsed_query.0,
2095                &parsed_query.1,
2096                principal.as_ref(),
2097                allow,
2098            )
2099        },
2100    )
2101    .await
2102}
2103
2104#[allow(clippy::too_many_arguments)]
2105async fn execute_wire_query_sql<R>(
2106    engine: Arc<RwLock<Engine>>,
2107    tx_gate: TxGate,
2108    tx_permit: &mut Option<OwnedSemaphorePermit>,
2109    query: String,
2110    result_mode: WireResultMode,
2111    principal: Option<Principal>,
2112    query_timeout: Duration,
2113    tx_wait_timeout: Duration,
2114    metrics: &Arc<Metrics>,
2115    reader: &mut BufReader<R>,
2116    wire_read_buffer: &mut Vec<u8>,
2117    pending_messages: &mut InFlightReadAhead,
2118) -> (
2119    Message,
2120    Option<PendingDurability>,
2121    Option<ConnectionTermination>,
2122)
2123where
2124    R: AsyncRead + Unpin,
2125{
2126    let query_deadline = Instant::now() + query_timeout;
2127    let stmt_result = sql::parse_sql(&query).map_err(|e| e.to_string());
2128    let parsed_query = Arc::new((query, stmt_result));
2129    let tx_control = parsed_transaction_control(&parsed_query.1);
2130    let autocommit_admission = parsed_query
2131        .1
2132        .as_ref()
2133        .map(statement_admission)
2134        .unwrap_or(AdmissionMode::Writer);
2135    run_wire_query_state_machine(
2136        engine,
2137        tx_gate,
2138        tx_permit,
2139        parsed_query,
2140        tx_control,
2141        autocommit_admission,
2142        result_mode,
2143        principal,
2144        query_timeout,
2145        query_deadline,
2146        tx_wait_timeout,
2147        metrics,
2148        reader,
2149        wire_read_buffer,
2150        pending_messages,
2151        |engine,
2152         parsed_query: Arc<(String, Result<powdb_query::ast::Statement, String>)>,
2153         principal: Option<Principal>,
2154         allow| {
2155            dispatch_sql_query_parsed(
2156                &engine,
2157                &parsed_query.0,
2158                &parsed_query.1,
2159                principal.as_ref(),
2160                allow,
2161            )
2162        },
2163    )
2164    .await
2165}
2166
2167// One over clippy's default arg limit: the metrics handle was threaded through
2168// to instrument the typed query result. Bundling these into a struct would add
2169// more noise than it removes for an internal dispatcher.
2170#[allow(clippy::too_many_arguments)]
2171async fn execute_wire_query_with_params<R>(
2172    engine: Arc<RwLock<Engine>>,
2173    tx_gate: TxGate,
2174    tx_permit: &mut Option<OwnedSemaphorePermit>,
2175    query: String,
2176    params: Vec<WireParam>,
2177    result_mode: WireResultMode,
2178    principal: Option<Principal>,
2179    query_timeout: Duration,
2180    tx_wait_timeout: Duration,
2181    metrics: &Arc<Metrics>,
2182    reader: &mut BufReader<R>,
2183    wire_read_buffer: &mut Vec<u8>,
2184    pending_messages: &mut InFlightReadAhead,
2185) -> (
2186    Message,
2187    Option<PendingDurability>,
2188    Option<ConnectionTermination>,
2189)
2190where
2191    R: AsyncRead + Unpin,
2192{
2193    let query_deadline = Instant::now() + query_timeout;
2194    let bound: Vec<powdb_query::ast::ParamValue> = params.iter().map(wire_param_to_value).collect();
2195    let stmt_result = parser::parse_with_params(&query, &bound).map_err(|e| e.to_string());
2196    let parsed_query = Arc::new((query, bound, stmt_result));
2197    let tx_control = parsed_transaction_control(&parsed_query.2);
2198    let autocommit_admission = parsed_query
2199        .2
2200        .as_ref()
2201        .map(statement_admission)
2202        .unwrap_or(AdmissionMode::Writer);
2203    run_wire_query_state_machine(
2204        engine,
2205        tx_gate,
2206        tx_permit,
2207        parsed_query,
2208        tx_control,
2209        autocommit_admission,
2210        result_mode,
2211        principal,
2212        query_timeout,
2213        query_deadline,
2214        tx_wait_timeout,
2215        metrics,
2216        reader,
2217        wire_read_buffer,
2218        pending_messages,
2219        |engine,
2220         parsed_query: Arc<(
2221            String,
2222            Vec<powdb_query::ast::ParamValue>,
2223            Result<powdb_query::ast::Statement, String>,
2224        )>,
2225         principal: Option<Principal>,
2226         allow| {
2227            dispatch_query_with_bound_params_parsed(
2228                &engine,
2229                &parsed_query.0,
2230                &parsed_query.1,
2231                &parsed_query.2,
2232                principal.as_ref(),
2233                allow,
2234            )
2235        },
2236    )
2237    .await
2238}
2239
2240/// A statement's metric sample whose recording is deferred until its WAL
2241/// durability obligation settles: a Full-mode fsync failure downgrades the
2242/// client's success reply to an error, and the metrics must tell the same
2243/// story (and the latency must include the wait the client observed).
2244struct DeferredQueryMetric {
2245    start: Instant,
2246    outcome: QueryOutcome,
2247    exceeded_timeout: bool,
2248}
2249
2250/// Durability ticket + the deferred metric of the statement that produced it.
2251type PendingDurability = (WalDurabilityTicket, DeferredQueryMetric);
2252
2253#[allow(clippy::too_many_arguments)]
2254async fn run_blocking_query<T, F, R>(
2255    engine: Arc<RwLock<Engine>>,
2256    input: T,
2257    principal: Option<Principal>,
2258    result_mode: WireResultMode,
2259    query_timeout: Duration,
2260    query_deadline: Instant,
2261    metrics: &Arc<Metrics>,
2262    reader: &mut BufReader<R>,
2263    wire_read_buffer: &mut Vec<u8>,
2264    pending_messages: &mut InFlightReadAhead,
2265    f: F,
2266) -> (
2267    Message,
2268    Option<PendingDurability>,
2269    Option<ConnectionTermination>,
2270    bool,
2271)
2272where
2273    T: Send + 'static,
2274    F: FnOnce(Arc<RwLock<Engine>>, T, Option<Principal>) -> DispatchOutcome + Send + 'static,
2275    R: AsyncRead + Unpin,
2276{
2277    let _in_flight = metrics.in_flight_guard();
2278    let start = Instant::now();
2279
2280    // Cooperative cancellation. The blocking closure installs this token for the
2281    // executor thread; every unbounded executor loop polls it. We give it a
2282    // deadline of `now + query_timeout` so the query self-terminates even if the
2283    // async timeout arm below is slow to be scheduled — that is what actually
2284    // makes the timeout enforceable (a `spawn_blocking` thread cannot be
2285    // aborted, so before this the timeout arm just awaited the runaway query to
2286    // completion while it held the engine lock / tx-gate permit).
2287    let timeout_ms = query_timeout.as_millis().min(u128::from(u64::MAX)) as u64;
2288    let cancel = Arc::new(powdb_query::cancel::ExecCancel::with_deadline(
2289        query_deadline,
2290        timeout_ms,
2291    ));
2292    let cancel_task = Arc::clone(&cancel);
2293    let mut handle = tokio::task::spawn_blocking(move || {
2294        let _cancel_guard = powdb_query::cancel::install(cancel_task);
2295        f(engine, input, principal)
2296    });
2297    let mut exceeded_timeout = false;
2298    let mut termination = None;
2299    let timeout = tokio::time::sleep(query_deadline.saturating_duration_since(Instant::now()));
2300    tokio::pin!(timeout);
2301    let join_result = loop {
2302        tokio::select! {
2303            result = &mut handle => break result,
2304            _ = &mut timeout => {
2305                exceeded_timeout = true;
2306                // Signal the executor to stop at its next cancellation checkpoint,
2307                // then await the (now promptly returning) handle. The closure
2308                // returns a typed timeout error and releases the engine lock /
2309                // tx-gate permit as it unwinds.
2310                cancel.cancel(powdb_query::cancel::CancelReason::Timeout);
2311                break handle.await;
2312            }
2313            read = read_message_cancel_safe(
2314                reader,
2315                wire_read_buffer,
2316                pending_messages.remaining_bytes(),
2317            ) => {
2318                match read {
2319                    Ok(Some(DecodedWireMessage { message: Message::Disconnect, .. })) => {
2320                        cancel.cancel(powdb_query::cancel::CancelReason::Disconnect);
2321                        termination = Some(ConnectionTermination::Closed);
2322                        break handle.await;
2323                    }
2324                    Ok(Some(frame)) => {
2325                        if pending_messages.len() + 1 >= MAX_IN_FLIGHT_READ_AHEAD_FRAMES
2326                            || pending_messages.wire_bytes + frame.wire_len
2327                                >= MAX_IN_FLIGHT_READ_AHEAD_BYTES
2328                        {
2329                            // Never stop observing the socket behind a full
2330                            // queue. Reaching either hard cap cancels the query
2331                            // and closes this connection immediately.
2332                            cancel.cancel(powdb_query::cancel::CancelReason::Disconnect);
2333                            termination = Some(ConnectionTermination::ReadError);
2334                            break handle.await;
2335                        }
2336                        // Preserve frames that arrive while the blocking query
2337                        // runs. The normal batching/main-loop path consumes them
2338                        // in order once execution completes.
2339                        pending_messages.push_back(frame);
2340                    }
2341                    Ok(None) => {
2342                        cancel.cancel(powdb_query::cancel::CancelReason::Disconnect);
2343                        termination = Some(ConnectionTermination::Closed);
2344                        break handle.await;
2345                    }
2346                    Err(_) => {
2347                        cancel.cancel(powdb_query::cancel::CancelReason::Disconnect);
2348                        termination = Some(ConnectionTermination::ReadError);
2349                        break handle.await;
2350                    }
2351                }
2352            }
2353        }
2354    };
2355
2356    let (message, ticket, outcome, readonly_needs_write) = match join_result {
2357        Ok((Ok(result), ticket)) => match query_result_to_message(result, result_mode) {
2358            Ok(message) => (message, ticket, QueryOutcome::Ok, false),
2359            Err(e) => (
2360                Message::Error {
2361                    message: sanitize_error(&e.to_string()),
2362                },
2363                ticket,
2364                QueryOutcome::Error,
2365                false,
2366            ),
2367        },
2368        Ok((Err(QueryError::ReadonlyNeedsWrite), ticket)) => {
2369            if exceeded_timeout {
2370                (
2371                    Message::Error {
2372                        message: sanitize_error(&QueryError::Timeout { timeout_ms }.to_string()),
2373                    },
2374                    ticket,
2375                    QueryOutcome::Timeout,
2376                    true,
2377                )
2378            } else {
2379                (
2380                    Message::Error {
2381                        message: "query execution error".into(),
2382                    },
2383                    ticket,
2384                    QueryOutcome::Error,
2385                    true,
2386                )
2387            }
2388        }
2389        Ok((Err(e), ticket)) => {
2390            // A deadline-driven cancellation returns Timeout even when the async
2391            // timeout arm has not fired yet (the executor self-cancels): treat it
2392            // as a timeout for metrics either way.
2393            if matches!(e, QueryError::Timeout { .. }) {
2394                exceeded_timeout = true;
2395            }
2396            let outcome = if matches!(e, QueryError::MemoryLimitExceeded { .. }) {
2397                QueryOutcome::MemoryLimit
2398            } else if matches!(e, QueryError::Timeout { .. }) {
2399                QueryOutcome::Timeout
2400            } else {
2401                QueryOutcome::Error
2402            };
2403            (
2404                Message::Error {
2405                    message: sanitize_error(&e.to_string()),
2406                },
2407                ticket,
2408                outcome,
2409                false,
2410            )
2411        }
2412        Err(e) => (
2413            Message::Error {
2414                message: format!("internal error: {e}"),
2415            },
2416            None,
2417            QueryOutcome::Error,
2418            false,
2419        ),
2420    };
2421    let readonly_retry = readonly_needs_write && !exceeded_timeout && termination.is_none();
2422    match ticket {
2423        // The statement's durability (and thus its true outcome and the
2424        // latency the client observes) settles at batch end — defer the
2425        // metric to the settlement site instead of recording a success that
2426        // a failed fsync would falsify.
2427        Some(ticket) => (
2428            message,
2429            Some((
2430                ticket,
2431                DeferredQueryMetric {
2432                    start,
2433                    outcome,
2434                    exceeded_timeout,
2435                },
2436            )),
2437            termination,
2438            readonly_retry,
2439        ),
2440        None => {
2441            if !readonly_retry {
2442                if exceeded_timeout {
2443                    metrics.record_query(start.elapsed(), QueryOutcome::Timeout);
2444                } else {
2445                    metrics.record_query(start.elapsed(), outcome);
2446                }
2447            }
2448            (message, None, termination, readonly_retry)
2449        }
2450    }
2451}
2452
2453/// Settle a WAL durability ticket off the async path, AFTER the TxGate
2454/// permit has been dropped — that ordering is what lets committers on other
2455/// connections append (and share the fsync) while this one waits.
2456///
2457/// Returns `None` when the covering fsync succeeded, or `Some(client-facing
2458/// error message)` when it failed — in which case no statement the ticket
2459/// covers may be acknowledged as durable (it executed in memory only).
2460async fn settle_durability_ticket(ticket: WalDurabilityTicket) -> Option<String> {
2461    match tokio::task::spawn_blocking(move || ticket.wait()).await {
2462        Ok(Ok(())) => None,
2463        Ok(Err(e)) => Some(sanitize_error(&format!("WAL durability sync failed: {e}"))),
2464        Err(e) => Some(format!("internal error: {e}")),
2465    }
2466}
2467
2468fn is_success_response(msg: &Message) -> bool {
2469    matches!(
2470        msg,
2471        Message::ResultRows { .. }
2472            | Message::ResultScalar { .. }
2473            | Message::ResultRowsNative { .. }
2474            | Message::ResultScalarNative { .. }
2475            | Message::ResultOk { .. }
2476            | Message::ResultMessage { .. }
2477    )
2478}
2479
2480fn rollback_open_transaction(engine: Arc<RwLock<Engine>>, principal: Option<Principal>) {
2481    let (res, ticket) = dispatch_query(&engine, "rollback", principal.as_ref(), true);
2482    let _ = res;
2483    // Rollback takes the sync-preserving path (no ticket), but settle one
2484    // defensively if it ever appears so the durability watermark stays honest.
2485    if let Some(ticket) = ticket {
2486        let _ = ticket.wait();
2487    }
2488}
2489
2490fn is_query_cancellation_response(message: &Message) -> bool {
2491    matches!(
2492        message,
2493        Message::Error { message }
2494            if message.starts_with("query timeout after")
2495                || message == "query cancelled by client disconnect"
2496    )
2497}
2498
2499/// Roll back this connection's explicit transaction while it still owns the
2500/// transaction-gate permit, then release the permit. A timed-out/cancelled
2501/// statement cannot leave an ambiguous transaction open and block every later
2502/// writer; releasing first would let another connection enter the engine before
2503/// this rollback has restored the prior snapshot.
2504async fn rollback_connection_transaction(
2505    engine: Arc<RwLock<Engine>>,
2506    principal: Option<Principal>,
2507    tx_permit: &mut Option<OwnedSemaphorePermit>,
2508) {
2509    if tx_permit.is_none() {
2510        return;
2511    }
2512    let _ = tokio::task::spawn_blocking(move || rollback_open_transaction(engine, principal)).await;
2513    tx_permit.take();
2514}
2515
2516/// Read one post-auth wire frame without losing partially-read bytes when the
2517/// future is cancelled by `tokio::select!`.
2518///
2519/// `Message::read_from` uses `read_exact`, whose future is not cancellation
2520/// safe: racing it against query completion can consume part of the next frame
2521/// and then drop those bytes. This reader stores every completed `read` in a
2522/// connection-owned buffer before awaiting again, so a query may safely race
2523/// socket EOF / `DISCONNECT` while preserving ordinary pipelined frames.
2524struct DecodedWireMessage {
2525    message: Message,
2526    wire_len: usize,
2527}
2528
2529#[derive(Default)]
2530struct InFlightReadAhead {
2531    frames: VecDeque<DecodedWireMessage>,
2532    wire_bytes: usize,
2533}
2534
2535impl InFlightReadAhead {
2536    fn len(&self) -> usize {
2537        self.frames.len()
2538    }
2539
2540    fn is_empty(&self) -> bool {
2541        self.frames.is_empty()
2542    }
2543
2544    fn remaining_bytes(&self) -> usize {
2545        MAX_IN_FLIGHT_READ_AHEAD_BYTES.saturating_sub(self.wire_bytes)
2546    }
2547
2548    fn push_back(&mut self, frame: DecodedWireMessage) {
2549        self.wire_bytes += frame.wire_len;
2550        self.frames.push_back(frame);
2551    }
2552
2553    fn pop_front(&mut self) -> Option<Message> {
2554        let frame = self.frames.pop_front()?;
2555        self.wire_bytes -= frame.wire_len;
2556        Some(frame.message)
2557    }
2558}
2559
2560async fn read_message_cancel_safe<R>(
2561    reader: &mut BufReader<R>,
2562    buffered: &mut Vec<u8>,
2563    max_frame_len: usize,
2564) -> std::io::Result<Option<DecodedWireMessage>>
2565where
2566    R: AsyncRead + Unpin,
2567{
2568    loop {
2569        if buffered.len() >= 6 {
2570            let payload_len = u32::from_le_bytes(
2571                buffered[2..6]
2572                    .try_into()
2573                    .expect("four-byte wire payload length"),
2574            ) as usize;
2575            if payload_len > MAX_WIRE_PAYLOAD_SIZE {
2576                return Err(std::io::Error::new(
2577                    std::io::ErrorKind::InvalidData,
2578                    format!("payload too large: {payload_len} bytes (max {MAX_WIRE_PAYLOAD_SIZE})"),
2579                ));
2580            }
2581            let frame_len = 6usize.checked_add(payload_len).ok_or_else(|| {
2582                std::io::Error::new(
2583                    std::io::ErrorKind::InvalidData,
2584                    "wire frame length overflow",
2585                )
2586            })?;
2587            if frame_len > max_frame_len {
2588                return Err(std::io::Error::new(
2589                    std::io::ErrorKind::InvalidData,
2590                    format!(
2591                        "wire frame exceeds the available in-flight read-ahead budget: \
2592                         {frame_len} bytes (available {max_frame_len})"
2593                    ),
2594                ));
2595            }
2596            if buffered.len() >= frame_len {
2597                let frame: Vec<u8> = buffered.drain(..frame_len).collect();
2598                return Message::decode(&frame)
2599                    .map(|message| {
2600                        Some(DecodedWireMessage {
2601                            message,
2602                            wire_len: frame_len,
2603                        })
2604                    })
2605                    .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e));
2606            }
2607        }
2608
2609        let mut chunk = [0u8; 8192];
2610        // Read only the bytes needed for this frame stage. Besides preserving
2611        // cancellation safety, this prevents a large pipelined payload from
2612        // overshooting the in-flight byte budget in one buffered read.
2613        let wanted = if buffered.len() < 6 {
2614            6 - buffered.len()
2615        } else {
2616            let payload_len = u32::from_le_bytes(
2617                buffered[2..6]
2618                    .try_into()
2619                    .expect("four-byte wire payload length"),
2620            ) as usize;
2621            6usize
2622                .checked_add(payload_len)
2623                .and_then(|frame_len| frame_len.checked_sub(buffered.len()))
2624                .ok_or_else(|| {
2625                    std::io::Error::new(
2626                        std::io::ErrorKind::InvalidData,
2627                        "invalid buffered wire frame length",
2628                    )
2629                })?
2630        };
2631        let read_limit = wanted.min(chunk.len());
2632        let read = reader.read(&mut chunk[..read_limit]).await?;
2633        if read == 0 {
2634            if buffered.len() < 6 {
2635                // Match the existing protocol behavior: EOF before a complete
2636                // header is a clean connection close.
2637                buffered.clear();
2638                return Ok(None);
2639            }
2640            return Err(std::io::Error::new(
2641                std::io::ErrorKind::UnexpectedEof,
2642                "connection closed in the middle of a wire frame",
2643            ));
2644        }
2645        buffered.extend_from_slice(&chunk[..read]);
2646    }
2647}
2648
2649#[derive(Clone, Copy, Debug)]
2650enum ConnectionTermination {
2651    Closed,
2652    ReadError,
2653}
2654
2655pub async fn handle_connection<S>(stream: S, opts: ConnOpts<'_>)
2656where
2657    S: AsyncRead + AsyncWrite + Unpin,
2658{
2659    let ConnOpts {
2660        engine,
2661        tx_gate,
2662        expected_password,
2663        users,
2664        shutdown_rx,
2665        idle_timeout,
2666        query_timeout,
2667        rate_limiter,
2668        peer_addr,
2669        metrics,
2670        tx_wait_timeout,
2671        db_name: server_db_name,
2672    } = opts;
2673
2674    let peer = peer_addr
2675        .map(|a| a.to_string())
2676        .unwrap_or_else(|| "unknown".into());
2677    let peer_ip = peer_addr.map(|a| a.ip());
2678
2679    let (reader, writer) = tokio::io::split(stream);
2680    let mut reader = BufReader::new(reader);
2681    let mut writer = BufWriter::new(writer);
2682
2683    // Wait for Connect message (with idle timeout).
2684    // Accept Ping messages before authentication so load balancers can
2685    // health-check without completing a full CONNECT handshake.
2686    // Uses the smaller pre-auth payload limit (4 KB) to prevent memory abuse.
2687    let connect_msg = loop {
2688        match tokio::time::timeout(idle_timeout, Message::read_from_preauth(&mut reader)).await {
2689            Ok(Ok(Some(Message::Ping))) => {
2690                debug!(peer = %peer, "pre-auth ping");
2691                if !write_msg(&mut writer, &Message::Pong).await {
2692                    return;
2693                }
2694                continue;
2695            }
2696            Ok(Ok(Some(msg))) => break msg,
2697            Ok(Ok(None)) => {
2698                debug!(peer = %peer, "client closed before CONNECT");
2699                return;
2700            }
2701            Ok(Err(e)) => {
2702                error!(peer = %peer, error = %e, "error reading CONNECT");
2703                return;
2704            }
2705            Err(_) => {
2706                warn!(peer = %peer, "idle timeout waiting for CONNECT");
2707                return;
2708            }
2709        }
2710    };
2711
2712    // The authenticated identity for this connection. Bound at connect time
2713    // and enforced on every query by `dispatch_query`.
2714    let principal: Option<Principal>;
2715    let credential_auth_configured = !users.is_empty() || expected_password.is_some();
2716    match connect_msg {
2717        Message::Connect {
2718            db_name,
2719            password,
2720            username,
2721        } => {
2722            // Check rate limiting before verifying credentials.
2723            if let (Some(limiter), Some(ip)) = (rate_limiter, peer_ip) {
2724                if is_rate_limited(limiter, ip) {
2725                    warn!(peer = %peer, "rate limited: too many auth failures");
2726                    let err = Message::Error {
2727                        message: "too many auth failures, try again later".into(),
2728                    };
2729                    write_msg(&mut writer, &err).await;
2730                    return;
2731                }
2732            }
2733
2734            let outcome = authenticate_connect(
2735                &users,
2736                expected_password.as_ref().map(|p| p.as_str()),
2737                username.as_deref(),
2738                password.as_ref().map(|p| p.as_str()),
2739            );
2740
2741            match outcome {
2742                AuthOutcome::Rejected => {
2743                    warn!(peer = %peer, db = %db_name, "auth rejected");
2744                    metrics.inc_auth_failure();
2745                    // Record the failure for rate limiting.
2746                    if let (Some(limiter), Some(ip)) = (rate_limiter, peer_ip) {
2747                        record_auth_failure(limiter, ip);
2748                    }
2749                    let err = Message::Error {
2750                        message: "authentication failed".into(),
2751                    };
2752                    write_msg(&mut writer, &err).await;
2753                    return;
2754                }
2755                AuthOutcome::Authenticated {
2756                    principal: auth_principal,
2757                } => {
2758                    // Auth succeeded — clear any prior failure count.
2759                    if let (Some(limiter), Some(ip)) = (rate_limiter, peer_ip) {
2760                        clear_auth_failures(limiter, ip);
2761                    }
2762                    match &auth_principal {
2763                        Some(p) => {
2764                            info!(peer = %peer, db = %db_name, user = %p.name, role = %p.role, "authenticated");
2765                        }
2766                        None => {
2767                            info!(peer = %peer, db = %db_name, "client connected");
2768                        }
2769                    }
2770                    principal = auth_principal;
2771                }
2772            }
2773
2774            // One process serves one database. When pinned to a name, reject a
2775            // CONNECT that explicitly asks for a different one (checked after
2776            // auth so db existence never leaks to unauthenticated clients).
2777            // When unpinned, accept anything but warn once per process so the
2778            // silent one-db-per-process mismatch is visible without spamming
2779            // the log on every pooled connection.
2780            match check_db_name(server_db_name.as_deref(), &db_name) {
2781                Ok(()) => {
2782                    if server_db_name.is_none() && !db_name.is_empty() && db_name != DEFAULT_DB_NAME
2783                    {
2784                        static NAMED_DB_WARNED: std::sync::atomic::AtomicBool =
2785                            std::sync::atomic::AtomicBool::new(false);
2786                        if !NAMED_DB_WARNED.swap(true, std::sync::atomic::Ordering::Relaxed) {
2787                            warn!(
2788                                peer = %peer, db = %db_name,
2789                                "client requested a named database but this server serves a single global database; name ignored"
2790                            );
2791                        }
2792                    }
2793                }
2794                Err(msg) => {
2795                    warn!(peer = %peer, db = %db_name, "rejected: unknown database");
2796                    let err = Message::Error { message: msg };
2797                    write_msg(&mut writer, &err).await;
2798                    return;
2799                }
2800            }
2801
2802            let ok = Message::ConnectOk {
2803                version: env!("CARGO_PKG_VERSION").into(),
2804            };
2805            if !write_msg(&mut writer, &ok).await {
2806                return;
2807            }
2808        }
2809        _ => {
2810            warn!(peer = %peer, "first message was not CONNECT");
2811            let err = Message::Error {
2812                message: "expected CONNECT".into(),
2813            };
2814            write_msg(&mut writer, &err).await;
2815            return;
2816        }
2817    }
2818
2819    let mut tx_permit: Option<OwnedSemaphorePermit> = None;
2820    // Persistent framing state makes reads cancellation-safe while they race
2821    // an in-flight blocking query. Frames decoded during execution retain
2822    // their original order here for normal pipelined processing afterwards.
2823    let mut wire_read_buffer = Vec::new();
2824    let mut pending_messages = InFlightReadAhead::default();
2825    // A non-query frame decoded during read-ahead batching, carried over to
2826    // the next iteration of the main loop.
2827    let mut carry: Option<Message> = None;
2828
2829    // Main query loop with idle timeout and shutdown awareness.
2830    'conn: loop {
2831        let msg = if let Some(m) = carry.take() {
2832            m
2833        } else if let Some(m) = pending_messages.pop_front() {
2834            m
2835        } else {
2836            tokio::select! {
2837                // Read next message with idle timeout.
2838                result = tokio::time::timeout(
2839                    idle_timeout,
2840                    read_message_cancel_safe(
2841                        &mut reader,
2842                        &mut wire_read_buffer,
2843                        MAX_WIRE_PAYLOAD_SIZE + 6,
2844                    ),
2845                ) => {
2846                    match result {
2847                        Ok(Ok(Some(frame))) => frame.message,
2848                        Ok(Ok(None)) => break,
2849                        Ok(Err(e)) => {
2850                            error!(peer = %peer, error = %e, "read error");
2851                            break;
2852                        }
2853                        Err(_) => {
2854                            info!(peer = %peer, "idle timeout, closing connection");
2855                            let err = Message::Error { message: "idle timeout".into() };
2856                            write_msg(&mut writer, &err).await;
2857                            break;
2858                        }
2859                    }
2860                }
2861                // If server is shutting down, notify client and close.
2862                _ = shutdown_rx.changed() => {
2863                    if *shutdown_rx.borrow() {
2864                        info!(peer = %peer, "server shutting down, closing connection");
2865                        let err = Message::Error { message: "server shutting down".into() };
2866                        write_msg(&mut writer, &err).await;
2867                        break;
2868                    }
2869                    continue;
2870                }
2871            }
2872        };
2873
2874        // Plain query frames take the batched path: a pipelining client's
2875        // whole burst is executed with ONE durability wait at the end
2876        // (durability generations are cumulative, so the newest statement's
2877        // ticket covers every earlier one). Everything else is handled one
2878        // frame at a time exactly as before.
2879        if matches!(
2880            msg,
2881            Message::Query { .. }
2882                | Message::QuerySql { .. }
2883                | Message::QueryWithParams { .. }
2884                | Message::QueryNative { .. }
2885                | Message::QuerySqlNative { .. }
2886                | Message::QueryWithParamsNative { .. }
2887        ) {
2888            /// Read-ahead cap per batch: bounds unflushed responses and keeps
2889            /// the reply latency of the first statement bounded.
2890            const MAX_PIPELINE_BATCH: usize = 128;
2891            /// Byte cap on retained (unflushed) response payloads per batch:
2892            /// large row results stop read-ahead, so one connection can never
2893            /// hold gigabytes of replies hostage to the batch's durability
2894            /// wait.
2895            const MAX_PIPELINE_BATCH_BYTES: usize = 4 << 20;
2896
2897            /// Approximate encoded size of a response, for the batch byte
2898            /// cap. Counts the dominant string payloads; exact per-frame
2899            /// overhead is irrelevant at the 4 MiB cap.
2900            fn approx_response_bytes(msg: &Message) -> usize {
2901                match msg {
2902                    Message::ResultRows { columns, rows } => {
2903                        columns.iter().map(|c| c.len() + 4).sum::<usize>()
2904                            + rows
2905                                .iter()
2906                                .map(|r| r.iter().map(|v| v.len() + 4).sum::<usize>())
2907                                .sum::<usize>()
2908                    }
2909                    Message::ResultScalar { value } => value.len(),
2910                    Message::ResultRowsNative { columns, rows } => {
2911                        columns.iter().map(|c| c.len() + 4).sum::<usize>()
2912                            + rows
2913                                .iter()
2914                                .map(|row| {
2915                                    row.iter()
2916                                        .map(|value| 5 + native_value_body_len(value))
2917                                        .sum::<usize>()
2918                                })
2919                                .sum::<usize>()
2920                    }
2921                    Message::ResultScalarNative { value } => 5 + native_value_body_len(value),
2922                    Message::ResultMessage { message } | Message::Error { message } => {
2923                        message.len()
2924                    }
2925                    _ => 16,
2926                }
2927            }
2928
2929            /// Whether the reader's buffered bytes hold at least one COMPLETE
2930            /// frame (6-byte header + payload). Read-ahead must never await
2931            /// the socket: blocking on a partial frame would hold the batch's
2932            /// durability settlement and unflushed replies hostage to a slow
2933            /// (or malicious) client, up to the idle timeout.
2934            fn complete_frame_buffered(buf: &[u8]) -> bool {
2935                buf.len() >= 6 && {
2936                    let payload_len =
2937                        u32::from_le_bytes(buf[2..6].try_into().expect("4-byte slice")) as usize;
2938                    buf.len() - 6 >= payload_len
2939                }
2940            }
2941
2942            let mut responses: Vec<Message> = Vec::new();
2943            let mut response_bytes: usize = 0;
2944            let mut last_ticket: Option<WalDurabilityTicket> = None;
2945            let mut deferred_metrics: Vec<DeferredQueryMetric> = Vec::new();
2946            let mut fatal: Option<ConnectionTermination> = None;
2947            let mut current = msg;
2948            loop {
2949                let (response, ticket, termination) = match current {
2950                    Message::Query { query } => {
2951                        if query.len() > MAX_QUERY_LENGTH {
2952                            (
2953                                Message::Error {
2954                                    message: format!(
2955                                        "query too large: {} bytes (max {})",
2956                                        query.len(),
2957                                        MAX_QUERY_LENGTH
2958                                    ),
2959                                },
2960                                None,
2961                                None,
2962                            )
2963                        } else {
2964                            debug!(peer = %peer, query = %query, "received query");
2965                            execute_wire_query(
2966                                engine.clone(),
2967                                tx_gate.clone(),
2968                                &mut tx_permit,
2969                                query,
2970                                WireResultMode::LegacyText,
2971                                principal.clone(),
2972                                query_timeout,
2973                                tx_wait_timeout,
2974                                &metrics,
2975                                &mut reader,
2976                                &mut wire_read_buffer,
2977                                &mut pending_messages,
2978                            )
2979                            .await
2980                        }
2981                    }
2982                    Message::QuerySql { query } => {
2983                        if query.len() > MAX_QUERY_LENGTH {
2984                            (
2985                                Message::Error {
2986                                    message: format!(
2987                                        "query too large: {} bytes (max {})",
2988                                        query.len(),
2989                                        MAX_QUERY_LENGTH
2990                                    ),
2991                                },
2992                                None,
2993                                None,
2994                            )
2995                        } else {
2996                            debug!(peer = %peer, query = %query, "received SQL query");
2997                            execute_wire_query_sql(
2998                                engine.clone(),
2999                                tx_gate.clone(),
3000                                &mut tx_permit,
3001                                query,
3002                                WireResultMode::LegacyText,
3003                                principal.clone(),
3004                                query_timeout,
3005                                tx_wait_timeout,
3006                                &metrics,
3007                                &mut reader,
3008                                &mut wire_read_buffer,
3009                                &mut pending_messages,
3010                            )
3011                            .await
3012                        }
3013                    }
3014                    Message::QueryWithParams { query, params } => {
3015                        if query.len() > MAX_QUERY_LENGTH {
3016                            (
3017                                Message::Error {
3018                                    message: format!(
3019                                        "query too large: {} bytes (max {})",
3020                                        query.len(),
3021                                        MAX_QUERY_LENGTH
3022                                    ),
3023                                },
3024                                None,
3025                                None,
3026                            )
3027                        } else {
3028                            debug!(peer = %peer, query = %query, n_params = params.len(), "received parameterized query");
3029                            execute_wire_query_with_params(
3030                                engine.clone(),
3031                                tx_gate.clone(),
3032                                &mut tx_permit,
3033                                query,
3034                                params,
3035                                WireResultMode::LegacyText,
3036                                principal.clone(),
3037                                query_timeout,
3038                                tx_wait_timeout,
3039                                &metrics,
3040                                &mut reader,
3041                                &mut wire_read_buffer,
3042                                &mut pending_messages,
3043                            )
3044                            .await
3045                        }
3046                    }
3047                    Message::QueryNative { query } => {
3048                        if query.len() > MAX_QUERY_LENGTH {
3049                            (
3050                                Message::Error {
3051                                    message: format!(
3052                                        "query too large: {} bytes (max {})",
3053                                        query.len(),
3054                                        MAX_QUERY_LENGTH
3055                                    ),
3056                                },
3057                                None,
3058                                None,
3059                            )
3060                        } else {
3061                            debug!(peer = %peer, query = %query, "received native query");
3062                            execute_wire_query(
3063                                engine.clone(),
3064                                tx_gate.clone(),
3065                                &mut tx_permit,
3066                                query,
3067                                WireResultMode::Native,
3068                                principal.clone(),
3069                                query_timeout,
3070                                tx_wait_timeout,
3071                                &metrics,
3072                                &mut reader,
3073                                &mut wire_read_buffer,
3074                                &mut pending_messages,
3075                            )
3076                            .await
3077                        }
3078                    }
3079                    Message::QuerySqlNative { query } => {
3080                        if query.len() > MAX_QUERY_LENGTH {
3081                            (
3082                                Message::Error {
3083                                    message: format!(
3084                                        "query too large: {} bytes (max {})",
3085                                        query.len(),
3086                                        MAX_QUERY_LENGTH
3087                                    ),
3088                                },
3089                                None,
3090                                None,
3091                            )
3092                        } else {
3093                            debug!(peer = %peer, query = %query, "received native SQL query");
3094                            execute_wire_query_sql(
3095                                engine.clone(),
3096                                tx_gate.clone(),
3097                                &mut tx_permit,
3098                                query,
3099                                WireResultMode::Native,
3100                                principal.clone(),
3101                                query_timeout,
3102                                tx_wait_timeout,
3103                                &metrics,
3104                                &mut reader,
3105                                &mut wire_read_buffer,
3106                                &mut pending_messages,
3107                            )
3108                            .await
3109                        }
3110                    }
3111                    Message::QueryWithParamsNative { query, params } => {
3112                        if query.len() > MAX_QUERY_LENGTH {
3113                            (
3114                                Message::Error {
3115                                    message: format!(
3116                                        "query too large: {} bytes (max {})",
3117                                        query.len(),
3118                                        MAX_QUERY_LENGTH
3119                                    ),
3120                                },
3121                                None,
3122                                None,
3123                            )
3124                        } else {
3125                            debug!(peer = %peer, query = %query, n_params = params.len(), "received native parameterized query");
3126                            execute_wire_query_with_params(
3127                                engine.clone(),
3128                                tx_gate.clone(),
3129                                &mut tx_permit,
3130                                query,
3131                                params,
3132                                WireResultMode::Native,
3133                                principal.clone(),
3134                                query_timeout,
3135                                tx_wait_timeout,
3136                                &metrics,
3137                                &mut reader,
3138                                &mut wire_read_buffer,
3139                                &mut pending_messages,
3140                            )
3141                            .await
3142                        }
3143                    }
3144                    _ => unreachable!("batch loop only receives plain query frames"),
3145                };
3146                if let Some((t, m)) = ticket {
3147                    // Later tickets cover earlier generations — keep only the
3148                    // newest; the batch-end wait settles them all. Every
3149                    // deferred metric is kept: each records after settlement.
3150                    last_ticket = Some(t);
3151                    deferred_metrics.push(m);
3152                }
3153                response_bytes += approx_response_bytes(&response);
3154                responses.push(response);
3155                if let Some(reason) = termination {
3156                    fatal = Some(reason);
3157                    break;
3158                }
3159
3160                // Read ahead only when a COMPLETE next frame is already
3161                // buffered (never await the socket mid-batch) and the
3162                // retained replies stay small. While an explicit transaction
3163                // is open the connection holds the TxGate, so batching would
3164                // only extend the exclusive window — flush instead.
3165                if tx_permit.is_some()
3166                    || responses.len() >= MAX_PIPELINE_BATCH
3167                    || response_bytes >= MAX_PIPELINE_BATCH_BYTES
3168                    || (pending_messages.is_empty()
3169                        && !complete_frame_buffered(&wire_read_buffer)
3170                        && !complete_frame_buffered(reader.buffer()))
3171                {
3172                    break;
3173                }
3174                // The full frame is buffered, so this returns without socket
3175                // I/O; the timeout is a defensive backstop only.
3176                let next_message = if let Some(message) = pending_messages.pop_front() {
3177                    Ok(Some(message))
3178                } else {
3179                    tokio::time::timeout(
3180                        idle_timeout,
3181                        read_message_cancel_safe(
3182                            &mut reader,
3183                            &mut wire_read_buffer,
3184                            MAX_WIRE_PAYLOAD_SIZE + 6,
3185                        ),
3186                    )
3187                    .await
3188                    .map(|result| result.map(|frame| frame.map(|frame| frame.message)))
3189                    .unwrap_or_else(|_| {
3190                        Err(std::io::Error::new(
3191                            std::io::ErrorKind::TimedOut,
3192                            "timeout decoding fully-buffered frame",
3193                        ))
3194                    })
3195                };
3196                match next_message {
3197                    Ok(Some(
3198                        next @ (Message::Query { .. }
3199                        | Message::QuerySql { .. }
3200                        | Message::QueryWithParams { .. }
3201                        | Message::QueryNative { .. }
3202                        | Message::QuerySqlNative { .. }
3203                        | Message::QueryWithParamsNative { .. }),
3204                    )) => {
3205                        // If another connection currently holds the TxGate,
3206                        // the next statement would block on the gate with
3207                        // this batch's replies still unflushed (pre-batching,
3208                        // they'd already have been written). Flush first and
3209                        // handle the frame on the next main-loop iteration.
3210                        // Benign TOCTOU: worst case is one early flush or one
3211                        // gate wait with an empty reply queue.
3212                        if tx_gate.available_permits() == 0 {
3213                            carry = Some(next);
3214                            break;
3215                        }
3216                        current = next;
3217                    }
3218                    Ok(Some(other)) => {
3219                        // Not a plain query — flush this batch, then handle
3220                        // the frame on the next main-loop iteration.
3221                        carry = Some(other);
3222                        break;
3223                    }
3224                    Ok(None) => {
3225                        fatal = Some(ConnectionTermination::Closed);
3226                        break;
3227                    }
3228                    Err(e) => {
3229                        error!(peer = %peer, error = %e, "read error");
3230                        fatal = Some(ConnectionTermination::ReadError);
3231                        break;
3232                    }
3233                }
3234            }
3235
3236            // ONE durability wait for the whole batch, then the deferred
3237            // metrics: a durability failure records Ok statements as errors,
3238            // and latency includes the settlement wait the client observed.
3239            let mut durability_failed = false;
3240            if let Some(ticket) = last_ticket {
3241                if let Some(message) = settle_durability_ticket(ticket).await {
3242                    // The covering fsync failed: nothing in this batch may be
3243                    // acknowledged as durable.
3244                    durability_failed = true;
3245                    for r in responses.iter_mut() {
3246                        if is_success_response(r) {
3247                            *r = Message::Error {
3248                                message: message.clone(),
3249                            };
3250                        }
3251                    }
3252                }
3253            }
3254            for m in deferred_metrics.drain(..) {
3255                let outcome = if m.exceeded_timeout {
3256                    QueryOutcome::Timeout
3257                } else if durability_failed && matches!(m.outcome, QueryOutcome::Ok) {
3258                    QueryOutcome::Error
3259                } else {
3260                    m.outcome
3261                };
3262                metrics.record_query(m.start.elapsed(), outcome);
3263            }
3264
3265            for r in &responses {
3266                if !write_msg(&mut writer, r).await {
3267                    break 'conn;
3268                }
3269            }
3270            match fatal {
3271                None => continue,
3272                Some(ConnectionTermination::Closed | ConnectionTermination::ReadError) => break,
3273            }
3274        }
3275
3276        let response = match msg {
3277            Message::Ping => {
3278                debug!(peer = %peer, "ping");
3279                Message::Pong
3280            }
3281            Message::SyncStatus { replica_id } => {
3282                let engine = engine.clone();
3283                let principal = principal.clone();
3284                let log_context = SyncLogContext::status(&replica_id);
3285                execute_gated_sync(
3286                    SyncExecutionContext {
3287                        tx_gate: tx_gate.clone(),
3288                        connection_has_transaction: tx_permit.is_some(),
3289                        operation: SyncOperation::Status,
3290                        log_context,
3291                        metrics: &metrics,
3292                        query_timeout,
3293                    },
3294                    (engine, replica_id, credential_auth_configured, principal),
3295                    |(engine, replica_id, credential_authenticated, principal)| {
3296                        dispatch_sync_status_decision(
3297                            &engine,
3298                            replica_id,
3299                            credential_authenticated,
3300                            principal.as_ref(),
3301                        )
3302                    },
3303                )
3304                .await
3305            }
3306            Message::SyncPull {
3307                replica_id,
3308                since_lsn,
3309                max_units,
3310                max_bytes,
3311                database_id,
3312                primary_generation,
3313                wal_format_version,
3314                catalog_version,
3315                segment_format_version,
3316            } => {
3317                let engine = engine.clone();
3318                let principal = principal.clone();
3319                let request = SyncPullRequest {
3320                    replica_id,
3321                    since_lsn,
3322                    max_units,
3323                    max_bytes,
3324                    database_id,
3325                    primary_generation,
3326                    wal_format_version,
3327                    catalog_version,
3328                    segment_format_version,
3329                };
3330                let log_context = SyncLogContext::pull(&request);
3331                execute_gated_sync(
3332                    SyncExecutionContext {
3333                        tx_gate: tx_gate.clone(),
3334                        connection_has_transaction: tx_permit.is_some(),
3335                        operation: SyncOperation::Pull,
3336                        log_context,
3337                        metrics: &metrics,
3338                        query_timeout,
3339                    },
3340                    (engine, request, credential_auth_configured, principal),
3341                    |(engine, request, credential_authenticated, principal)| {
3342                        dispatch_sync_pull_decision(
3343                            &engine,
3344                            request,
3345                            credential_authenticated,
3346                            principal.as_ref(),
3347                        )
3348                    },
3349                )
3350                .await
3351            }
3352            Message::SyncAck {
3353                replica_id,
3354                applied_lsn,
3355                remote_lsn,
3356            } => {
3357                let engine = engine.clone();
3358                let principal = principal.clone();
3359                let log_context = SyncLogContext::ack(&replica_id, applied_lsn, remote_lsn);
3360                execute_gated_sync(
3361                    SyncExecutionContext {
3362                        tx_gate: tx_gate.clone(),
3363                        connection_has_transaction: tx_permit.is_some(),
3364                        operation: SyncOperation::Ack,
3365                        log_context,
3366                        metrics: &metrics,
3367                        query_timeout,
3368                    },
3369                    (
3370                        engine,
3371                        replica_id,
3372                        applied_lsn,
3373                        remote_lsn,
3374                        credential_auth_configured,
3375                        principal,
3376                    ),
3377                    |(
3378                        engine,
3379                        replica_id,
3380                        applied_lsn,
3381                        observed_remote_lsn,
3382                        credential_authenticated,
3383                        principal,
3384                    )| {
3385                        dispatch_sync_ack_decision(
3386                            &engine,
3387                            replica_id,
3388                            applied_lsn,
3389                            observed_remote_lsn,
3390                            credential_authenticated,
3391                            principal.as_ref(),
3392                        )
3393                    },
3394                )
3395                .await
3396            }
3397            Message::Disconnect => {
3398                debug!(peer = %peer, "received DISCONNECT");
3399                break;
3400            }
3401            _ => Message::Error {
3402                message: "unexpected message type".into(),
3403            },
3404        };
3405
3406        if !write_msg(&mut writer, &response).await {
3407            break;
3408        }
3409    }
3410
3411    // Roll back any open transaction the client left behind on disconnect.
3412    // The permit must stay alive in `tx_permit` for the duration of the awaited
3413    // rollback and be released only afterwards — mirroring the query-timeout
3414    // path above. Using `tx_permit.take().is_some()` here would drop the permit
3415    // (freeing the TxGate) *before* the rollback runs, letting another
3416    // connection BEGIN a transaction that this stale rollback would then clobber.
3417    if tx_permit.is_some() {
3418        let engine = engine.clone();
3419        let principal = principal.clone();
3420        let _ =
3421            tokio::task::spawn_blocking(move || rollback_open_transaction(engine, principal)).await;
3422    }
3423    tx_permit.take();
3424
3425    info!(peer = %peer, "client disconnected");
3426}
3427
3428fn charge_response_bytes(total: &mut usize, bytes: usize) -> Result<(), QueryError> {
3429    *total = total.saturating_add(bytes);
3430    if *total > MAX_RESPONSE_PAYLOAD_SIZE {
3431        return Err(QueryError::Execution(format!(
3432            "result too large: encoded response exceeds {} bytes; add a limit or narrower projection",
3433            MAX_RESPONSE_PAYLOAD_SIZE
3434        )));
3435    }
3436    Ok(())
3437}
3438
3439fn native_value_body_len(value: &Value) -> usize {
3440    match value {
3441        Value::Empty => 0,
3442        Value::Int(_) | Value::Float(_) | Value::DateTime(_) => 8,
3443        Value::Bool(_) => 1,
3444        Value::Str(value) => value.len(),
3445        Value::Uuid(_) => 16,
3446        Value::Bytes(value) => value.len(),
3447        Value::Json(value) => value.len(),
3448    }
3449}
3450
3451fn query_result_to_message(
3452    result: QueryResult,
3453    result_mode: WireResultMode,
3454) -> Result<Message, QueryError> {
3455    match result {
3456        QueryResult::Rows { columns, rows } => {
3457            let mut encoded_bytes = 2usize; // column count
3458            for col in &columns {
3459                charge_response_bytes(&mut encoded_bytes, 4 + col.len())?;
3460            }
3461            charge_response_bytes(&mut encoded_bytes, 4)?; // row count
3462
3463            match result_mode {
3464                WireResultMode::Native => {
3465                    for row in &rows {
3466                        for value in row {
3467                            charge_response_bytes(
3468                                &mut encoded_bytes,
3469                                5 + native_value_body_len(value),
3470                            )?;
3471                        }
3472                    }
3473                    Ok(Message::ResultRowsNative { columns, rows })
3474                }
3475                WireResultMode::LegacyText => {
3476                    let mut str_rows = Vec::with_capacity(rows.len());
3477                    for row in rows {
3478                        let mut str_row = Vec::with_capacity(row.len());
3479                        for value in row {
3480                            let display = value_to_display(&value);
3481                            charge_response_bytes(&mut encoded_bytes, 4 + display.len())?;
3482                            str_row.push(display);
3483                        }
3484                        str_rows.push(str_row);
3485                    }
3486                    Ok(Message::ResultRows {
3487                        columns,
3488                        rows: str_rows,
3489                    })
3490                }
3491            }
3492        }
3493        QueryResult::Scalar(value) => match result_mode {
3494            WireResultMode::Native => {
3495                let mut encoded_bytes = 0;
3496                charge_response_bytes(&mut encoded_bytes, 5 + native_value_body_len(&value))?;
3497                Ok(Message::ResultScalarNative { value })
3498            }
3499            WireResultMode::LegacyText => Ok(Message::ResultScalar {
3500                value: value_to_display(&value),
3501            }),
3502        },
3503        QueryResult::Modified(n) => Ok(Message::ResultOk { affected: n }),
3504        QueryResult::Created(name) => Ok(Message::ResultMessage {
3505            message: format!("type {name} created"),
3506        }),
3507        QueryResult::Executed { message } => Ok(Message::ResultMessage { message }),
3508    }
3509}
3510
3511// Canonical wire rendering lives on `Value` (`powdb_storage`) so the server,
3512// CLI, and embedded bindings render results identically. Kept as a thin alias
3513// to minimize churn at the call sites in this module.
3514fn value_to_display(v: &Value) -> String {
3515    v.to_wire_string()
3516}
3517
3518#[cfg(test)]
3519mod tests {
3520    use super::*;
3521    use powdb_storage::wal::WalRecordType;
3522    use powdb_sync::{
3523        write_identity_snapshot, write_segment_atomic, DatabaseIdentity, IdentitySnapshot,
3524        ReplicaCursor, RetainedSegment, RetainedUnit,
3525    };
3526
3527    // ---- Wire NULL rendering (Fix: remote protocol rendered NULL as `{}`) ----
3528
3529    #[test]
3530    fn null_serializes_as_null_bareword_on_wire() {
3531        assert_eq!(value_to_display(&Value::Empty), "null");
3532    }
3533
3534    // ---- Error sanitization allowlist ----
3535
3536    #[test]
3537    fn unique_violation_error_surfaces_to_remote_clients() {
3538        // The storage layer reports the actionable message; the server must
3539        // not replace it with the generic "query execution error".
3540        assert_eq!(
3541            sanitize_error("unique constraint violation on User.email"),
3542            "unique constraint violation on User.email"
3543        );
3544    }
3545
3546    #[test]
3547    fn internal_errors_stay_generic() {
3548        assert_eq!(
3549            sanitize_error("some internal io panic detail"),
3550            "query execution error"
3551        );
3552    }
3553
3554    #[test]
3555    fn cancellation_errors_surface_to_remote_clients() {
3556        // A cancelled/timed-out query must reach the client with its real
3557        // message (both are derived from the configured timeout or a client
3558        // disconnect and leak no internal state) rather than the generic mask.
3559        for msg in [
3560            &QueryError::Timeout { timeout_ms: 2000 }.to_string(),
3561            &QueryError::Cancelled.to_string(),
3562        ] {
3563            assert_eq!(sanitize_error(msg), *msg, "should pass through verbatim");
3564        }
3565        // Sanity-check the exact wording the executor emits.
3566        assert_eq!(
3567            QueryError::Timeout { timeout_ms: 2000 }.to_string(),
3568            "query timeout after 2000ms"
3569        );
3570        assert_eq!(
3571            QueryError::Cancelled.to_string(),
3572            "query cancelled by client disconnect"
3573        );
3574    }
3575
3576    // ---- JSON (v0.12): canonical-text wire rendering + parse-error passthrough ----
3577
3578    #[test]
3579    fn json_cell_renders_canonical_text_on_wire() {
3580        // A Json value flows through the same string-cell path as every other
3581        // value (value_to_display -> Value::to_wire_string). PJ1 is canonical,
3582        // so keys come back sorted bytewise regardless of input order and the
3583        // client receives parseable JSON text with no protocol change.
3584        let pj1 = powdb_storage::pj1::parse_json_text(r#"{"b":2,"a":1,"nested":{"z":true}}"#)
3585            .expect("valid JSON");
3586        let result = QueryResult::Rows {
3587            columns: vec!["doc".into()],
3588            rows: vec![vec![Value::Json(pj1.into())]],
3589        };
3590        match query_result_to_message(result, WireResultMode::LegacyText).expect("encodes") {
3591            Message::ResultRows { columns, rows } => {
3592                assert_eq!(columns, vec!["doc"]);
3593                assert_eq!(
3594                    rows,
3595                    vec![vec![r#"{"a":1,"b":2,"nested":{"z":true}}"#.to_string()]]
3596                );
3597            }
3598            other => panic!("expected ResultRows, got {other:?}"),
3599        }
3600    }
3601
3602    #[test]
3603    fn json_parse_error_surfaces_to_remote_clients() {
3604        // Lane B rejects invalid JSON on insert as QueryError::TypeError, whose
3605        // Display is "type mismatch: <detail>" (crates/query/src/result.rs).
3606        // That prefix is allowlisted, so the actionable detail reaches the
3607        // client instead of the generic "query execution error". The raw
3608        // storage-layer phrasing ("invalid JSON: ...") is also allowlisted as
3609        // defense-in-depth. Internal PJ1 corruption ("malformed PJ1: ...") is
3610        // deliberately NOT allowlisted: it leaks storage internals and never
3611        // occurs on the client-driven insert path.
3612        for msg in [
3613            "type mismatch: invalid JSON: unexpected character 'x' at position 3",
3614            "invalid JSON: nesting exceeds depth cap 128",
3615        ] {
3616            assert_eq!(sanitize_error(msg), msg, "should pass through verbatim");
3617        }
3618        assert_eq!(
3619            sanitize_error("malformed PJ1: truncated"),
3620            "query execution error",
3621            "internal storage corruption must stay masked"
3622        );
3623    }
3624
3625    // `describe <Type>` renders a json column's type as the bareword "json"
3626    // over the wire. introspect_describe emits type_id_to_name(TypeId::Json) =
3627    // "json" (crates/query/src/executor/compiled.rs) as a Str cell, which flows
3628    // through value_to_display unchanged; Lane B's DDL keyword makes `type Doc
3629    // { body: json }` accepted, so this runs end to end (v0.12, Lane D).
3630    #[test]
3631    fn describe_shows_json_type_over_the_wire() {
3632        let dir = tempfile::tempdir().unwrap();
3633        let mut engine = Engine::new(dir.path()).unwrap();
3634        engine
3635            .execute_powql("type Doc { required id: int, body: json }")
3636            .expect("json column DDL should be accepted once Lane B lands");
3637        let result = engine.execute_powql("describe Doc").expect("describe runs");
3638        let msg = query_result_to_message(result, WireResultMode::LegacyText).expect("encodes");
3639        match msg {
3640            Message::ResultRows { columns, rows } => {
3641                assert_eq!(columns[1], "type");
3642                // The `body` column's type cell must be the bareword "json".
3643                let body = rows
3644                    .iter()
3645                    .find(|r| r[0] == "body")
3646                    .expect("body column present");
3647                assert_eq!(body[1], "json");
3648            }
3649            other => panic!("expected ResultRows, got {other:?}"),
3650        }
3651    }
3652
3653    // ---- Named-database gate (P-10) ----
3654
3655    #[test]
3656    fn db_name_unpinned_accepts_any_name() {
3657        for requested in ["", "default", "prod", "anything"] {
3658            assert!(
3659                check_db_name(None, requested).is_ok(),
3660                "rejected {requested}"
3661            );
3662        }
3663    }
3664
3665    #[test]
3666    fn db_name_pinned_accepts_match_empty_and_default_sentinel() {
3667        // The configured name, the empty name, and the client default sentinel
3668        // are all "no foreign database explicitly requested".
3669        assert!(check_db_name(Some("prod"), "prod").is_ok());
3670        assert!(check_db_name(Some("prod"), "").is_ok());
3671        assert!(check_db_name(Some("prod"), DEFAULT_DB_NAME).is_ok());
3672    }
3673
3674    #[test]
3675    fn db_name_pinned_rejects_foreign_with_clear_message() {
3676        let err = check_db_name(Some("prod"), "staging").unwrap_err();
3677        assert_eq!(err, "unknown database 'staging'; this server serves 'prod'");
3678    }
3679
3680    // ---- Explicit-transaction gate wait timeout (P-4) ----
3681
3682    #[tokio::test]
3683    async fn begin_permit_acquires_when_gate_is_free() {
3684        let gate = new_tx_gate();
3685        let metrics = Arc::new(Metrics::new());
3686        let permit = acquire_begin_permit(&gate, Duration::from_secs(5), &metrics)
3687            .await
3688            .expect("should acquire a free gate");
3689        assert_eq!(gate.available_permits(), 0, "permit must be held");
3690        drop(permit);
3691        assert_eq!(
3692            gate.available_permits(),
3693            DEFAULT_TX_GATE_READER_PERMITS as usize,
3694            "permit pool must release on drop"
3695        );
3696    }
3697
3698    #[tokio::test]
3699    async fn begin_permit_times_out_with_clear_error_and_truthful_metric() {
3700        let gate = new_tx_gate();
3701        let metrics = Arc::new(Metrics::new());
3702        // Hold the full writer admission so the next acquire must time out.
3703        let _held = gate
3704            .clone()
3705            .acquire_many_owned(DEFAULT_TX_GATE_READER_PERMITS)
3706            .await
3707            .unwrap();
3708        let err = acquire_begin_permit(&gate, Duration::from_millis(25), &metrics)
3709            .await
3710            .expect_err("must time out while the gate is held");
3711        match err {
3712            Message::Error { message } => {
3713                assert!(
3714                    message.contains("transaction gate timeout after 25ms"),
3715                    "unexpected message: {message}"
3716                );
3717                assert!(
3718                    message.contains("waiting for concurrent transaction"),
3719                    "unexpected message: {message}"
3720                );
3721            }
3722            other => panic!("expected Error, got {other:?}"),
3723        }
3724        let rendered = metrics.render();
3725        assert!(rendered.contains("powdb_tx_gate_timeouts_total 1"));
3726        // A timed-out begin is a failed statement from the client's view.
3727        assert!(rendered.contains("powdb_queries_total{result=\"error\"} 1"));
3728    }
3729
3730    #[test]
3731    fn admission_classification_has_query_shape_parity_and_fails_closed() {
3732        assert_eq!(
3733            classify_query_admission("User filter .id = 1"),
3734            AdmissionMode::Reader
3735        );
3736        assert_eq!(
3737            classify_sql_admission("SELECT * FROM User WHERE id = 1"),
3738            AdmissionMode::Reader
3739        );
3740        assert_eq!(
3741            classify_params_admission("User filter .id = $1", &[WireParam::Int(1)]),
3742            AdmissionMode::Reader
3743        );
3744
3745        assert_eq!(
3746            classify_query_admission("insert User { id := 1 }"),
3747            AdmissionMode::Writer
3748        );
3749        assert_eq!(
3750            classify_sql_admission("INSERT INTO User (id) VALUES (1)"),
3751            AdmissionMode::Writer
3752        );
3753        assert_eq!(
3754            classify_params_admission("insert User { id := $1 }", &[WireParam::Int(1)]),
3755            AdmissionMode::Writer
3756        );
3757        assert_eq!(
3758            classify_query_admission("this is not valid PowQL"),
3759            AdmissionMode::Writer,
3760            "uncertain statements must never enter through reader admission"
3761        );
3762    }
3763
3764    #[tokio::test]
3765    async fn writer_admission_excludes_readers() {
3766        let gate = new_tx_gate();
3767        let metrics = Arc::new(Metrics::new());
3768        let writer = acquire_begin_permit(&gate, Duration::from_secs(1), &metrics)
3769            .await
3770            .expect("writer admission");
3771        let blocked = acquire_autocommit_permit(
3772            &gate,
3773            AdmissionMode::Reader,
3774            Duration::from_millis(25),
3775            &metrics,
3776        )
3777        .await;
3778        assert!(blocked.is_err(), "reader must wait behind writer admission");
3779        drop(writer);
3780        let _reader = acquire_autocommit_permit(
3781            &gate,
3782            AdmissionMode::Reader,
3783            Duration::from_secs(1),
3784            &metrics,
3785        )
3786        .await
3787        .expect("reader must proceed after writer releases");
3788    }
3789
3790    #[tokio::test]
3791    async fn queued_writer_is_not_starved_by_later_readers() {
3792        let gate = new_tx_gate();
3793        let metrics = Arc::new(Metrics::new());
3794        let first_reader = acquire_autocommit_permit(
3795            &gate,
3796            AdmissionMode::Reader,
3797            Duration::from_secs(1),
3798            &metrics,
3799        )
3800        .await
3801        .expect("first reader admission");
3802
3803        let writer_gate = gate.clone();
3804        let writer_metrics = metrics.clone();
3805        let mut writer = tokio::spawn(async move {
3806            acquire_begin_permit(&writer_gate, Duration::from_secs(1), &writer_metrics).await
3807        });
3808        tokio::time::sleep(Duration::from_millis(10)).await;
3809
3810        let late_gate = gate.clone();
3811        let late_metrics = metrics.clone();
3812        let mut late_reader = tokio::spawn(async move {
3813            acquire_autocommit_permit(
3814                &late_gate,
3815                AdmissionMode::Reader,
3816                Duration::from_secs(1),
3817                &late_metrics,
3818            )
3819            .await
3820        });
3821        assert!(
3822            tokio::time::timeout(Duration::from_millis(25), &mut late_reader)
3823                .await
3824                .is_err(),
3825            "a later reader must queue behind the waiting writer"
3826        );
3827
3828        drop(first_reader);
3829        let writer_permit = tokio::time::timeout(Duration::from_secs(1), &mut writer)
3830            .await
3831            .expect("writer must acquire once prior readers drain")
3832            .expect("writer task")
3833            .expect("writer admission");
3834        assert!(
3835            tokio::time::timeout(Duration::from_millis(25), &mut late_reader)
3836                .await
3837                .is_err(),
3838            "writer must hold exclusive admission before the late reader"
3839        );
3840        drop(writer_permit);
3841        let _late_reader = tokio::time::timeout(Duration::from_secs(1), late_reader)
3842            .await
3843            .expect("late reader must eventually acquire")
3844            .expect("late reader task")
3845            .expect("late reader admission");
3846    }
3847
3848    fn dirty_view_engine() -> (tempfile::TempDir, Arc<RwLock<Engine>>) {
3849        let dir = tempfile::tempdir().unwrap();
3850        let mut engine = Engine::new(dir.path()).unwrap();
3851        engine
3852            .execute_powql("type Source { required id: int }")
3853            .unwrap();
3854        engine.execute_powql("insert Source { id := 1 }").unwrap();
3855        engine
3856            .execute_powql("materialize Snapshot as Source")
3857            .unwrap();
3858        engine.execute_powql("insert Source { id := 2 }").unwrap();
3859        (dir, Arc::new(RwLock::new(engine)))
3860    }
3861
3862    #[test]
3863    fn dirty_view_requests_explicit_escalation_on_every_frontend() {
3864        let (_dir, engine) = dirty_view_engine();
3865
3866        assert!(matches!(
3867            dispatch_query(&engine, "Snapshot", None, false).0,
3868            Err(QueryError::ReadonlyNeedsWrite)
3869        ));
3870        assert!(matches!(
3871            dispatch_sql_query(&engine, "SELECT * FROM Snapshot", None, false).0,
3872            Err(QueryError::ReadonlyNeedsWrite)
3873        ));
3874        assert!(matches!(
3875            dispatch_query_with_params(
3876                &engine,
3877                "Snapshot filter .id = $1",
3878                &[WireParam::Int(1)],
3879                None,
3880                false,
3881            )
3882            .0,
3883            Err(QueryError::ReadonlyNeedsWrite)
3884        ));
3885    }
3886
3887    #[tokio::test]
3888    async fn dirty_view_upgrade_waits_for_held_reader_then_records_once() {
3889        let (_dir, engine) = dirty_view_engine();
3890        let gate = new_tx_gate_with_permits(2);
3891        let metrics = Arc::new(Metrics::new());
3892        let held_reader = acquire_autocommit_permit(
3893            &gate,
3894            AdmissionMode::Reader,
3895            Duration::from_secs(1),
3896            &metrics,
3897        )
3898        .await
3899        .expect("held reader admission");
3900
3901        // Keep the peer open so the query monitor waits instead of treating
3902        // EOF as a client disconnect while the admission upgrade is blocked.
3903        let (_client, server) = tokio::io::duplex(1024);
3904        let task_gate = gate.clone();
3905        let task_metrics = metrics.clone();
3906        let mut task = tokio::spawn(async move {
3907            let mut reader = BufReader::new(server);
3908            let mut wire_read_buffer = Vec::new();
3909            let mut pending_messages = InFlightReadAhead::default();
3910            let mut tx_permit = None;
3911            execute_wire_query(
3912                engine,
3913                task_gate,
3914                &mut tx_permit,
3915                "Snapshot".into(),
3916                WireResultMode::Native,
3917                None,
3918                Duration::from_secs(2),
3919                Duration::from_secs(1),
3920                &task_metrics,
3921                &mut reader,
3922                &mut wire_read_buffer,
3923                &mut pending_messages,
3924            )
3925            .await
3926        });
3927
3928        assert!(
3929            tokio::time::timeout(Duration::from_millis(100), &mut task)
3930                .await
3931                .is_err(),
3932            "dirty-view retry must wait for exclusive admission while another reader is held"
3933        );
3934        drop(held_reader);
3935
3936        let (message, ticket, termination) = tokio::time::timeout(Duration::from_secs(2), task)
3937            .await
3938            .expect("upgrade must finish after the held reader releases")
3939            .expect("query task");
3940        assert!(matches!(message, Message::ResultRowsNative { .. }));
3941        assert!(termination.is_none());
3942
3943        let (ticket, metric) = ticket.expect("view refresh must defer its WAL metric");
3944        drop(ticket);
3945        metrics.record_query(metric.start.elapsed(), metric.outcome);
3946
3947        let rendered = metrics.render();
3948        assert!(rendered.contains("powdb_queries_total{result=\"ok\"} 1"));
3949        assert!(rendered.contains("powdb_queries_total{result=\"error\"} 0"));
3950    }
3951
3952    #[tokio::test]
3953    async fn timed_out_readonly_escalation_is_not_retried_or_reported_as_generic_error() {
3954        let (_dir, engine) = dirty_view_engine();
3955        let metrics = Arc::new(Metrics::new());
3956        // Keep the peer open so the socket monitor does not turn this into a
3957        // disconnect before the deadline fires.
3958        let (_client, server) = tokio::io::duplex(1024);
3959        let mut reader = BufReader::new(server);
3960        let mut wire_read_buffer = Vec::new();
3961        let mut pending_messages = InFlightReadAhead::default();
3962        let query_timeout = Duration::from_millis(20);
3963        let query_deadline = Instant::now() + query_timeout;
3964
3965        let (message, ticket, termination, retry) = run_blocking_query(
3966            engine,
3967            (),
3968            None,
3969            WireResultMode::Native,
3970            query_timeout,
3971            query_deadline,
3972            &metrics,
3973            &mut reader,
3974            &mut wire_read_buffer,
3975            &mut pending_messages,
3976            |_engine, (), _principal| {
3977                // Ignore the token deliberately: this reproduces the race in
3978                // which the async deadline wins but the joined task's final
3979                // result is the internal dirty-view escalation sentinel.
3980                std::thread::sleep(Duration::from_millis(50));
3981                (Err(QueryError::ReadonlyNeedsWrite), None)
3982            },
3983        )
3984        .await;
3985
3986        assert!(ticket.is_none());
3987        assert!(termination.is_none());
3988        assert!(!retry, "a timed-out query must never be resurrected");
3989        match message {
3990            Message::Error { message } => assert!(
3991                message.contains("query timeout after 20ms"),
3992                "timeout must remain client-visible, got {message}"
3993            ),
3994            other => panic!("expected timeout error, got {other:?}"),
3995        }
3996        let rendered = metrics.render();
3997        assert!(rendered.contains("powdb_query_timeouts_total 1"));
3998        assert!(rendered.contains("powdb_queries_total{result=\"error\"} 1"));
3999    }
4000
4001    #[test]
4002    fn resource_limit_errors_surface_actionable_hints() {
4003        // These carry user-actionable guidance and leak no internal state, so
4004        // they must reach the client verbatim — not be masked to the generic
4005        // message. The exact strings come from QueryError's Display impl
4006        // (crates/query/src/result.rs).
4007        for msg in [
4008            "sort input exceeds row limit — add a LIMIT clause",
4009            "join result exceeds row limit",
4010            "query exceeded memory budget: requested 100 bytes, limit 50 bytes",
4011            "result too large: encoded response exceeds 1024 bytes; add a limit or narrower projection",
4012        ] {
4013            assert_eq!(sanitize_error(msg), msg, "should pass through verbatim");
4014        }
4015    }
4016
4017    #[test]
4018    fn oversized_result_is_rejected_before_wire_encoding() {
4019        let long = "x".repeat(MAX_RESPONSE_PAYLOAD_SIZE);
4020        let result = QueryResult::Rows {
4021            columns: vec!["payload".into()],
4022            rows: vec![vec![Value::Str(long)]],
4023        };
4024        let err = query_result_to_message(result, WireResultMode::LegacyText).unwrap_err();
4025        assert!(
4026            err.to_string().starts_with("result too large"),
4027            "unexpected error: {err}"
4028        );
4029    }
4030
4031    // ---- Role enforcement (Fix: readonly role was not enforced) ----
4032
4033    fn parsed(q: &str) -> powdb_query::ast::Statement {
4034        parser::parse(q).unwrap()
4035    }
4036
4037    fn principal(role: &str) -> Option<Principal> {
4038        Some(Principal {
4039            name: "u".into(),
4040            role: role.into(),
4041        })
4042    }
4043
4044    #[test]
4045    fn readonly_can_read_but_not_write() {
4046        let p = principal("readonly");
4047        // Reads pass.
4048        assert!(check_statement_permitted(p.as_ref(), &parsed("User")).is_ok());
4049        assert!(check_statement_permitted(p.as_ref(), &parsed("count(User)")).is_ok());
4050        assert!(check_statement_permitted(p.as_ref(), &parsed("explain User")).is_ok());
4051        // Writes, DDL, and transaction control are denied.
4052        for q in [
4053            r#"insert User { name := "x" }"#,
4054            "User filter .id = 1 update { age := 2 }",
4055            "User filter .id = 1 delete",
4056            "drop User",
4057            "alter User add column c: str",
4058            "type T { required id: int }",
4059            "begin",
4060            "commit",
4061            "rollback",
4062        ] {
4063            let err = check_statement_permitted(p.as_ref(), &parsed(q))
4064                .expect_err(&format!("must deny: {q}"));
4065            assert!(
4066                err.to_string().contains("permission denied"),
4067                "unexpected error for {q}: {err}"
4068            );
4069        }
4070    }
4071
4072    #[test]
4073    fn readwrite_and_admin_have_full_query_access() {
4074        for role in ["readwrite", "admin"] {
4075            let p = principal(role);
4076            assert!(check_statement_permitted(p.as_ref(), &parsed("User")).is_ok());
4077            assert!(check_statement_permitted(
4078                p.as_ref(),
4079                &parsed(r#"insert User { name := "x" }"#)
4080            )
4081            .is_ok());
4082            assert!(check_statement_permitted(p.as_ref(), &parsed("drop User")).is_ok());
4083        }
4084    }
4085
4086    #[test]
4087    fn unknown_role_fails_closed_for_writes() {
4088        let p = principal("mystery");
4089        assert!(check_statement_permitted(p.as_ref(), &parsed("User")).is_ok());
4090        assert!(
4091            check_statement_permitted(p.as_ref(), &parsed(r#"insert User { name := "x" }"#))
4092                .is_err()
4093        );
4094    }
4095
4096    #[test]
4097    fn no_principal_means_full_access() {
4098        // Shared-password / open mode: no per-user identity, no restriction.
4099        assert!(check_statement_permitted(None, &parsed("drop User")).is_ok());
4100        assert!(check_statement_permitted(None, &parsed(r#"insert User { name := "x" }"#)).is_ok());
4101    }
4102
4103    fn store_with_alice() -> UserStore {
4104        let mut s = UserStore::new();
4105        s.create_user("alice", "pw", "readwrite").unwrap();
4106        s
4107    }
4108
4109    // ---- Empty store: legacy shared-password fallback ----
4110
4111    #[test]
4112    fn empty_store_no_password_is_open() {
4113        let s = UserStore::new();
4114        assert_eq!(
4115            authenticate_connect(&s, None, None, None),
4116            AuthOutcome::Authenticated { principal: None }
4117        );
4118        // Even a stray username/password is accepted (legacy open behavior).
4119        assert_eq!(
4120            authenticate_connect(&s, None, Some("x"), Some("y")),
4121            AuthOutcome::Authenticated { principal: None }
4122        );
4123    }
4124
4125    #[test]
4126    fn empty_store_correct_shared_password_succeeds() {
4127        let s = UserStore::new();
4128        assert_eq!(
4129            authenticate_connect(&s, Some("pw"), None, Some("pw")),
4130            AuthOutcome::Authenticated { principal: None }
4131        );
4132    }
4133
4134    #[test]
4135    fn empty_store_wrong_shared_password_rejected() {
4136        let s = UserStore::new();
4137        assert_eq!(
4138            authenticate_connect(&s, Some("pw"), None, Some("bad")),
4139            AuthOutcome::Rejected
4140        );
4141    }
4142
4143    #[test]
4144    fn empty_store_missing_password_rejected_when_expected() {
4145        let s = UserStore::new();
4146        assert_eq!(
4147            authenticate_connect(&s, Some("pw"), None, None),
4148            AuthOutcome::Rejected
4149        );
4150    }
4151
4152    #[test]
4153    fn empty_store_ignores_username_for_shared_password() {
4154        // A new client may send a username even against a shared-password
4155        // server; the username is ignored and the password still governs.
4156        let s = UserStore::new();
4157        assert_eq!(
4158            authenticate_connect(&s, Some("pw"), Some("whoever"), Some("pw")),
4159            AuthOutcome::Authenticated { principal: None }
4160        );
4161    }
4162
4163    // ---- Populated store: multi-user auth ----
4164
4165    #[test]
4166    fn user_auth_success_binds_principal() {
4167        let s = store_with_alice();
4168        assert_eq!(
4169            authenticate_connect(&s, None, Some("alice"), Some("pw")),
4170            AuthOutcome::Authenticated {
4171                principal: Some(Principal {
4172                    name: "alice".into(),
4173                    role: "readwrite".into(),
4174                })
4175            }
4176        );
4177    }
4178
4179    #[test]
4180    fn user_auth_wrong_password_rejected() {
4181        let s = store_with_alice();
4182        assert_eq!(
4183            authenticate_connect(&s, None, Some("alice"), Some("bad")),
4184            AuthOutcome::Rejected
4185        );
4186    }
4187
4188    #[test]
4189    fn user_auth_unknown_user_rejected() {
4190        let s = store_with_alice();
4191        assert_eq!(
4192            authenticate_connect(&s, None, Some("mallory"), Some("pw")),
4193            AuthOutcome::Rejected
4194        );
4195    }
4196
4197    #[test]
4198    fn user_auth_missing_username_rejected() {
4199        let s = store_with_alice();
4200        assert_eq!(
4201            authenticate_connect(&s, None, None, Some("pw")),
4202            AuthOutcome::Rejected
4203        );
4204    }
4205
4206    #[test]
4207    fn user_auth_missing_password_rejected() {
4208        let s = store_with_alice();
4209        assert_eq!(
4210            authenticate_connect(&s, Some("pw"), Some("alice"), None),
4211            AuthOutcome::Rejected
4212        );
4213    }
4214
4215    #[test]
4216    fn user_auth_ignores_shared_password_when_users_present() {
4217        // With users present, the shared password is irrelevant: supplying it as
4218        // the password without a valid user must NOT authenticate.
4219        let s = store_with_alice();
4220        assert_eq!(
4221            authenticate_connect(&s, Some("shared"), None, Some("shared")),
4222            AuthOutcome::Rejected
4223        );
4224    }
4225
4226    #[test]
4227    fn replica_fingerprint_is_stable_and_redacted() {
4228        let replica_id = "customer-prod-replica-a";
4229        let fingerprint = replica_fingerprint(replica_id);
4230        assert_eq!(fingerprint, replica_fingerprint(replica_id));
4231        assert_eq!(fingerprint, log_replica_fingerprint(replica_id));
4232        assert_ne!(fingerprint, replica_fingerprint("customer-prod-replica-b"));
4233        assert_eq!(fingerprint.len(), 16);
4234        assert!(fingerprint.chars().all(|c| c.is_ascii_hexdigit()));
4235        assert!(!fingerprint.contains("customer"));
4236        assert!(!fingerprint.contains("replica"));
4237        assert!(!fingerprint.contains(replica_id));
4238    }
4239
4240    #[test]
4241    fn invalid_replica_ids_use_fixed_log_fingerprint() {
4242        assert_eq!(log_replica_fingerprint(""), INVALID_REPLICA_FINGERPRINT);
4243        assert_eq!(
4244            log_replica_fingerprint("customer/prod/replica"),
4245            INVALID_REPLICA_FINGERPRINT
4246        );
4247        assert_eq!(
4248            log_replica_fingerprint(&"a".repeat(4096)),
4249            INVALID_REPLICA_FINGERPRINT
4250        );
4251    }
4252
4253    #[test]
4254    fn sync_error_classes_use_bounded_labels() {
4255        assert_eq!(SyncErrorClass::AuthRequired.as_label(), "auth_required");
4256        assert_eq!(
4257            SyncErrorClass::PermissionDenied.as_label(),
4258            "permission_denied"
4259        );
4260        assert_eq!(
4261            SyncErrorClass::IdentityOrFormatMismatch.as_label(),
4262            "identity_or_format_mismatch"
4263        );
4264        assert_eq!(SyncErrorClass::AckValidation.as_label(), "ack_validation");
4265        assert_eq!(SyncErrorClass::Internal.as_label(), "internal");
4266    }
4267
4268    fn sync_identity() -> DatabaseIdentity {
4269        DatabaseIdentity {
4270            database_id: *b"server-sync-test",
4271            primary_generation: 1,
4272        }
4273    }
4274
4275    fn retained_unit(lsn: u64) -> RetainedUnit {
4276        RetainedUnit {
4277            tx_id: 1,
4278            record_type: 4,
4279            lsn,
4280            data: lsn.to_le_bytes().to_vec(),
4281        }
4282    }
4283
4284    fn retained_unit_with(tx_id: u64, record_type: WalRecordType, lsn: u64) -> RetainedUnit {
4285        RetainedUnit {
4286            tx_id,
4287            record_type: record_type as u8,
4288            lsn,
4289            data: lsn.to_le_bytes().to_vec(),
4290        }
4291    }
4292
4293    fn write_sync_identity_and_tail(data_dir: &std::path::Path, through_lsn: u64) {
4294        let identity = sync_identity();
4295        write_identity_snapshot(data_dir, &IdentitySnapshot::from_identity(identity, 1)).unwrap();
4296        let units = (1..=through_lsn).map(retained_unit).collect();
4297        let segment = RetainedSegment::new(identity.segment_identity(), units).unwrap();
4298        write_segment_atomic(&retained_segments_dir(data_dir), &segment).unwrap();
4299    }
4300
4301    fn write_sync_identity_and_units(data_dir: &std::path::Path, units: Vec<RetainedUnit>) {
4302        let identity = sync_identity();
4303        write_identity_snapshot(data_dir, &IdentitySnapshot::from_identity(identity, 1)).unwrap();
4304        let segment = RetainedSegment::new(identity.segment_identity(), units).unwrap();
4305        write_segment_atomic(&retained_segments_dir(data_dir), &segment).unwrap();
4306    }
4307
4308    fn write_sync_identity_only(data_dir: &std::path::Path) {
4309        let identity = sync_identity();
4310        write_identity_snapshot(data_dir, &IdentitySnapshot::from_identity(identity, 1)).unwrap();
4311    }
4312
4313    fn admin_principal() -> Principal {
4314        Principal {
4315            name: "admin".into(),
4316            role: "admin".into(),
4317        }
4318    }
4319
4320    #[test]
4321    fn sync_protocol_requires_credential_auth_and_rejects_readonly() {
4322        let dir = tempfile::tempdir().unwrap();
4323        let engine = Arc::new(RwLock::new(Engine::new(dir.path()).unwrap()));
4324
4325        match dispatch_sync_status(&engine, "replica-a".into(), false, None) {
4326            Message::Error { message } => {
4327                assert!(message.contains("requires authentication"));
4328            }
4329            other => panic!("expected auth error, got {other:?}"),
4330        }
4331
4332        let readonly = Principal {
4333            name: "reader".into(),
4334            role: "readonly".into(),
4335        };
4336        match dispatch_sync_status(&engine, "replica-a".into(), true, Some(&readonly)) {
4337            Message::Error { message } => {
4338                assert!(message.contains("permission denied"));
4339            }
4340            other => panic!("expected permission error, got {other:?}"),
4341        }
4342    }
4343
4344    #[test]
4345    fn sync_status_pull_and_ack_use_server_remote_lsn() {
4346        let dir = tempfile::tempdir().unwrap();
4347        let mut engine = Engine::new(dir.path()).unwrap();
4348        engine
4349            .execute_powql("type SyncT { required id: int, v: str }")
4350            .unwrap();
4351        engine
4352            .execute_powql(r#"insert SyncT { id := 1, v := "one" }"#)
4353            .unwrap();
4354        let remote_lsn = engine.catalog().max_lsn();
4355        assert!(remote_lsn > 0);
4356        write_sync_identity_and_tail(dir.path(), remote_lsn);
4357        powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
4358            .unwrap();
4359
4360        let engine = Arc::new(RwLock::new(engine));
4361        let principal = admin_principal();
4362        let status = match dispatch_sync_status(&engine, "replica-a".into(), true, Some(&principal))
4363        {
4364            Message::SyncStatusResult { status } => status,
4365            other => panic!("expected sync status, got {other:?}"),
4366        };
4367        assert_eq!(status.remote_lsn, remote_lsn);
4368        assert_eq!(status.servable_lsn, Some(remote_lsn));
4369        assert_eq!(status.unarchived_lsn, Some(0));
4370        assert_eq!(status.last_applied_lsn, Some(0));
4371        assert_eq!(status.repair_action, WireSyncRepairAction::Pull);
4372        assert!(status.stale);
4373
4374        let identity = sync_identity().segment_identity();
4375        let pull = SyncPullRequest {
4376            replica_id: "replica-a".into(),
4377            since_lsn: 0,
4378            max_units: MAX_SYNC_PULL_UNITS,
4379            max_bytes: MAX_SYNC_PULL_BYTES,
4380            database_id: identity.database_id,
4381            primary_generation: identity.primary_generation,
4382            wal_format_version: identity.wal_format_version,
4383            catalog_version: identity.catalog_version,
4384            segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
4385        };
4386        let units = match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
4387            Message::SyncPullResult {
4388                status,
4389                units,
4390                has_more,
4391            } => {
4392                assert_eq!(status.repair_action, WireSyncRepairAction::Pull);
4393                assert!(!has_more);
4394                units
4395            }
4396            other => panic!("expected sync pull result, got {other:?}"),
4397        };
4398        assert_eq!(units.len() as u64, remote_lsn);
4399        assert_eq!(units.last().unwrap().lsn, remote_lsn);
4400
4401        let ack = match dispatch_sync_ack(
4402            &engine,
4403            "replica-a".into(),
4404            remote_lsn,
4405            remote_lsn,
4406            true,
4407            Some(&principal),
4408        ) {
4409            Message::SyncAckResult {
4410                previous_applied_lsn,
4411                applied_lsn,
4412                remote_lsn: ack_remote_lsn,
4413                advanced,
4414                status,
4415            } => {
4416                assert_eq!(previous_applied_lsn, 0);
4417                assert_eq!(applied_lsn, remote_lsn);
4418                assert_eq!(ack_remote_lsn, remote_lsn);
4419                assert!(advanced);
4420                status
4421            }
4422            other => panic!("expected sync ack result, got {other:?}"),
4423        };
4424        assert_eq!(ack.repair_action, WireSyncRepairAction::None);
4425        assert!(!ack.stale);
4426        assert_eq!(ack.lag_lsn, Some(0));
4427    }
4428
4429    fn seed_pullable_replica(engine: &mut Engine) -> u64 {
4430        let data_dir = engine.catalog().data_dir().to_path_buf();
4431        let remote_lsn = engine.catalog().max_lsn();
4432        assert!(remote_lsn > 0);
4433        write_sync_identity_and_tail(&data_dir, remote_lsn);
4434        powdb_sync::upsert_replica_cursor(&data_dir, ReplicaCursor::active("replica-a", 0))
4435            .unwrap();
4436        remote_lsn
4437    }
4438
4439    fn pull_request_with_catalog_version(catalog_version: u16) -> SyncPullRequest {
4440        let identity = sync_identity().segment_identity();
4441        SyncPullRequest {
4442            replica_id: "replica-a".into(),
4443            since_lsn: 0,
4444            max_units: MAX_SYNC_PULL_UNITS,
4445            max_bytes: MAX_SYNC_PULL_BYTES,
4446            database_id: identity.database_id,
4447            primary_generation: identity.primary_generation,
4448            wal_format_version: identity.wal_format_version,
4449            catalog_version,
4450            segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
4451        }
4452    }
4453
4454    #[test]
4455    fn fresh_database_expects_legacy_catalog_version_and_accepts_v5_replica() {
4456        use powdb_storage::catalog::{CATALOG_VERSION, LEGACY_CATALOG_VERSION};
4457
4458        let dir = tempfile::tempdir().unwrap();
4459        let mut engine = Engine::new(dir.path()).unwrap();
4460        engine
4461            .execute_powql("type Doc { required id: int, data: json }")
4462            .unwrap();
4463        engine
4464            .execute_powql(r#"insert Doc { id := 1, data := "{\"score\":20}" }"#)
4465            .unwrap();
4466        // No expression index created yet: the database stays at the legacy
4467        // catalog format, exactly as a v0.12 database on disk.
4468        assert_eq!(
4469            engine.catalog().active_catalog_version(),
4470            LEGACY_CATALOG_VERSION
4471        );
4472        let remote_lsn = seed_pullable_replica(&mut engine);
4473
4474        let engine = Arc::new(RwLock::new(engine));
4475        let principal = admin_principal();
4476
4477        // A replica whose maximum is the legacy version (as v0.12 clients state)
4478        // is accepted against a legacy-active server.
4479        let pull = pull_request_with_catalog_version(LEGACY_CATALOG_VERSION);
4480        match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
4481            Message::SyncPullResult { units, .. } => {
4482                assert_eq!(units.len() as u64, remote_lsn);
4483            }
4484            other => panic!("expected sync pull result, got {other:?}"),
4485        }
4486
4487        // A newer replica (states this binary's max) is also accepted.
4488        let pull = pull_request_with_catalog_version(CATALOG_VERSION);
4489        assert!(matches!(
4490            dispatch_sync_pull(&engine, pull, true, Some(&principal)),
4491            Message::SyncPullResult { .. }
4492        ));
4493
4494        // A replica whose maximum is older than the active format is rejected
4495        // with a message naming both versions.
4496        let pull = pull_request_with_catalog_version(LEGACY_CATALOG_VERSION - 1);
4497        match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
4498            Message::Error { message } => {
4499                assert!(message.contains("v4"), "message: {message}");
4500                assert!(message.contains("v5"), "message: {message}");
4501                assert!(
4502                    message.contains("rebootstrap with an upgraded replica required"),
4503                    "message: {message}"
4504                );
4505            }
4506            other => panic!("expected identity mismatch error, got {other:?}"),
4507        }
4508    }
4509
4510    #[test]
4511    fn activated_database_expects_v6_and_rejects_v5_replica() {
4512        use powdb_storage::catalog::{CATALOG_VERSION, LEGACY_CATALOG_VERSION};
4513
4514        let dir = tempfile::tempdir().unwrap();
4515        let mut engine = Engine::new(dir.path()).unwrap();
4516        engine
4517            .execute_powql("type Doc { required id: int, data: json }")
4518            .unwrap();
4519        engine
4520            .execute_powql(r#"insert Doc { id := 1, data := "{\"score\":20}" }"#)
4521            .unwrap();
4522        // Creating a JSON-path expression index activates the v6 catalog format.
4523        engine
4524            .execute_powql("alter Doc add index (.data->score)")
4525            .unwrap();
4526        assert_eq!(engine.catalog().active_catalog_version(), CATALOG_VERSION);
4527        let remote_lsn = seed_pullable_replica(&mut engine);
4528
4529        let engine = Arc::new(RwLock::new(engine));
4530        let principal = admin_principal();
4531
4532        // A v0.12 replica (states catalog_version 5) genuinely cannot read the
4533        // now-activated v6 data and is rejected with the targeted message.
4534        let pull = pull_request_with_catalog_version(LEGACY_CATALOG_VERSION);
4535        match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
4536            Message::Error { message } => {
4537                assert!(message.contains("v5"), "message: {message}");
4538                assert!(message.contains("v6"), "message: {message}");
4539                assert!(
4540                    message.contains("rebootstrap with an upgraded replica required"),
4541                    "message: {message}"
4542                );
4543            }
4544            other => panic!("expected identity mismatch error, got {other:?}"),
4545        }
4546
4547        // A v6-capable replica is accepted.
4548        let pull = pull_request_with_catalog_version(CATALOG_VERSION);
4549        match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
4550            Message::SyncPullResult { units, .. } => {
4551                assert_eq!(units.len() as u64, remote_lsn);
4552            }
4553            other => panic!("expected sync pull result, got {other:?}"),
4554        }
4555    }
4556
4557    #[test]
4558    fn sync_pull_and_ack_reject_transaction_cut_boundaries() {
4559        let dir = tempfile::tempdir().unwrap();
4560        let mut engine = Engine::new(dir.path()).unwrap();
4561        engine
4562            .execute_powql("type SyncT { required id: int }")
4563            .unwrap();
4564        for id in 1..=3 {
4565            engine
4566                .execute_powql(&format!("insert SyncT {{ id := {id} }}"))
4567                .unwrap();
4568        }
4569        let remote_lsn = engine.catalog().max_lsn();
4570        assert!(remote_lsn >= 3);
4571        write_sync_identity_and_units(
4572            dir.path(),
4573            vec![
4574                retained_unit_with(77, WalRecordType::Begin, 1),
4575                retained_unit_with(77, WalRecordType::Insert, 2),
4576                retained_unit_with(77, WalRecordType::Commit, 3),
4577            ],
4578        );
4579        powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
4580            .unwrap();
4581
4582        let engine = Arc::new(RwLock::new(engine));
4583        let principal = admin_principal();
4584        let identity = sync_identity().segment_identity();
4585        let cut_pull = SyncPullRequest {
4586            replica_id: "replica-a".into(),
4587            since_lsn: 0,
4588            max_units: 2,
4589            max_bytes: MAX_SYNC_PULL_BYTES,
4590            database_id: identity.database_id,
4591            primary_generation: identity.primary_generation,
4592            wal_format_version: identity.wal_format_version,
4593            catalog_version: identity.catalog_version,
4594            segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
4595        };
4596        match dispatch_sync_pull(&engine, cut_pull, true, Some(&principal)) {
4597            Message::Error { message } => assert!(message.contains("cuts through transaction")),
4598            other => panic!("expected transaction-cut pull error, got {other:?}"),
4599        }
4600
4601        let cut_bytes_pull = SyncPullRequest {
4602            replica_id: "replica-a".into(),
4603            since_lsn: 0,
4604            max_units: 3,
4605            max_bytes: 58,
4606            database_id: identity.database_id,
4607            primary_generation: identity.primary_generation,
4608            wal_format_version: identity.wal_format_version,
4609            catalog_version: identity.catalog_version,
4610            segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
4611        };
4612        match dispatch_sync_pull(&engine, cut_bytes_pull, true, Some(&principal)) {
4613            Message::Error { message } => assert!(message.contains("cuts through transaction")),
4614            other => panic!("expected byte-capped transaction-cut pull error, got {other:?}"),
4615        }
4616
4617        let full_pull = SyncPullRequest {
4618            replica_id: "replica-a".into(),
4619            since_lsn: 0,
4620            max_units: 3,
4621            max_bytes: MAX_SYNC_PULL_BYTES,
4622            database_id: identity.database_id,
4623            primary_generation: identity.primary_generation,
4624            wal_format_version: identity.wal_format_version,
4625            catalog_version: identity.catalog_version,
4626            segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
4627        };
4628        match dispatch_sync_pull(&engine, full_pull, true, Some(&principal)) {
4629            Message::SyncPullResult { units, .. } => {
4630                assert_eq!(units.len(), 3);
4631                assert_eq!(units.last().unwrap().lsn, 3);
4632            }
4633            other => panic!("expected complete transaction pull, got {other:?}"),
4634        }
4635
4636        match dispatch_sync_ack(
4637            &engine,
4638            "replica-a".into(),
4639            2,
4640            remote_lsn,
4641            true,
4642            Some(&principal),
4643        ) {
4644            Message::Error { message } => assert!(message.contains("cuts through transaction")),
4645            other => panic!("expected transaction-cut ack error, got {other:?}"),
4646        }
4647        let cursor = powdb_sync::read_replica_cursors(dir.path()).unwrap();
4648        assert_eq!(cursor[0].applied_lsn, 0);
4649
4650        match dispatch_sync_ack(
4651            &engine,
4652            "replica-a".into(),
4653            3,
4654            remote_lsn,
4655            true,
4656            Some(&principal),
4657        ) {
4658            Message::SyncAckResult {
4659                previous_applied_lsn,
4660                applied_lsn,
4661                advanced,
4662                ..
4663            } => {
4664                assert_eq!(previous_applied_lsn, 0);
4665                assert_eq!(applied_lsn, 3);
4666                assert!(advanced);
4667            }
4668            other => panic!("expected complete transaction ack, got {other:?}"),
4669        }
4670    }
4671
4672    #[test]
4673    fn sync_pull_byte_cap_returns_applyable_prefix_with_reused_tx_id() {
4674        let dir = tempfile::tempdir().unwrap();
4675        let mut engine = Engine::new(dir.path()).unwrap();
4676        engine
4677            .execute_powql("type SyncT { required id: int }")
4678            .unwrap();
4679        for id in 1..=6 {
4680            engine
4681                .execute_powql(&format!("insert SyncT {{ id := {id} }}"))
4682                .unwrap();
4683        }
4684        let remote_lsn = engine.catalog().max_lsn();
4685        assert!(remote_lsn >= 6);
4686        write_sync_identity_and_units(
4687            dir.path(),
4688            vec![
4689                retained_unit_with(1, WalRecordType::Begin, 1),
4690                retained_unit_with(1, WalRecordType::Insert, 2),
4691                retained_unit_with(1, WalRecordType::Commit, 3),
4692                retained_unit_with(1, WalRecordType::Begin, 4),
4693                retained_unit_with(1, WalRecordType::Insert, 5),
4694                retained_unit_with(1, WalRecordType::Commit, 6),
4695            ],
4696        );
4697        powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
4698            .unwrap();
4699
4700        let engine = Arc::new(RwLock::new(engine));
4701        let principal = admin_principal();
4702        let identity = sync_identity().segment_identity();
4703        let pull = SyncPullRequest {
4704            replica_id: "replica-a".into(),
4705            since_lsn: 0,
4706            max_units: 6,
4707            max_bytes: 100,
4708            database_id: identity.database_id,
4709            primary_generation: identity.primary_generation,
4710            wal_format_version: identity.wal_format_version,
4711            catalog_version: identity.catalog_version,
4712            segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
4713        };
4714
4715        match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
4716            Message::SyncPullResult {
4717                status,
4718                units,
4719                has_more,
4720            } => {
4721                assert_eq!(status.repair_action, WireSyncRepairAction::Pull);
4722                assert_eq!(units.len(), 3);
4723                assert_eq!(units.last().unwrap().lsn, 3);
4724                assert!(has_more);
4725            }
4726            other => panic!("expected byte-capped applyable prefix, got {other:?}"),
4727        }
4728    }
4729
4730    #[test]
4731    fn sync_pull_never_serves_units_beyond_server_remote_lsn() {
4732        let dir = tempfile::tempdir().unwrap();
4733        let mut engine = Engine::new(dir.path()).unwrap();
4734        engine
4735            .execute_powql("type SyncT { required id: int }")
4736            .unwrap();
4737        engine.execute_powql("insert SyncT { id := 1 }").unwrap();
4738        let remote_lsn = engine.catalog().max_lsn();
4739        assert!(remote_lsn > 0);
4740        write_sync_identity_and_tail(dir.path(), remote_lsn + 2);
4741        powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
4742            .unwrap();
4743
4744        let engine = Arc::new(RwLock::new(engine));
4745        let principal = admin_principal();
4746        let identity = sync_identity().segment_identity();
4747        let pull = SyncPullRequest {
4748            replica_id: "replica-a".into(),
4749            since_lsn: 0,
4750            max_units: MAX_SYNC_PULL_UNITS,
4751            max_bytes: MAX_SYNC_PULL_BYTES,
4752            database_id: identity.database_id,
4753            primary_generation: identity.primary_generation,
4754            wal_format_version: identity.wal_format_version,
4755            catalog_version: identity.catalog_version,
4756            segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
4757        };
4758
4759        match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
4760            Message::SyncPullResult {
4761                status,
4762                units,
4763                has_more,
4764            } => {
4765                assert_eq!(status.remote_lsn, remote_lsn);
4766                assert_eq!(status.servable_lsn, Some(remote_lsn));
4767                assert_eq!(status.unarchived_lsn, Some(0));
4768                assert_eq!(status.repair_action, WireSyncRepairAction::Pull);
4769                assert!(!has_more);
4770                assert_eq!(units.len() as u64, remote_lsn);
4771                assert_eq!(units.last().unwrap().lsn, remote_lsn);
4772                assert!(units.iter().all(|unit| unit.lsn <= remote_lsn));
4773            }
4774            other => panic!("expected capped sync pull result, got {other:?}"),
4775        }
4776    }
4777
4778    #[test]
4779    fn sync_status_reports_await_archive_when_primary_outruns_retained_tail() {
4780        let dir = tempfile::tempdir().unwrap();
4781        let mut engine = Engine::new(dir.path()).unwrap();
4782        engine
4783            .execute_powql("type SyncT { required id: int }")
4784            .unwrap();
4785        engine.execute_powql("insert SyncT { id := 1 }").unwrap();
4786        let remote_lsn = engine.catalog().max_lsn();
4787        assert!(remote_lsn > 0);
4788        write_sync_identity_only(dir.path());
4789        powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
4790            .unwrap();
4791
4792        let engine = Arc::new(RwLock::new(engine));
4793        let principal = admin_principal();
4794        let identity = sync_identity().segment_identity();
4795        let status = match dispatch_sync_status(&engine, "replica-a".into(), true, Some(&principal))
4796        {
4797            Message::SyncStatusResult { status } => status,
4798            other => panic!("expected sync status, got {other:?}"),
4799        };
4800        assert_eq!(status.remote_lsn, remote_lsn);
4801        assert_eq!(status.servable_lsn, Some(0));
4802        assert_eq!(status.unarchived_lsn, Some(remote_lsn));
4803        assert_eq!(status.repair_action, WireSyncRepairAction::AwaitArchive);
4804        assert!(status
4805            .last_sync_error
4806            .as_deref()
4807            .unwrap()
4808            .contains("not yet archived"));
4809
4810        let pull = SyncPullRequest {
4811            replica_id: "replica-a".into(),
4812            since_lsn: 0,
4813            max_units: MAX_SYNC_PULL_UNITS,
4814            max_bytes: MAX_SYNC_PULL_BYTES,
4815            database_id: identity.database_id,
4816            primary_generation: identity.primary_generation,
4817            wal_format_version: identity.wal_format_version,
4818            catalog_version: identity.catalog_version,
4819            segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
4820        };
4821        match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
4822            Message::SyncPullResult {
4823                status,
4824                units,
4825                has_more,
4826            } => {
4827                assert_eq!(status.repair_action, WireSyncRepairAction::AwaitArchive);
4828                assert!(units.is_empty());
4829                assert!(!has_more);
4830            }
4831            other => panic!("expected await-archive sync pull result, got {other:?}"),
4832        }
4833    }
4834
4835    #[test]
4836    fn sync_pull_serves_partial_retained_prefix_when_archive_lags_remote_lsn() {
4837        let dir = tempfile::tempdir().unwrap();
4838        let mut engine = Engine::new(dir.path()).unwrap();
4839        engine
4840            .execute_powql("type SyncT { required id: int }")
4841            .unwrap();
4842        engine.execute_powql("insert SyncT { id := 1 }").unwrap();
4843        engine.execute_powql("insert SyncT { id := 2 }").unwrap();
4844        let remote_lsn = engine.catalog().max_lsn();
4845        assert!(remote_lsn > 1);
4846        let servable_lsn = remote_lsn - 1;
4847        write_sync_identity_and_tail(dir.path(), servable_lsn);
4848        powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
4849            .unwrap();
4850
4851        let engine = Arc::new(RwLock::new(engine));
4852        let principal = admin_principal();
4853        let identity = sync_identity().segment_identity();
4854        let pull = SyncPullRequest {
4855            replica_id: "replica-a".into(),
4856            since_lsn: 0,
4857            max_units: MAX_SYNC_PULL_UNITS,
4858            max_bytes: MAX_SYNC_PULL_BYTES,
4859            database_id: identity.database_id,
4860            primary_generation: identity.primary_generation,
4861            wal_format_version: identity.wal_format_version,
4862            catalog_version: identity.catalog_version,
4863            segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
4864        };
4865
4866        match dispatch_sync_pull(&engine, pull, true, Some(&principal)) {
4867            Message::SyncPullResult {
4868                status,
4869                units,
4870                has_more,
4871            } => {
4872                assert_eq!(status.remote_lsn, remote_lsn);
4873                assert_eq!(status.servable_lsn, Some(servable_lsn));
4874                assert_eq!(status.unarchived_lsn, Some(1));
4875                assert_eq!(status.repair_action, WireSyncRepairAction::Pull);
4876                assert!(!has_more);
4877                assert_eq!(units.len() as u64, servable_lsn);
4878                assert_eq!(units.last().unwrap().lsn, servable_lsn);
4879                assert!(units.iter().all(|unit| unit.lsn <= servable_lsn));
4880            }
4881            other => panic!("expected partial sync pull result, got {other:?}"),
4882        }
4883    }
4884
4885    #[test]
4886    fn sync_pull_rejects_cursor_or_format_mismatch() {
4887        let dir = tempfile::tempdir().unwrap();
4888        let mut engine = Engine::new(dir.path()).unwrap();
4889        engine
4890            .execute_powql("type SyncT { required id: int }")
4891            .unwrap();
4892        engine.execute_powql("insert SyncT { id := 1 }").unwrap();
4893        let remote_lsn = engine.catalog().max_lsn();
4894        write_sync_identity_and_tail(dir.path(), remote_lsn);
4895        powdb_sync::upsert_replica_cursor(dir.path(), ReplicaCursor::active("replica-a", 0))
4896            .unwrap();
4897        let engine = Arc::new(RwLock::new(engine));
4898        let principal = admin_principal();
4899        let identity = sync_identity().segment_identity();
4900
4901        let wrong_cursor = SyncPullRequest {
4902            replica_id: "replica-a".into(),
4903            since_lsn: 1,
4904            max_units: 10,
4905            max_bytes: MAX_SYNC_PULL_BYTES,
4906            database_id: identity.database_id,
4907            primary_generation: identity.primary_generation,
4908            wal_format_version: identity.wal_format_version,
4909            catalog_version: identity.catalog_version,
4910            segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION,
4911        };
4912        match dispatch_sync_pull(&engine, wrong_cursor, true, Some(&principal)) {
4913            Message::Error { message } => assert!(message.contains("does not match")),
4914            other => panic!("expected cursor mismatch error, got {other:?}"),
4915        }
4916
4917        let wrong_format = SyncPullRequest {
4918            replica_id: "replica-a".into(),
4919            since_lsn: 0,
4920            max_units: 10,
4921            max_bytes: MAX_SYNC_PULL_BYTES,
4922            database_id: identity.database_id,
4923            primary_generation: identity.primary_generation,
4924            wal_format_version: identity.wal_format_version,
4925            catalog_version: identity.catalog_version,
4926            segment_format_version: RETAINED_SEGMENT_FORMAT_VERSION + 1,
4927        };
4928        match dispatch_sync_pull(&engine, wrong_format, true, Some(&principal)) {
4929            Message::Error { message } => assert!(message.contains("rebootstrap required")),
4930            other => panic!("expected format mismatch error, got {other:?}"),
4931        }
4932    }
4933}