use std::collections::BTreeMap; use std::fmt; use std::sync::Arc; use bytes::{Bytes, BytesMut}; use futures_util::TryStreamExt; use object_store::aws::AmazonS3Builder; use object_store::path::Path; use object_store::{ClientOptions, ObjectStore, ObjectStoreExt, PutMode, PutOptions}; use reqwest::header::HeaderValue; use tokio::sync::RwLock; use super::config::S3BackupConfig; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) enum BackupObjectCreateResult { Created, AlreadyExists, } #[async_trait::async_trait] pub(crate) trait BackupObjectStore: Send + Sync { async fn put_object(&self, key: &str, bytes: Bytes) -> Result<(), BackupStoreError>; async fn put_object_if_absent( &self, key: &str, bytes: Bytes, ) -> Result; async fn get_object_limited( &self, key: &str, max_bytes: usize, ) -> Result; async fn delete_object(&self, key: &str) -> Result<(), BackupStoreError>; async fn list_keys_limited( &self, prefix: &str, max_objects: usize, ) -> Result, BackupStoreError>; } #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) struct BackupStoreError { message: String, } impl BackupStoreError { fn new(message: impl Into) -> Self { Self { message: message.into(), } } fn object_store(operation: &str, key: &str, error: impl fmt::Display) -> Self { Self::new(format!( "S3 backup object store {operation} failed for `{key}`: {error}" )) } } impl fmt::Display for BackupStoreError { fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { formatter.write_str(&self.message) } } impl std::error::Error for BackupStoreError {} #[derive(Debug, Default, Clone)] pub(crate) struct FakeBackupObjectStore { objects: Arc>>, } #[async_trait::async_trait] impl BackupObjectStore for FakeBackupObjectStore { async fn put_object(&self, key: &str, bytes: Bytes) -> Result<(), BackupStoreError> { self.objects.write().await.insert(key.to_string(), bytes); Ok(()) } async fn put_object_if_absent( &self, key: &str, bytes: Bytes, ) -> Result { let mut objects = self.objects.write().await; if objects.contains_key(key) { Ok(BackupObjectCreateResult::AlreadyExists) } else { objects.insert(key.to_string(), bytes); Ok(BackupObjectCreateResult::Created) } } async fn get_object_limited( &self, key: &str, max_bytes: usize, ) -> Result { let bytes = self .objects .read() .await .get(key) .cloned() .ok_or_else(|| BackupStoreError::new(format!("backup object `{key}` not found")))?; if bytes.len() > max_bytes { return Err(BackupStoreError::new(format!( "backup object `{key}` exceeds the configured {max_bytes} byte read limit" ))); } Ok(bytes) } async fn delete_object(&self, key: &str) -> Result<(), BackupStoreError> { self.objects.write().await.remove(key); Ok(()) } async fn list_keys_limited( &self, prefix: &str, max_objects: usize, ) -> Result, BackupStoreError> { let prefix = directory_list_prefix(prefix); let keys: Vec<_> = self .objects .read() .await .keys() .filter(|key| key.starts_with(&prefix)) .cloned() .collect(); if keys.len() > max_objects { return Err(BackupStoreError::new(format!( "backup object listing exceeds the configured {max_objects} object limit" ))); } Ok(keys) } } #[cfg(test)] impl FakeBackupObjectStore { pub(crate) async fn object_bytes(&self, key: &str) -> Option { self.objects.read().await.get(key).cloned() } } #[derive(Debug)] pub(crate) struct ObjectStoreS3BackupStore { store: object_store::aws::AmazonS3, } impl ObjectStoreS3BackupStore { pub(crate) fn from_config(config: &S3BackupConfig) -> Result { // 部分 S3 兼容网关(如中国科技云 s3.cstcloud.cn)按 User-Agent 放行请求, // object_store 默认 UA 会被拒,因此允许自定义 User-Agent。 let client_options = if config.user_agent.trim().is_empty() { ClientOptions::new() } else { ClientOptions::new().with_user_agent( HeaderValue::from_str(config.user_agent.trim()).map_err(|error| { BackupStoreError::new(format!("S3 备份 User-Agent 配置无效: {error}")) })?, ) }; let store = AmazonS3Builder::new() .with_client_options(client_options) .with_endpoint(config.endpoint.clone()) .with_region(config.region.clone()) .with_bucket_name(config.bucket.clone()) .with_access_key_id(config.access_key_id.clone()) .with_secret_access_key(config.secret_access_key.clone()) .with_virtual_hosted_style_request(!config.path_style) .build() .map_err(|error| { BackupStoreError::new(format!( "S3 backup object store configuration failed: {error}" )) })?; Ok(Self { store }) } } #[async_trait::async_trait] impl BackupObjectStore for ObjectStoreS3BackupStore { async fn put_object(&self, key: &str, bytes: Bytes) -> Result<(), BackupStoreError> { self.store .put(&Path::from(key), bytes.into()) .await .map(|_| ()) .map_err(|error| BackupStoreError::object_store("put", key, error)) } async fn put_object_if_absent( &self, key: &str, bytes: Bytes, ) -> Result { let options = PutOptions { mode: PutMode::Create, ..PutOptions::default() }; match self .store .put_opts(&Path::from(key), bytes.into(), options) .await { Ok(_) => Ok(BackupObjectCreateResult::Created), Err(object_store::Error::AlreadyExists { .. }) => { Ok(BackupObjectCreateResult::AlreadyExists) } Err(error) => Err(BackupStoreError::object_store( "conditional put", key, error, )), } } async fn get_object_limited( &self, key: &str, max_bytes: usize, ) -> Result { let result = self .store .get(&Path::from(key)) .await .map_err(|error| BackupStoreError::object_store("get", key, error))?; if result.meta.size > u64::try_from(max_bytes).unwrap_or(u64::MAX) { return Err(BackupStoreError::new(format!( "backup object `{key}` exceeds the configured {max_bytes} byte read limit" ))); } let object_size = result.meta.size; let mut stream = result.into_stream(); let mut bytes = BytesMut::with_capacity( usize::try_from(object_size) .unwrap_or(max_bytes) .min(max_bytes) .min(8 * 1024 * 1024), ); while let Some(chunk) = stream .try_next() .await .map_err(|error| BackupStoreError::object_store("read", key, error))? { if bytes.len().saturating_add(chunk.len()) > max_bytes { return Err(BackupStoreError::new(format!( "backup object `{key}` exceeds the configured {max_bytes} byte read limit" ))); } bytes.extend_from_slice(&chunk); } Ok(bytes.freeze()) } async fn delete_object(&self, key: &str) -> Result<(), BackupStoreError> { self.store .delete(&Path::from(key)) .await .map_err(|error| BackupStoreError::object_store("delete", key, error)) } async fn list_keys_limited( &self, prefix: &str, max_objects: usize, ) -> Result, BackupStoreError> { let prefix_path = list_prefix_path(prefix); let mut objects = self.store.list(prefix_path.as_ref()); let mut keys = Vec::new(); while let Some(meta) = objects .try_next() .await .map_err(|error| BackupStoreError::object_store("list", prefix, error))? { if keys.len() >= max_objects { return Err(BackupStoreError::new(format!( "backup object listing exceeds the configured {max_objects} object limit" ))); } keys.push(meta.location.to_string()); } keys.sort(); Ok(keys) } } fn directory_list_prefix(prefix: &str) -> String { let prefix = prefix.trim_end_matches('/'); if prefix.is_empty() { String::new() } else { format!("{prefix}/") } } fn list_prefix_path(prefix: &str) -> Option { let prefix = prefix.trim_end_matches('/'); if prefix.is_empty() { None } else { Some(Path::from(prefix)) } } #[cfg(test)] mod tests { use super::{ list_prefix_path, BackupObjectCreateResult, BackupObjectStore, FakeBackupObjectStore, }; #[tokio::test] async fn fake_backup_object_store_puts_and_lists() { let store = FakeBackupObjectStore::default(); store .put_object( "prod/aether-data-backup-20260524-010000.json.zst", bytes::Bytes::from_static(b"one"), ) .await .unwrap(); store .put_object( "prod/aether-data-backup-20260524-020000.json.zst", bytes::Bytes::from_static(b"two"), ) .await .unwrap(); let keys = store.list_keys_limited("prod/", 2).await.unwrap(); assert_eq!( keys, vec![ "prod/aether-data-backup-20260524-010000.json.zst", "prod/aether-data-backup-20260524-020000.json.zst", ] ); } #[tokio::test] async fn fake_backup_object_store_enforces_read_and_listing_limits() { let store = FakeBackupObjectStore::default(); store .put_object("prod/one", bytes::Bytes::from_static(b"1234")) .await .unwrap(); store .put_object("prod/two", bytes::Bytes::from_static(b"5678")) .await .unwrap(); assert!(store.get_object_limited("prod/one", 3).await.is_err()); assert_eq!( store.get_object_limited("prod/one", 4).await.unwrap(), bytes::Bytes::from_static(b"1234") ); assert!(store.list_keys_limited("prod/", 1).await.is_err()); assert_eq!(store.list_keys_limited("prod/", 2).await.unwrap().len(), 2); } #[tokio::test] async fn fake_backup_object_store_conditional_put_never_overwrites() { let store = FakeBackupObjectStore::default(); let key = "prod/aether-data-backup-20260524-010000.json.zst.aes256gcm"; assert_eq!( store .put_object_if_absent(key, bytes::Bytes::from_static(b"first")) .await .unwrap(), BackupObjectCreateResult::Created ); assert_eq!( store .put_object_if_absent(key, bytes::Bytes::from_static(b"second")) .await .unwrap(), BackupObjectCreateResult::AlreadyExists ); assert_eq!( store.object_bytes(key).await.as_deref(), Some(b"first".as_slice()) ); } #[tokio::test] async fn fake_backup_object_store_lists_normalized_directory_prefixes() { let store = FakeBackupObjectStore::default(); store .put_object( "prod/aether-data-backup-20260524-010000.json.zst", bytes::Bytes::from_static(b"one"), ) .await .unwrap(); store .put_object( "prod-backups/aether-data-backup-20260524-010000.json.zst", bytes::Bytes::from_static(b"two"), ) .await .unwrap(); let keys = store.list_keys_limited("prod", 10).await.unwrap(); assert_eq!( keys, vec!["prod/aether-data-backup-20260524-010000.json.zst"] ); } #[test] fn s3_list_prefix_path_lets_object_store_add_directory_delimiter() { assert_eq!( list_prefix_path("prod/") .as_ref() .map(std::string::ToString::to_string) .as_deref(), Some("prod") ); assert!(list_prefix_path("").is_none()); } }