1use crate::settings::SettingsReadGuard;
2use aws_sdk_s3::{
3 Client as S3Client,
4 config::{Config as S3Config, Credentials, Region, retry::RetryConfig, timeout::TimeoutConfig},
5 primitives::ByteStream,
6 types::{CompletedMultipartUpload, CompletedPart},
7};
8use compact_str::ToCompactString;
9use serde::{Deserialize, Serialize};
10use std::{path::Path, sync::Arc};
11use tokio::io::{AsyncReadExt, AsyncWriteExt};
12use tokio_util::bytes::{Bytes, BytesMut};
13use utoipa::ToSchema;
14
15#[derive(ToSchema, Deserialize, Serialize)]
16pub struct StorageAsset {
17 pub name: compact_str::CompactString,
18 pub url: String,
19 pub size: u64,
20 pub is_directory: bool,
21 pub created: chrono::DateTime<chrono::Utc>,
22}
23
24fn get_s3_client(
25 access_key: &str,
26 secret_key: &str,
27 region: &str,
28 endpoint: &str,
29 path_style: bool,
30) -> Result<S3Client, anyhow::Error> {
31 let credentials = Credentials::new(access_key, secret_key, None, None, "calagopus-static");
32
33 let timeout_config = TimeoutConfig::builder()
34 .connect_timeout(std::time::Duration::from_secs(10))
35 .build();
36
37 let config = S3Config::builder()
38 .behavior_version(aws_sdk_s3::config::BehaviorVersion::latest())
39 .credentials_provider(credentials)
40 .region(Region::new(region.to_string()))
41 .endpoint_url(endpoint)
42 .force_path_style(path_style)
43 .timeout_config(timeout_config)
44 .retry_config(RetryConfig::standard())
45 .build();
46
47 Ok(S3Client::from_conf(config))
48}
49
50pub struct StorageUrlRetriever<'a> {
51 settings: SettingsReadGuard<'a>,
52}
53
54impl<'a> StorageUrlRetriever<'a> {
55 pub fn new(settings: SettingsReadGuard<'a>) -> Self {
56 Self { settings }
57 }
58
59 pub fn get_settings(&self) -> &super::settings::AppSettings {
60 &self.settings
61 }
62
63 pub fn get_url(&self, path: impl AsRef<str>) -> String {
64 match &self.settings.storage_driver {
65 super::settings::StorageDriver::Filesystem { .. } => {
66 format!(
67 "{}/{}",
68 self.settings.app.url.trim_end_matches('/'),
69 path.as_ref()
70 )
71 }
72 super::settings::StorageDriver::S3 { public_url, .. } => {
73 format!("{}/{}", public_url.trim_end_matches('/'), path.as_ref())
74 }
75 }
76 }
77}
78
79pub struct Storage {
80 settings: Arc<super::settings::Settings>,
81}
82
83impl Storage {
84 pub fn new(settings: Arc<super::settings::Settings>) -> Self {
85 Self { settings }
86 }
87
88 pub async fn retrieve_urls(&self) -> Result<StorageUrlRetriever<'_>, anyhow::Error> {
89 let settings = self.settings.get().await?;
90
91 Ok(StorageUrlRetriever::new(settings))
92 }
93
94 pub async fn remove(&self, path: Option<impl AsRef<str>>) -> Result<(), anyhow::Error> {
95 let path = match path {
96 Some(path) => path,
97 None => return Ok(()),
98 };
99 let path = path.as_ref();
100
101 if path.is_empty() || path.contains("..") || path.starts_with('/') {
102 return Err(anyhow::anyhow!("invalid path"));
103 }
104
105 let settings = self.settings.get().await?;
106
107 tracing::debug!(path, "removing file");
108
109 match &settings.storage_driver {
110 super::settings::StorageDriver::Filesystem { path: base_path } => {
111 let base_filesystem =
112 match crate::cap::CapFilesystem::async_new(base_path.into()).await {
113 Ok(base_filesystem) => base_filesystem,
114 Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(()),
115 Err(err) => return Err(err.into()),
116 };
117 drop(settings);
118
119 if let Err(err) = base_filesystem.async_remove_file(&path).await
120 && err
121 .downcast_ref::<std::io::Error>()
122 .is_none_or(|e| e.kind() != std::io::ErrorKind::NotFound)
123 {
124 return Err(err);
125 }
126
127 if let Some(parent) = Path::new(path).parent().map(|p| p.to_path_buf()) {
128 tokio::spawn(async move {
129 tokio::time::sleep(std::time::Duration::from_secs(10)).await;
130
131 let mut directory = match base_filesystem.async_read_dir(&parent).await {
132 Ok(directory) => directory,
133 Err(_) => return,
134 };
135
136 if directory.next_entry().await.is_none() {
137 base_filesystem.async_remove_dir(parent).await.ok();
138 }
139 });
140 }
141 }
142 super::settings::StorageDriver::S3 {
143 access_key,
144 secret_key,
145 bucket,
146 region,
147 endpoint,
148 path_style,
149 ..
150 } => {
151 let s3_client =
152 get_s3_client(access_key, secret_key, region, endpoint, *path_style)?;
153 let bucket = bucket.clone();
154 drop(settings);
155
156 s3_client
157 .delete_object()
158 .bucket(bucket)
159 .key(path)
160 .send()
161 .await?;
162 }
163 }
164
165 Ok(())
166 }
167
168 pub async fn store(
169 &self,
170 path: impl AsRef<str>,
171 data: impl tokio::io::AsyncRead + Unpin,
172 content_type: impl AsRef<str>,
173 ) -> Result<u64, anyhow::Error> {
174 let path = path.as_ref();
175 let content_type = content_type.as_ref();
176
177 if path.is_empty() || path.contains("..") || path.starts_with('/') {
178 return Err(anyhow::anyhow!("invalid path"));
179 }
180
181 let settings = self.settings.get().await?;
182
183 tracing::debug!(path, content_type, "storing file");
184
185 match &settings.storage_driver {
186 super::settings::StorageDriver::Filesystem { path: base_path } => {
187 tokio::fs::create_dir_all(base_path).await?;
188
189 let base_filesystem =
190 crate::cap::CapFilesystem::async_new(base_path.into()).await?;
191 drop(settings);
192
193 if let Some(parent) = Path::new(path).parent() {
194 base_filesystem.async_create_dir_all(parent).await?;
195 }
196
197 let mut file = base_filesystem.async_create(path).await?;
198 let mut data = data;
199 let bytes = tokio::io::copy(&mut data, &mut file).await?;
200
201 file.shutdown().await?;
202 Ok(bytes)
203 }
204 super::settings::StorageDriver::S3 {
205 access_key,
206 secret_key,
207 bucket,
208 region,
209 endpoint,
210 path_style,
211 ..
212 } => {
213 let s3_client =
214 get_s3_client(access_key, secret_key, region, endpoint, *path_style)?;
215 let bucket = bucket.clone();
216 drop(settings);
217
218 upload_multipart(&s3_client, &bucket, path, content_type, data).await
219 }
220 }
221 }
222
223 pub async fn list(
224 &self,
225 base: impl AsRef<str>,
226 directory: impl AsRef<str>,
227 page: usize,
228 per_page: usize,
229 ) -> Result<crate::models::Pagination<StorageAsset>, anyhow::Error> {
230 let base = base.as_ref();
231 let directory = directory.as_ref();
232
233 if base.is_empty() || base.contains("..") || base.starts_with('/') {
234 return Err(anyhow::anyhow!("invalid base path"));
235 }
236 if !directory.is_empty()
237 && (directory.contains("..") || directory.starts_with('/') || directory.ends_with('/'))
238 {
239 return Err(anyhow::anyhow!("invalid directory path"));
240 }
241
242 let settings = self.settings.get().await?;
243
244 match &settings.storage_driver {
245 super::settings::StorageDriver::Filesystem { path: base_path } => {
246 let dir_path = if directory.is_empty() {
247 Path::new(base_path).join(base)
248 } else {
249 Path::new(base_path).join(base).join(directory)
250 };
251
252 let base_filesystem = match crate::cap::CapFilesystem::async_new(dir_path).await {
253 Ok(base_filesystem) => base_filesystem,
254 Err(err) if err.kind() == std::io::ErrorKind::NotFound => {
255 return Ok(crate::models::Pagination {
256 total: 0,
257 per_page: per_page as i64,
258 page: page as i64,
259 data: Vec::new(),
260 });
261 }
262 Err(err) => return Err(err.into()),
263 };
264 drop(settings);
265
266 let mut dir_reader = base_filesystem.async_read_dir("").await?;
267 let mut raw_dirs: Vec<String> = Vec::new();
268 let mut raw_files: Vec<String> = Vec::new();
269
270 while let Some(Ok((is_dir, name))) = dir_reader.next_entry().await {
271 if is_dir {
272 raw_dirs.push(name);
273 } else {
274 raw_files.push(name);
275 }
276 }
277
278 raw_dirs.sort_unstable();
279 raw_files.sort_unstable();
280
281 let total = (raw_dirs.len() + raw_files.len()) as i64;
282 let start = (page - 1) * per_page;
283
284 let storage_url_retriever = self.retrieve_urls().await?;
285
286 let mut entries = Vec::new();
287
288 for (is_dir, name) in raw_dirs
289 .into_iter()
290 .map(|n| (true, n))
291 .chain(raw_files.into_iter().map(|n| (false, n)))
292 .skip(start)
293 .take(per_page)
294 {
295 let full_name = if directory.is_empty() {
296 name.clone()
297 } else {
298 format!("{directory}/{name}")
299 };
300
301 let (size, created) = if is_dir {
302 (0u64, chrono::DateTime::<chrono::Utc>::default())
303 } else {
304 let metadata = match base_filesystem.async_metadata(&name).await {
305 Ok(m) => m,
306 Err(_) => continue,
307 };
308 let created = metadata
309 .created()
310 .or_else(|_| metadata.modified())?
311 .into_std()
312 .into();
313 (metadata.len(), created)
314 };
315
316 entries.push(StorageAsset {
317 url: storage_url_retriever.get_url(format!("{base}/{full_name}")),
318 name: full_name.to_compact_string(),
319 size,
320 is_directory: is_dir,
321 created,
322 });
323 }
324
325 Ok(crate::models::Pagination {
326 total,
327 per_page: per_page as i64,
328 page: page as i64,
329 data: entries,
330 })
331 }
332 super::settings::StorageDriver::S3 {
333 access_key,
334 secret_key,
335 bucket,
336 region,
337 endpoint,
338 path_style,
339 ..
340 } => {
341 let s3_client =
342 get_s3_client(access_key, secret_key, region, endpoint, *path_style)?;
343 let bucket_name = bucket.clone();
344 drop(settings);
345
346 let s3_prefix = if directory.is_empty() {
347 format!("{base}/")
348 } else {
349 format!("{base}/{directory}/")
350 };
351 let strip_prefix = format!("{base}/");
352
353 let storage_url_retriever = self.retrieve_urls().await?;
354
355 let mut dirs = Vec::new();
356 let mut files = Vec::new();
357
358 let mut paginator = s3_client
359 .list_objects_v2()
360 .bucket(&*bucket_name)
361 .prefix(&s3_prefix)
362 .delimiter("/")
363 .into_paginator()
364 .send();
365
366 while let Some(result) = paginator.next().await {
367 let page = result?;
368
369 for cp in page.common_prefixes() {
370 let Some(prefix) = cp.prefix() else { continue };
371 let name = prefix
372 .trim_start_matches(&strip_prefix)
373 .trim_end_matches('/')
374 .to_compact_string();
375 dirs.push(StorageAsset {
376 url: storage_url_retriever.get_url(prefix),
377 name,
378 size: 0,
379 is_directory: true,
380 created: chrono::DateTime::<chrono::Utc>::default(),
381 });
382 }
383
384 for entry in page.contents() {
385 let Some(key) = entry.key() else { continue };
386 if key == s3_prefix {
387 continue;
388 }
389 let name = key.trim_start_matches(&strip_prefix).to_compact_string();
390 let size = entry.size().unwrap_or(0).max(0) as u64;
391 let created = entry
392 .last_modified()
393 .and_then(|dt| {
394 chrono::DateTime::<chrono::Utc>::from_timestamp(
395 dt.secs(),
396 dt.subsec_nanos(),
397 )
398 })
399 .unwrap_or_default();
400
401 files.push(StorageAsset {
402 url: storage_url_retriever.get_url(key),
403 name,
404 size,
405 is_directory: false,
406 created,
407 });
408 }
409 }
410
411 let total = (dirs.len() + files.len()) as i64;
412 let start = (page - 1) * per_page;
413
414 Ok(crate::models::Pagination {
415 total,
416 per_page: per_page as i64,
417 page: page as i64,
418 data: dirs
419 .into_iter()
420 .chain(files)
421 .skip(start)
422 .take(per_page)
423 .collect(),
424 })
425 }
426 }
427 }
428}
429
430const PART_SIZE: usize = 16 * 1024 * 1024;
431
432async fn upload_multipart(
433 client: &S3Client,
434 bucket: &str,
435 key: &str,
436 content_type: &str,
437 mut data: impl tokio::io::AsyncRead + Unpin,
438) -> Result<u64, anyhow::Error> {
439 let first_part = read_part(&mut data, PART_SIZE).await?;
440
441 if first_part.len() < PART_SIZE {
442 let total = first_part.len() as u64;
443 client
444 .put_object()
445 .bucket(bucket)
446 .key(key)
447 .content_type(content_type)
448 .body(ByteStream::from(first_part))
449 .send()
450 .await?;
451 return Ok(total);
452 }
453
454 let create = client
455 .create_multipart_upload()
456 .bucket(bucket)
457 .key(key)
458 .content_type(content_type)
459 .send()
460 .await?;
461
462 let upload_id = create
463 .upload_id()
464 .ok_or_else(|| anyhow::anyhow!("S3 did not return an upload_id"))?
465 .to_string();
466
467 let result = run_multipart(client, bucket, key, &upload_id, &mut data, first_part).await;
468
469 match result {
470 Ok(total) => Ok(total),
471 Err(err) => {
472 if let Err(abort_err) = client
473 .abort_multipart_upload()
474 .bucket(bucket)
475 .key(key)
476 .upload_id(&upload_id)
477 .send()
478 .await
479 {
480 tracing::warn!(
481 bucket,
482 key,
483 upload_id,
484 "failed to abort multipart upload after error: {:#?}",
485 abort_err
486 );
487 }
488 Err(err)
489 }
490 }
491}
492
493async fn run_multipart(
494 client: &S3Client,
495 bucket: &str,
496 key: &str,
497 upload_id: &str,
498 data: &mut (impl tokio::io::AsyncRead + Unpin),
499 first_part: Bytes,
500) -> Result<u64, anyhow::Error> {
501 let mut completed = Vec::new();
502 let mut total: u64 = 0;
503 let mut part_number: i32 = 1;
504 let mut current = first_part;
505
506 loop {
507 let part_len = current.len() as u64;
508 let resp = client
509 .upload_part()
510 .bucket(bucket)
511 .key(key)
512 .upload_id(upload_id)
513 .part_number(part_number)
514 .body(ByteStream::from(current))
515 .send()
516 .await?;
517
518 completed.push(
519 CompletedPart::builder()
520 .part_number(part_number)
521 .set_e_tag(resp.e_tag().map(|s| s.to_string()))
522 .build(),
523 );
524 total += part_len;
525 part_number += 1;
526
527 let next = read_part(data, PART_SIZE).await?;
528 if next.is_empty() {
529 break;
530 }
531 current = next;
532 }
533
534 let completed_upload = CompletedMultipartUpload::builder()
535 .set_parts(Some(completed))
536 .build();
537
538 client
539 .complete_multipart_upload()
540 .bucket(bucket)
541 .key(key)
542 .upload_id(upload_id)
543 .multipart_upload(completed_upload)
544 .send()
545 .await?;
546
547 Ok(total)
548}
549
550async fn read_part(
551 data: &mut (impl tokio::io::AsyncRead + Unpin),
552 cap: usize,
553) -> Result<Bytes, std::io::Error> {
554 let mut buf = BytesMut::with_capacity(cap);
555 while buf.len() < cap {
556 let n = data.read_buf(&mut buf).await?;
557 if n == 0 {
558 break;
559 }
560 }
561 Ok(buf.freeze())
562}