Add support for Cloudflare R2 Storage

crates/aws_s3/src/storage.rs

@@ -9,9 +9,9 @@ use async_trait::async_trait;
 use aws_config::retry::RetryConfig;
 use aws_sdk_s3::{
     operation::{
-        get_object_attributes::{
-            GetObjectAttributesError,
-            GetObjectAttributesOutput,
+        head_object::{
+            HeadObjectOutput,
+            HeadObjectError,
         },
         upload_part::builders::UploadPartFluentBuilder,
     },
@@ -26,7 +26,7 @@ use aws_sdk_s3::{
     Client,
 };
 use aws_utils::{
-    must_config_from_env,
+    must_s3_config_from_env,
     s3::S3Client,
 };
 use bytes::Bytes;
@@ -132,7 +132,7 @@ impl<RT: Runtime> S3Storage<RT> {
         key_prefix: String,
         runtime: RT,
     ) -> anyhow::Result<Self> {
-        let config = must_config_from_env()
+        let config = must_s3_config_from_env()
             .context("AWS env variables are required when using S3 storage")?
             .retry_config(RetryConfig::standard())
             .load()
@@ -434,26 +434,24 @@ impl<RT: Runtime> Storage for S3Storage<RT> {
             .split_once('/')
             .with_context(|| format!("Invalid fully qualified S3 key {:?}", key))?;
         let result: Result<
-            GetObjectAttributesOutput,
-            aws_sdk_s3::error::SdkError<GetObjectAttributesError>,
+            HeadObjectOutput,
+            aws_sdk_s3::error::SdkError<HeadObjectError>,
         > = self
             .client
-            .get_object_attributes()
+            .head_object()
             .bucket(bucket)
             .key(s3_key)
-            .object_attributes(aws_sdk_s3::types::ObjectAttributes::Checksum)
-            .object_attributes(aws_sdk_s3::types::ObjectAttributes::ObjectSize)
             .send()
             .await;
         match result {
-            Ok(object_attributes) => {
-                let size = object_attributes
-                    .object_size
+            Ok(head_attributes) => {
+                let size = head_attributes
+                    .content_length
                     .context("Object is missing size")? as u64;
                 Ok(Some(ObjectAttributes { size }))
             },
             Err(aws_sdk_s3::error::SdkError::ServiceError(err)) => match err.err() {
-                GetObjectAttributesError::NoSuchKey(_) => Ok(None),
+                HeadObjectError::NotFound(_) => Ok(None),
                 // Other service errors from S3
                 _ => Err(err.into_err().into()),
             },

crates/aws_utils/src/lib.rs

@@ -14,6 +14,9 @@ use aws_types::region::Region;
 
 pub mod s3;
 
+static S3_ENDPOINT_URL: LazyLock<Option<String>> =
+    LazyLock::new(|| env::var("S3_ENDPOINT_URL").ok());
+
 static AWS_ACCESS_KEY_ID: LazyLock<Option<String>> =
     LazyLock::new(|| env::var("AWS_ACCESS_KEY_ID").ok());
 
@@ -41,3 +44,11 @@ pub fn must_config_from_env() -> anyhow::Result<ConfigLoader> {
         .region(region)
         .credentials_provider(credentials))
 }
+
+pub fn must_s3_config_from_env() -> anyhow::Result<ConfigLoader> {
+    let mut config_loader = must_config_from_env()?;
+    if let Some(s3_endpoint_url) = S3_ENDPOINT_URL.clone() {
+        config_loader = config_loader.endpoint_url(s3_endpoint_url);
+    }
+    Ok(config_loader)
+}

crates/aws_utils/src/s3.rs

@@ -8,7 +8,7 @@ use aws_sdk_s3::{
 };
 use futures_async_stream::try_stream;
 
-use crate::must_config_from_env;
+use crate::must_s3_config_from_env;
 
 #[derive(Clone, Debug)]
 pub struct S3Client(pub Client);
@@ -21,7 +21,7 @@ impl S3Client {
             true => RetryConfig::standard(),
             false => RetryConfig::disabled(),
         };
-        let config = must_config_from_env()
+        let config = must_s3_config_from_env()
             .context("AWS env variables are required when using AWS Lambda")?
             .retry_config(retry_config)
             .load()