From 9b374e4ba3ff053da08396bc9a830e25a1bb03eb Mon Sep 17 00:00:00 2001 From: Beinan Date: Wed, 15 Jul 2026 01:07:39 +0000 Subject: [PATCH] fix(storage): harden URI and store name handling --- crates/lance-context-api/src/lib.rs | 4 + crates/lance-context-core/src/lib.rs | 2 + crates/lance-context-core/src/namespace.rs | 9 +- crates/lance-context-core/src/storage.rs | 140 ++++++++++++++++++ crates/lance-context-master/src/discovery.rs | 24 +-- crates/lance-context-master/src/main.rs | 5 +- crates/lance-context-master/src/state.rs | 24 +-- crates/lance-context-server/src/main.rs | 6 +- .../src/routes/contexts.rs | 45 ++++++ .../src/routes/rollouts.rs | 47 ++++++ crates/lance-context-server/src/state.rs | 41 +++-- 11 files changed, 274 insertions(+), 73 deletions(-) create mode 100644 crates/lance-context-core/src/storage.rs diff --git a/crates/lance-context-api/src/lib.rs b/crates/lance-context-api/src/lib.rs index 03530c73..d5ad06ab 100644 --- a/crates/lance-context-api/src/lib.rs +++ b/crates/lance-context-api/src/lib.rs @@ -148,6 +148,8 @@ pub trait RolloutStoreApi { #[derive(Debug, Serialize, Deserialize)] pub struct CreateContextRequest { + /// Portable dataset name: 1-128 ASCII characters matching + /// `[A-Za-z0-9_][A-Za-z0-9._-]*`; `_registry` and `_stats` are reserved. pub name: String, #[serde(default)] pub storage_options: Option>, @@ -571,6 +573,8 @@ pub struct CompactStatsResponse { #[derive(Debug, Serialize, Deserialize)] pub struct CreateRolloutStoreRequest { + /// Portable dataset name: 1-128 ASCII characters matching + /// `[A-Za-z0-9_][A-Za-z0-9._-]*`; `_registry` and `_stats` are reserved. pub name: String, #[serde(default)] pub storage_options: Option>, diff --git a/crates/lance-context-core/src/lib.rs b/crates/lance-context-core/src/lib.rs index 9e0d6559..70321269 100644 --- a/crates/lance-context-core/src/lib.rs +++ b/crates/lance-context-core/src/lib.rs @@ -12,6 +12,7 @@ mod registry; mod rollout; mod rollout_store; pub mod serde; +mod storage; mod store; pub use context::{Context, ContextEntry, Snapshot}; @@ -35,6 +36,7 @@ pub use record::{ pub use registry::{RegistryEntry, RolloutRegistry}; pub use rollout::{RolloutRecord, ROLE_ARTIFACT, ROLE_ASSISTANT, ROLE_GRADE, ROLE_TOOL}; pub use rollout_store::{rollout_schema, RolloutObservation, RolloutStore, RolloutStoreOptions}; +pub use storage::{create_local_dir_if_needed, join_uri, validate_store_name, MAX_STORE_NAME_LEN}; pub use store::{ CompactionConfig, CompactionStats, ContextStore, ContextStoreOptions, DistanceMetric, IdIndexType, ReadProjection, diff --git a/crates/lance-context-core/src/namespace.rs b/crates/lance-context-core/src/namespace.rs index 50d678af..c03ca12f 100644 --- a/crates/lance-context-core/src/namespace.rs +++ b/crates/lance-context-core/src/namespace.rs @@ -11,6 +11,7 @@ use lance::io::{ObjectStoreParams, StorageOptionsAccessor}; use lance::{Error as LanceError, Result as LanceResult}; use uuid::Uuid; +use crate::storage::join_uri; use crate::store::{ContextStore, ContextStoreOptions}; const MANIFEST_TABLE_NAME: &str = "__manifest"; @@ -432,14 +433,6 @@ fn normalize_root_uri(root_uri: String) -> LanceResult { Ok(root_uri.trim_end_matches('/').to_string()) } -fn join_uri(root_uri: &str, child: &str) -> String { - if root_uri == "/" { - format!("/{child}") - } else { - format!("{root_uri}/{child}") - } -} - fn invalid_input(message: impl Into) -> LanceError { LanceError::from(ArrowError::InvalidArgumentError(message.into())) } diff --git a/crates/lance-context-core/src/storage.rs b/crates/lance-context-core/src/storage.rs new file mode 100644 index 00000000..17f4d355 --- /dev/null +++ b/crates/lance-context-core/src/storage.rs @@ -0,0 +1,140 @@ +//! Shared helpers for dataset roots that may be local paths or object-store URIs. + +use std::io; +use std::path::Path; + +/// Maximum accepted length for a context or rollout store name. +pub const MAX_STORE_NAME_LEN: usize = 128; + +/// Join one dataset child onto a local directory or object-store URI. +#[must_use] +pub fn join_uri(base: &str, child: &str) -> String { + if has_uri_scheme(base) { + format!("{}/{}", base.trim_end_matches('/'), child) + } else { + Path::new(base).join(child).to_string_lossy().to_string() + } +} + +/// Create `data_dir` when it is a local filesystem path. +/// +/// Object-store roots such as `s3://`, `gs://`, and `az://` must never be +/// passed to `std::fs`, where they would be interpreted as relative paths. +pub fn create_local_dir_if_needed(data_dir: &str) -> io::Result<()> { + if has_uri_scheme(data_dir) { + return Ok(()); + } + std::fs::create_dir_all(data_dir) +} + +/// Validate a logical context/rollout store name before using it in a URI. +/// +/// Names are intentionally limited to a portable single-segment form so they +/// cannot escape `data_dir` or acquire platform-specific path semantics. +pub fn validate_store_name(name: &str) -> Result<(), String> { + if name.is_empty() { + return Err("store name must not be empty".to_string()); + } + + let mut chars = name.chars(); + let first = chars.next().expect("non-empty checked above"); + if !(first.is_ascii_alphanumeric() || first == '_') { + return Err("store name must start with an ASCII letter, digit, or '_'".to_string()); + } + if !chars.all(|ch| ch.is_ascii_alphanumeric() || matches!(ch, '-' | '_' | '.')) { + return Err( + "store name may contain only ASCII letters, digits, '.', '-', and '_'".to_string(), + ); + } + if matches!(name, "_registry" | "_stats") { + return Err(format!("store name '{name}' is reserved")); + } + if name.len() > MAX_STORE_NAME_LEN { + return Err(format!( + "store name must be at most {MAX_STORE_NAME_LEN} characters" + )); + } + Ok(()) +} + +fn has_uri_scheme(value: &str) -> bool { + let Some((scheme, _)) = value.split_once("://") else { + return false; + }; + let mut chars = scheme.chars(); + matches!(chars.next(), Some(ch) if ch.is_ascii_alphabetic()) + && chars.all(|ch| ch.is_ascii_alphanumeric() || matches!(ch, '+' | '-' | '.')) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn joins_local_paths_and_object_store_uris() { + assert_eq!( + join_uri("/data/rollouts", "exp.rollout.lance"), + "/data/rollouts/exp.rollout.lance" + ); + assert_eq!( + join_uri("relative/data/", "exp.rollout.lance"), + "relative/data/exp.rollout.lance" + ); + assert_eq!( + join_uri("s3://bucket/prefix/", "exp.rollout.lance"), + "s3://bucket/prefix/exp.rollout.lance" + ); + assert_eq!( + join_uri("gs://bucket", "_registry.rollout.lance"), + "gs://bucket/_registry.rollout.lance" + ); + } + + #[test] + fn recognizes_uri_schemes_without_treating_local_paths_as_uris() { + assert!(has_uri_scheme("s3://bucket/prefix")); + assert!(has_uri_scheme("az://container/prefix")); + assert!(has_uri_scheme("file:///tmp/data")); + assert!(!has_uri_scheme("/tmp/data")); + assert!(!has_uri_scheme("./data")); + assert!(!has_uri_scheme("experiment:data")); + } + + #[test] + fn creates_only_local_directories() { + let root = tempfile::TempDir::new().unwrap(); + let nested = root.path().join("nested/data"); + create_local_dir_if_needed(nested.to_str().unwrap()).unwrap(); + assert!(nested.is_dir()); + + create_local_dir_if_needed("s3://bucket/prefix").unwrap(); + create_local_dir_if_needed("gs://bucket/prefix").unwrap(); + } + + #[test] + fn validates_portable_single_segment_store_names() { + for name in [ + "exp-alpha", + "conntest_lianghongxu", + "run.2026_07-15", + "_private", + "7", + ] { + validate_store_name(name).unwrap(); + } + + for name in [ + "", + "../outside", + "nested/store", + r"nested\store", + "_registry", + "has space", + "experiment:name", + ] { + assert!(validate_store_name(name).is_err(), "{name}"); + } + assert!(validate_store_name(&"a".repeat(MAX_STORE_NAME_LEN)).is_ok()); + assert!(validate_store_name(&"a".repeat(MAX_STORE_NAME_LEN + 1)).is_err()); + } +} diff --git a/crates/lance-context-master/src/discovery.rs b/crates/lance-context-master/src/discovery.rs index 4f77af8e..c4713fed 100644 --- a/crates/lance-context-master/src/discovery.rs +++ b/crates/lance-context-master/src/discovery.rs @@ -1,9 +1,7 @@ //! Startup discovery for rollout datasets created before the registry existed. -use std::path::Path; - use lance::io::ObjectStore; -use lance_context_core::RolloutRegistry; +use lance_context_core::{join_uri, RolloutRegistry}; const ROLLOUT_SUFFIX: &str = ".rollout.lance"; @@ -31,14 +29,6 @@ pub async fn backfill_registry( registry.insert_missing(&entries).await } -fn join_uri(base: &str, child: &str) -> String { - if base.contains("://") { - format!("{}/{}", base.trim_end_matches('/'), child) - } else { - Path::new(base).join(child).to_string_lossy().to_string() - } -} - #[cfg(test)] mod tests { use super::*; @@ -95,16 +85,4 @@ mod tests { names.sort(); assert_eq!(names, vec!["existing", "legacy"]); } - - #[test] - fn joins_local_paths_and_object_store_uris() { - assert_eq!( - join_uri("/data/rollouts", "exp.rollout.lance"), - "/data/rollouts/exp.rollout.lance" - ); - assert_eq!( - join_uri("s3://bucket/prefix/", "exp.rollout.lance"), - "s3://bucket/prefix/exp.rollout.lance" - ); - } } diff --git a/crates/lance-context-master/src/main.rs b/crates/lance-context-master/src/main.rs index 1bc95ee1..54983048 100644 --- a/crates/lance-context-master/src/main.rs +++ b/crates/lance-context-master/src/main.rs @@ -6,6 +6,7 @@ use tower_http::services::ServeDir; use tower_http::trace::TraceLayer; use tracing_subscriber::EnvFilter; +use lance_context_core::create_local_dir_if_needed; use lance_context_master::config::MasterConfig; use lance_context_master::routes; use lance_context_master::scanner; @@ -21,9 +22,9 @@ async fn main() { let config = MasterConfig::parse(); let addr = format!("{}:{}", config.host, config.port); - if let Err(e) = std::fs::create_dir_all(&config.data_dir) { + if let Err(e) = create_local_dir_if_needed(&config.data_dir) { tracing::error!( - "Failed to create data directory '{}': {}", + "Failed to create local data directory '{}': {}", config.data_dir, e ); diff --git a/crates/lance-context-master/src/state.rs b/crates/lance-context-master/src/state.rs index d0393470..c3bc0f93 100644 --- a/crates/lance-context-master/src/state.rs +++ b/crates/lance-context-master/src/state.rs @@ -1,11 +1,10 @@ //! Shared master state: durable registry + stats table. use std::collections::HashMap; -use std::path::PathBuf; use std::sync::Arc; use lance_context_api::CompactJobStatus; -use lance_context_core::RolloutRegistry; +use lance_context_core::{join_uri, RolloutRegistry}; use tokio::sync::{mpsc, Mutex, RwLock}; use crate::config::MasterConfig; @@ -25,7 +24,7 @@ pub struct MasterState { /// Periodically-refreshed per-experiment metrics (master-owned). pub stats: Mutex, /// Shared data directory / object-store prefix. - pub base_path: PathBuf, + pub base_uri: String, /// Effective configuration. pub config: MasterConfig, /// Enqueues an experiment name for (serial) compaction. Both automatic @@ -41,15 +40,9 @@ pub struct MasterState { impl MasterState { /// Open (or create) the registry and stats datasets under `data_dir`. pub async fn new(config: MasterConfig) -> lance::Result> { - let base_path = PathBuf::from(&config.data_dir); - let registry_uri = base_path - .join("_registry.rollout.lance") - .to_string_lossy() - .to_string(); - let stats_uri = base_path - .join("_stats.rollout.lance") - .to_string_lossy() - .to_string(); + let base_uri = config.data_dir.clone(); + let registry_uri = join_uri(&base_uri, "_registry.rollout.lance"); + let stats_uri = join_uri(&base_uri, "_stats.rollout.lance"); let mut registry = RolloutRegistry::open_or_create(®istry_uri, None).await?; let backfilled = discovery::backfill_registry(&config.data_dir, &mut registry).await?; if backfilled > 0 { @@ -63,7 +56,7 @@ impl MasterState { Ok(Arc::new(Self { registry: RwLock::new(registry), stats: Mutex::new(stats), - base_path, + base_uri, config, compact_tx, compact_rx: Mutex::new(Some(compact_rx)), @@ -74,10 +67,7 @@ impl MasterState { /// Physical rollout dataset URI for `name`, matching the data-plane's /// `rollout_uri` convention (`{name}.rollout.lance`). pub fn rollout_uri(&self, name: &str) -> String { - self.base_path - .join(format!("{}.rollout.lance", name)) - .to_string_lossy() - .to_string() + join_uri(&self.base_uri, &format!("{}.rollout.lance", name)) } } diff --git a/crates/lance-context-server/src/main.rs b/crates/lance-context-server/src/main.rs index 6e300b5a..f76f29a6 100644 --- a/crates/lance-context-server/src/main.rs +++ b/crates/lance-context-server/src/main.rs @@ -11,6 +11,8 @@ use tower_http::cors::CorsLayer; use tower_http::trace::TraceLayer; use tracing_subscriber::EnvFilter; +use lance_context_core::create_local_dir_if_needed; + use crate::config::ServerConfig; use crate::state::AppState; @@ -23,9 +25,9 @@ async fn main() { let config = ServerConfig::parse(); let addr = format!("{}:{}", config.host, config.port); - if let Err(e) = std::fs::create_dir_all(&config.data_dir) { + if let Err(e) = create_local_dir_if_needed(&config.data_dir) { tracing::error!( - "Failed to create data directory '{}': {}", + "Failed to create local data directory '{}': {}", config.data_dir, e ); diff --git a/crates/lance-context-server/src/routes/contexts.rs b/crates/lance-context-server/src/routes/contexts.rs index fbc45389..bbfb8315 100644 --- a/crates/lance-context-server/src/routes/contexts.rs +++ b/crates/lance-context-server/src/routes/contexts.rs @@ -14,6 +14,7 @@ pub async fn create_context( State(state): State>, Json(req): Json, ) -> Result<(axum::http::StatusCode, Json), AppError> { + AppState::validate_name(&req.name)?; let stores = state.stores.read().await; if stores.contains_key(&req.name) { return Err(AppError::AlreadyExists(format!( @@ -107,6 +108,7 @@ pub async fn delete_context( State(state): State>, Path(name): Path, ) -> Result { + AppState::validate_name(&name)?; let mut stores = state.stores.write().await; if stores.remove(&name).is_none() { return Err(AppError::NotFound(format!( @@ -122,3 +124,46 @@ pub async fn delete_context( Ok(axum::http::StatusCode::NO_CONTENT) } + +#[cfg(test)] +mod tests { + use super::*; + use tempfile::TempDir; + + fn create_request(name: &str) -> CreateContextRequest { + CreateContextRequest { + name: name.to_string(), + storage_options: None, + id_index_type: None, + blob_columns: None, + embedding_dim: None, + distance_metric: None, + } + } + + #[tokio::test] + async fn create_rejects_names_that_are_not_portable_path_segments() { + let dir = TempDir::new().unwrap(); + let state = Arc::new(AppState::new_for_test(dir.path().to_path_buf()).await); + + for name in ["../escape", "nested/store", r"nested\store", "_registry"] { + let err = create_context(State(state.clone()), Json(create_request(name))) + .await + .unwrap_err(); + assert!(matches!(err, AppError::InvalidRequest(_)), "{name}"); + } + assert!(state.stores.read().await.is_empty()); + } + + #[tokio::test] + async fn open_rejects_invalid_context_names_before_storage_access() { + let dir = TempDir::new().unwrap(); + let state = AppState::new_for_test(dir.path().to_path_buf()).await; + + let err = match state.get_or_open_context_store("../escape").await { + Ok(_) => panic!("invalid name unexpectedly opened"), + Err(err) => err, + }; + assert!(matches!(err, AppError::InvalidRequest(_))); + } +} diff --git a/crates/lance-context-server/src/routes/rollouts.rs b/crates/lance-context-server/src/routes/rollouts.rs index ccdae359..e0288da9 100644 --- a/crates/lance-context-server/src/routes/rollouts.rs +++ b/crates/lance-context-server/src/routes/rollouts.rs @@ -30,6 +30,7 @@ pub async fn create_rollout_store( State(state): State>, Json(req): Json, ) -> Result<(StatusCode, Json), AppError> { + AppState::validate_name(&req.name)?; // Existence is tracked durably in the registry, not by cache membership. if state .rollout_registry @@ -582,6 +583,52 @@ mod tests { (state, dir) } + #[tokio::test] + async fn create_rejects_names_that_are_not_portable_path_segments() { + let dir = TempDir::new().unwrap(); + let state = Arc::new(AppState::new_for_test(dir.path().to_path_buf()).await); + + for name in [ + "../escape", + "nested/store", + r"nested\store", + "_registry", + &"a".repeat(lance_context_core::MAX_STORE_NAME_LEN + 1), + ] { + let err = create_rollout_store( + State(state.clone()), + Json(CreateRolloutStoreRequest { + name: name.to_string(), + storage_options: None, + }), + ) + .await + .unwrap_err(); + assert!(matches!(err, AppError::InvalidRequest(_)), "{name}"); + } + assert!(state.rollout_stores.lock().await.is_empty()); + assert!(state + .rollout_registry + .write() + .await + .list() + .await + .unwrap() + .is_empty()); + } + + #[tokio::test] + async fn open_rejects_invalid_rollout_names_before_registry_access() { + let dir = TempDir::new().unwrap(); + let state = AppState::new_for_test(dir.path().to_path_buf()).await; + + let err = match state.get_or_open_rollout_store("../escape").await { + Ok(_) => panic!("invalid name unexpectedly opened"), + Err(err) => err, + }; + assert!(matches!(err, AppError::InvalidRequest(_))); + } + fn record_with_size(id: &str, payload_size: Option) -> AddRolloutRequest { AddRolloutRequest { id: id.to_string(), diff --git a/crates/lance-context-server/src/state.rs b/crates/lance-context-server/src/state.rs index 913dd8c6..ba0508d8 100644 --- a/crates/lance-context-server/src/state.rs +++ b/crates/lance-context-server/src/state.rs @@ -1,10 +1,12 @@ use std::num::NonZeroUsize; +#[cfg(test)] use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; use lance_context_core::{ - ContextStore, ContextStoreOptions, RolloutRegistry, RolloutStore, RolloutStoreOptions, + join_uri, validate_store_name, ContextStore, ContextStoreOptions, RolloutRegistry, + RolloutStore, RolloutStoreOptions, }; use lru::LruCache; use tokio::sync::{Mutex, RwLock}; @@ -34,7 +36,7 @@ pub struct AppState { /// because every operation refreshes the snapshot and therefore takes /// `&mut`. pub rollout_registry: RwLock, - pub base_path: PathBuf, + pub base_uri: String, /// Stable identity of this server instance, used as the MemWAL shard key for /// rollout writes so each instance owns exactly one shard. `None` falls back /// to a single shared shard (single-instance deployments only). @@ -55,11 +57,8 @@ impl AppState { /// under `data_dir`. Async because opening the registry touches storage. pub async fn new(config: ServerConfig) -> Result { let instance_id = config.resolved_instance_id(); - let base_path = PathBuf::from(&config.data_dir); - let registry_uri = base_path - .join("_registry.rollout.lance") - .to_string_lossy() - .to_string(); + let base_uri = config.data_dir.clone(); + let registry_uri = join_uri(&base_uri, "_registry.rollout.lance"); let registry = RolloutRegistry::open_or_create(®istry_uri, None) .await .map_err(AppError::from_lance)?; @@ -69,7 +68,7 @@ impl AppState { stores: RwLock::new(std::collections::HashMap::new()), rollout_stores: Mutex::new(LruCache::new(capacity)), rollout_registry: RwLock::new(registry), - base_path, + base_uri, instance_id, rollout_merge_after_generations: config.rollout_merge_after_generations, rollout_cleanup_interval_secs: config.rollout_cleanup_interval_secs, @@ -78,10 +77,7 @@ impl AppState { } pub fn context_uri(&self, name: &str) -> String { - self.base_path - .join(format!("{}.lance", name)) - .to_string_lossy() - .to_string() + join_uri(&self.base_uri, &format!("{}.lance", name)) } /// Build a default-configured `AppState` rooted at `base_path`, for tests. @@ -98,10 +94,8 @@ impl AppState { base_path: PathBuf, instance_id: Option, ) -> Self { - let registry_uri = base_path - .join("_registry.rollout.lance") - .to_string_lossy() - .to_string(); + let base_uri = base_path.to_string_lossy().to_string(); + let registry_uri = join_uri(&base_uri, "_registry.rollout.lance"); let registry = RolloutRegistry::open_or_create(®istry_uri, None) .await .expect("open test registry"); @@ -111,7 +105,7 @@ impl AppState { NonZeroUsize::new(DEFAULT_ROLLOUT_CACHE_CAPACITY).unwrap(), )), rollout_registry: RwLock::new(registry), - base_path, + base_uri, instance_id, rollout_merge_after_generations: 0, rollout_cleanup_interval_secs: 0, @@ -123,10 +117,7 @@ impl AppState { /// rollout store and a context store may share the same name without /// colliding on disk. pub fn rollout_uri(&self, name: &str) -> String { - self.base_path - .join(format!("{}.rollout.lance", name)) - .to_string_lossy() - .to_string() + join_uri(&self.base_uri, &format!("{}.rollout.lance", name)) } fn rollout_store_options(&self) -> RolloutStoreOptions { @@ -151,6 +142,7 @@ impl AppState { uri: &str, store: Arc>, ) -> Result<(), AppError> { + Self::validate_name(name)?; self.rollout_registry .write() .await @@ -170,6 +162,7 @@ impl AppState { /// Remove a rollout store from the durable registry and evict any resident /// handle. Returns whether the store existed. pub async fn unregister_rollout(&self, name: &str) -> Result { + Self::validate_name(name)?; let existed = self .rollout_registry .write() @@ -203,6 +196,7 @@ impl AppState { &self, name: &str, ) -> Result>, AppError> { + Self::validate_name(name)?; // Fast path: resident in this process's LRU (updates recency). if let Some(store) = self.rollout_stores.lock().await.get(name) { metrics::counter!("rollout_store_cache_hits_total").increment(1); @@ -252,6 +246,7 @@ impl AppState { &self, name: &str, ) -> Result>, AppError> { + Self::validate_name(name)?; if let Some(store) = self.stores.read().await.get(name) { return Ok(store.clone()); } @@ -270,6 +265,10 @@ impl AppState { Ok(opened) } + pub fn validate_name(name: &str) -> Result<(), AppError> { + validate_store_name(name).map_err(AppError::InvalidRequest) + } + /// Spawn the single, process-wide WAL-cleanup sweeper. /// /// This replaces the former one-timer-per-store model, which does not scale