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