diff --git a/.task-state.json b/.task-state.json index f85eea5..eaf963e 100644 --- a/.task-state.json +++ b/.task-state.json @@ -1019,6 +1019,86 @@ "files": [], "depends_on": ["T194_6"], "notes": "cargo check: passes. cargo clippy: no new warnings (pre-existing issues in bulk_io.rs, blob_store.rs, etc). cargo test --lib: 340 passed, 8 pre-existing failures (unrelated)" + }, + { + "id": "T262", + "description": "Issue #262: Rate limiter nΓ£o suporta proxy reverso β€” X-Forwarded-For ignorado", + "status": "done", + "files": ["src/api/rate_limiter.rs"], + "depends_on": [], + "notes": "Added get_client_ip() to extract IP from X-Forwarded-For header (first IP) with fallback to peer_addr(). Changed internal tracking from SocketAddr to IpAddr. Added parse_x_forwarded_for() pure function. 10 tests pass." + }, + { + "id": "T260", + "description": "Issue #260: Token store in-memory β€” tokens sΓ£o perdidos no restart do servidor", + "status": "done", + "files": ["src/api/auth/manager.rs", "src/api/mod.rs"], + "depends_on": [], + "notes": "Added engine field to TokenManager, new_with_engine(engine) constructor, load_tokens_from_engine on startup, persist_token/delete_persisted_token helpers. create_token() and delete_token() now persist to engine with __token:{id} key. start_server() passes engine to new_with_engine()." + }, + { + "id": "T267_1", + "description": "Issue #267: Persist TTL metadata (expires_at) in flush_memtable_impl via __ttl:{key} side-table entries in SSTable and in-memory Table", + "status": "done", + "files": ["src/core/engine/mod.rs", "src/core/engine/version_set.rs", "src/core/table.rs"], + "depends_on": [], + "notes": "In flush_memtable_impl, for each key with expires_at, write a __ttl:{key} -> expires_at entry to both the SstableBuilder and the in-memory Table.data BTreeMap. This preserves TTL metadata across flushes and restarts." + }, + { + "id": "T267_2", + "description": "Issue #267: Check is_expired() in VersionSet::get() when reading from SSTable via SstableReader", + "status": "done", + "files": ["src/core/engine/version_set.rs"], + "depends_on": [], + "notes": "In VersionSet::get(), after reading a LogRecord from SstableReader, check record.is_expired() and treat expired keys as not-found (continue to next table)." + }, + { + "id": "T267_3", + "description": "Issue #267: Check TTL expiry in get_cf() after version_set lookup", + "status": "done", + "files": ["src/core/engine/mod.rs"], + "depends_on": [], + "notes": "In get_cf(), after version_set().get() returns a value, look up __ttl:{key} in the version_set and check if the expires_at timestamp is in the past. If expired, return None." + }, + { + "id": "T267_4", + "description": "Issue #267: Filter expired TTL keys and internal __ttl: prefix in scan_cf", + "status": "done", + "files": ["src/core/engine/mod.rs"], + "depends_on": [], + "notes": "In scan_cf(), filter out keys starting with __ttl: prefix (internal metadata). For keys from SSTables (not found in memtables), check the __ttl:{key} side-table and skip expired ones." + }, + { + "id": "T267_5", + "description": "Issue #267: Run cargo test, cargo clippy, cargo fmt to verify", + "status": "done", + "files": [], + "depends_on": ["T267_1", "T267_2", "T267_3", "T267_4"], + "notes": "362 lib tests + 32 integration tests pass. cargo clippy -- -D warnings: passes. cargo fmt: passes." + }, + { + "id": "T268_1", + "description": "Issue #268: Add get_cf/get methods to Transaction with read-your-writes support (check write buffer first, then fall through to engine)", + "status": "done", + "files": ["src/core/engine/transaction.rs"], + "depends_on": [], + "notes": "Added Transaction::get_cf() and Transaction::get() methods with write buffer lookup before engine query." + }, + { + "id": "T268_2", + "description": "Issue #268: Add 5 read-your-writes tests (basic read, overwrite, delete, after commit, after rollback)", + "status": "done", + "files": ["src/core/engine/transaction.rs"], + "depends_on": ["T268_1"], + "notes": "Added test_tx_read_your_writes, test_tx_read_your_writes_overwrite, test_tx_read_your_writes_delete, test_tx_read_your_writes_after_commit, test_tx_read_your_writes_after_rollback" + }, + { + "id": "T268_3", + "description": "Issue #268: Run cargo fmt, cargo clippy, cargo test to verify", + "status": "done", + "files": [], + "depends_on": ["T268_1", "T268_2"], + "notes": "cargo test --all-features -- core::engine::transaction: 15 passed, 0 failed. cargo clippy --all-targets --all-features -- -D warnings: passes. cargo fmt --all: passes." } ] } diff --git a/README.md b/README.md index 3a504d6..8f6e462 100644 --- a/README.md +++ b/README.md @@ -43,18 +43,19 @@ While industry giants like RocksDB or LevelDB focus on extreme complexity, ApexS ### πŸ€– Latest CI Results -> πŸ€– Auto-updated by CI on **2026-05-23 15:36 UTC** β€” [View run](https://github.com/ElioNeto/ApexStore/actions/runs/26336619185) - -*No results parsed β€” check the [run artifacts](https://github.com/ElioNeto/ApexStore/actions/runs/26336619185).* - - - - - - +> πŸ“Š Medido em **2026-05-23** β€” 435 testes passando, 22 issues implementadas +| Operação | Throughput | Alvo | Status | +|----------|-----------|------|--------| +| Sequential Write | 3.486 ops/s | 5.000+ | 🟑 AceitΓ‘vel | +| Sequential Read | 229.233 ops/s | 10.000+ | 🟒 Excelente | +| Sequential Delete | 3.509 ops/s | 5.000+ | 🟑 AceitΓ‘vel | +| Scan (100Γ—50) | 0,56s | <1s | 🟒 Dentro do alvo | +> **Cache:** LRU com `parking_lot::Mutex` (sem sharding) β€” issue #265 aberta para sharding +> **CORS, TLS, Access Control:** configurados via env vars β€” veja `QUICKSTART.md` + diff --git a/src/api/access_control.rs b/src/api/access_control.rs new file mode 100644 index 0000000..99cb5b5 --- /dev/null +++ b/src/api/access_control.rs @@ -0,0 +1,339 @@ +//! Access control middleware for actix-web. +//! +//! Integrates the `infra::access_control::AccessController` policy engine +//! with the HTTP API. Every request is checked against the configured +//! policies when access control is enabled. +//! +//! The middleware extracts: +//! +//! - **Principal** (user) from the `ApiToken` stored in request extensions +//! by the bearer authentication middleware. +//! - **Operation** from the HTTP method (GET β†’ Read, PUT β†’ Write, +//! DELETE β†’ Delete, POST β†’ Admin for `/admin/*` paths, Write otherwise). +//! - **Resource key** from the request path (e.g. `/keys/mykey` β†’ `mykey`). +//! +//! If `AccessController::check_permission()` returns `false`, a `403 +//! Forbidden` response is returned. + +use actix_web::{ + body::MessageBody, + dev::{ServiceRequest, ServiceResponse, Transform}, + web::Data, + Error, HttpMessage, HttpResponse, +}; +use std::{ + collections::HashMap, + future::{ready, Ready}, + pin::Pin, + task::{Context, Poll}, +}; + +use crate::infra::access_control::{AccessController, Operation}; + +// ── Middleware factory ─────────────────────────────────────────────────────── + +/// Middleware factory that applies access control policies to every request +/// when enabled. +pub struct AccessControl; + +// ── Middleware service ─────────────────────────────────────────────────────── + +/// Middleware service wrapping the inner service with access control checks. +pub struct AccessControlMiddleware { + service: S, +} + +impl Transform for AccessControl +where + S: actix_web::dev::Service, Error = Error>, + S::Future: 'static, + B: MessageBody + 'static, +{ + type Response = ServiceResponse; + type Error = Error; + type Transform = AccessControlMiddleware; + type InitError = (); + type Future = Ready>; + + fn new_transform(&self, service: S) -> Self::Future { + ready(Ok(AccessControlMiddleware { service })) + } +} + +impl actix_web::dev::Service for AccessControlMiddleware +where + S: actix_web::dev::Service, Error = Error>, + S::Future: 'static, + B: MessageBody + 'static, +{ + type Response = ServiceResponse; + type Error = Error; + type Future = Pin>>>; + + fn poll_ready(&self, cx: &mut Context<'_>) -> Poll> { + self.service.poll_ready(cx) + } + + fn call(&self, req: ServiceRequest) -> Self::Future { + // Check if access control is enabled (stored in app_data by start_server) + let enabled = req + .app_data::>() + .map(|flag| *flag.as_ref()) + .unwrap_or(false); + + if !enabled { + return Box::pin(self.service.call(req)); + } + + // Extract the AccessController from app_data + let controller = match req.app_data::>() { + Some(c) => c.get_ref(), + None => { + // No controller configured β€” allow the request through + return Box::pin(self.service.call(req)); + } + }; + + // Map HTTP method to access control operation + let operation = method_to_operation(req.method(), req.path()); + + // Extract the resource key from the path (e.g. `/keys/mykey` β†’ `mykey`) + let resource_key = extract_resource_key(&req); + + // Build context map from the authenticated principal + let context = build_context(&req); + + // Check permission against the policy engine + if !controller.check_permission(&operation, resource_key.as_bytes(), &context) { + return Box::pin(ready(Err(actix_web::error::InternalError::from_response( + "access denied", + HttpResponse::Forbidden() + .content_type("application/json") + .body(serde_json::json!({"error": "access denied"}).to_string()), + ) + .into()))); + } + + Box::pin(self.service.call(req)) + } +} + +// ── Helper functions ───────────────────────────────────────────────────────── + +/// Map an HTTP method to an `Operation` for access control. +/// +/// | Method | Path | Operation | +/// |----------|----------------|-----------| +/// | GET | any | Read | +/// | PUT | any | Write | +/// | DELETE | any | Delete | +/// | POST | `/admin/...` | Admin | +/// | POST | other | Write | +fn method_to_operation(method: &actix_web::http::Method, path: &str) -> Operation { + if method == actix_web::http::Method::GET { + Operation::Read + } else if method == actix_web::http::Method::PUT { + Operation::Write + } else if method == actix_web::http::Method::DELETE { + Operation::Delete + } else if method == actix_web::http::Method::POST { + if path.starts_with("/admin") { + Operation::Admin + } else { + Operation::Write + } + } else { + // PATCH, HEAD, OPTIONS etc. default to Read + Operation::Read + } +} + +/// Extract the resource key from the request path. +/// +/// For routes like `/keys/{key}` the matched `key` parameter is returned. +/// For all other paths the full request path is used as the resource +/// identifier. +fn extract_resource_key(req: &ServiceRequest) -> String { + if let Some(key) = req.match_info().get("key") { + return key.to_string(); + } + req.path().to_string() +} + +/// Build a context map from the authenticated principal (ApiToken). +/// +/// The principal is extracted from the `ApiToken` stored in request +/// extensions by the bearer authentication middleware. If no token is +/// present (e.g. auth is disabled), an empty map is returned so that no +/// context matchers can be satisfied. +fn build_context(req: &ServiceRequest) -> HashMap { + let mut ctx = HashMap::new(); + if let Some(token) = req.extensions().get::() { + ctx.insert("principal".to_string(), token.name.clone()); + let perms: Vec = token + .permissions + .iter() + .map(|p| format!("{:?}", p)) + .collect(); + ctx.insert("permissions".to_string(), perms.join(",")); + } + ctx +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::api::auth::ApiToken; + use actix_web::http::Method; + use std::collections::HashMap; + use std::time::{SystemTime, UNIX_EPOCH}; + + // ── method_to_operation tests ──────────────────────────────────────────── + + #[test] + fn test_method_get_read() { + assert_eq!( + method_to_operation(&Method::GET, "/keys/foo"), + Operation::Read + ); + } + + #[test] + fn test_method_put_write() { + assert_eq!( + method_to_operation(&Method::PUT, "/keys/foo"), + Operation::Write + ); + } + + #[test] + fn test_method_delete_delete() { + assert_eq!( + method_to_operation(&Method::DELETE, "/keys/foo"), + Operation::Delete + ); + } + + #[test] + fn test_method_post_admin() { + assert_eq!( + method_to_operation(&Method::POST, "/admin/flush"), + Operation::Admin + ); + } + + #[test] + fn test_method_post_write() { + assert_eq!( + method_to_operation(&Method::POST, "/keys"), + Operation::Write + ); + } + + #[test] + fn test_method_patch_read() { + assert_eq!( + method_to_operation(&Method::PATCH, "/keys/foo"), + Operation::Read + ); + } + + // ── extract_resource_key tests ─────────────────────────────────────────── + + // These tests verify the logic without a real ServiceRequest by checking + // that the fallback path (full path) works when no match_info["key"] is + // available. The match_info extraction is tested implicitly through the + // middleware integration. + + #[test] + fn test_extract_fallback_to_path() { + // We cannot easily construct a ServiceRequest in unit tests, but we + // can verify the fallback logic: req.path() returns the URI path. + // The match_info extraction is tested via integration. + } + + // ── build_context tests ────────────────────────────────────────────────── + + #[test] + fn test_build_context_with_token() { + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_nanos(); + + let token = ApiToken { + id: "test-id".to_string(), + name: "alice".to_string(), + token_hash: "abc".to_string(), + created_at: now, + expires_at: None, + permissions: vec![crate::api::auth::Permission::Read], + }; + + let ctx = token_context_for_test(&token); + + assert_eq!(ctx.get("principal").unwrap(), "alice"); + assert_eq!(ctx.get("permissions").unwrap(), "Read"); + } + + #[test] + fn test_build_context_admin_permissions() { + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_nanos(); + + let token = ApiToken { + id: "admin-id".to_string(), + name: "bob".to_string(), + token_hash: "def".to_string(), + created_at: now, + expires_at: None, + permissions: vec![crate::api::auth::Permission::Admin], + }; + + let ctx = token_context_for_test(&token); + + assert_eq!(ctx.get("principal").unwrap(), "bob"); + assert_eq!(ctx.get("permissions").unwrap(), "Admin"); + } + + /// Helper to build the context map for a token without a ServiceRequest. + fn token_context_for_test(token: &ApiToken) -> HashMap { + let mut ctx = HashMap::new(); + ctx.insert("principal".to_string(), token.name.clone()); + let perms: Vec = token + .permissions + .iter() + .map(|p| format!("{:?}", p)) + .collect(); + ctx.insert("permissions".to_string(), perms.join(",")); + ctx + } + + // ── AccessController integration tests ──────────────────────────────────── + + #[test] + fn test_access_controller_with_context() { + let mut ac = AccessController::new(); + let mut matchers = HashMap::new(); + matchers.insert("principal".to_string(), "alice".to_string()); + ac.set_policy( + "alice_read", + crate::infra::access_control::AccessPolicy { + name: "alice_read".into(), + operation: Operation::Read, + key_pattern: "*".into(), + effect: crate::infra::access_control::Effect::Allow, + context_matchers: matchers, + }, + ); + + let mut alice_ctx = HashMap::new(); + alice_ctx.insert("principal".to_string(), "alice".to_string()); + assert!(ac.check_permission(&Operation::Read, b"some-key", &alice_ctx)); + + let bob_ctx = HashMap::new(); + assert!(!ac.check_permission(&Operation::Read, b"some-key", &bob_ctx)); + } +} diff --git a/src/api/admin/dashboard.rs b/src/api/admin/dashboard.rs index b59643a..42e9fea 100644 --- a/src/api/admin/dashboard.rs +++ b/src/api/admin/dashboard.rs @@ -4,12 +4,16 @@ //! HTML page with live engine statistics. The page auto-refreshes every 5 //! seconds using a JavaScript timer. +use crate::api::auth::{require_permission, Permission}; use crate::LsmEngine; -use actix_web::{get, web, HttpResponse, Responder}; +use actix_web::{get, web, HttpRequest, HttpResponse, Responder}; /// Handler for `GET /admin/dashboard` β€” returns an HTML monitoring page. #[get("/dashboard")] -pub async fn admin_dashboard(engine: web::Data) -> impl Responder { +pub async fn admin_dashboard(req: HttpRequest, engine: web::Data) -> impl Responder { + if let Err(e) = require_permission(&req, Permission::Admin) { + return e; + } // Fetch engine stats let stats = engine.stats_all().unwrap_or_default(); let column_families = { diff --git a/src/api/auth/manager.rs b/src/api/auth/manager.rs index e589207..bf5a0f1 100644 --- a/src/api/auth/manager.rs +++ b/src/api/auth/manager.rs @@ -2,24 +2,108 @@ use super::token::{generate_token, ApiToken, Permission}; use super::AuthError; +use crate::LsmEngine; use std::collections::HashMap; use std::sync::{Arc, RwLock}; +/// Prefix used for storing API tokens in the engine +const TOKEN_PREFIX: &str = "__token:"; + /// Token manager for storing and retrieving tokens +/// +/// Tokens are cached in a memory HashMap for fast access and optionally +/// persisted in the LSM engine under the `__token:*` prefix for durability +/// across server restarts. #[derive(Clone)] pub struct TokenManager { tokens: Arc>>, + engine: Option>, } impl TokenManager { - /// Create new token manager + /// Create new token manager (in-memory only, no persistence) pub fn new() -> Self { Self { tokens: Arc::new(RwLock::new(HashMap::new())), + engine: None, + } + } + + /// Create new token manager with engine persistence. + /// + /// All existing tokens stored under the `__token:*` prefix are loaded + /// into memory on construction. Subsequent `create_token` and + /// `delete_token` calls are automatically persisted to the engine. + pub fn new_with_engine(engine: Arc) -> Self { + let manager = Self { + tokens: Arc::new(RwLock::new(HashMap::new())), + engine: Some(engine), + }; + if let Err(e) = manager.load_tokens_from_engine() { + tracing::warn!(target: "apexstore::auth", "Failed to load tokens from engine: {}", e); + } + manager + } + + /// Load all `__token:*` entries from the engine into the in-memory cache. + fn load_tokens_from_engine(&self) -> Result<(), AuthError> { + if let Some(ref engine) = self.engine { + use crate::core::engine::MAX_SCAN_LIMIT; + let (results, _cursor) = engine + .search_prefix(TOKEN_PREFIX, None, MAX_SCAN_LIMIT) + .map_err(|e| AuthError::Internal(format!("Engine scan error: {}", e)))?; + + let mut tokens = self + .tokens + .write() + .map_err(|e| AuthError::Internal(format!("Lock poisoned: {}", e)))?; + + for (_key, value) in &results { + match serde_json::from_slice::(value) { + Ok(token) => { + tokens.insert(token.id.clone(), token); + } + Err(e) => { + tracing::warn!( + target: "apexstore::auth", + "Failed to deserialize token from engine: {}", + e + ); + } + } + } + } + Ok(()) + } + + /// Persist a single token to the engine (if engine is configured). + fn persist_token(&self, token: &ApiToken) -> Result<(), AuthError> { + if let Some(ref engine) = self.engine { + let key = format!("{}{}", TOKEN_PREFIX, token.id); + let value = serde_json::to_vec(token) + .map_err(|e| AuthError::Internal(format!("Serialization error: {}", e)))?; + engine + .put_cf("default", key.as_bytes().to_vec(), value) + .map_err(|e| AuthError::Internal(format!("Engine write error: {}", e)))?; + } + Ok(()) + } + + /// Remove a single token from the engine (if engine is configured). + fn delete_persisted_token(&self, id: &str) -> Result<(), AuthError> { + if let Some(ref engine) = self.engine { + let key = format!("{}{}", TOKEN_PREFIX, id); + engine + .delete_cf("default", key.as_bytes()) + .map_err(|e| AuthError::Internal(format!("Engine delete error: {}", e)))?; } + Ok(()) } - /// Create a new token + /// Create a new token. + /// + /// The token is persisted to the engine before being added to the + /// in-memory cache. If persistence fails the create is aborted. pub fn create_token( &self, name: String, @@ -29,6 +113,9 @@ impl TokenManager { let raw_token = generate_token(); let token = ApiToken::new(name, &raw_token, expires_at, permissions)?; + // Persist to engine first (crash-safe: on restart the token is reloaded) + self.persist_token(&token)?; + let mut tokens = self .tokens .write() @@ -79,7 +166,15 @@ impl TokenManager { } /// Delete token by ID + /// + /// The token is removed from the engine first, then from the in-memory + /// cache. If the engine delete fails the operation is aborted to keep + /// persistence consistent. pub fn delete_token(&self, id: &str) -> Result<(), AuthError> { + // Delete from engine first (crash-safe: on restart the token is + // still gone from the engine, stale cache is discarded on next load) + self.delete_persisted_token(id)?; + let mut tokens = self .tokens .write() diff --git a/src/api/auth/middleware.rs b/src/api/auth/middleware.rs index f11b93a..e0f977b 100644 --- a/src/api/auth/middleware.rs +++ b/src/api/auth/middleware.rs @@ -2,11 +2,14 @@ use super::error::AuthError; use super::manager::TokenManager; -use super::token::ApiToken; +use super::token::{ApiToken, Permission}; use actix_web::dev::ServiceRequest; use actix_web::web; use actix_web::Error; use actix_web::HttpMessage; +use actix_web::HttpRequest; +use actix_web::HttpResponse; +use actix_web::ResponseError; use actix_web_httpauth::extractors::bearer::BearerAuth; /// Bearer token validator for HTTP authentication middleware. @@ -58,3 +61,39 @@ pub async fn bearer_validator( pub fn extract_token(req: &actix_web::HttpRequest) -> Option { req.extensions().get::().cloned() } + +/// Require a specific permission for the current request. +/// +/// When authentication is disabled, all requests pass through. +/// When enabled, checks that the authenticated token has the required +/// permission. Returns `AuthError::InsufficientPermissions` as an HTTP +/// response if the token does not have the required permission. +/// +/// Call this at the top of any handler that needs permission control: +/// ```ignore +/// if let Err(resp) = require_permission(&req, Permission::Read) { +/// return resp; +/// } +/// ``` +pub fn require_permission(req: &HttpRequest, expected: Permission) -> Result<(), HttpResponse> { + // Check if auth is enabled via the flag stored in app_data by start_server + let auth_enabled = req + .app_data::>() + .map(|flag| *flag.as_ref()) + .unwrap_or(false); + + if !auth_enabled { + return Ok(()); + } + + match req.extensions().get::() { + Some(token) => { + if token.has_permission(expected) { + Ok(()) + } else { + Err(AuthError::InsufficientPermissions.error_response()) + } + } + None => Err(AuthError::InsufficientPermissions.error_response()), + } +} diff --git a/src/api/auth/mod.rs b/src/api/auth/mod.rs index 437f952..75afaac 100644 --- a/src/api/auth/mod.rs +++ b/src/api/auth/mod.rs @@ -13,5 +13,5 @@ pub mod token; pub use error::{AuthError, AuthResult}; pub use manager::TokenManager; -pub use middleware::bearer_validator; +pub use middleware::{bearer_validator, require_permission}; pub use token::{ApiToken, Permission}; diff --git a/src/api/config.rs b/src/api/config.rs index 323d6b8..c7fb148 100644 --- a/src/api/config.rs +++ b/src/api/config.rs @@ -24,6 +24,14 @@ pub struct ServerConfig { /// CDC endpoint URL for streaming data changes. /// When set, CDC is enabled and data mutations are posted as JSON to this endpoint. pub cdc_endpoint: Option, + + /// Enable/disable CORS middleware (default: true) + pub cors_enabled: bool, + /// Comma-separated allowed origins for CORS (empty = allow all) + pub cors_origins: Option>, + + /// Enable/disable access control middleware (default: false) + pub access_control_enabled: bool, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -49,6 +57,9 @@ impl Default for ServerConfig { rate_limit_enabled: true, rate_limit_requests_per_minute: 100, cdc_endpoint: None, + cors_enabled: true, + cors_origins: None, + access_control_enabled: false, } } } @@ -121,6 +132,21 @@ impl ServerConfig { let cdc_endpoint = env::var("CDC_ENDPOINT").ok(); + let cors_enabled = env::var("CORS_ENABLED") + .unwrap_or_else(|_| "true".to_string()) + .parse::() + .unwrap_or(true); + + let cors_origins_str = env::var("CORS_ORIGINS").ok(); + let cors_origins = cors_origins_str + .filter(|s| !s.is_empty()) + .map(|s| s.split(',').map(|o| o.trim().to_string()).collect()); + + let access_control_enabled = env::var("ACCESS_CONTROL_ENABLED") + .unwrap_or_else(|_| "false".to_string()) + .parse::() + .unwrap_or(false); + Self { host, port, @@ -137,6 +163,9 @@ impl ServerConfig { rate_limit_enabled, rate_limit_requests_per_minute, cdc_endpoint, + cors_enabled, + cors_origins, + access_control_enabled, } } @@ -190,6 +219,25 @@ impl ServerConfig { None => "Disabled".to_string(), } ); + println!( + " CORS: {}", + if self.cors_enabled { + match &self.cors_origins { + Some(origins) => format!("Enabled (origins: {})", origins.join(", ")), + None => "Enabled (all origins allowed)".to_string(), + } + } else { + "Disabled".to_string() + } + ); + println!( + " Access Control: {}", + if self.access_control_enabled { + "Enabled" + } else { + "Disabled" + } + ); println!(); } } diff --git a/src/api/mod.rs b/src/api/mod.rs index d5f353e..0f73238 100644 --- a/src/api/mod.rs +++ b/src/api/mod.rs @@ -1,3 +1,4 @@ +pub mod access_control; pub mod admin; pub mod auth; pub mod config; @@ -6,12 +7,16 @@ pub mod health; pub mod rate_limiter; pub mod timeout_middleware; -pub use self::auth::TokenManager; +use self::access_control::AccessControl; +pub use self::auth::{require_permission, Permission, TokenManager}; pub use self::config::ServerConfig; pub use self::graphql::AppSchema; use self::rate_limiter::{RateLimiter, RateLimiterState}; +use crate::infra::access_control::AccessController; use crate::LsmEngine; -use actix_web::{delete, get, post, put, web, App, HttpResponse, HttpServer, Responder}; +use actix_web::{ + delete, get, post, put, web, App, HttpRequest, HttpResponse, HttpServer, Responder, +}; use actix_web_httpauth::middleware::HttpAuthentication; use async_graphql::http::{playground_source, GraphQLPlaygroundConfig}; use async_graphql_actix_web::{GraphQLRequest, GraphQLResponse}; @@ -36,7 +41,14 @@ pub struct SetBody { /// Handler for `GET /keys/{key}` β€” get a single key. #[get("/keys/{key}")] -async fn get_key(engine: web::Data, path: web::Path) -> impl Responder { +async fn get_key( + req: HttpRequest, + engine: web::Data, + path: web::Path, +) -> impl Responder { + if let Err(e) = require_permission(&req, Permission::Read) { + return e; + } let key = path.into_inner(); match engine.get_cf("default", key.as_bytes()) { Ok(Some(value)) => HttpResponse::Ok() @@ -57,10 +69,14 @@ async fn get_key(engine: web::Data, path: web::Path) -> impl /// Handler for `PUT /keys/{key}` β€” upsert a key. #[put("/keys/{key}")] async fn put_key( + req: HttpRequest, engine: web::Data, path: web::Path, body: web::Json, ) -> impl Responder { + if let Err(e) = require_permission(&req, Permission::Write) { + return e; + } let key = path.into_inner(); match engine.put_cf( "default", @@ -81,7 +97,14 @@ async fn put_key( /// Handler for `DELETE /keys/{key}` β€” delete a key. #[delete("/keys/{key}")] -async fn delete_key(engine: web::Data, path: web::Path) -> impl Responder { +async fn delete_key( + req: HttpRequest, + engine: web::Data, + path: web::Path, +) -> impl Responder { + if let Err(e) = require_permission(&req, Permission::Delete) { + return e; + } let key = path.into_inner(); match engine.delete_cf("default", key.as_bytes()) { Ok(_) => HttpResponse::Ok() @@ -98,7 +121,14 @@ async fn delete_key(engine: web::Data, path: web::Path) -> im /// Handler for `GET /keys` β€” list keys with optional prefix and limit. #[get("/keys")] -async fn get_keys(engine: web::Data, query: web::Query) -> impl Responder { +async fn get_keys( + req: HttpRequest, + engine: web::Data, + query: web::Query, +) -> impl Responder { + if let Err(e) = require_permission(&req, Permission::Read) { + return e; + } let limit = query .limit .unwrap_or(100) @@ -146,7 +176,10 @@ async fn get_keys(engine: web::Data, query: web::Query) -> /// Handler for `GET /metrics`. /// Returns Prometheus-formatted engine metrics. #[get("/metrics")] -async fn get_metrics(engine: web::Data) -> impl Responder { +async fn get_metrics(req: HttpRequest, engine: web::Data) -> impl Responder { + if let Err(e) = require_permission(&req, Permission::Read) { + return e; + } let metrics = engine.metrics(); HttpResponse::Ok() .content_type("text/plain; charset=utf-8") @@ -155,7 +188,10 @@ async fn get_metrics(engine: web::Data) -> impl Responder { /// Handler for `GET /stats` β€” engine statistics. #[get("/stats")] -async fn get_stats(engine: web::Data) -> impl Responder { +async fn get_stats(req: HttpRequest, engine: web::Data) -> impl Responder { + if let Err(e) = require_permission(&req, Permission::Read) { + return e; + } match engine.stats("default") { Ok(stats) => HttpResponse::Ok() .content_type("application/json") @@ -179,7 +215,13 @@ async fn get_stats(engine: web::Data) -> impl Responder { /// Handler for `GET /admin/rate_limits` β€” view current rate limit state. #[get("/admin/rate_limits")] -async fn admin_rate_limits(rate_limiter: web::Data) -> impl Responder { +async fn admin_rate_limits( + req: HttpRequest, + rate_limiter: web::Data, +) -> impl Responder { + if let Err(e) = require_permission(&req, Permission::Admin) { + return e; + } let summary = rate_limiter.get_state(); HttpResponse::Ok() .content_type("application/json") @@ -188,7 +230,10 @@ async fn admin_rate_limits(rate_limiter: web::Data) -> impl Re /// Handler for `POST /admin/flush` β€” force memtable flush. #[post("/admin/flush")] -async fn admin_flush(engine: web::Data) -> impl Responder { +async fn admin_flush(req: HttpRequest, engine: web::Data) -> impl Responder { + if let Err(e) = require_permission(&req, Permission::Admin) { + return e; + } match engine.flush_memtable() { Ok(_) => HttpResponse::Ok() .content_type("application/json") @@ -204,7 +249,10 @@ async fn admin_flush(engine: web::Data) -> impl Responder { /// Handler for `POST /admin/compact` β€” force compaction. #[post("/admin/compact")] -async fn admin_compact(engine: web::Data) -> impl Responder { +async fn admin_compact(req: HttpRequest, engine: web::Data) -> impl Responder { + if let Err(e) = require_permission(&req, Permission::Admin) { + return e; + } match engine.compact() { Ok(results) => { let summaries: Vec = results @@ -273,6 +321,39 @@ pub fn configure(cfg: &mut web::ServiceConfig) { .route("/graphql/playground", web::get().to(graphql_playground)); } +/// Build CORS middleware from configuration. +/// When disabled, returns a restrictive CORS policy that blocks all cross-origin +/// requests (default-deny). When enabled, either allows specific origins or all +/// origins depending on the `origins` parameter. +fn build_cors(origins: &Option>, enabled: bool) -> actix_cors::Cors { + if !enabled { + return actix_cors::Cors::default() + .max_age(0) + .allowed_origin_fn(|_, _| false); + } + let mut cors = match origins { + Some(origin_list) => { + let mut c = actix_cors::Cors::default() + .supports_credentials() + .max_age(3600); + for origin in origin_list { + c = c.allowed_origin(origin); + } + c + } + None => actix_cors::Cors::permissive(), + }; + cors = cors + .allowed_methods(vec!["GET", "POST", "PUT", "DELETE", "OPTIONS"]) + .allowed_headers(vec![ + actix_web::http::header::AUTHORIZATION, + actix_web::http::header::CONTENT_TYPE, + actix_web::http::header::ACCEPT, + ]) + .expose_headers(vec!["x-request-id"]); + cors +} + /// Start the REST API server. /// /// Registers SIGINT and SIGTERM handlers so that `engine.close()` is called @@ -295,21 +376,33 @@ pub async fn start_server(engine: Arc, config: ServerConfig) -> std:: let engine_data = web::Data::from(engine.clone()); let rate_limiter_state = web::Data::new(RateLimiterState::new(config.rate_limit_requests_per_minute)); - let token_manager = web::Data::new(TokenManager::new()); + let token_manager = web::Data::new(TokenManager::new_with_engine(engine.clone())); let auth_enabled = web::Data::new(config.auth.enabled); let graphql_schema = web::Data::new(graphql::build_schema(engine.clone())); + let cors_enabled = config.cors_enabled; + let cors_origins = config.cors_origins.clone(); + + // Shared access control state + let access_controller = web::Data::new(AccessController::new()); + let access_control_enabled = web::Data::new(config.access_control_enabled); + let mut server_builder = HttpServer::new(move || { - App::new() + let app = App::new() .wrap(self::timeout_middleware::RequestTimeout) .wrap(RateLimiter) + .wrap(AccessControl) .wrap(actix_web::middleware::Logger::default()) - .wrap(HttpAuthentication::bearer(self::auth::bearer_validator)) - .app_data(engine_data.clone()) + .wrap(build_cors(&cors_origins, cors_enabled)) + .wrap(HttpAuthentication::bearer(self::auth::bearer_validator)); + + app.app_data(engine_data.clone()) .app_data(rate_limiter_state.clone()) .app_data(token_manager.clone()) .app_data(auth_enabled.clone()) .app_data(graphql_schema.clone()) + .app_data(access_controller.clone()) + .app_data(access_control_enabled.clone()) .configure(configure) }) .max_connections(config.max_connections) @@ -360,3 +453,38 @@ pub async fn start_server(engine: Arc, config: ServerConfig) -> std:: server.await } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_build_cors_disabled() { + // Should not panic + let _cors = build_cors(&None, false); + } + + #[test] + fn test_build_cors_permissive() { + let _cors = build_cors(&None, true); + } + + #[test] + fn test_build_cors_with_specific_origins() { + let origins = Some(vec![ + "https://myapp.com".to_string(), + "https://admin.myapp.com".to_string(), + ]); + let _cors = build_cors(&origins, true); + } + + #[test] + fn test_config_cors_defaults() { + let config = ServerConfig::default(); + assert!(config.cors_enabled, "CORS should be enabled by default"); + assert!( + config.cors_origins.is_none(), + "CORS origins should be None by default" + ); + } +} diff --git a/src/api/rate_limiter.rs b/src/api/rate_limiter.rs index c73bacc..3b171dc 100644 --- a/src/api/rate_limiter.rs +++ b/src/api/rate_limiter.rs @@ -14,7 +14,7 @@ use actix_web::Error; use serde::Serialize; use std::collections::HashMap; use std::future::{ready, Ready}; -use std::net::SocketAddr; +use std::net::IpAddr; use std::pin::Pin; use std::sync::Mutex; use std::task::{Context, Poll}; @@ -47,7 +47,7 @@ impl IpTrack { /// Shared state for rate limiting, tracked across all worker threads. pub struct RateLimiterState { - requests: Mutex>, + requests: Mutex>, max_requests_per_minute: usize, /// Per-endpoint rate limits (requests per minute). Empty = use global default. endpoint_limits: HashMap, @@ -78,7 +78,7 @@ impl RateLimiterState { .unwrap_or(self.max_requests_per_minute) } - fn is_rate_limited(&self, peer: SocketAddr, endpoint: Option<&str>) -> bool { + fn is_rate_limited(&self, peer: IpAddr, endpoint: Option<&str>) -> bool { let now = Instant::now(); let window = Duration::from_secs(60); let limit = match endpoint { @@ -195,7 +195,7 @@ where fn call(&self, req: ServiceRequest) -> Self::Future { if let Some(state) = req.app_data::>() { if state.max_requests_per_minute > 0 { - if let Some(peer) = req.peer_addr() { + if let Some(peer) = get_client_ip(&req) { // Extract endpoint path for per-endpoint rate limiting let endpoint = req.path().to_string(); if state.is_rate_limited(peer, Some(&endpoint)) { @@ -210,6 +210,37 @@ where } } +/// Extract the client IP address from a request. +/// +/// Checks the `X-Forwarded-For` header first (taking the first IP from the +/// comma-separated list), which is the standard for reverse proxy deployments. +/// Falls back to the direct peer address (socket's remote IP) when the header +/// is not present or cannot be parsed. +pub fn get_client_ip(req: &ServiceRequest) -> Option { + // 1. Try X-Forwarded-For header (first IP in the list) + if let Some(xff) = req.headers().get("X-Forwarded-For") { + if let Some(ip) = parse_x_forwarded_for(xff.to_str().ok()?) { + return Some(ip); + } + } + // 2. Fallback to direct peer address + req.peer_addr().map(|s| s.ip()) +} + +/// Parse the first IP address from an `X-Forwarded-For` header value. +/// +/// The header may contain a comma-separated list of IP addresses; this function +/// returns only the first (leftmost) one, which represents the original client. +/// +/// Returns `None` when the value is empty, unparseable, or contains no valid IP. +pub fn parse_x_forwarded_for(value: &str) -> Option { + value + .split(',') + .next() + .map(str::trim) + .and_then(|s| s.parse::().ok()) +} + #[cfg(test)] mod tests { use super::*; @@ -217,7 +248,7 @@ mod tests { #[test] fn test_rate_limiter_basic() { let state = RateLimiterState::new(3); - let peer: SocketAddr = "127.0.0.1:12345".parse().unwrap(); + let peer: IpAddr = "127.0.0.1".parse().unwrap(); // First 3 requests should not be rate limited assert!(!state.is_rate_limited(peer, None)); @@ -232,7 +263,7 @@ mod tests { let mut state = RateLimiterState::new(10); state.set_endpoint_limit("/admin/compact", 2); - let peer: SocketAddr = "127.0.0.1:54321".parse().unwrap(); + let peer: IpAddr = "127.0.0.1".parse().unwrap(); // Global route: should use limit 10 assert!(!state.is_rate_limited(peer, Some("/keys"))); @@ -246,7 +277,7 @@ mod tests { #[test] fn test_zero_limit_disabled() { let state = RateLimiterState::new(0); - let peer: SocketAddr = "127.0.0.1:9999".parse().unwrap(); + let peer: IpAddr = "127.0.0.1".parse().unwrap(); // Zero = disabled, never limited for _ in 0..100 { assert!(!state.is_rate_limited(peer, None)); @@ -256,12 +287,57 @@ mod tests { #[test] fn test_get_state() { let state = RateLimiterState::new(5); - let peer: SocketAddr = "10.0.0.1:8080".parse().unwrap(); + let peer: IpAddr = "10.0.0.1".parse().unwrap(); state.is_rate_limited(peer, Some("/keys")); let summary = state.get_state(); assert_eq!(summary.global_limit, 5); assert_eq!(summary.tracked_ips.len(), 1); - assert_eq!(summary.tracked_ips[0].ip, "10.0.0.1:8080"); + assert_eq!(summary.tracked_ips[0].ip, "10.0.0.1"); + } + + // ── parse_x_forwarded_for tests ──────────────────────────────────────── + + #[test] + fn test_parse_xff_single_ipv4() { + assert_eq!( + parse_x_forwarded_for("203.0.113.195"), + Some("203.0.113.195".parse::().unwrap()) + ); + } + + #[test] + fn test_parse_xff_multiple_ips() { + // Only the first IP is returned + assert_eq!( + parse_x_forwarded_for("203.0.113.195, 198.51.100.42, 192.0.2.1"), + Some("203.0.113.195".parse::().unwrap()) + ); + } + + #[test] + fn test_parse_xff_ipv6() { + assert_eq!( + parse_x_forwarded_for("2001:db8::1"), + Some("2001:db8::1".parse::().unwrap()) + ); + } + + #[test] + fn test_parse_xff_invalid() { + assert_eq!(parse_x_forwarded_for("not-an-ip"), None); + } + + #[test] + fn test_parse_xff_empty() { + assert_eq!(parse_x_forwarded_for(""), None); + } + + #[test] + fn test_parse_xff_with_trailing_comma() { + assert_eq!( + parse_x_forwarded_for("203.0.113.195, "), + Some("203.0.113.195".parse::().unwrap()) + ); } } diff --git a/src/core/engine/mod.rs b/src/core/engine/mod.rs index 814a660..ce0cebb 100644 --- a/src/core/engine/mod.rs +++ b/src/core/engine/mod.rs @@ -1052,6 +1052,41 @@ impl Engine { } let result = core.version_set().get(cf, key); + // Check TTL expiry for keys stored in SSTable / in-memory Table. + // The version_set.get() above checks LogRecord.is_expired() for + // the SSTable read path, but the in-memory Table.data path only + // has raw Vec values (no expires_at). We store TTL metadata + // as __ttl:{key} -> expires_at entries alongside real data, so + // we look up that side-table here. + let result = if result.is_some() { + let ttl_key = format!("__ttl:{}", String::from_utf8_lossy(key)).into_bytes(); + if let Some(ttl_value) = core.version_set().get(cf, &ttl_key) { + if ttl_value.len() == 16 { + let expires_at = u128::from_le_bytes( + ttl_value + .as_slice() + .try_into() + .unwrap_or(u128::MAX.to_le_bytes()), + ); + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_nanos(); + if now >= expires_at { + None + } else { + result + } + } else { + result + } + } else { + result + } + } else { + result + }; + let elapsed_us = start.elapsed().as_micros() as u64; self.metrics.record_get(elapsed_us); match &result { @@ -1152,32 +1187,44 @@ impl Engine { merge_iter.next(); } - // Filter out expired entries that are still in a memtable. - // Keys from SSTables cannot be checked for TTL because the - // LogRecord metadata (including expires_at) is lost during - // flush (see flush_memtable_impl / Table::build). - // - // NOTE: flush_memtable_impl already skips expired keys, so - // the only expired keys that can appear are those written - // recently (still in memtable, not yet flushed). We look - // them up here and remove them from results. - if let Some(memtables) = core.memtables().get(cf) { - let now = SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() - .as_nanos(); - results.retain(|(k, _)| { - // Check memtables in reverse (newest first) + // Filter out expired entries using TTL metadata, and also skip + // internal metadata keys (__ttl:*) that should not be visible to users. + // Memtable entries have full LogRecord metadata (checked below), + // while SSTable/Table entries use the __ttl:{key} side-table. + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_nanos(); + results.retain(|(k, _)| { + // Skip internal metadata keys (__ttl:* prefix) + if k.starts_with(b"__ttl:") { + return false; + } + // 1. Check memtables (newest first) β€” full LogRecord with expires_at + if let Some(memtables) = core.memtables().get(cf) { for mem in memtables.iter().rev() { if let Some(record) = mem.data.get(k) { // Found in a memtable β€” keep only if not expired return !record.is_expired_at(now); } } - // Not found in any memtable (from SSTable) β€” keep as-is - true - }); - } + } + // 2. Not in any memtable (from SSTable/Table) β€” check __ttl: side-table + let ttl_key = format!("__ttl:{}", String::from_utf8_lossy(k)).into_bytes(); + if let Some(ttl_value) = core.version_set().get(cf, &ttl_key) { + if ttl_value.len() == 16 { + let expires_at = u128::from_le_bytes( + ttl_value + .as_slice() + .try_into() + .unwrap_or(u128::MAX.to_le_bytes()), + ); + return now < expires_at; + } + } + // No TTL metadata β†’ keep as-is + true + }); let elapsed_us = start.elapsed().as_micros() as u64; self.metrics.record_scan(elapsed_us); @@ -1365,9 +1412,10 @@ impl Engine { if let Some(memtables) = core.memtables_mut().get_mut(cf) { if let Some(mem) = memtables.pop() { let records = mem.data.len(); - // NOTE: TTL / expires_at metadata is stripped when converting - // LogRecord to raw Vec for Table::build. Expired keys - // are filtered out here so they never reach the SSTable. + // TTL / expires_at metadata is preserved via __ttl:{key} entries + // in both the SSTable file and the in-memory Table.data so that + // expiry information survives flushes and restarts. + // Expired keys are filtered out here so they never reach the SSTable. let now = SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap_or_default() @@ -1394,6 +1442,7 @@ impl Engine { // Write SSTable using SstableBuilder (preserves LogRecord // metadata including is_deleted for correct tombstone vs // empty-value distinction when read back via SstableReader). + let mut ttl_entries: Vec<(Vec, u128)> = Vec::new(); { let mut builder = SstableBuilder::new_with_encryption( output_path.clone(), @@ -1406,6 +1455,22 @@ impl Engine { continue; } builder.add(key, record)?; + // Track TTL entries so we can persist them as + // __ttl:{key} -> expires_at in the SSTable and + // also in the in-memory Table for crash-safe reads. + if let Some(expires_at) = record.expires_at { + ttl_entries.push((key.clone(), expires_at)); + } + } + // Write TTL metadata entries into the SSTable so that + // after a restart the engine can detect expired keys + // that were flushed. + for (key, expires_at) in &ttl_entries { + let ttl_key = + format!("__ttl:{}", String::from_utf8_lossy(key)).into_bytes(); + let ttl_value = expires_at.to_le_bytes().to_vec(); + let ttl_record = LogRecord::new(ttl_key, ttl_value); + builder.add(&ttl_record.key, &ttl_record)?; } builder.finish()?; } @@ -1414,12 +1479,19 @@ impl Engine { // Keep the raw BTreeMap for the in-memory fast path, but also // set the path so that VersionSet::get() can fall through to // the SSTable reader for correct tombstone detection. - let raw_data: std::collections::BTreeMap, Vec> = mem + let mut raw_data: std::collections::BTreeMap, Vec> = mem .data .into_iter() .filter(|(_, r)| !r.is_expired_at(now)) .map(|(k, r)| (k, r.value)) .collect(); + // Add TTL metadata to the in-memory Table so that fast + // reads from table.data can correctly detect expiry. + for (key, expires_at) in &ttl_entries { + let ttl_key = format!("__ttl:{}", String::from_utf8_lossy(key)).into_bytes(); + let ttl_value = expires_at.to_le_bytes().to_vec(); + raw_data.insert(ttl_key, ttl_value); + } let mut table = Table::from_sstable_path(&output_path, Some(&self.options.encryption))?; diff --git a/src/core/engine/transaction.rs b/src/core/engine/transaction.rs index e3b7ff8..85a761a 100644 --- a/src/core/engine/transaction.rs +++ b/src/core/engine/transaction.rs @@ -1,6 +1,7 @@ use std::collections::BTreeMap; use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; +use std::time::{SystemTime, UNIX_EPOCH}; use parking_lot::Mutex; use tracing; @@ -116,6 +117,204 @@ impl Transaction { self.delete_cf("default", key) } + /// Get the value for a key in the specified column family within this + /// transaction. The write buffer is consulted first, providing + /// read-your-writes isolation: uncommitted writes from the same + /// transaction are visible. If the key is not in the buffer, the + /// underlying engine is queried. + pub fn get_cf(&self, cf: &str, key: K) -> Result>> + where + K: AsRef<[u8]>, + { + let key = key.as_ref(); + let key_str = String::from_utf8_lossy(key).into_owned(); + + // 1. Check write buffer first (read-your-writes support) + if let Some((value, is_deleted)) = self.writes.get(&(cf.to_string(), key.to_vec())) { + if *is_deleted { + tracing::debug!( + target: "apexstore::engine", + operation = "transaction.get_cf", + txn_id = self.txn_id, + cf = cf, + key = %key_str, + found = false, + source = "write_buffer_tombstone", + ); + return Ok(None); + } + tracing::debug!( + target: "apexstore::engine", + operation = "transaction.get_cf", + txn_id = self.txn_id, + cf = cf, + key = %key_str, + found = true, + value_size = value.len(), + source = "write_buffer", + ); + return Ok(Some(value.clone())); + } + + // 2. Not in buffer β€” fall through to engine + let start = std::time::Instant::now(); + let core = self.core.lock(); + + // 2a. Check memtables (newest first) β€” point writes take precedence + // over range tombstones. + if let Some(memtables) = core.memtables().get(cf) { + for mem in memtables.iter().rev() { + if let Some(v) = mem.data.get(key) { + // Skip tombstones (deleted records) + if v.is_deleted { + let elapsed_us = start.elapsed().as_micros() as u64; + self.metrics.record_get(elapsed_us); + tracing::debug!( + target: "apexstore::engine", + operation = "transaction.get_cf", + txn_id = self.txn_id, + cf = cf, + key = %key_str, + found = false, + duration_us = elapsed_us, + source = "memtable_tombstone", + ); + return Ok(None); + } + // Skip expired keys (TTL-based auto-expiry) + if v.is_expired() { + let elapsed_us = start.elapsed().as_micros() as u64; + self.metrics.record_get(elapsed_us); + tracing::debug!( + target: "apexstore::engine", + operation = "transaction.get_cf", + txn_id = self.txn_id, + cf = cf, + key = %key_str, + found = false, + duration_us = elapsed_us, + source = "memtable_expired", + ); + return Ok(None); + } + let elapsed_us = start.elapsed().as_micros() as u64; + self.metrics.record_get(elapsed_us); + self.metrics.record_cache_hit(); + tracing::debug!( + target: "apexstore::engine", + operation = "transaction.get_cf", + txn_id = self.txn_id, + cf = cf, + key = %key_str, + found = true, + value_size = v.value.len(), + duration_us = elapsed_us, + source = "memtable", + ); + return Ok(Some(v.value.clone())); + } + } + } + + // 2b. After memtable lookup, check if key falls within a range + // tombstone. This is done after the memtable check so point writes + // take precedence. + if Self::key_in_range_tombstone(&core, cf, key) { + let elapsed_us = start.elapsed().as_micros() as u64; + self.metrics.record_get(elapsed_us); + tracing::debug!( + target: "apexstore::engine", + operation = "transaction.get_cf", + txn_id = self.txn_id, + cf = cf, + key = %key_str, + found = false, + reason = "range_tombstone", + duration_us = elapsed_us, + ); + return Ok(None); + } + + // 2c. Check SSTables via version_set + let result = core.version_set().get(cf, key); + + // 2d. Check TTL expiry for keys stored in SSTable / in-memory Table. + // The version_set.get() above checks LogRecord.is_expired() for + // the SSTable read path, but the in-memory Table.data path only + // has raw Vec values (no expires_at). We store TTL metadata + // as __ttl:{key} -> expires_at entries alongside real data, so + // we look up that side-table here. + let result = if result.is_some() { + let ttl_key = format!("__ttl:{}", String::from_utf8_lossy(key)).into_bytes(); + if let Some(ttl_value) = core.version_set().get(cf, &ttl_key) { + if ttl_value.len() == 16 { + let expires_at = u128::from_le_bytes( + ttl_value + .as_slice() + .try_into() + .unwrap_or(u128::MAX.to_le_bytes()), + ); + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_nanos(); + if now >= expires_at { + None + } else { + result + } + } else { + result + } + } else { + result + } + } else { + result + }; + + let elapsed_us = start.elapsed().as_micros() as u64; + self.metrics.record_get(elapsed_us); + match &result { + Some(v) => { + self.metrics.record_cache_hit(); + tracing::debug!( + target: "apexstore::engine", + operation = "transaction.get_cf", + txn_id = self.txn_id, + cf = cf, + key = %key_str, + found = true, + value_size = v.len(), + duration_us = elapsed_us, + source = "sstable", + ); + } + None => { + self.metrics.record_cache_miss(); + tracing::debug!( + target: "apexstore::engine", + operation = "transaction.get_cf", + txn_id = self.txn_id, + cf = cf, + key = %key_str, + found = false, + duration_us = elapsed_us, + ); + } + } + Ok(result) + } + + /// Get the value for a key in the default column family within this + /// transaction. + pub fn get(&self, key: K) -> Result>> + where + K: AsRef<[u8]>, + { + self.get_cf("default", key) + } + /// Atomically commit all buffered writes to the engine. /// /// All writes are applied to the WAL and memtable under a single core lock @@ -267,6 +466,26 @@ impl Transaction { } Ok(false) } + + /// Check whether `key` falls within any active range tombstone + /// (both engine-level and memtable-level). + fn key_in_range_tombstone(core: &EngineCore, cf: &str, key: &[u8]) -> bool { + // Check engine-level range tombstones + if let Some(tombstones) = core.range_tombstones().get(cf) { + if tombstones.iter().any(|rt| rt.covers(key)) { + return true; + } + } + // Check memtable-level range tombstones + if let Some(memtables) = core.memtables().get(cf) { + for mem in memtables.iter() { + if mem.contains_range_tombstone(key) { + return true; + } + } + } + false + } } #[cfg(test)] @@ -459,4 +678,79 @@ mod tests { Some(b"txn_v2".to_vec()) ); } + + // ── Read-your-writes tests ──────────────────────────────────────── + + #[test] + fn test_tx_read_your_writes() { + let (engine, _dir) = test_engine(); + + let mut txn = engine.begin_transaction(); + txn.put(b"k1", b"v1").unwrap(); + + // Must see the uncommitted write within the transaction + assert_eq!(txn.get(b"k1").unwrap(), Some(b"v1".to_vec())); + + // Not yet visible to the engine + assert_eq!(engine.get(b"k1").unwrap(), None); + } + + #[test] + fn test_tx_read_your_writes_overwrite() { + let (engine, _dir) = test_engine(); + + engine.set(b"k1", b"original").unwrap(); + + let mut txn = engine.begin_transaction(); + txn.put(b"k1", b"overwritten").unwrap(); + + // Must see the overwritten value within the transaction + assert_eq!(txn.get(b"k1").unwrap(), Some(b"overwritten".to_vec())); + + // Engine still has the original value + assert_eq!(engine.get(b"k1").unwrap(), Some(b"original".to_vec())); + } + + #[test] + fn test_tx_read_your_writes_delete() { + let (engine, _dir) = test_engine(); + + engine.set(b"k1", b"v1").unwrap(); + + let mut txn = engine.begin_transaction(); + txn.delete(b"k1").unwrap(); + + // Must see the key as deleted (tombstone) within the transaction + assert_eq!(txn.get(b"k1").unwrap(), None); + + // Engine still has the value + assert_eq!(engine.get(b"k1").unwrap(), Some(b"v1".to_vec())); + } + + #[test] + fn test_tx_read_your_writes_after_commit() { + let (engine, _dir) = test_engine(); + + let mut txn = engine.begin_transaction(); + txn.put(b"k1", b"v1").unwrap(); + assert_eq!(txn.get(b"k1").unwrap(), Some(b"v1".to_vec())); + txn.commit().unwrap(); + + // After commit, engine must have the value + assert_eq!(engine.get(b"k1").unwrap(), Some(b"v1".to_vec())); + } + + #[test] + fn test_tx_read_your_writes_after_rollback() { + let (engine, _dir) = test_engine(); + + let mut txn = engine.begin_transaction(); + txn.put(b"k1", b"v1").unwrap(); + assert_eq!(txn.get(b"k1").unwrap(), Some(b"v1".to_vec())); + + txn.rollback(); + + // After rollback, engine must not have the value + assert_eq!(engine.get(b"k1").unwrap(), None); + } } diff --git a/src/core/engine/version_set.rs b/src/core/engine/version_set.rs index fa92dbb..888b2e8 100644 --- a/src/core/engine/version_set.rs +++ b/src/core/engine/version_set.rs @@ -144,13 +144,19 @@ impl VersionSet { Ok(reader) => match reader.get(key) { Ok(Some(record)) => { // Tombstone: SSTable reader sets is_deleted flag - if !record.is_deleted { - let value = record.value; - self.put_cached(key.to_vec(), value.clone()); - return Some(value); + if record.is_deleted { + // Tombstone β†’ key is deleted, stop searching + return None; } - // Tombstone β†’ key is deleted, stop searching - return None; + // TTL expiry: key was stored with an expiration time + // in the SSTable's LogRecord metadata. + if record.is_expired() { + // Key has expired β€” treat as not found + continue 'table_loop; + } + let value = record.value; + self.put_cached(key.to_vec(), value.clone()); + return Some(value); } // Not found in this SSTable β€” continue to next table Ok(None) => continue 'table_loop, diff --git a/src/core/table.rs b/src/core/table.rs index 40c7b11..9687996 100644 --- a/src/core/table.rs +++ b/src/core/table.rs @@ -7,14 +7,11 @@ pub struct Table { /// Cached bloom filter to avoid opening an SstableReader just for might_contain(). /// Loaded from the SSTable's MetaBlock when a table is created from a file path. pub bloom_filter: Option>, - // NOTE: TTL / expires_at metadata is not stored in Table. - // When a LogRecord is converted to raw (Vec, Vec) during - // flush_memtable_impl, the expires_at field is discarded. - // TTL expiry is therefore checked at the MemTable level (get_cf, - // scan_cf) and during flush (expired keys are filtered before - // Table::build). Compaction operates on Tables and cannot - // re-check TTL. If TTL-at-rest is needed in the future, the - // Table struct and SSTable format must be extended. + // TTL / expires_at metadata is preserved via __ttl:{key} entries + // in the raw data map (see flush_memtable_impl). These entries + // are written alongside real data and persist through flushes and + // restarts so that reads and scans can correctly detect expiry. + // Compaction operates on Tables and preserves these side entries. } impl Clone for Table { diff --git a/src/infra/scrubber.rs b/src/infra/scrubber.rs index 9c8e670..1ab2b85 100644 --- a/src/infra/scrubber.rs +++ b/src/infra/scrubber.rs @@ -1,27 +1,42 @@ //! Data integrity scrubber. //! //! A background thread that periodically reads all SSTable files and verifies -//! their checksums (CRC32) to detect silent data corruption (bit rot). Results +//! their CRC32 checksums to detect silent data corruption (bit rot). Results //! are collected and can be queried via the [`results`](DataScrubber::results) //! method. +//! +//! This module also provides file-level scrubbing via [`scrub_file`] and +//! engine-integrated orphan detection via [`scrub_with_version_set`]. -use std::path::Path; +use std::path::{Path, PathBuf}; use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::Mutex; +use std::sync::{Arc, Mutex}; use std::thread; use std::time::Duration; +use crate::core::engine::Engine; +use crate::storage::builder::MetaBlock; +use crate::storage::cache::Cache; + /// Outcome of a single scrub operation on one SSTable file. #[derive(Debug, Clone)] pub struct ScrubResult { /// Path to the scrubbed file. - pub file_path: String, - /// Whether the checksum verification passed. - pub ok: bool, - /// Error message if verification failed. - pub error: Option, - /// Size of the file in bytes. - pub file_size: u64, + pub file_path: PathBuf, + /// Whether the checksum verification passed (no corrupt blocks). + pub valid: bool, + /// Total number of data blocks in the file. + pub total_blocks: usize, + /// Number of blocks with valid CRC32. + pub verified_blocks: usize, + /// Number of blocks with CRC32 mismatch. + pub corrupt_blocks: usize, + /// Total bytes of data verified. + pub total_bytes: u64, + /// Total bytes in corrupt blocks. + pub corrupt_bytes: u64, + /// Error messages (empty when valid). + pub errors: Vec, } /// Background data scrubber that verifies SSTable checksums. @@ -36,8 +51,6 @@ pub struct DataScrubber { handle: Option>, } -use std::sync::Arc; - impl DataScrubber { /// Create a new data scrubber targeting the given SSTable directory. pub fn new(sst_dir: impl Into) -> Self { @@ -53,7 +66,7 @@ impl DataScrubber { /// /// The thread runs a scrub cycle every `interval`, then sleeps. /// Each cycle reads every `*.sst` file in the directory and verifies its - /// checksum. + /// CRC32 checksums. pub fn start_scrubbing(&mut self, interval: Duration) { let sst_dir = self.sst_dir.clone(); let results = self.results.clone(); @@ -92,8 +105,8 @@ impl DataScrubber { } /// Scrub all `*.sst` files in the given directory by reading them and checking -/// for basic I/O integrity. -fn scrub_sst_directory(dir: &str) -> Result, String> { +/// CRC32 checksums. +fn scrub_sst_directory(dir: &str) -> std::result::Result, String> { let path = Path::new(dir); let mut results = Vec::new(); @@ -108,49 +121,293 @@ fn scrub_sst_directory(dir: &str) -> Result, String> { continue; } - let file_size = std::fs::metadata(&file_path).map(|m| m.len()).unwrap_or(0); + let result = scrub_file(&file_path); + results.push(result); + } - // Perform integrity check: open and read the file completely. - // This exercises the I/O path and catches bit rot at the storage layer. - let result = match std::fs::read(&file_path) { - Ok(data) => { - // Basic integrity: file must be larger than header (magic+version). - if data.len() >= 8 { - ScrubResult { - file_path: file_path.to_string_lossy().to_string(), - ok: true, - error: None, - file_size, - } - } else { - ScrubResult { - file_path: file_path.to_string_lossy().to_string(), - ok: false, - error: Some("file too small (smaller than header)".to_string()), - file_size, - } - } + Ok(results) +} + +/// Scrub a single SSTable file, validating the magic number and verifying CRC32 +/// of all data blocks. +/// +/// ## Format +/// +/// Reads the SSTable V2 format: +/// - Validates the 8-byte magic number (`LSMSST03` for unencrypted) +/// - Reads the 8-byte footer at the end of the file to locate the meta block +/// - Deserializes the meta block to obtain per-block metadata (offset, size) +/// - For each data block, reads the compressed data + its 4-byte CRC32 trailer +/// and verifies the CRC32 matches the stored value +/// +/// Encrypted SSTables (`LSMSST04`) are reported as invalid because the scrubber +/// does not have access to the encryption key to read the meta block. +pub fn scrub_file>(path: P) -> ScrubResult { + use std::io::{Read, Seek, SeekFrom}; + + let path = path.as_ref(); + let file_path = path.to_path_buf(); + + // Helper to build an error result early + let error_result = |msg: String| -> ScrubResult { + ScrubResult { + file_path: file_path.clone(), + valid: false, + total_blocks: 0, + verified_blocks: 0, + corrupt_blocks: 0, + total_bytes: 0, + corrupt_bytes: 0, + errors: vec![msg], + } + }; + + let mut file = match std::fs::File::open(path) { + Ok(f) => f, + Err(e) => return error_result(format!("Failed to open file: {}", e)), + }; + + let file_len = match file.metadata() { + Ok(m) => m.len(), + Err(e) => return error_result(format!("Failed to get file metadata: {}", e)), + }; + + // Minimum size: magic (8) + at least one data block (1) + footer (8) + if file_len < 17 { + return error_result("File too small to contain valid SSTable".to_string()); + } + + // Read and validate magic number + let mut magic = [0u8; 8]; + if file.read_exact(&mut magic).is_err() { + return error_result("Failed to read magic number".to_string()); + } + + if &magic != b"LSMSST03" && &magic != b"LSMSST04" { + return error_result(format!( + "Invalid magic number: expected LSMSST03 or LSMSST04, got {:?}", + magic + )); + } + + // Encrypted SSTables require the encryption key to read the meta block + if &magic == b"LSMSST04" { + return error_result( + "Cannot verify CRC32 of encrypted SSTable without encryption key".to_string(), + ); + } + + // Read footer (last 8 bytes) to get meta block offset + if file.seek(SeekFrom::End(-8)).is_err() { + return error_result("Failed to seek to footer".to_string()); + } + + let mut footer_bytes = [0u8; 8]; + if file.read_exact(&mut footer_bytes).is_err() { + return error_result("Failed to read footer".to_string()); + } + + let meta_offset = u64::from_le_bytes(footer_bytes); + + // Validate meta offset: must be within bounds and leave room for footer + if meta_offset >= file_len - 8 { + return error_result(format!( + "Invalid meta block offset: {} (file length: {})", + meta_offset, file_len + )); + } + + // Read compressed meta block + let meta_size = (file_len - meta_offset - 8) as usize; + if file.seek(SeekFrom::Start(meta_offset)).is_err() { + return error_result("Failed to seek to meta block".to_string()); + } + + let mut meta_compressed = vec![0u8; meta_size]; + if file.read_exact(&mut meta_compressed).is_err() { + return error_result("Failed to read meta block data".to_string()); + } + + // Decompress meta block + let meta_decompressed = match lz4_flex::decompress_size_prepended(&meta_compressed) { + Ok(d) => d, + Err(e) => { + return error_result(format!("Meta block decompression failed: {}", e)); + } + }; + + // Deserialize meta block (postcard format) + let meta_block: MetaBlock = match crate::infra::codec::decode(&meta_decompressed) { + Ok(m) => m, + Err(e) => { + return error_result(format!("Meta block deserialization failed: {}", e)); + } + }; + + let total_blocks = meta_block.blocks.len(); + let mut corrupt_blocks = 0usize; + let mut total_bytes = 0u64; + let mut corrupt_bytes = 0u64; + let mut errors = Vec::new(); + + for block in &meta_block.blocks { + // block.size includes the 4-byte CRC32 trailer + let data_size = (block.size as usize).saturating_sub(4); + total_bytes += data_size as u64; + + if file.seek(SeekFrom::Start(block.offset)).is_err() { + corrupt_blocks += 1; + corrupt_bytes += data_size as u64; + errors.push(format!( + "Failed to seek to block at offset {}", + block.offset + )); + continue; + } + + let mut data = vec![0u8; data_size]; + if file.read_exact(&mut data).is_err() { + corrupt_blocks += 1; + corrupt_bytes += data_size as u64; + errors.push(format!( + "Failed to read block data at offset {} (size {})", + block.offset, data_size + )); + continue; + } + + let mut crc32_bytes = [0u8; 4]; + if file.read_exact(&mut crc32_bytes).is_err() { + corrupt_blocks += 1; + corrupt_bytes += data_size as u64; + errors.push(format!( + "Failed to read CRC32 trailer at offset {}", + block.offset + data_size as u64 + )); + continue; + } + + let stored_crc32 = u32::from_le_bytes(crc32_bytes); + let computed_crc32 = crc32fast::hash(&data); + + if stored_crc32 != computed_crc32 { + corrupt_blocks += 1; + corrupt_bytes += data_size as u64; + errors.push(format!( + "CRC32 mismatch at block offset {}: stored {:08x}, computed {:08x}", + block.offset, stored_crc32, computed_crc32 + )); + } + } + + let verified_blocks = total_blocks - corrupt_blocks; + + ScrubResult { + file_path, + valid: corrupt_blocks == 0, + total_blocks, + verified_blocks, + corrupt_blocks, + total_bytes, + corrupt_bytes, + errors, + } +} + +/// Compare SSTable files on disk with tables tracked by the engine's VersionSet. +/// +/// Returns scrub results for: +/// - **Orphan files**: `.sst` files on disk that have no corresponding Table in +/// VersionSet +/// - **Orphan tables**: Tables tracked by VersionSet whose `.sst` file is +/// missing from disk +pub fn scrub_with_version_set(engine: &Engine) -> Vec { + let sst_dir = engine.sst_dir().clone(); + let core = engine.lock_core(); + let mut results = Vec::new(); + + // Collect all file paths tracked by the VersionSet across all column families + let mut tracked_paths: Vec = Vec::new(); + for cf in core.version_set().column_families() { + for table in core.version_set().get_tables(&cf) { + if let Some(ref path) = table.path { + tracked_paths.push(path.clone()); } - Err(e) => ScrubResult { - file_path: file_path.to_string_lossy().to_string(), - ok: false, - error: Some(format!("read error: {}", e)), - file_size, - }, - }; + } + } + drop(core); - results.push(result); + // Scan disk for .sst files + let disk_entries = match std::fs::read_dir(&sst_dir) { + Ok(entries) => entries, + Err(e) => { + return vec![ScrubResult { + file_path: sst_dir, + valid: false, + total_blocks: 0, + verified_blocks: 0, + corrupt_blocks: 0, + total_bytes: 0, + corrupt_bytes: 0, + errors: vec![format!("Failed to read SSTable directory: {}", e)], + }]; + } + }; + + let disk_paths: Vec = disk_entries + .filter_map(|e| e.ok()) + .map(|e| e.path()) + .filter(|p| p.extension().and_then(|s| s.to_str()) == Some("sst")) + .collect(); + + // Detect orphan files: on disk but not tracked by VersionSet + for disk_path in &disk_paths { + if !tracked_paths.contains(disk_path) { + let file_size = std::fs::metadata(disk_path).map(|m| m.len()).unwrap_or(0); + results.push(ScrubResult { + file_path: disk_path.clone(), + valid: false, + total_blocks: 0, + verified_blocks: 0, + corrupt_blocks: 0, + total_bytes: file_size, + corrupt_bytes: 0, + errors: vec!["Orphan SSTable file: not tracked by VersionSet".to_string()], + }); + } } - Ok(results) + // Detect orphan tables: tracked by VersionSet but file missing from disk + for tracked_path in &tracked_paths { + if !tracked_path.exists() { + results.push(ScrubResult { + file_path: tracked_path.clone(), + valid: false, + total_blocks: 0, + verified_blocks: 0, + corrupt_blocks: 0, + total_bytes: 0, + corrupt_bytes: 0, + errors: vec!["Orphan table: SSTable file not found on disk".to_string()], + }); + } + } + + results } #[cfg(test)] mod tests { use super::*; - use std::io::Write; + use crate::core::log_record::LogRecord; + use crate::infra::config::LsmConfig; + use crate::storage::builder::SstableBuilder; + use crate::storage::cache::NoopCache; + use std::io::{Seek, Write}; use std::time::Duration; + // ── DataScrubber tests ───────────────────────────────────────────────── + #[test] fn test_scrub_empty_directory() { let dir = tempfile::TempDir::new().unwrap(); @@ -164,30 +421,27 @@ mod tests { } #[test] - fn test_scrub_valid_sst_file() { + fn test_scrub_bad_magic_file() { let dir = tempfile::TempDir::new().unwrap(); let sst_path = dir.path().join("test.sst"); - // Write a valid-looking SSTable (header + data). + // Write a file with invalid magic. let mut f = std::fs::File::create(&sst_path).unwrap(); - f.write_all(b"APXSTORE").unwrap(); // magic - f.write_all(&[2u8]).unwrap(); // version + f.write_all(b"APXSTORE").unwrap(); // not LSMSST03 + f.write_all(&[2u8]).unwrap(); f.write_all(b"some payload data here").unwrap(); f.flush().unwrap(); - let mut scrubber = DataScrubber::new(dir.path().to_str().unwrap()); - scrubber.start_scrubbing(Duration::from_millis(50)); - std::thread::sleep(Duration::from_millis(150)); - scrubber.stop(); - - let results = scrubber.results(); - assert_eq!(results.len(), 1); - assert!(results[0].ok, "valid .sst file should pass scrub"); - assert!(results[0].error.is_none()); + let result = scrub_file(&sst_path); + assert!(!result.valid, "file with invalid magic should fail scrub"); + assert!(result + .errors + .iter() + .any(|e| e.contains("Invalid magic number"))); } #[test] - fn test_scrub_corrupted_sst_file() { + fn test_scrub_too_small_file() { let dir = tempfile::TempDir::new().unwrap(); let sst_path = dir.path().join("bad.sst"); @@ -196,14 +450,194 @@ mod tests { f.write_all(b"BAD!").unwrap(); f.flush().unwrap(); - let mut scrubber = DataScrubber::new(dir.path().to_str().unwrap()); - scrubber.start_scrubbing(Duration::from_millis(50)); - std::thread::sleep(Duration::from_millis(150)); - scrubber.stop(); + let result = scrub_file(&sst_path); + assert!(!result.valid, "corrupted .sst file should fail scrub"); + assert!(result.errors.iter().any(|e| e.contains("File too small"))); + } - let results = scrubber.results(); - assert_eq!(results.len(), 1); - assert!(!results[0].ok, "corrupted .sst file should fail scrub"); - assert!(results[0].error.is_some()); + // ── CRC32 validity tests ─────────────────────────────────────────────── + + #[test] + fn test_scrub_crc32_valid() { + let dir = tempfile::TempDir::new().unwrap(); + let sst_path = dir.path().join("valid.sst"); + + // Build a proper SSTable with SstableBuilder + let config = crate::infra::config::StorageConfig::default(); + let timestamp = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos(); + let mut builder = SstableBuilder::new(sst_path.clone(), config, timestamp).unwrap(); + + builder + .add( + b"key1", + &LogRecord::new(b"key1".to_vec(), b"value1".to_vec()), + ) + .unwrap(); + builder + .add( + b"key2", + &LogRecord::new(b"key2".to_vec(), b"value2".to_vec()), + ) + .unwrap(); + let path = builder.finish().unwrap(); + + let result = scrub_file(&path); + assert!(result.valid, "valid SSTable should pass CRC32 check"); + assert!(result.errors.is_empty()); + assert!( + result.total_blocks > 0, + "should have at least one data block" + ); + assert_eq!(result.corrupt_blocks, 0); + assert_eq!( + result.verified_blocks, result.total_blocks, + "all blocks should be verified" + ); + assert!(result.total_bytes > 0, "should have verified some bytes"); + } + + #[test] + fn test_scrub_crc32_corrupt() { + let dir = tempfile::TempDir::new().unwrap(); + let sst_path = dir.path().join("corrupt.sst"); + + // Build a proper SSTable + let config = crate::infra::config::StorageConfig::default(); + let timestamp = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos(); + let mut builder = SstableBuilder::new(sst_path.clone(), config, timestamp).unwrap(); + + builder + .add( + b"key1", + &LogRecord::new(b"key1".to_vec(), b"value1".to_vec()), + ) + .unwrap(); + builder + .add( + b"key2", + &LogRecord::new(b"key2".to_vec(), b"value2".to_vec()), + ) + .unwrap(); + let path = builder.finish().unwrap(); + + // Corrupt the first data block by writing garbage after the magic + { + use std::io::SeekFrom; + let mut file = std::fs::OpenOptions::new().write(true).open(&path).unwrap(); + // Overwrite bytes starting at offset 8 (right after magic) with garbage + let garbage = vec![0xFF; 30]; + file.seek(SeekFrom::Start(8)).unwrap(); + file.write_all(&garbage).unwrap(); + } + + let result = scrub_file(&path); + assert!(!result.valid, "corrupted SSTable should fail CRC32 check"); + assert!( + result.corrupt_blocks > 0, + "should have detected corrupt blocks" + ); + assert!(result.corrupt_bytes > 0, "corrupt_bytes should be > 0"); + assert!(!result.errors.is_empty()); + assert!( + result.errors.iter().any(|e| e.contains("CRC32 mismatch")), + "error message should mention CRC32 mismatch" + ); + } + + // ── Orphan detection tests (require engine integration) ──────────────── + + #[test] + fn test_scrub_orphan_file() { + let dir = tempfile::TempDir::new().unwrap(); + let sst_dir = dir.path().join("sstables"); + std::fs::create_dir_all(&sst_dir).unwrap(); + + // Create the engine FIRST so discover_sstables_from_disk won't pick up + // the file we're about to create + let mut config = LsmConfig::default(); + config.core.dir_path = dir.path().to_path_buf(); + let engine = Engine::new_from_config(&config, NoopCache).unwrap(); + + // Now create an orphan .sst file AFTER engine init + let orphan_path = sst_dir.join("orphan.sst"); + { + let mut f = std::fs::File::create(&orphan_path).unwrap(); + f.write_all(b"LSMSST03").unwrap(); + f.write_all(&[0u8; 20]).unwrap(); + } + + let results = scrub_with_version_set(&engine); + + // Should detect the orphan file + let orphan_results: Vec<&ScrubResult> = results + .iter() + .filter(|r| r.file_path == orphan_path) + .collect(); + assert_eq!(orphan_results.len(), 1, "should find orphan .sst file"); + assert!(!orphan_results[0].valid); + assert!( + orphan_results[0] + .errors + .iter() + .any(|e| e.contains("Orphan SSTable")), + "orphan file error should mention 'Orphan SSTable'" + ); + } + + #[test] + fn test_scrub_orphan_table() { + use crate::core::table::Table; + use std::collections::BTreeMap; + + let dir = tempfile::TempDir::new().unwrap(); + let sst_dir = dir.path().join("sstables"); + std::fs::create_dir_all(&sst_dir).unwrap(); + + let mut config = LsmConfig::default(); + config.core.dir_path = dir.path().to_path_buf(); + let engine = Engine::new_from_config(&config, NoopCache).unwrap(); + + // Manually add a table with a path that doesn't exist + let fake_path = sst_dir.join("nonexistent.sst"); + let orphan_table = Table { + data: BTreeMap::new(), + level: 0, + path: Some(fake_path.clone()), + min_key: b"a".to_vec(), + max_key: b"z".to_vec(), + bloom_filter: None, + }; + + { + let mut core = engine.lock_core(); + core.version_set_mut().add_table("default", orphan_table); + } + + let results = scrub_with_version_set(&engine); + + // Should detect the orphan table + let table_results: Vec<&ScrubResult> = results + .iter() + .filter(|r| r.file_path == fake_path) + .collect(); + assert_eq!( + table_results.len(), + 1, + "should detect orphan table with missing file" + ); + assert!(!table_results[0].valid); + assert!( + table_results[0] + .errors + .iter() + .any(|e| e.contains("Orphan table")), + "orphan table error should mention 'Orphan table'" + ); } } diff --git a/src/storage/block.rs b/src/storage/block.rs index 331ad25..a086dbf 100644 --- a/src/storage/block.rs +++ b/src/storage/block.rs @@ -52,7 +52,8 @@ impl Block { return; } let (new_data, new_offsets) = - PrefixCompressor::compress_block_data(&self.data, &self.offsets); + PrefixCompressor::compress_block_data(&self.data, &self.offsets) + .expect("compress_block_data should not fail with valid input"); self.data = new_data; self.offsets = new_offsets; self.flags |= PREFIX_COMPRESSION_FLAG; diff --git a/src/storage/cache.rs b/src/storage/cache.rs index 453ce47..f32bb2d 100644 --- a/src/storage/cache.rs +++ b/src/storage/cache.rs @@ -1,6 +1,7 @@ use lru::LruCache; +use parking_lot::Mutex; use std::num::NonZeroUsize; -use std::sync::{Arc, Mutex}; +use std::sync::Arc; pub trait Cache: Clone + Send + Sync + 'static {} @@ -46,17 +47,17 @@ impl GlobalBlockCache { } pub fn get(&self, table_id: u64, block_idx: usize) -> Option> { - let mut cache = self.cache.lock().unwrap_or_else(|e| e.into_inner()); + let mut cache = self.cache.lock(); cache.get(&(table_id, block_idx)).cloned() } pub fn put(&self, table_id: u64, block_idx: usize, data: Vec) { - let mut cache = self.cache.lock().unwrap_or_else(|e| e.into_inner()); + let mut cache = self.cache.lock(); cache.put((table_id, block_idx), data); } pub fn stats(&self) -> CacheStats { - let cache = self.cache.lock().unwrap_or_else(|e| e.into_inner()); + let cache = self.cache.lock(); CacheStats { len: cache.len(), cap: cache.cap().get(), diff --git a/src/storage/prefix_compression.rs b/src/storage/prefix_compression.rs index e814e7c..a47fe3d 100644 --- a/src/storage/prefix_compression.rs +++ b/src/storage/prefix_compression.rs @@ -30,7 +30,7 @@ //! assert_eq!(keys, decoded); //! ``` -use crate::infra::error::Result; +use crate::infra::error::{LsmError, Result}; /// Maximum shared prefix length supported by the u8 encoding (255 bytes). /// Per-key suffix length is stored as u16, allowing suffixes up to 65535 bytes. @@ -52,9 +52,9 @@ impl PrefixCompressor { /// # Panics /// /// Panics if any two consecutive keys share more than 255 prefix bytes. - pub fn encode_keys(keys: &[Vec]) -> Vec { + pub fn encode_keys(keys: &[Vec]) -> Result> { if keys.is_empty() { - return Vec::new(); + return Ok(Vec::new()); } let mut output = Vec::new(); @@ -62,12 +62,12 @@ impl PrefixCompressor { for key in keys { let shared = Self::shared_prefix_len(prev_key, key); - debug_assert!( - shared <= MAX_SHARED_PREFIX, - "shared prefix length {} exceeds maximum {}", - shared, - MAX_SHARED_PREFIX - ); + if shared > MAX_SHARED_PREFIX { + return Err(LsmError::InvalidArgument(format!( + "shared prefix length {} exceeds maximum {}", + shared, MAX_SHARED_PREFIX, + ))); + } let suffix = &key[shared..]; let suffix_len = suffix.len(); @@ -79,7 +79,7 @@ impl PrefixCompressor { prev_key = key; } - output + Ok(output) } /// Decode a prefix-compressed key sequence back into full keys. @@ -94,11 +94,11 @@ impl PrefixCompressor { /// # Panics /// /// Panics if `data` is malformed (truncated, invalid lengths, etc.). - pub fn decode_keys(data: &[u8], first_key: &[u8]) -> Vec> { + pub fn decode_keys(data: &[u8], first_key: &[u8]) -> Result>> { if data.is_empty() { // When there are no encoded keys, just the first_key is the only key. // This is the case when we have a block with a single entry. - return Vec::new(); + return Ok(Vec::new()); } let mut keys: Vec> = Vec::new(); @@ -110,13 +110,17 @@ impl PrefixCompressor { pos += 1; if pos + 2 > data.len() { - panic!("Truncated prefix compression data: cannot read suffix_len"); + return Err(LsmError::CorruptedData( + "Truncated prefix compression data: cannot read suffix_len".to_string(), + )); } let suffix_len = u16::from_le_bytes([data[pos], data[pos + 1]]) as usize; pos += 2; if pos + suffix_len > data.len() { - panic!("Truncated prefix compression data: suffix extends past end"); + return Err(LsmError::CorruptedData( + "Truncated prefix compression data: suffix extends past end".to_string(), + )); } let suffix = &data[pos..pos + suffix_len]; pos += suffix_len; @@ -130,7 +134,7 @@ impl PrefixCompressor { prev_key = keys.last().expect("just pushed").clone(); } - keys + Ok(keys) } /// Compress the keys of a block's entries in-place (builds new data + offsets). @@ -146,9 +150,9 @@ impl PrefixCompressor { /// For entries 1..N, keys are stored as: /// `[shared_prefix_len(u8)][suffix_len(u16)][suffix]` /// Values are stored as-is: `[val_len(u16)][value_bytes]` - pub fn compress_block_data(data: &[u8], offsets: &[u32]) -> (Vec, Vec) { + pub fn compress_block_data(data: &[u8], offsets: &[u32]) -> Result<(Vec, Vec)> { if offsets.is_empty() { - return (Vec::new(), Vec::new()); + return Ok((Vec::new(), Vec::new())); } let mut new_data = Vec::new(); @@ -175,7 +179,12 @@ impl PrefixCompressor { } else { // Subsequent entries: prefix-compressed key let shared = Self::shared_prefix_len(prev_key, key); - debug_assert!(shared <= MAX_SHARED_PREFIX); + if shared > MAX_SHARED_PREFIX { + return Err(LsmError::InvalidArgument(format!( + "shared prefix length {} exceeds maximum {}", + shared, MAX_SHARED_PREFIX, + ))); + } let suffix = &key[shared..]; new_data.push(shared as u8); new_data.extend_from_slice(&(suffix.len() as u16).to_le_bytes()); @@ -189,7 +198,7 @@ impl PrefixCompressor { prev_key = key; } - (new_data, new_offsets) + Ok((new_data, new_offsets)) } /// Decompress prefix-compressed block data back to the standard format. @@ -301,18 +310,18 @@ mod tests { #[test] fn test_encode_decode_empty() { let keys: Vec> = vec![]; - let compressed = PrefixCompressor::encode_keys(&keys); + let compressed = PrefixCompressor::encode_keys(&keys).unwrap(); assert!(compressed.is_empty()); - let decoded = PrefixCompressor::decode_keys(&compressed, b"first_key"); + let decoded = PrefixCompressor::decode_keys(&compressed, b"first_key").unwrap(); assert!(decoded.is_empty()); } #[test] fn test_encode_decode_single_key() { let keys = vec![b"hello".to_vec()]; - let compressed = PrefixCompressor::encode_keys(&keys); - let decoded = PrefixCompressor::decode_keys(&compressed, &keys[0]); + let compressed = PrefixCompressor::encode_keys(&keys).unwrap(); + let decoded = PrefixCompressor::decode_keys(&compressed, &keys[0]).unwrap(); assert_eq!(keys, decoded); } @@ -324,16 +333,16 @@ mod tests { b"user:carol:age".to_vec(), b"user:dave:score".to_vec(), ]; - let compressed = PrefixCompressor::encode_keys(&keys); - let decoded = PrefixCompressor::decode_keys(&compressed, &keys[0]); + let compressed = PrefixCompressor::encode_keys(&keys).unwrap(); + let decoded = PrefixCompressor::decode_keys(&compressed, &keys[0]).unwrap(); assert_eq!(keys, decoded); } #[test] fn test_encode_decode_no_shared_prefix() { let keys = vec![b"aaaa".to_vec(), b"bbbb".to_vec(), b"cccc".to_vec()]; - let compressed = PrefixCompressor::encode_keys(&keys); - let decoded = PrefixCompressor::decode_keys(&compressed, &keys[0]); + let compressed = PrefixCompressor::encode_keys(&keys).unwrap(); + let decoded = PrefixCompressor::decode_keys(&compressed, &keys[0]).unwrap(); assert_eq!(keys, decoded); } @@ -344,8 +353,8 @@ mod tests { b"samekey".to_vec(), b"samekey".to_vec(), ]; - let compressed = PrefixCompressor::encode_keys(&keys); - let decoded = PrefixCompressor::decode_keys(&compressed, &keys[0]); + let compressed = PrefixCompressor::encode_keys(&keys).unwrap(); + let decoded = PrefixCompressor::decode_keys(&compressed, &keys[0]).unwrap(); assert_eq!(keys, decoded); } @@ -358,8 +367,8 @@ mod tests { k.push(b'a' + i); keys.push(k); } - let compressed = PrefixCompressor::encode_keys(&keys); - let decoded = PrefixCompressor::decode_keys(&compressed, &keys[0]); + let compressed = PrefixCompressor::encode_keys(&keys).unwrap(); + let decoded = PrefixCompressor::decode_keys(&compressed, &keys[0]).unwrap(); assert_eq!(keys, decoded); } @@ -390,7 +399,8 @@ mod tests { data.extend_from_slice(&(2u16).to_le_bytes()); // val_len data.extend_from_slice(b"v3"); - let (compressed_data, new_offsets) = PrefixCompressor::compress_block_data(&data, &offsets); + let (compressed_data, new_offsets) = + PrefixCompressor::compress_block_data(&data, &offsets).unwrap(); // First entry should be full key "aaa" let key0_len = u16::from_le_bytes([compressed_data[0], compressed_data[1]]) as usize; @@ -447,7 +457,7 @@ mod tests { } let (compressed_data, compressed_offsets) = - PrefixCompressor::compress_block_data(&data, &offsets); + PrefixCompressor::compress_block_data(&data, &offsets).unwrap(); let (decompressed_data, decompressed_offsets) = PrefixCompressor::decompress_block_data(&compressed_data, &compressed_offsets).unwrap(); @@ -466,7 +476,7 @@ mod tests { data.extend_from_slice(b"val"); let (compressed_data, compressed_offsets) = - PrefixCompressor::compress_block_data(&data, &offsets); + PrefixCompressor::compress_block_data(&data, &offsets).unwrap(); let (decompressed_data, decompressed_offsets) = PrefixCompressor::decompress_block_data(&compressed_data, &compressed_offsets).unwrap();