Something went wrong. Try again.
A downstream fork of Tranquil PDS for use in margin.cafe
Something went wrong. Try again.
23 kB · 711 lines
Rust
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712pub use tranquil_infra::{BlobStorage, StorageError, StreamUploadResult};
use async_trait::async_trait;use bytes::Bytes;use futures::Stream;use sha2::{Digest, Sha256};use std::path::{Path, PathBuf};use std::pin::Pin;use std::sync::Arc;
const EXDEV: i32 = 18;const CID_SHARD_PREFIX_LEN: usize = 9;
fn split_cid_path(key: &str) -> Option<(&str, &str)> { let is_cid = key.get(..3).is_some_and(|p| p.eq_ignore_ascii_case("baf")); (key.len() > CID_SHARD_PREFIX_LEN && is_cid).then(|| key.split_at(CID_SHARD_PREFIX_LEN))}
fn validate_key(key: &str) -> Result<(), StorageError> { let dominated_by_traversal = key .split('/') .filter(|seg| !seg.is_empty()) .try_fold(0i32, |depth, segment| match segment { ".." => { let new_depth = depth - 1; (new_depth >= 0).then_some(new_depth) } "." => Some(depth), _ => Some(depth + 1), }) .is_none();
let has_null = key.contains('\0'); let is_absolute = key.starts_with('/');
match (dominated_by_traversal, has_null, is_absolute) { (true, _, _) => Err(StorageError::Other(format!( "Path traversal detected in key: {}", key ))), (_, true, _) => Err(StorageError::Other(format!( "Null byte in key: {}", key.replace('\0', "\\0") ))), (_, _, true) => Err(StorageError::Other(format!( "Absolute path not allowed: {}", key ))), _ => Ok(()), }}
async fn cleanup_orphaned_tmp_files(tmp_path: &Path) { let tmp_path = tmp_path.to_path_buf(); let cleaned = tokio::task::spawn_blocking(move || { std::fs::read_dir(&tmp_path) .into_iter() .flatten() .filter_map(Result::ok) .filter(|e| e.path().is_file()) .filter_map(|entry| std::fs::remove_file(entry.path()).ok()) .count() }) .await .unwrap_or(0);
if cleaned > 0 { tracing::info!( count = cleaned, "Cleaned orphaned tmp files from previous run" ); }}
async fn rename_with_fallback(src: &Path, dst: &Path) -> Result<(), StorageError> { match tokio::fs::rename(src, dst).await { Ok(()) => Ok(()), Err(e) if e.raw_os_error() == Some(EXDEV) => { tokio::fs::copy(src, dst).await?; tokio::fs::File::open(dst).await?.sync_all().await?; let _ = tokio::fs::remove_file(src).await; Ok(()) } Err(e) => Err(StorageError::Io(e)), }}
async fn ensure_parent_dir(path: &Path) -> Result<(), StorageError> { if let Some(parent) = path.parent() { tokio::fs::create_dir_all(parent).await?; } Ok(())}
fn map_io_not_found(key: &str) -> impl FnOnce(std::io::Error) -> StorageError + '_ { |e| match e.kind() { std::io::ErrorKind::NotFound => StorageError::NotFound(key.to_string()), _ => StorageError::Io(e), }}
#[cfg(feature = "s3")]mod s3 { use super::*; use aws_config::BehaviorVersion; use aws_config::meta::region::RegionProviderChain; use aws_sdk_s3::Client; use aws_sdk_s3::primitives::ByteStream; use aws_sdk_s3::types::CompletedMultipartUpload; use aws_sdk_s3::types::CompletedPart;
const MIN_PART_SIZE: usize = 5 * 1024 * 1024;
pub struct S3BlobStorage { client: Client, bucket: String, }
impl S3BlobStorage { pub async fn new() -> Self { let cfg = tranquil_config::get(); let bucket = cfg .storage .s3_bucket .clone() .expect("storage.s3_bucket (S3_BUCKET) must be set"); let client = create_s3_client().await; Self { client, bucket } }
pub async fn with_bucket(bucket: String) -> Self { let client = create_s3_client().await; Self { client, bucket } } }
async fn create_s3_client() -> Client { let region_provider = RegionProviderChain::default_provider().or_else("us-east-1");
let config = aws_config::defaults(BehaviorVersion::latest()) .region(region_provider) .load() .await;
tranquil_config::get() .storage .s3_endpoint .as_deref() .map_or_else( || Client::new(&config), |endpoint| { let s3_config = aws_sdk_s3::config::Builder::from(&config) .endpoint_url(endpoint) .force_path_style(true) .build(); Client::from_conf(s3_config) }, ) }
#[async_trait] impl BlobStorage for S3BlobStorage { async fn put(&self, key: &str, data: &[u8]) -> Result<(), StorageError> { self.put_bytes(key, Bytes::copy_from_slice(data)).await }
async fn put_bytes(&self, key: &str, data: Bytes) -> Result<(), StorageError> { self.client .put_object() .bucket(&self.bucket) .key(key) .body(ByteStream::from(data)) .send() .await .map_err(|e| StorageError::Backend(e.to_string()))?;
Ok(()) }
async fn get(&self, key: &str) -> Result<Vec<u8>, StorageError> { self.get_bytes(key).await.map(|b| b.to_vec()) }
async fn get_bytes(&self, key: &str) -> Result<Bytes, StorageError> { let resp = self .client .get_object() .bucket(&self.bucket) .key(key) .send() .await .map_err(|e| StorageError::Backend(e.to_string()))?;
resp.body .collect() .await .map(|agg| agg.into_bytes()) .map_err(|e| StorageError::Backend(e.to_string())) }
async fn get_head(&self, key: &str, size: usize) -> Result<Bytes, StorageError> { let range = format!("bytes=0-{}", size.saturating_sub(1)); let resp = self .client .get_object() .bucket(&self.bucket) .key(key) .range(range) .send() .await .map_err(|e| StorageError::Backend(e.to_string()))?;
resp.body .collect() .await .map(|agg| agg.into_bytes()) .map_err(|e| StorageError::Backend(e.to_string())) }
async fn delete(&self, key: &str) -> Result<(), StorageError> { self.client .delete_object() .bucket(&self.bucket) .key(key) .send() .await .map_err(|e| StorageError::Backend(e.to_string()))?;
Ok(()) }
async fn put_stream( &self, key: &str, stream: Pin<Box<dyn Stream<Item = Result<Bytes, std::io::Error>> + Send>>, ) -> Result<StreamUploadResult, StorageError> { use futures::StreamExt;
let create_resp = self .client .create_multipart_upload() .bucket(&self.bucket) .key(key) .send() .await .map_err(|e| { StorageError::Backend(format!("Failed to create multipart upload: {}", e)) })?;
let upload_id = create_resp .upload_id() .ok_or_else(|| StorageError::Backend("No upload ID returned".to_string()))? .to_string();
let upload_part = |client: &Client, bucket: &str, key: &str, upload_id: &str, part_num: i32, data: Vec<u8>| -> std::pin::Pin< Box<dyn std::future::Future<Output = Result<CompletedPart, StorageError>> + Send>, > { let client = client.clone(); let bucket = bucket.to_string(); let key = key.to_string(); let upload_id = upload_id.to_string(); Box::pin(async move { let resp = client .upload_part() .bucket(&bucket) .key(&key) .upload_id(&upload_id) .part_number(part_num) .body(ByteStream::from(data)) .send() .await .map_err(|e| { StorageError::Backend(format!("Failed to upload part: {}", e)) })?;
let etag = resp .e_tag() .ok_or_else(|| { StorageError::Backend("No ETag returned for part".to_string()) })? .to_string();
Ok(CompletedPart::builder() .part_number(part_num) .e_tag(etag) .build()) }) };
struct UploadState { hasher: Sha256, total_size: u64, part_number: i32, completed_parts: Vec<CompletedPart>, buffer: Vec<u8>, }
let initial_state = UploadState { hasher: Sha256::new(), total_size: 0, part_number: 1, completed_parts: Vec::new(), buffer: Vec::with_capacity(MIN_PART_SIZE), };
let abort_upload = || async { let _ = self .client .abort_multipart_upload() .bucket(&self.bucket) .key(key) .upload_id(&upload_id) .send() .await; };
let result: Result<UploadState, StorageError> = { let mut state = initial_state;
let chunk_results: Vec<Result<Bytes, std::io::Error>> = stream.collect().await;
for chunk_result in chunk_results { match chunk_result { Ok(chunk) => { state.hasher.update(&chunk); state.total_size += chunk.len() as u64; state.buffer.extend_from_slice(&chunk);
if state.buffer.len() >= MIN_PART_SIZE { let part_data = std::mem::replace( &mut state.buffer, Vec::with_capacity(MIN_PART_SIZE), ); let part = upload_part( &self.client, &self.bucket, key, &upload_id, state.part_number, part_data, ) .await?; state.completed_parts.push(part); state.part_number += 1; } } Err(e) => { abort_upload().await; return Err(StorageError::Io(e)); } } }
Ok(state) };
let mut state = result?;
if !state.buffer.is_empty() { let part = upload_part( &self.client, &self.bucket, key, &upload_id, state.part_number, std::mem::take(&mut state.buffer), ) .await?; state.completed_parts.push(part); }
if state.completed_parts.is_empty() { abort_upload().await; return Err(StorageError::Other("Empty upload".to_string())); }
let completed_upload = CompletedMultipartUpload::builder() .set_parts(Some(state.completed_parts)) .build();
self.client .complete_multipart_upload() .bucket(&self.bucket) .key(key) .upload_id(&upload_id) .multipart_upload(completed_upload) .send() .await .map_err(|e| { StorageError::Backend(format!("Failed to complete multipart upload: {}", e)) })?;
let hash: [u8; 32] = state.hasher.finalize().into(); Ok(StreamUploadResult { sha256_hash: hash, size: state.total_size, }) }
async fn copy(&self, src_key: &str, dst_key: &str) -> Result<(), StorageError> { let copy_source = format!("{}/{}", self.bucket, src_key);
self.client .copy_object() .bucket(&self.bucket) .copy_source(©_source) .key(dst_key) .send() .await .map_err(|e| StorageError::Backend(format!("Failed to copy object: {}", e)))?;
Ok(()) } }}
#[cfg(feature = "s3")]pub use s3::S3BlobStorage;
pub struct FilesystemBlobStorage { base_path: PathBuf, tmp_path: PathBuf,}
impl FilesystemBlobStorage { pub async fn new(base_path: impl Into<PathBuf>) -> Result<Self, StorageError> { let base_path = base_path.into(); let tmp_path = base_path.join(".tmp"); tokio::fs::create_dir_all(&base_path).await?; tokio::fs::create_dir_all(&tmp_path).await?; cleanup_orphaned_tmp_files(&tmp_path).await; Ok(Self { base_path, tmp_path, }) }
fn resolve_path(&self, key: &str) -> Result<PathBuf, StorageError> { validate_key(key)?; Ok(split_cid_path(key).map_or_else( || self.base_path.join(key), |(dir, file)| self.base_path.join(dir).join(file), )) }
async fn atomic_write(&self, path: &Path, data: &[u8]) -> Result<(), StorageError> { use tokio::io::AsyncWriteExt;
let tmp_file_name = uuid::Uuid::new_v4().to_string(); let tmp_path = self.tmp_path.join(&tmp_file_name);
let mut file = tokio::fs::File::create(&tmp_path).await?; file.write_all(data).await?; file.sync_all().await?; drop(file);
rename_with_fallback(&tmp_path, path).await }}
#[async_trait]impl BlobStorage for FilesystemBlobStorage { async fn put(&self, key: &str, data: &[u8]) -> Result<(), StorageError> { let path = self.resolve_path(key)?; ensure_parent_dir(&path).await?; self.atomic_write(&path, data).await }
async fn put_bytes(&self, key: &str, data: Bytes) -> Result<(), StorageError> { self.put(key, &data).await }
async fn get(&self, key: &str) -> Result<Vec<u8>, StorageError> { let path = self.resolve_path(key)?; tokio::fs::read(&path).await.map_err(map_io_not_found(key)) }
async fn get_bytes(&self, key: &str) -> Result<Bytes, StorageError> { self.get(key).await.map(Bytes::from) }
async fn get_head(&self, key: &str, size: usize) -> Result<Bytes, StorageError> { use tokio::io::AsyncReadExt; let path = self.resolve_path(key)?; let mut file = tokio::fs::File::open(&path) .await .map_err(map_io_not_found(key))?; let mut buffer = vec![0u8; size]; let n = file.read(&mut buffer).await?; buffer.truncate(n); Ok(Bytes::from(buffer)) }
async fn delete(&self, key: &str) -> Result<(), StorageError> { let path = self.resolve_path(key)?; tokio::fs::remove_file(&path).await.or_else(|e| { (e.kind() == std::io::ErrorKind::NotFound) .then_some(()) .ok_or(StorageError::Io(e)) }) }
async fn put_stream( &self, key: &str, stream: Pin<Box<dyn Stream<Item = Result<Bytes, std::io::Error>> + Send>>, ) -> Result<StreamUploadResult, StorageError> { use futures::TryStreamExt; use tokio::io::AsyncWriteExt;
let tmp_file_name = uuid::Uuid::new_v4().to_string(); let tmp_path = self.tmp_path.join(&tmp_file_name); let final_path = self.resolve_path(key)?; ensure_parent_dir(&final_path).await?;
let file = tokio::fs::File::create(&tmp_path).await?;
struct StreamState { file: tokio::fs::File, hasher: Sha256, total_size: u64, }
let initial = StreamState { file, hasher: Sha256::new(), total_size: 0, };
let final_state = stream .map_err(StorageError::Io) .try_fold(initial, |mut state, chunk| async move { state.hasher.update(&chunk); state.total_size += chunk.len() as u64; state.file.write_all(&chunk).await?; Ok(state) }) .await?;
final_state.file.sync_all().await?; drop(final_state.file);
rename_with_fallback(&tmp_path, &final_path).await?;
let hash: [u8; 32] = final_state.hasher.finalize().into(); Ok(StreamUploadResult { sha256_hash: hash, size: final_state.total_size, }) }
async fn copy(&self, src_key: &str, dst_key: &str) -> Result<(), StorageError> { let src_path = self.resolve_path(src_key)?; let dst_path = self.resolve_path(dst_key)?; ensure_parent_dir(&dst_path).await?; tokio::fs::copy(&src_path, &dst_path) .await .map_err(map_io_not_found(src_key))?; tokio::fs::File::open(&dst_path).await?.sync_all().await?; Ok(()) }}
pub async fn create_blob_storage() -> Arc<dyn BlobStorage> { let cfg = tranquil_config::get(); let backend = &cfg.storage.backend;
match backend.as_str() { #[cfg(feature = "s3")] "s3" => { tracing::info!("Initializing S3 blob storage"); Arc::new(S3BlobStorage::new().await) } #[cfg(not(feature = "s3"))] "s3" => { panic!( "BLOB_STORAGE_BACKEND=s3 but binary was compiled without s3 feature. \ Rebuild with --features s3 to enable S3 storage." ); } _ => { tracing::info!("Initializing filesystem blob storage"); let path = cfg.storage.path.clone(); FilesystemBlobStorage::new(path) .await .unwrap_or_else(|e| { panic!( "Failed to initialize filesystem blob storage: {}. \ Set BLOB_STORAGE_PATH to a valid directory path.", e ); }) .pipe(Arc::new) } }}
trait Pipe: Sized { fn pipe<F, R>(self, f: F) -> R where F: FnOnce(Self) -> R, { f(self) }}
impl<T> Pipe for T {}
#[cfg(test)]mod tests { use super::*;
#[test] fn split_path_from_raw_blob_cid() { let cid = "bafkreihdwdcefgh4dqkjv67uzcmw7ojee6xedzdetojuzjevtenxquvyku"; assert_eq!( split_cid_path(cid), Some(( "bafkreihd", "wdcefgh4dqkjv67uzcmw7ojee6xedzdetojuzjevtenxquvyku" )) ); }
#[test] fn split_path_from_dag_cbor_cid() { let cid = "bafyreigdmqpykrgxyaxtlafqpqhzrb7qy2rh75nldvfd4tucqmqqme5yje"; assert_eq!( split_cid_path(cid), Some(( "bafyreigd", "mqpykrgxyaxtlafqpqhzrb7qy2rh75nldvfd4tucqmqqme5yje" )) ); }
#[test] fn no_split_for_temp_keys() { assert_eq!(split_cid_path("temp/abc123"), None); }
#[test] fn no_split_for_short_keys() { assert_eq!(split_cid_path("bafkreihd"), None); assert_eq!(split_cid_path("bafkrei"), None); assert_eq!(split_cid_path("baf"), None); assert_eq!(split_cid_path("ba"), None); assert_eq!(split_cid_path(""), None); }
#[test] fn no_split_for_non_cid_keys() { assert_eq!(split_cid_path("something/else/entirely"), None); assert_eq!(split_cid_path("Qmabcdefghijklmnop"), None); }
#[test] fn split_cid_case_insensitive() { let upper = "BAFKREIHDWDCEFGH4DQKJV67UZCMW7OJEE6XEDZDETOJUZJEVTENXQUVYKU"; let mixed = "BaFkReIhDwDcEfGh4DqKjV67UzCmW7OjEe6XeDzDeTojUzJevTeNxQuVyKu"; assert_eq!( split_cid_path(upper), Some(( "BAFKREIHD", "WDCEFGH4DQKJV67UZCMW7OJEE6XEDZDETOJUZJEVTENXQUVYKU" )) ); assert_eq!( split_cid_path(mixed), Some(( "BaFkReIhD", "wDcEfGh4DqKjV67UzCmW7OjEe6XeDzDeTojUzJevTeNxQuVyKu" )) ); }
#[test] fn split_at_minimum_length() { let cid = "bafkreihdx"; assert_eq!(split_cid_path(cid), Some(("bafkreihd", "x"))); }
#[test] fn resolve_path_shards_cid_keys() { let base = PathBuf::from("/blobs"); let cid = "bafkreihdwdcefgh4dqkjv67uzcmw7ojee6xedzdetojuzjevtenxquvyku";
let expected = PathBuf::from("/blobs/bafkreihd/wdcefgh4dqkjv67uzcmw7ojee6xedzdetojuzjevtenxquvyku"); let result = split_cid_path(cid) .map_or_else(|| base.join(cid), |(dir, file)| base.join(dir).join(file)); assert_eq!(result, expected); }
#[test] fn resolve_path_no_shard_for_temp() { let base = PathBuf::from("/blobs"); let key = "temp/abc123";
let expected = PathBuf::from("/blobs/temp/abc123"); let result = split_cid_path(key) .map_or_else(|| base.join(key), |(dir, file)| base.join(dir).join(file)); assert_eq!(result, expected); }}