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
24pub type AuthRateLimiter = Arc<Mutex<HashMap<IpAddr, (u32, Instant)>>>;
26
27pub type TxGate = Arc<Semaphore>;
32
33pub fn new_tx_gate() -> TxGate {
35 Arc::new(Semaphore::new(1))
36}
37
38const MAX_QUERY_LENGTH: usize = 1024 * 1024;
40
41const MAX_SYNC_PULL_UNITS: u32 = 4096;
43
44const MAX_SYNC_PULL_BYTES: u64 = 16 * 1024 * 1024;
46
47#[cfg(not(test))]
51const MAX_RESPONSE_PAYLOAD_SIZE: usize = 64 * 1024 * 1024;
52#[cfg(test)]
53const MAX_RESPONSE_PAYLOAD_SIZE: usize = 1024;
54
55const WRITE_TIMEOUT: Duration = Duration::from_secs(30);
58
59const MAX_AUTH_FAILURES: u32 = 5;
61
62const AUTH_FAILURE_WINDOW: Duration = Duration::from_secs(60);
64
65pub fn new_rate_limiter() -> AuthRateLimiter {
67 Arc::new(Mutex::new(HashMap::new()))
68}
69
70fn is_rate_limited(limiter: &AuthRateLimiter, ip: IpAddr) -> bool {
73 let mut map = limiter.lock().unwrap_or_else(|e| e.into_inner());
74 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
85fn 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 if now.duration_since(entry.1) >= AUTH_FAILURE_WINDOW {
92 *entry = (1, now);
93 } else {
94 entry.0 += 1;
95 }
96}
97
98fn 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
104fn 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#[derive(Debug, Clone, PartialEq, Eq)]
121pub struct Principal {
122 pub name: String,
123 pub role: String,
124}
125
126fn 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
147fn 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
162fn check_statement_permitted(
174 principal: Option<&Principal>,
175 stmt: &powdb_query::ast::Statement,
176) -> Result<(), QueryError> {
177 let Some(p) = principal else {
178 return Ok(());
181 };
182 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#[derive(Debug, Clone, PartialEq, Eq)]
204pub enum AuthOutcome {
205 Authenticated { principal: Option<Principal> },
209 Rejected,
212}
213
214pub 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 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 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
263const 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 "sort input exceeds",
289 "join result exceeds",
290 "query exceeded memory budget",
291 "result too large",
292];
293
294fn 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
307async 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
321pub struct ConnOpts<'a> {
324 pub engine: Arc<RwLock<Engine>>,
325 pub tx_gate: TxGate,
326 pub expected_password: Option<Zeroizing<String>>,
329 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 pub metrics: Arc<Metrics>,
340}
341
342fn 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 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 }
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
431fn 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#[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 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 }
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#[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, ¶ms) {
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, ¶ms, 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, ¶ms, 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, ¶ms, 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, ¶ms, 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 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 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 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 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 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 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 loop {
1907 let msg = tokio::select! {
1908 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 _ = 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 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; 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)?; 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
2210fn 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 #[test]
2229 fn null_serializes_as_null_bareword_on_wire() {
2230 assert_eq!(value_to_display(&Value::Empty), "null");
2231 }
2232
2233 #[test]
2236 fn unique_violation_error_surfaces_to_remote_clients() {
2237 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 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 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 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 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 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 #[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 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 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 #[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 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}