8923bc0
use std::pin::Pin;
use config::s3::S3Config;
use futures::StreamExt as _;
use http::HeaderMap;
use s3::Bucket;
use s3::bucket::{CHUNK_SIZE, Part};
use s3::creds::Credentials;
use s3::region::Region;
use s3::serde_types::{DateTime as S3DateTime, Object as S3Object};
use tokio::io::{AsyncRead, AsyncReadExt};
use crate::error::StorageError;
pub const MAX_PRESIGN_EXPIRY_SECS: u32 = 604_800;
#[derive(Debug, Clone)]
pub struct ObjectInfo {
pub key: String,
pub size: u64,
pub last_modified: Option<S3DateTime>,
pub etag: Option<String>,
}
impl From<S3Object> for ObjectInfo {
fn from(obj: S3Object) -> Self {
Self {
key: obj.key,
size: obj.size,
last_modified: Some(obj.last_modified),
etag: obj.e_tag,
}
}
}
#[derive(Debug, Clone)]
pub struct ObjectPage {
pub objects: Vec<ObjectInfo>,
pub next_continuation_token: Option<String>,
}
#[derive(Clone)]
pub struct S3Client {
bucket: Bucket,
/// Same bucket, but pointed at `s3.public_endpoint`. Presigned URLs carry
/// the endpoint host inside their signature, so they must be produced by a
/// client whose region points at the address the *caller* can reach (eg.
/// https://s3.gitdata.ai) rather than the in-cluster service DNS name.
presign_bucket: Bucket,
bucket_name: String,
}
impl S3Client {
pub fn new(config: &S3Config) -> Result<Self, StorageError> {
let (bucket_name, region) = config.validate().map_err(StorageError::InvalidConfig)?;
let region = {
match config.endpoint.as_deref().filter(|s| !s.is_empty()) {
Some(endpoint) => Ok(Region::Custom {
region: region.to_string(),
endpoint: endpoint.to_string(),
}),
None => region.parse::<Region>().map_err(|err| {
StorageError::InvalidConfig(format!("region {region:?}: {err}"))
}),
}
}?;
let credentials = {
match (
config.access_key_id.as_deref().filter(|s| !s.is_empty()),
config
.secret_access_key
.as_deref()
.filter(|s| !s.is_empty()),
) {
(Some(access), Some(secret)) => Ok(Credentials::new(
Some(access),
Some(secret),
None,
None,
None,
)?),
(Some(_), None) => Err(StorageError::InvalidConfig(
"access_key_id provided without secret_access_key".into(),
)),
(None, Some(_)) => Err(StorageError::InvalidConfig(
"secret_access_key provided without access_key_id".into(),
)),
(None, None) => Ok(Credentials::default()?),
}
}?;
let mut bucket = Bucket::new(bucket_name, region, credentials)?;
if config.use_path_style() {
bucket = bucket.with_path_style();
}
if let Some(timeout) = config.request_timeout() {
bucket = bucket.with_request_timeout(timeout);
}
let mut presign_bucket = bucket.clone();
if let Some(public_endpoint) = config.public_endpoint.as_deref().filter(|s| !s.is_empty()) {
presign_bucket.region = Region::Custom {
region: config.region().to_string(),
endpoint: public_endpoint.to_string(),
};
}
Ok(Self {
bucket,
presign_bucket,
bucket_name: bucket_name.to_string(),
})
}
pub fn bucket_name(&self) -> &str {
&self.bucket_name
}
pub async fn put(&self, key: &str, body: &[u8]) -> Result<(), StorageError> {
self.put_with_content_type(key, body, "application/octet-stream")
.await
}
pub async fn put_with_content_type(
&self,
key: &str,
body: &[u8],
content_type: &str,
) -> Result<(), StorageError> {
let response = self
.bucket
.put_object_with_content_type(key, body, content_type)
.await?;
let status = response.status_code();
if (200..300).contains(&status) {
Ok(())
} else {
Err(StorageError::UnexpectedStatus {
operation: "put_object",
status,
})
}
}
/// True streaming upload with bounded memory: consumes the reader in
/// 8 MiB chunks and uploads them serially (single PUT below one chunk,
/// multipart above), hashing nothing but counting bytes through a tee.
///
/// Returns the number of bytes actually read from the reader.
pub async fn put_stream<R: AsyncRead + Unpin>(
&self,
key: &str,
reader: &mut R,
content_type: &str,
) -> Result<u64, StorageError> {
let first = read_chunk(reader).await?;
if first.len() < CHUNK_SIZE {
self.put_with_content_type(key, &first, content_type)
.await?;
return Ok(first.len() as u64);
}
let msg = self
.bucket
.initiate_multipart_upload(key, content_type)
.await?;
let path = msg.key;
let upload_id = msg.upload_id;
let mut chunk = first;
let mut part_number: u32 = 0;
let mut parts: Vec<Part> = Vec::new();
let mut total: u64 = 0;
loop {
total += chunk.len() as u64;
part_number += 1;
match self
.bucket
.put_multipart_chunk(chunk, &path, part_number, &upload_id, content_type)
.await
{
Ok(part) => parts.push(part),
Err(e) => {
abort_upload(&self.bucket, &path, &upload_id).await;
return Err(StorageError::from(e));
}
}
chunk = match read_chunk(reader).await {
Ok(c) => c,
Err(e) => {
abort_upload(&self.bucket, &path, &upload_id).await;
return Err(e);
}
};
if chunk.is_empty() {
abort_upload(&self.bucket, &path, &upload_id).await;
return Err(StorageError::UnexpectedStatus {
operation: "put_stream",
status: 400,
});
}
if chunk.len() < CHUNK_SIZE {
total += chunk.len() as u64;
part_number += 1;
match self
.bucket
.put_multipart_chunk(chunk, &path, part_number, &upload_id, content_type)
.await
{
Ok(part) => parts.push(part),
Err(e) => {
abort_upload(&self.bucket, &path, &upload_id).await;
return Err(StorageError::from(e));
}
}
break;
}
}
match self
.bucket
.complete_multipart_upload(&path, &upload_id, parts)
.await
{
Ok(r) if (200..300).contains(&r.status_code()) => Ok(total),
Ok(r) => {
abort_upload(&self.bucket, &path, &upload_id).await;
Err(StorageError::UnexpectedStatus {
operation: "complete_multipart_upload",
status: r.status_code(),
})
}
Err(e) => {
abort_upload(&self.bucket, &path, &upload_id).await;
Err(StorageError::from(e))
}
}
}
pub async fn get(&self, key: &str) -> Result<Vec<u8>, StorageError> {
let response = self.bucket.get_object(key).await?;
let status = response.status_code();
if !(200..300).contains(&status) {
return Err(StorageError::UnexpectedStatus {
operation: "get_object",
status,
});
}
Ok(response.as_slice().to_vec())
}
/// Stream an object out as a byte stream (ready for an HTTP response
/// body). The returned stream yields `Result<Bytes, StorageError>`.
pub async fn get_stream(
&self,
key: &str,
) -> Result<
Pin<Box<dyn futures::Stream<Item = Result<bytes::Bytes, StorageError>> + Send>>,
StorageError,
> {
let stream = self.bucket.get_object_stream(key).await?;
if !(200..300).contains(&stream.status_code) {
return Err(StorageError::UnexpectedStatus {
operation: "get_object",
status: stream.status_code,
});
}
Ok(Box::pin(
stream
.bytes
.map(move |item| item.map_err(StorageError::from)),
))
}
pub async fn delete(&self, key: &str) -> Result<(), StorageError> {
let response = self.bucket.delete_object(key).await?;
let status = response.status_code();
if (200..300).contains(&status) || status == 404 {
Ok(())
} else {
Err(StorageError::UnexpectedStatus {
operation: "delete_object",
status,
})
}
}
/// Whether a key is present. A missing key is a `false`, not an error: the
/// underlying `HEAD` answers 404, which rust-s3 surfaces as a failure, and
/// callers use this to choose between "create" and "already there".
pub async fn exists(&self, key: &str) -> Result<bool, StorageError> {
match self.bucket.head_object(key).await {
Ok((_, status)) => Ok((200..300).contains(&status)),
Err(s3::error::S3Error::HttpFailWithBody(404, _))
| Err(s3::error::S3Error::HttpFailWithBody(403, _)) => Ok(false),
Err(err) => Err(StorageError::S3(err)),
}
}
pub async fn head(&self, key: &str) -> Result<Option<(u64, String)>, StorageError> {
let (meta, status) = self.bucket.head_object(key).await?;
if !(200..300).contains(&status) {
return Ok(None);
}
Ok(Some((
meta.content_length.unwrap_or(0).max(0) as u64,
meta.content_type.unwrap_or_default(),
)))
}
pub async fn list_page(
&self,
prefix: &str,
continuation_token: Option<&str>,
max_keys: Option<usize>,
) -> Result<ObjectPage, StorageError> {
let (page, status) = self
.bucket
.list_page(
prefix.to_string(),
None,
continuation_token.map(str::to_owned),
None,
max_keys,
)
.await?;
if !(200..300).contains(&status) {
return Err(StorageError::UnexpectedStatus {
operation: "list",
status,
});
}
Ok(ObjectPage {
objects: page.contents.into_iter().map(ObjectInfo::from).collect(),
next_continuation_token: page.next_continuation_token,
})
}
pub async fn list(&self, prefix: &str) -> Result<Vec<ObjectInfo>, StorageError> {
let mut out = Vec::new();
let mut token: Option<String> = None;
loop {
let page = self.list_page(prefix, token.as_deref(), None).await?;
let more = page.next_continuation_token.is_some();
out.extend(page.objects);
match page.next_continuation_token {
Some(next) => token = Some(next),
None => {
if !more {
break;
}
}
}
if token.is_none() {
break;
}
}
Ok(out)
}
pub fn bucket(&self) -> &Bucket {
&self.bucket
}
pub fn presign_get(&self, key: &str, expiry_secs: u32) -> Result<String, StorageError> {
check_expiry(expiry_secs)?;
Ok(self.presign_bucket.presign_get(key, expiry_secs, None)?)
}
pub fn presign_put(&self, key: &str, expiry_secs: u32) -> Result<String, StorageError> {
check_expiry(expiry_secs)?;
Ok(self.presign_bucket.presign_put(key, expiry_secs, None)?)
}
pub fn presign_put_with_headers(
&self,
key: &str,
expiry_secs: u32,
custom_headers: HeaderMap,
) -> Result<String, StorageError> {
check_expiry(expiry_secs)?;
Ok(self
.presign_bucket
.presign_put(key, expiry_secs, Some(custom_headers))?)
}
pub fn presign_delete(&self, key: &str, expiry_secs: u32) -> Result<String, StorageError> {
check_expiry(expiry_secs)?;
Ok(self.presign_bucket.presign_delete(key, expiry_secs)?)
}
pub fn presign_post(
&self,
key: &str,
expiry_secs: u32,
post_policy: String,
) -> Result<String, StorageError> {
check_expiry(expiry_secs)?;
Ok(self
.presign_bucket
.presign_post(key, expiry_secs, post_policy)?)
}
}
fn check_expiry(expiry_secs: u32) -> Result<(), StorageError> {
if expiry_secs > MAX_PRESIGN_EXPIRY_SECS {
return Err(StorageError::InvalidConfig(format!(
"presign expiry must be <= {MAX_PRESIGN_EXPIRY_SECS} seconds (got {expiry_secs})"
)));
}
Ok(())
}
/// Read up to one `CHUNK_SIZE` block from the reader (EOF-tolerant).
async fn read_chunk<R: AsyncRead + Unpin>(reader: &mut R) -> Result<Vec<u8>, StorageError> {
let mut chunk = Vec::with_capacity(CHUNK_SIZE);
let mut take = reader.take(CHUNK_SIZE as u64);
let read = AsyncReadExt::read_to_end(&mut take, &mut chunk).await?;
chunk.truncate(read);
Ok(chunk)
}
/// Fire-and-forget multipart abort (errors are intentionally swallowed: the
/// original failure is what the caller cares about).
async fn abort_upload(bucket: &Bucket, path: &str, upload_id: &str) {
if let Err(error) = bucket.abort_upload(path, upload_id).await {
tracing::warn!(%error, "failed to abort multipart upload");
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::io;
#[tokio::test]
async fn read_chunk_respects_chunk_size() {
let mut data = (0..CHUNK_SIZE as u32)
.map(|i| (i % 251) as u8)
.collect::<Vec<_>>();
let mut cursor = io::Cursor::new(&data);
let first = read_chunk(&mut cursor).await.expect("read");
assert_eq!(first.len(), CHUNK_SIZE);
let second = read_chunk(&mut cursor).await.expect("read");
assert!(second.is_empty());
data.truncate(CHUNK_SIZE - 1);
let mut cursor = io::Cursor::new(&data);
let first = read_chunk(&mut cursor).await.expect("read");
assert_eq!(first.len(), CHUNK_SIZE - 1);
}
}