feat: add configurable S3 object storage

This commit is contained in:
237899745
2026-07-25 13:23:11 +08:00
parent 61fa9cb820
commit d1f093685d
29 changed files with 3703 additions and 284 deletions

View File

@@ -5,3 +5,4 @@ pub mod idempotency;
pub mod mail;
pub mod quota;
pub mod settings;
pub mod storage;

718
src/services/storage.rs Normal file
View File

@@ -0,0 +1,718 @@
use crate::error::{AppError, ErrorCode};
use crate::services::settings;
use crate::state::AppState;
use aws_sdk_s3::config::{
BehaviorVersion, Credentials, Region, RequestChecksumCalculation, ResponseChecksumValidation,
};
use aws_sdk_s3::presigning::PresigningConfig;
use aws_sdk_s3::primitives::ByteStream;
use aws_sdk_s3::types::{CompletedMultipartUpload, CompletedPart};
use aws_sdk_s3::Client;
use chrono::{DateTime, Datelike, Utc};
use sqlx::FromRow;
use std::path::{Path, PathBuf};
use std::time::Duration;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use uuid::Uuid;
const MULTIPART_THRESHOLD: u64 = 64 * 1024 * 1024;
const MULTIPART_PART_SIZE: usize = 16 * 1024 * 1024;
#[derive(Debug, Clone, FromRow)]
pub struct StorageEndpoint {
pub id: Uuid,
pub name: String,
pub internal_endpoint: String,
pub public_endpoint: String,
pub bucket: String,
pub region: String,
pub access_key_encrypted: String,
pub secret_key_encrypted: String,
pub access_key_hint: String,
pub force_path_style: bool,
pub presign_ttl_seconds: i32,
pub is_active: bool,
pub last_test_at: Option<DateTime<Utc>>,
pub last_test_ok: Option<bool>,
pub last_test_error: Option<String>,
pub created_at: DateTime<Utc>,
pub updated_at: DateTime<Utc>,
pub deleted_at: Option<DateTime<Utc>>,
}
#[derive(Debug, Clone)]
pub struct StoredObject {
pub backend: String,
pub endpoint_id: Option<Uuid>,
pub key: String,
pub etag: Option<String>,
pub size: u64,
}
#[derive(Debug, Clone)]
pub struct ObjectLocator {
pub backend: String,
pub endpoint_id: Option<Uuid>,
pub key: String,
}
pub async fn list_endpoints(state: &AppState) -> Result<Vec<StorageEndpoint>, AppError> {
sqlx::query_as::<_, StorageEndpoint>(
r#"
SELECT id, name, internal_endpoint, public_endpoint, bucket, region,
access_key_encrypted, secret_key_encrypted, access_key_hint,
force_path_style, presign_ttl_seconds, is_active,
last_test_at, last_test_ok, last_test_error,
created_at, updated_at, deleted_at
FROM storage_endpoints
WHERE deleted_at IS NULL
ORDER BY is_active DESC, updated_at DESC
"#,
)
.fetch_all(&state.db)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "查询存储端点失败").with_source(err))
}
pub async fn get_endpoint(
state: &AppState,
endpoint_id: Uuid,
) -> Result<StorageEndpoint, AppError> {
sqlx::query_as::<_, StorageEndpoint>(
r#"
SELECT id, name, internal_endpoint, public_endpoint, bucket, region,
access_key_encrypted, secret_key_encrypted, access_key_hint,
force_path_style, presign_ttl_seconds, is_active,
last_test_at, last_test_ok, last_test_error,
created_at, updated_at, deleted_at
FROM storage_endpoints
WHERE id = $1
"#,
)
.bind(endpoint_id)
.fetch_optional(&state.db)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "查询存储端点失败").with_source(err))?
.ok_or_else(|| AppError::new(ErrorCode::NotFound, "存储端点不存在"))
}
pub async fn active_endpoint(state: &AppState) -> Result<Option<StorageEndpoint>, AppError> {
sqlx::query_as::<_, StorageEndpoint>(
r#"
SELECT id, name, internal_endpoint, public_endpoint, bucket, region,
access_key_encrypted, secret_key_encrypted, access_key_hint,
force_path_style, presign_ttl_seconds, is_active,
last_test_at, last_test_ok, last_test_error,
created_at, updated_at, deleted_at
FROM storage_endpoints
WHERE is_active = true AND deleted_at IS NULL
LIMIT 1
"#,
)
.fetch_optional(&state.db)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "查询活动存储端点失败").with_source(err))
}
pub fn result_key(retention_hours: i64, task_id: Uuid, file_id: Uuid, extension: &str) -> String {
let now = Utc::now();
format!(
"results/{}/{:04}/{:02}/{task_id}/{file_id}.{}",
retention_prefix(retention_hours),
now.year(),
now.month(),
extension.trim_start_matches('.')
)
}
pub fn archive_key(retention_hours: i64, task_id: Uuid) -> String {
let now = Utc::now();
format!(
"archives/{}/{:04}/{:02}/{task_id}.zip",
retention_prefix(retention_hours),
now.year(),
now.month()
)
}
fn retention_prefix(hours: i64) -> String {
match hours {
0..=24 => "1d".to_string(),
25..=168 => "7d".to_string(),
169..=360 => "15d".to_string(),
value => format!("custom-{}h", value.max(1)),
}
}
pub async fn store_bytes(
state: &AppState,
key: &str,
bytes: Vec<u8>,
content_type: &str,
) -> Result<StoredObject, AppError> {
if let Some(endpoint) = active_endpoint(state).await? {
let client = client_for(state, &endpoint, EndpointKind::Internal)?;
let size = bytes.len() as u64;
let output = client
.put_object()
.bucket(&endpoint.bucket)
.key(key)
.content_type(content_type)
.body(ByteStream::from(bytes))
.send()
.await
.map_err(|err| storage_error("上传 S3 对象失败", err))?;
return Ok(StoredObject {
backend: "s3".to_string(),
endpoint_id: Some(endpoint.id),
key: key.to_string(),
etag: output.e_tag().map(ToOwned::to_owned),
size,
});
}
let path = local_path(state, key)?;
if let Some(parent) = path.parent() {
tokio::fs::create_dir_all(parent).await.map_err(|err| {
AppError::new(ErrorCode::StorageUnavailable, "创建本地存储目录失败").with_source(err)
})?;
}
tokio::fs::write(&path, &bytes).await.map_err(|err| {
AppError::new(ErrorCode::StorageUnavailable, "写入本地存储失败").with_source(err)
})?;
Ok(StoredObject {
backend: "local".to_string(),
endpoint_id: None,
key: path.to_string_lossy().to_string(),
etag: None,
size: bytes.len() as u64,
})
}
pub async fn store_file(
state: &AppState,
key: &str,
path: &Path,
content_type: &str,
) -> Result<StoredObject, AppError> {
let metadata = tokio::fs::metadata(path).await.map_err(|err| {
AppError::new(ErrorCode::StorageUnavailable, "读取待上传文件失败").with_source(err)
})?;
if let Some(endpoint) = active_endpoint(state).await? {
let client = client_for(state, &endpoint, EndpointKind::Internal)?;
let etag = if metadata.len() >= MULTIPART_THRESHOLD {
multipart_upload(&client, &endpoint.bucket, key, path, content_type).await?
} else {
let body = ByteStream::from_path(path).await.map_err(|err| {
AppError::new(ErrorCode::StorageUnavailable, "读取待上传文件失败").with_source(err)
})?;
client
.put_object()
.bucket(&endpoint.bucket)
.key(key)
.content_type(content_type)
.body(body)
.send()
.await
.map_err(|err| storage_error("上传 S3 对象失败", err))?
.e_tag()
.map(ToOwned::to_owned)
};
return Ok(StoredObject {
backend: "s3".to_string(),
endpoint_id: Some(endpoint.id),
key: key.to_string(),
etag,
size: metadata.len(),
});
}
let destination = local_path(state, key)?;
if let Some(parent) = destination.parent() {
tokio::fs::create_dir_all(parent).await.map_err(|err| {
AppError::new(ErrorCode::StorageUnavailable, "创建本地存储目录失败").with_source(err)
})?;
}
tokio::fs::copy(path, &destination).await.map_err(|err| {
AppError::new(ErrorCode::StorageUnavailable, "保存本地文件失败").with_source(err)
})?;
Ok(StoredObject {
backend: "local".to_string(),
endpoint_id: None,
key: destination.to_string_lossy().to_string(),
etag: None,
size: metadata.len(),
})
}
pub async fn read_bytes(state: &AppState, object: &ObjectLocator) -> Result<Vec<u8>, AppError> {
if object.backend == "local" {
return tokio::fs::read(&object.key).await.map_err(|err| {
AppError::new(ErrorCode::StorageUnavailable, "读取本地文件失败").with_source(err)
});
}
let endpoint = endpoint_for_object(state, object).await?;
let client = client_for(state, &endpoint, EndpointKind::Internal)?;
let response = client
.get_object()
.bucket(&endpoint.bucket)
.key(&object.key)
.send()
.await
.map_err(|err| storage_error("读取 S3 对象失败", err))?;
let bytes = response
.body
.collect()
.await
.map_err(|err| storage_error("接收 S3 对象失败", err))?;
Ok(bytes.into_bytes().to_vec())
}
pub async fn download_to_file(
state: &AppState,
object: &ObjectLocator,
destination: &Path,
) -> Result<(), AppError> {
if object.backend == "local" {
tokio::fs::copy(&object.key, destination)
.await
.map_err(|err| {
AppError::new(ErrorCode::StorageUnavailable, "复制本地文件失败").with_source(err)
})?;
return Ok(());
}
let endpoint = endpoint_for_object(state, object).await?;
let client = client_for(state, &endpoint, EndpointKind::Internal)?;
let response = client
.get_object()
.bucket(&endpoint.bucket)
.key(&object.key)
.send()
.await
.map_err(|err| storage_error("读取 S3 对象失败", err))?;
let mut source = response.body.into_async_read();
let mut output = tokio::fs::File::create(destination).await.map_err(|err| {
AppError::new(ErrorCode::StorageUnavailable, "创建临时文件失败").with_source(err)
})?;
tokio::io::copy(&mut source, &mut output)
.await
.map_err(|err| {
AppError::new(ErrorCode::StorageUnavailable, "保存临时文件失败").with_source(err)
})?;
output.flush().await.map_err(|err| {
AppError::new(ErrorCode::StorageUnavailable, "刷新临时文件失败").with_source(err)
})?;
Ok(())
}
pub async fn presign_download(
state: &AppState,
object: &ObjectLocator,
file_name: &str,
content_type: &str,
) -> Result<Option<String>, AppError> {
if object.backend == "local" {
return Ok(None);
}
let endpoint = endpoint_for_object(state, object).await?;
let client = client_for(state, &endpoint, EndpointKind::Public)?;
let ttl = Duration::from_secs(endpoint.presign_ttl_seconds as u64);
let config = PresigningConfig::expires_in(ttl).map_err(|err| {
AppError::new(ErrorCode::Internal, "生成下载签名配置失败").with_source(err)
})?;
let disposition = content_disposition(file_name);
let request = client
.get_object()
.bucket(&endpoint.bucket)
.key(&object.key)
.response_content_disposition(disposition)
.response_content_type(content_type)
.presigned(config)
.await
.map_err(|err| storage_error("生成 S3 下载地址失败", err))?;
Ok(Some(request.uri().to_string()))
}
pub async fn delete_object(state: &AppState, object: &ObjectLocator) -> Result<(), AppError> {
if object.backend == "local" {
match tokio::fs::remove_file(&object.key).await {
Ok(()) => return Ok(()),
Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(()),
Err(err) => {
return Err(
AppError::new(ErrorCode::StorageUnavailable, "删除本地文件失败")
.with_source(err),
)
}
}
}
let endpoint = endpoint_for_object(state, object).await?;
let client = client_for(state, &endpoint, EndpointKind::Internal)?;
client
.delete_object()
.bucket(&endpoint.bucket)
.key(&object.key)
.send()
.await
.map_err(|err| storage_error("删除 S3 对象失败", err))?;
Ok(())
}
pub async fn test_endpoint(state: &AppState, endpoint: &StorageEndpoint) -> Result<(), AppError> {
let client = client_for(state, endpoint, EndpointKind::Internal)?;
client
.head_bucket()
.bucket(&endpoint.bucket)
.send()
.await
.map_err(|err| storage_error("S3 Bucket 访问失败", err))?;
let key = format!(".imageforge-health/{}.txt", Uuid::new_v4());
let expected = b"imageforge-storage-check".to_vec();
client
.put_object()
.bucket(&endpoint.bucket)
.key(&key)
.content_type("text/plain")
.body(ByteStream::from(expected.clone()))
.send()
.await
.map_err(|err| storage_error("S3 写入测试失败", err))?;
let result = async {
let response = client
.get_object()
.bucket(&endpoint.bucket)
.key(&key)
.send()
.await
.map_err(|err| storage_error("S3 读取测试失败", err))?;
let actual = response
.body
.collect()
.await
.map_err(|err| storage_error("S3 测试对象接收失败", err))?
.into_bytes();
if actual.as_ref() != expected.as_slice() {
return Err(AppError::new(
ErrorCode::StorageUnavailable,
"S3 读写校验内容不一致",
));
}
let public_client = client_for(state, endpoint, EndpointKind::Public)?;
let signed = public_client
.get_object()
.bucket(&endpoint.bucket)
.key(&key)
.presigned(
PresigningConfig::expires_in(Duration::from_secs(60)).map_err(|err| {
AppError::new(ErrorCode::Internal, "生成公网下载测试签名失败").with_source(err)
})?,
)
.await
.map_err(|err| storage_error("生成公网下载测试地址失败", err))?;
let response = reqwest::Client::builder()
.timeout(Duration::from_secs(15))
.build()
.map_err(|err| storage_error("创建公网下载测试客户端失败", err))?
.get(signed.uri().to_string())
.header(reqwest::header::USER_AGENT, "ImageForge-Storage-Check/1.0")
.send()
.await
.map_err(|err| public_test_request_error(endpoint, err))?;
let status = response.status();
if !status.is_success() {
return Err(AppError::new(
ErrorCode::StorageUnavailable,
format!("公网 Endpoint 下载测试返回 HTTP {status}"),
));
}
let public_bytes = response.bytes().await.map_err(|err| {
AppError::new(ErrorCode::StorageUnavailable, "读取公网下载测试响应失败")
.with_source(format!("response body error: {}", err.is_timeout()))
})?;
if public_bytes.as_ref() != expected.as_slice() {
return Err(AppError::new(
ErrorCode::StorageUnavailable,
"公网 Endpoint 下载内容校验不一致",
));
}
Ok(())
}
.await;
let delete_result = client
.delete_object()
.bucket(&endpoint.bucket)
.key(&key)
.send()
.await;
result?;
delete_result.map_err(|err| storage_error("S3 删除测试失败", err))?;
Ok(())
}
fn client_for(
state: &AppState,
endpoint: &StorageEndpoint,
kind: EndpointKind,
) -> Result<Client, AppError> {
let access_key =
settings::decrypt_secret(state, &endpoint.access_key_encrypted).map_err(|err| {
AppError::new(ErrorCode::StorageUnavailable, "存储端点凭据不可用").with_source(err)
})?;
let secret_key =
settings::decrypt_secret(state, &endpoint.secret_key_encrypted).map_err(|err| {
AppError::new(ErrorCode::StorageUnavailable, "存储端点凭据不可用").with_source(err)
})?;
let endpoint_url = match kind {
EndpointKind::Internal => &endpoint.internal_endpoint,
EndpointKind::Public => &endpoint.public_endpoint,
};
Ok(build_client(
access_key,
secret_key,
endpoint_url,
&endpoint.region,
endpoint.force_path_style,
))
}
fn build_client(
access_key: String,
secret_key: String,
endpoint_url: &str,
region: &str,
force_path_style: bool,
) -> Client {
let config = aws_sdk_s3::Config::builder()
.behavior_version(BehaviorVersion::latest())
.credentials_provider(Credentials::new(
access_key,
secret_key,
None,
None,
"imageforge-admin",
))
.region(Region::new(region.to_string()))
.endpoint_url(endpoint_url)
.force_path_style(force_path_style)
.request_checksum_calculation(RequestChecksumCalculation::WhenRequired)
.response_checksum_validation(ResponseChecksumValidation::WhenRequired)
.build();
Client::from_conf(config)
}
async fn endpoint_for_object(
state: &AppState,
object: &ObjectLocator,
) -> Result<StorageEndpoint, AppError> {
let endpoint_id = object
.endpoint_id
.ok_or_else(|| AppError::new(ErrorCode::StorageUnavailable, "S3 对象缺少存储端点标识"))?;
get_endpoint(state, endpoint_id).await
}
fn local_path(state: &AppState, key: &str) -> Result<PathBuf, AppError> {
if key.is_empty()
|| key.starts_with('/')
|| key.starts_with('\\')
|| key.split('/').any(|part| part == "..")
{
return Err(AppError::new(ErrorCode::InvalidRequest, "非法存储对象键"));
}
Ok(Path::new(&state.config.storage_path).join(key))
}
fn content_disposition(file_name: &str) -> String {
let ascii = file_name
.chars()
.map(|ch| {
if ch.is_ascii_alphanumeric() || matches!(ch, '.' | '-' | '_') {
ch
} else {
'_'
}
})
.collect::<String>();
let ascii = if ascii.is_empty() { "download" } else { &ascii };
let encoded =
percent_encoding::utf8_percent_encode(file_name, percent_encoding::NON_ALPHANUMERIC);
format!("attachment; filename=\"{ascii}\"; filename*=UTF-8''{encoded}")
}
async fn multipart_upload(
client: &Client,
bucket: &str,
key: &str,
path: &Path,
content_type: &str,
) -> Result<Option<String>, AppError> {
let created = client
.create_multipart_upload()
.bucket(bucket)
.key(key)
.content_type(content_type)
.send()
.await
.map_err(|err| storage_error("创建 S3 分片上传失败", err))?;
let upload_id = created
.upload_id()
.ok_or_else(|| AppError::new(ErrorCode::StorageUnavailable, "S3 未返回分片上传标识"))?;
let result = async {
let mut file = tokio::fs::File::open(path).await.map_err(|err| {
AppError::new(ErrorCode::StorageUnavailable, "读取待上传文件失败").with_source(err)
})?;
let mut parts = Vec::new();
let mut part_number = 1;
loop {
let mut buffer = vec![0u8; MULTIPART_PART_SIZE];
let mut filled = 0;
while filled < buffer.len() {
let read = file.read(&mut buffer[filled..]).await.map_err(|err| {
AppError::new(ErrorCode::StorageUnavailable, "读取上传分片失败")
.with_source(err)
})?;
if read == 0 {
break;
}
filled += read;
}
if filled == 0 {
break;
}
buffer.truncate(filled);
let uploaded = client
.upload_part()
.bucket(bucket)
.key(key)
.upload_id(upload_id)
.part_number(part_number)
.body(ByteStream::from(buffer))
.send()
.await
.map_err(|err| storage_error("上传 S3 分片失败", err))?;
let part = CompletedPart::builder()
.part_number(part_number)
.set_e_tag(uploaded.e_tag().map(ToOwned::to_owned))
.build();
parts.push(part);
part_number += 1;
}
let completed = CompletedMultipartUpload::builder()
.set_parts(Some(parts))
.build();
let response = client
.complete_multipart_upload()
.bucket(bucket)
.key(key)
.upload_id(upload_id)
.multipart_upload(completed)
.send()
.await
.map_err(|err| storage_error("完成 S3 分片上传失败", err))?;
Ok(response.e_tag().map(ToOwned::to_owned))
}
.await;
if result.is_err() {
let _ = client
.abort_multipart_upload()
.bucket(bucket)
.key(key)
.upload_id(upload_id)
.send()
.await;
}
result
}
fn storage_error(message: &'static str, err: impl std::fmt::Display) -> AppError {
AppError::new(ErrorCode::StorageUnavailable, message).with_source(err)
}
fn public_test_request_error(endpoint: &StorageEndpoint, err: reqwest::Error) -> AppError {
AppError::new(ErrorCode::StorageUnavailable, "公网 Endpoint 下载测试失败").with_source(format!(
"endpoint={}, connect={}, timeout={}",
endpoint.public_endpoint,
err.is_connect(),
err.is_timeout()
))
}
#[derive(Clone, Copy)]
enum EndpointKind {
Internal,
Public,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn retention_prefixes_match_product_policy() {
assert_eq!(retention_prefix(24), "1d");
assert_eq!(retention_prefix(168), "7d");
assert_eq!(retention_prefix(360), "15d");
assert_eq!(retention_prefix(720), "custom-720h");
}
#[test]
fn generated_keys_are_lifecycle_scoped() {
let task_id = Uuid::nil();
let file_id = Uuid::from_u128(1);
let key = result_key(168, task_id, file_id, "webp");
assert!(key.starts_with("results/7d/"));
assert!(key.ends_with("/00000000-0000-0000-0000-000000000001.webp"));
assert!(archive_key(360, task_id).starts_with("archives/15d/"));
}
#[test]
fn content_disposition_never_injects_headers() {
let value = content_disposition("a\r\nX-Test: yes.png");
assert!(!value.contains('\r'));
assert!(!value.contains('\n'));
assert!(value.contains("filename*=UTF-8''"));
}
#[tokio::test]
async fn presigning_uses_public_path_style_endpoint() {
let client = build_client(
"GKTESTACCESSKEY".to_string(),
"test-secret-key".to_string(),
"https://files.example.com",
"garage",
true,
);
let request = client
.get_object()
.bucket("imageforge-results")
.key("results/1d/test.webp")
.presigned(PresigningConfig::expires_in(Duration::from_secs(300)).unwrap())
.await
.unwrap();
let uri = request.uri().to_string();
assert!(
uri.starts_with("https://files.example.com/imageforge-results/results/1d/test.webp?")
);
assert!(uri.contains("X-Amz-Algorithm=AWS4-HMAC-SHA256"));
assert!(uri.contains("X-Amz-Expires=300"));
assert!(uri.contains("X-Amz-Signature="));
}
}