objectstore_service/backend/

s3_compatible.rs

//! S3-compatible backend with generic protocol support.

use std::time::SystemTime;
use std::{fmt, io};

use futures_util::{StreamExt, TryStreamExt};
use objectstore_types::metadata::{HEADER_SIZE, Metadata};
use objectstore_types::range::{ByteRange, ContentRange};
use reqwest::header::{HeaderMap, HeaderName};
use reqwest::{Body, IntoUrl, Method, RequestBuilder, Response, StatusCode};

use super::extensions::{ResponseExt, SendTraced};
use crate::backend::common::{
    self, Backend, DeleteResponse, GetResponse, MetadataResponse, PutResponse,
};
use crate::error::{Error, Result};
use crate::id::ObjectId;
use crate::stream::ClientStream;

/// Configuration for [`S3CompatibleBackend`].
///
/// Supports [Amazon S3] and other S3-compatible services. Authentication is handled via
/// environment variables (`AWS_ACCESS_KEY_ID` and `AWS_SECRET_ACCESS_KEY`) or IAM roles.
///
/// [Amazon S3]: https://aws.amazon.com/s3/
///
/// # Example
///
/// ```yaml
/// storage:
///   type: s3compatible
///   endpoint: https://s3.amazonaws.com
///   bucket: my-bucket
/// ```
#[derive(Debug, Clone, serde::Deserialize, serde::Serialize)]
pub struct S3CompatibleConfig {
    /// S3 endpoint URL.
    ///
    /// Examples: `https://s3.amazonaws.com`, `http://localhost:9000` (for MinIO)
    ///
    /// # Environment Variables
    ///
    /// - `OS__STORAGE__TYPE=s3compatible`
    /// - `OS__STORAGE__ENDPOINT=https://s3.amazonaws.com`
    pub endpoint: String,

    /// S3 bucket name.
    ///
    /// The bucket must exist before starting the server.
    ///
    /// # Environment Variables
    ///
    /// - `OS__STORAGE__BUCKET=my-bucket`
    pub bucket: String,
}

/// Prefix used for custom metadata in headers for the GCS backend.
///
/// See: <https://cloud.google.com/storage/docs/xml-api/reference-headers#xgoogmeta>
const GCS_CUSTOM_PREFIX: &str = "x-goog-meta-";
/// Header used to store the expiration time for GCS using the `daysSinceCustomTime` lifecycle
/// condition.
///
/// See: <https://cloud.google.com/storage/docs/xml-api/reference-headers#xgoogcustomtime>
const GCS_CUSTOM_TIME: &str = "x-goog-custom-time";

/// An authentication token that can be passed as a bearer credential.
pub trait Token: Send + Sync {
    /// Returns the token string.
    fn as_str(&self) -> &str;
}

/// Provides authentication tokens for S3-compatible requests.
pub trait TokenProvider: Send + Sync + 'static {
    /// Returns a fresh token, fetching or refreshing it as needed.
    fn get_token(&self) -> impl Future<Output = anyhow::Result<impl Token>> + Send;
}

/// Placeholder [`TokenProvider`] for unauthenticated backends.
#[derive(Debug)]
pub struct NoToken;

impl TokenProvider for NoToken {
    #[allow(refining_impl_trait)]
    async fn get_token(&self) -> anyhow::Result<NoToken> {
        unimplemented!()
    }
}
impl Token for NoToken {
    fn as_str(&self) -> &str {
        unimplemented!()
    }
}

/// S3-compatible storage backend with pluggable authentication.
pub struct S3CompatibleBackend<T> {
    client: reqwest::Client,

    endpoint: String,
    bucket: String,

    token_provider: Option<T>,
}

impl<T> S3CompatibleBackend<T> {
    /// Creates a new S3-compatible backend bound to the given bucket.
    pub fn new(endpoint: &str, bucket: &str, token_provider: T) -> Self {
        Self {
            client: common::reqwest_client(),
            endpoint: endpoint.into(),
            bucket: bucket.into(),
            token_provider: Some(token_provider),
        }
    }

    /// Formats the S3 object URL for the given key.
    fn object_url(&self, id: &ObjectId) -> String {
        format!("{}/{}/{}", self.endpoint, self.bucket, id.as_storage_path())
    }
}

/// Wraps [`Metadata::to_headers`] with GCS-specific concerns (tombstone + custom-time).
fn metadata_to_gcs_headers(
    metadata: &Metadata,
    prefix: &str,
) -> Result<HeaderMap, objectstore_types::metadata::Error> {
    let mut headers = metadata.to_headers(prefix)?;

    // The size is derived from the native `Content-Length` on every read, so it must not be
    // persisted: metadata updates rewrite *all* stored metadata, and a stored `x-sn-size` key
    // is rejected by the GCS JSON backend when it deserializes the object.
    let size = HeaderName::try_from(format!("{prefix}{HEADER_SIZE}"))?;
    headers.remove(&size);

    // GCS custom-time for lifecycle expiration
    if let Some(expires_at) = metadata.time_expires {
        let expires_at = humantime::format_rfc3339_seconds(expires_at);
        headers.append(GCS_CUSTOM_TIME, expires_at.to_string().parse()?);
    }
    Ok(headers)
}

impl<T> S3CompatibleBackend<T>
where
    T: TokenProvider,
{
    /// Creates a request builder with the appropriate authentication.
    async fn request(&self, method: Method, url: impl IntoUrl) -> Result<RequestBuilder> {
        let mut builder = self.client.request(method, url);
        if let Some(provider) = &self.token_provider {
            builder = builder.bearer_auth(
                provider
                    .get_token()
                    .await
                    .map_err(|err| Error::Generic {
                        context: "S3: failed to get authentication token".to_owned(),
                        cause: Some(err.into()),
                    })?
                    .as_str(),
            );
        }
        Ok(builder)
    }

    /// Fetches object metadata using the given HTTP method (GET or HEAD),
    /// bumps TTI if needed, and returns the parsed metadata along with the
    /// response (so `get_object` can read the body from a GET).
    async fn request_object(
        &self,
        method: Method,
        id: &ObjectId,
        range: Option<ByteRange>,
    ) -> Result<Option<(Metadata, Option<ContentRange>, Response)>> {
        let object_url = self.object_url(id);

        let mut builder = self.request(method, &object_url).await?;
        if let Some(r) = range {
            builder = builder.header(reqwest::header::RANGE, r.to_header_value());
        }
        let response = builder
            .send_traced()
            .await
            .map_err(|cause| Error::Reqwest {
                context: "S3: failed to send request".to_string(),
                cause,
            })?;

        if response.status() == StatusCode::NOT_FOUND {
            objectstore_log::debug!("Object not found");
            response.drain_body().await;
            return Ok(None);
        }

        if response.status() == StatusCode::RANGE_NOT_SATISFIABLE {
            let raw = response
                .headers()
                .get(reqwest::header::CONTENT_RANGE)
                .and_then(|v| v.to_str().ok());
            let total = raw.and_then(ContentRange::parse_unsatisfiable_total);
            let err = match total {
                Some(total) => Error::RangeNotSatisfiable { total },
                None => Error::generic(format!(
                    "S3: 416 response with invalid Content-Range: {raw:?}"
                )),
            };
            response.drain_body().await;
            return Err(err);
        }

        let response = response.check_error("S3: failed to get object").await?;

        let headers = response.headers();
        let mut metadata = Metadata::from_headers(headers, GCS_CUSTOM_PREFIX)?;

        let content_range = if response.status() == StatusCode::PARTIAL_CONTENT {
            let range = headers
                .get(reqwest::header::CONTENT_RANGE)
                .and_then(|v| v.to_str().ok())
                .and_then(|s| s.parse::<ContentRange>().ok())
                .ok_or_else(|| Error::Generic {
                    context: "S3: 206 response missing valid Content-Range header".to_owned(),
                    cause: None,
                })?;
            metadata.size = Some(range.total as usize);
            Some(range)
        } else {
            // NB: Read the header rather than `Response::content_length`, which reports the
            // length of the decoded body and is therefore always zero for a HEAD response.
            let size = headers
                .get(reqwest::header::CONTENT_LENGTH)
                .and_then(|value| value.to_str().ok())
                .map(|value| value.parse::<usize>())
                .transpose()
                .map_err(|cause| Error::Generic {
                    context: "S3: failed to parse Content-Length from object response".to_string(),
                    cause: Some(Box::new(cause)),
                })?;

            if let Some(size) = size {
                metadata.size = Some(size);
            } else {
                objectstore_log::warn!("S3: 200 response missing Content-Length header");
            }
            None
        };

        // TODO: Inject the access time from the request.
        let access_time = SystemTime::now();

        // Filter already expired objects but leave them to garbage collection
        if metadata.expiration_policy.is_timeout()
            && metadata.time_expires.is_some_and(|ts| ts < access_time)
        {
            objectstore_log::debug!("Object found but past expiry");
            response.drain_body().await;
            return Ok(None);
        }

        // TODO: extract into dedicated call from service
        // TODO: Schedule into background persistently so this doesn't get lost on restarts
        if let Some(new_expire_at) = metadata.check_tti_bump(access_time) {
            let mut bumped = metadata.clone();
            bumped.time_expires = Some(new_expire_at);
            self.update_metadata(id, &bumped).await?;
        }

        Ok(Some((metadata, content_range, response)))
    }

    /// Issues a request to update the metadata for the given object.
    async fn update_metadata(&self, id: &ObjectId, metadata: &Metadata) -> Result<()> {
        // NB: Meta updates require copy + REPLACE along with *all* metadata. See
        // https://cloud.google.com/storage/docs/xml-api/put-object-copy
        self.request(Method::PUT, self.object_url(id))
            .await?
            .header(
                "x-goog-copy-source",
                format!("/{}/{}", self.bucket, id.as_storage_path()),
            )
            .header("x-goog-metadata-directive", "REPLACE")
            .headers(metadata_to_gcs_headers(metadata, GCS_CUSTOM_PREFIX)?)
            .send_traced()
            .await
            .check_error("S3: update expiration time")
            .await?
            .drain_body()
            .await;

        Ok(())
    }
}

impl<T> fmt::Debug for S3CompatibleBackend<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("S3Compatible")
            .field("client", &self.client)
            .field("endpoint", &self.endpoint)
            .field("bucket", &self.bucket)
            .finish_non_exhaustive()
    }
}

impl S3CompatibleBackend<NoToken> {
    /// Creates a new S3-compatible backend that sends unauthenticated requests.
    pub fn without_token(config: S3CompatibleConfig) -> Self {
        Self {
            client: common::reqwest_client(),
            endpoint: config.endpoint,
            bucket: config.bucket,
            token_provider: None,
        }
    }
}

#[async_trait::async_trait]
impl<T: TokenProvider> Backend for S3CompatibleBackend<T> {
    fn name(&self) -> &'static str {
        "s3-compatible"
    }

    #[tracing::instrument(level = "debug", fields(?id), skip_all)]
    async fn put_object(
        &self,
        id: &ObjectId,
        metadata: &Metadata,
        stream: ClientStream,
    ) -> Result<PutResponse> {
        objectstore_log::debug!("Writing to s3_compatible backend");
        self.request(Method::PUT, self.object_url(id))
            .await?
            .headers(metadata_to_gcs_headers(metadata, GCS_CUSTOM_PREFIX)?)
            .body(Body::wrap_stream(stream))
            .send_traced()
            .await
            .check_error("S3: failed to put object")
            .await?
            .drain_body()
            .await;

        Ok(())
    }

    #[tracing::instrument(level = "debug", skip(self))]
    async fn get_object(&self, id: &ObjectId, range: Option<ByteRange>) -> Result<GetResponse> {
        objectstore_log::debug!("Reading from s3_compatible backend");

        let Some((metadata, content_range, response)) =
            self.request_object(Method::GET, id, range).await?
        else {
            return Ok(None);
        };

        let stream = response.bytes_stream().map_err(io::Error::other);
        Ok(Some((metadata, content_range, stream.boxed())))
    }

    #[tracing::instrument(level = "debug", skip(self))]
    async fn get_metadata(&self, id: &ObjectId) -> Result<MetadataResponse> {
        objectstore_log::debug!("Reading metadata from s3_compatible backend");
        let response = self.request_object(Method::HEAD, id, None).await?;
        Ok(response.map(|(metadata, _, _)| metadata))
    }

    #[tracing::instrument(level = "debug", skip(self))]
    async fn delete_object(&self, id: &ObjectId) -> Result<DeleteResponse> {
        objectstore_log::debug!("Deleting from s3_compatible backend");
        let response = self
            .request(Method::DELETE, self.object_url(id))
            .await?
            .send_traced()
            .await
            .map_err(|cause| Error::Reqwest {
                context: "S3: failed to send delete request".to_string(),
                cause,
            })?;

        // Do not error for objects that do not exist.
        if response.status() == StatusCode::NOT_FOUND {
            response.drain_body().await;
            return Ok(());
        }

        response
            .check_error("S3: failed to delete object")
            .await?
            .drain_body()
            .await;

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use std::collections::BTreeMap;
    use std::time::Duration;

    use anyhow::Result;
    use objectstore_types::metadata::ExpirationPolicy;
    use objectstore_types::scope::{Scope, Scopes};

    use super::*;
    use crate::backend::common::Backend;
    use crate::id::ObjectContext;
    use crate::stream;

    // NB: To run these tests, you need to have a MinIO server running. This is done
    // automatically in CI.
    //
    // Refer to the readme for how to set up MinIO via devservices.

    fn create_test_backend() -> S3CompatibleBackend<NoToken> {
        S3CompatibleBackend::without_token(S3CompatibleConfig {
            endpoint: "http://localhost:8089".into(),
            bucket: "test-bucket".into(),
        })
    }

    fn make_id() -> ObjectId {
        ObjectId::random(ObjectContext {
            usecase: "testing".into(),
            scopes: Scopes::from_iter([Scope::create("testing", "value").unwrap()]),
        })
    }

    #[test]
    fn metadata_to_gcs_headers_omits_size() {
        let metadata = Metadata {
            size: Some(4096),
            ..Default::default()
        };

        let headers = metadata_to_gcs_headers(&metadata, GCS_CUSTOM_PREFIX).unwrap();

        // Persisting the size would store a key that the GCS JSON backend rejects on read.
        assert!(headers.get("x-goog-meta-x-sn-size").is_none());
    }

    #[test]
    fn metadata_to_gcs_headers_uses_time_expires() {
        let expires = SystemTime::now() + Duration::from_hours(1);
        let metadata = Metadata {
            expiration_policy: ExpirationPolicy::TimeToLive(Duration::from_hours(1)),
            time_expires: Some(expires),
            ..Default::default()
        };

        let headers = metadata_to_gcs_headers(&metadata, GCS_CUSTOM_PREFIX).unwrap();

        // The lifecycle custom-time is the server-resolved expiry (second precision).
        let custom_time = headers.get(GCS_CUSTOM_TIME).unwrap().to_str().unwrap();
        let expected = humantime::format_rfc3339_seconds(expires).to_string();
        assert_eq!(custom_time, expected);
    }

    #[test]
    fn metadata_to_gcs_headers_escapes_unicode() {
        let metadata = Metadata {
            filename: Some("réport-📄.pdf".into()),
            custom: BTreeMap::from_iter([("release".into(), "vérsion-1.0-🚀".into())]),
            ..Default::default()
        };

        let headers = metadata_to_gcs_headers(&metadata, GCS_CUSTOM_PREFIX).unwrap();
        assert_eq!(
            headers.get("x-goog-meta-x-sn-filename").unwrap(),
            "r%C3%A9port-%F0%9F%93%84.pdf",
        );
        assert_eq!(
            headers.get("x-goog-meta-x-snme-release").unwrap(),
            "v%C3%A9rsion-1.0-%F0%9F%9A%80",
        );

        // The prefixed headers this backend writes are the ones it reads back.
        let roundtripped = Metadata::from_headers(&headers, GCS_CUSTOM_PREFIX).unwrap();
        assert_eq!(roundtripped.filename, metadata.filename);
        assert_eq!(roundtripped.custom, metadata.custom);
    }

    #[tokio::test]
    async fn test_get_metadata_nonexistent() -> Result<()> {
        let backend = create_test_backend();
        let id = make_id();
        let result = backend.get_metadata(&id).await?;
        assert!(result.is_none());
        Ok(())
    }

    #[tokio::test]
    #[ignore = "MinIO does not support streaming bodies (requires Content-Length)"]
    async fn test_get_metadata_reports_size() -> Result<()> {
        let backend = create_test_backend();
        let id = make_id();
        let payload = "hello, world";

        backend
            .put_object(&id, &Metadata::default(), stream::single(payload))
            .await?;

        // The size must come from the `Content-Length` header, not from the (empty) body of
        // the HEAD response.
        let metadata = backend.get_metadata(&id).await?.expect("object exists");
        assert_eq!(metadata.size, Some(payload.len()));

        Ok(())
    }

    #[tokio::test]
    #[ignore = "MinIO does not support streaming bodies (requires Content-Length)"]
    async fn test_ttl_immediate() -> Result<()> {
        let backend = create_test_backend();
        let id = make_id();
        let metadata = Metadata {
            expiration_policy: ExpirationPolicy::TimeToLive(Duration::from_secs(0)),
            time_expires: Some(SystemTime::now()),
            ..Default::default()
        };

        backend
            .put_object(&id, &metadata, stream::single("hello, world"))
            .await?;

        let get_result = backend.get_object(&id, None).await?;
        assert!(get_result.is_none());

        let head_result = backend.get_metadata(&id).await?;
        assert!(head_result.is_none());

        Ok(())
    }

    #[tokio::test]
    #[ignore = "MinIO does not support streaming bodies (requires Content-Length)"]
    async fn test_tti_immediate() -> Result<()> {
        let backend = create_test_backend();
        let id = make_id();
        let metadata = Metadata {
            expiration_policy: ExpirationPolicy::TimeToIdle(Duration::from_secs(0)),
            time_expires: Some(SystemTime::now()),
            ..Default::default()
        };

        backend
            .put_object(&id, &metadata, stream::single("hello, world"))
            .await?;

        let get_result = backend.get_object(&id, None).await?;
        assert!(get_result.is_none());

        let head_result = backend.get_metadata(&id).await?;
        assert!(head_result.is_none());

        Ok(())
    }
}