pub 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 ::s3::types::CompletedPart;
use ::s3::{AddressingStyle, Auth, Client};
const MIN_PART_SIZE: usize = 5 * 1024 * 1024;
pub struct S3BlobStorage {
client: Client,
bucket: String,
path: String,
}
impl S3BlobStorage {
pub async fn new(bucket: &str, endpoint: Option<&str>, path: &str) -> Self {
Self {
client: create_s3_client(endpoint),
bucket: bucket.to_string(),
path: path
.trim_start_matches("/")
.trim_end_matches("/")
.to_string(),
}
}
fn resolve_path(&self, key: &str) -> String {
if self.path.is_empty() {
return key.to_string();
}
format!("{}/{}", self.path, key)
}
}
fn create_s3_client(endpoint: Option<&str>) -> Client {
let region = std::env::var("AWS_REGION").unwrap_or_else(|_| "us-east-1".to_string());
let (url, style) = match endpoint {
Some(endpoint) => (endpoint.to_string(), AddressingStyle::Path),
None => (
format!("https://s3.{region}.amazonaws.com"),
AddressingStyle::Auto,
),
};
let auth =
Auth::from_env().expect("AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY must be set");
Client::builder(url)
.expect("storage.s3_endpoint (S3_ENDPOINT) must be a valid URL")
.region(region)
.addressing_style(style)
.auth(auth)
.base_retry_delay(std::time::Duration::from_secs(1))
.max_retry_delay(std::time::Duration::from_secs(20))
.build()
.expect("failed to build S3 client")
}
#[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> {
let path = self.resolve_path(key);
self.client
.objects()
.put(&self.bucket, path)
.body_bytes(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 path = self.resolve_path(key);
let resp = self
.client
.objects()
.get(&self.bucket, &path)
.send()
.await
.map_err(|e| StorageError::Backend(e.to_string()))?;
resp.bytes()
.await
.map_err(|e| StorageError::Backend(e.to_string()))
}
async fn get_head(&self, key: &str, size: usize) -> Result<Bytes, StorageError> {
let path = self.resolve_path(key);
let resp = self
.client
.objects()
.get(&self.bucket, &path)
.range_bytes(0, size.saturating_sub(1) as u64)
.map_err(|e| StorageError::Backend(e.to_string()))?
.send()
.await
.map_err(|e| StorageError::Backend(e.to_string()))?;
resp.bytes()
.await
.map_err(|e| StorageError::Backend(e.to_string()))
}
async fn delete(&self, key: &str) -> Result<(), StorageError> {
let path = self.resolve_path(key);
self.client
.objects()
.delete(&self.bucket, &path)
.send()
.await
.map_err(|e| StorageError::Backend(e.to_string()))?;
Ok(())
}
async fn put_stream(
&self,
key: &str,
mut stream: Pin<Box<dyn Stream<Item = Result<Bytes, std::io::Error>> + Send>>,
) -> Result<StreamUploadResult, StorageError> {
use futures::StreamExt;
let path = self.resolve_path(key);
let upload_id = self
.client
.objects()
.create_multipart_upload(&self.bucket, &path)
.send()
.await
.map_err(|e| {
StorageError::Backend(format!("Failed to create multipart upload: {}", e))
})?
.upload_id;
let upload_part = |client: &Client,
bucket: &str,
key: &str,
upload_id: &str,
part_num: u32,
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 path = self.resolve_path(key);
let upload_id = upload_id.to_string();
Box::pin(async move {
let resp = client
.objects()
.upload_part(&bucket, &path, &upload_id, part_num)
.body_bytes(data)
.send()
.await
.map_err(|e| {
StorageError::Backend(format!("Failed to upload part: {}", e))
})?;
let etag = resp.etag.ok_or_else(|| {
StorageError::Backend("No ETag returned for part".to_string())
})?;
CompletedPart::new(part_num, etag)
.map_err(|e| StorageError::Backend(e.to_string()))
})
};
struct UploadState {
hasher: Sha256,
total_size: u64,
part_number: u32,
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
.objects()
.abort_multipart_upload(&self.bucket, path, &upload_id)
.send()
.await;
};
let result: Result<StreamUploadResult, StorageError> = async {
let mut state = initial_state;
while let Some(chunk_result) = stream.next().await {
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) => return Err(StorageError::Io(e)),
}
}
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() {
return Err(StorageError::Other("Empty upload".to_string()));
}
let path = self.resolve_path(key);
self.client
.objects()
.complete_multipart_upload(&self.bucket, &path, &upload_id)
.parts(state.completed_parts)
.map_err(|e| StorageError::Backend(e.to_string()))?
.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,
})
}
.await;
if result.is_err() {
abort_upload().await;
}
result
}
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);
self.client
.objects()
.copy(&self.bucket, &src_path, &self.bucket, &dst_path)
.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");
let storage = &cfg.storage;
let bucket = storage
.s3_bucket
.as_deref()
.expect("storage.s3_bucket (S3_BUCKET) must be set");
Arc::new(
S3BlobStorage::new(bucket, storage.s3_endpoint.as_deref(), &storage.s3_path).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);
}
}