use std::time::Duration;
use axum::body::Bytes;
use futures_util::Stream;
use rusty_s3::actions::{
DeleteObject, GetObject, HeadBucket, HeadObject, ListObjectsV2, PutObject, S3Action,
};
use rusty_s3::{Bucket, Credentials, UrlStyle};
use crate::error::Error;
use crate::storage::s3::S3Config;
const COPY_SOURCE: &str = "x-amz-copy-source";
pub(crate) struct Entry {
pub(crate) key: String,
last_modified: String,
pub(crate) size: u64,
}
impl Entry {
pub(crate) fn age(&self) -> Option<Duration> {
let written = time::OffsetDateTime::parse(
&self.last_modified,
&time::format_description::well_known::Rfc3339,
)
.ok()?;
Duration::try_from(time::OffsetDateTime::now_utc() - written).ok()
}
}
pub struct Presigned {
pub href: String,
pub headers: Vec<(String, String)>,
}
async fn read_retrying(request: reqwest::RequestBuilder) -> Result<reqwest::Response, Error> {
let retry = request.try_clone();
match request.send().await {
Ok(response) => Ok(response),
Err(_) => match retry {
Some(retry) => retry.send().await.map_err(|_| unreachable_store()),
None => Err(unreachable_store()),
},
}
}
fn unreachable_store() -> Error {
Error::Storage(std::io::Error::other("the object store is unreachable"))
}
#[derive(Clone)]
pub struct Keyspace {
bucket: Bucket,
credentials: Credentials,
client: reqwest::Client,
lifetime: Duration,
}
impl Keyspace {
pub fn new(config: &S3Config) -> Result<Self, Error> {
crate::tls::install_crypto_provider();
let style = if config.path_style {
UrlStyle::Path
} else {
UrlStyle::VirtualHost
};
let bucket = Bucket::new(
config
.endpoint
.parse()
.map_err(|_| Error::Misconfigured("LFSX_S3_ENDPOINT is not a URL"))?,
style,
config.bucket.clone(),
config.region.clone(),
)
.map_err(|_| Error::Misconfigured("LFSX_S3_BUCKET is not a usable bucket name"))?;
Ok(Self {
bucket,
credentials: Credentials::new(config.access_key.clone(), config.secret_key.clone()),
client: reqwest::Client::new(),
lifetime: config.lifetime,
})
}
pub(crate) async fn reachable(&self) -> Result<(), Error> {
let action = HeadBucket::new(&self.bucket, Some(&self.credentials));
let response = read_retrying(self.client.head(action.sign(self.lifetime))).await?;
if !response.status().is_success() {
return Err(Error::Storage(std::io::Error::other(format!(
"the object store answered {} for the bucket",
response.status()
))));
}
Ok(())
}
pub(crate) fn signed_download(&self, key: &str) -> String {
GetObject::new(&self.bucket, Some(&self.credentials), key)
.sign(self.lifetime)
.to_string()
}
pub(crate) fn signed_upload(&self, key: &str, headers: Vec<(String, String)>) -> Presigned {
let mut action = PutObject::new(&self.bucket, Some(&self.credentials), key);
for (name, value) in &headers {
action
.headers_mut()
.insert(name.clone(), std::borrow::Cow::Owned(value.clone()));
}
Presigned {
href: action.sign(self.lifetime).to_string(),
headers,
}
}
pub(crate) async fn get_range(
&self,
key: &str,
start: u64,
length: u64,
) -> Result<impl Stream<Item = Result<Bytes, reqwest::Error>> + use<>, Error> {
let action = GetObject::new(&self.bucket, Some(&self.credentials), key);
let response = self
.client
.get(action.sign(self.lifetime))
.header(
reqwest::header::RANGE,
format!("bytes={start}-{}", start + length.saturating_sub(1)),
)
.send()
.await
.map_err(|_| unreachable_store())?;
if !response.status().is_success() {
return Err(Error::NotFound);
}
Ok(response.bytes_stream())
}
pub(crate) async fn head(&self, key: &str) -> Result<u64, Error> {
let action = HeadObject::new(&self.bucket, Some(&self.credentials), key);
let url = action.sign(self.lifetime);
let response = read_retrying(self.client.head(url)).await?;
if !response.status().is_success() {
return Err(Error::NotFound);
}
response
.headers()
.get(reqwest::header::CONTENT_LENGTH)
.and_then(|value| value.to_str().ok())
.and_then(|value| value.parse().ok())
.ok_or_else(|| {
Error::Storage(std::io::Error::other(
"the object store gave no object size",
))
})
}
pub(crate) async fn put(
&self,
key: &str,
body: reqwest::Body,
length: u64,
) -> Result<(), Error> {
let action = PutObject::new(&self.bucket, Some(&self.credentials), key);
let url = action.sign(self.lifetime);
let response = self
.client
.put(url)
.header(reqwest::header::CONTENT_LENGTH, length)
.body(body)
.send()
.await
.map_err(|_| {
Error::Storage(std::io::Error::other("the object store is unreachable"))
})?;
let status = response.status();
if !status.is_success() {
let detail = response.text().await.unwrap_or_default();
return Err(Error::Storage(std::io::Error::other(format!(
"the object store refused a write with {status}: {}",
detail.trim()
))));
}
Ok(())
}
pub(crate) async fn copy(&self, from: &str, to: &str) -> Result<(), Error> {
let source = format!("/{}/{from}", self.bucket.name());
let mut action = PutObject::new(&self.bucket, Some(&self.credentials), to);
action
.headers_mut()
.insert(COPY_SOURCE, std::borrow::Cow::Owned(source.clone()));
let response = self
.client
.put(action.sign(self.lifetime))
.header(COPY_SOURCE, source)
.header(reqwest::header::CONTENT_LENGTH, 0)
.send()
.await
.map_err(|_| unreachable_store())?;
self.expect_success(response, "copy").await?;
Ok(())
}
pub(crate) async fn put_if_absent(&self, key: &str, body: Vec<u8>) -> Result<bool, Error> {
let mut action = PutObject::new(&self.bucket, Some(&self.credentials), key);
action.headers_mut().insert("if-none-match", "*");
let url = action.sign(self.lifetime);
let length = body.len();
let response = self
.client
.put(url)
.header("if-none-match", "*")
.header(reqwest::header::CONTENT_LENGTH, length)
.body(body)
.send()
.await
.map_err(|_| unreachable_store())?;
if response.status() == reqwest::StatusCode::PRECONDITION_FAILED {
return Ok(false);
}
self.expect_success(response, "write").await?;
Ok(true)
}
pub(crate) async fn get_bytes(&self, key: &str) -> Result<Option<Vec<u8>>, Error> {
let action = GetObject::new(&self.bucket, Some(&self.credentials), key);
let response = read_retrying(self.client.get(action.sign(self.lifetime))).await?;
if response.status() == reqwest::StatusCode::NOT_FOUND {
return Ok(None);
}
let response = self.expect_success(response, "read").await?;
response
.bytes()
.await
.map(|bytes| Some(bytes.to_vec()))
.map_err(|_| unreachable_store())
}
pub(crate) async fn delete(&self, key: &str) -> Result<bool, Error> {
let existed = self.head(key).await.is_ok();
let action = DeleteObject::new(&self.bucket, Some(&self.credentials), key);
let response = self
.client
.delete(action.sign(self.lifetime))
.send()
.await
.map_err(|_| unreachable_store())?;
self.expect_success(response, "delete").await?;
Ok(existed)
}
pub(crate) async fn keys(&self, prefix: &str) -> Result<Vec<String>, Error> {
Ok(self
.entries(prefix)
.await?
.into_iter()
.map(|entry| entry.key)
.collect())
}
pub(crate) async fn entries(&self, prefix: &str) -> Result<Vec<Entry>, Error> {
let mut out = Vec::new();
let mut token: Option<String> = None;
loop {
let mut action = ListObjectsV2::new(&self.bucket, Some(&self.credentials));
action.with_prefix(prefix);
if let Some(token) = &token {
action.with_continuation_token(token);
}
let response = read_retrying(self.client.get(action.sign(self.lifetime))).await?;
let body = self
.expect_success(response, "list")
.await?
.text()
.await
.map_err(|_| unreachable_store())?;
let listing = ListObjectsV2::parse_response(&body).map_err(|error| {
Error::Storage(std::io::Error::other(format!(
"the object store sent a listing this server could not read: {error}"
)))
})?;
out.extend(listing.contents.into_iter().map(|object| Entry {
key: object.key,
last_modified: object.last_modified,
size: object.size,
}));
match listing.next_continuation_token {
Some(next) => token = Some(next),
None => break,
}
}
Ok(out)
}
async fn expect_success(
&self,
response: reqwest::Response,
what: &str,
) -> Result<reqwest::Response, Error> {
let status = response.status();
if status.is_success() {
return Ok(response);
}
let detail = response.text().await.unwrap_or_default();
Err(Error::Storage(std::io::Error::other(format!(
"the object store refused a {what} with {status}: {}",
detail.trim()
))))
}
}