Add S3 client auth check to cached object

Author: sd109

src/app.rs

@@ -224,7 +224,7 @@ async fn download_s3_object<'a>(
 /// * `chunk_cache`: ChunkCache object
 #[tracing::instrument(
     level = "DEBUG",
-    skip(client, request_data, resource_manager, chunk_cache)
+    skip(client, request_data, resource_manager, mem_permits, chunk_cache)
 )]
 async fn download_and_cache_s3_object<'a>(
     client: &s3_client::S3Client,
@@ -233,9 +233,32 @@ async fn download_and_cache_s3_object<'a>(
     mut mem_permits: Option<SemaphorePermit<'a>>,
     chunk_cache: &ChunkCache,
 ) -> Result<Bytes, ActiveStorageError> {
-    let key = format!("{},{:?}", client, request_data);
+    // We chose a cache key such that any changes to request data
+    // which may feasibly indicate a change to the upstream object
+    // lead to a new cache key.
+    let key = format!(
+        "{}-{}-{}-{}-{:?}-{:?}",
+        request_data.source.as_str(),
+        request_data.bucket,
+        request_data.object,
+        request_data.dtype,
+        request_data.byte_order,
+        request_data.compression,
+    );
 
     if let Some(metadata) = chunk_cache.get_metadata(&key).await {
+        // To avoid having to include the S3 client ID as part of the cache key
+        // (which means we'd have a separate cache for each authorised user and
+        // waste storage space) we instead make a lightweight check against the
+        // object store to ensure the user is authorised, even if the object data
+        // is already in the local cache.
+        let authorised = client
+            .is_authorised(&request_data.bucket, &request_data.object)
+            .await?;
+        if !authorised {
+            return Err(ActiveStorageError::Forbidden);
+        }
+
         // Update memory requested from resource manager to account for actual
         // size of data if we were previously unable to guess the size from request
         // data's size + offset parameters.
@@ -254,7 +277,10 @@ async fn download_and_cache_s3_object<'a>(
         // We only want to get chunks for which the metadata check succeeded too,
         // otherwise chunks which are missing metadata could bypass the resource
         // manager and exhaust system resources
-        let cache_value = chunk_cache.get(&key).await?;
+        let cache_value = chunk_cache
+            .get(&key)
+            .instrument(tracing::Span::current())
+            .await?;
         if let Some(bytes) = cache_value {
             return Ok(bytes);
         }

src/error.rs

@@ -3,6 +3,7 @@
 use aws_sdk_s3::error::ProvideErrorMetadata;
 use aws_sdk_s3::error::SdkError;
 use aws_sdk_s3::operation::get_object::GetObjectError;
+use aws_sdk_s3::operation::head_object::HeadObjectError;
 use aws_smithy_types::byte_stream::error::Error as ByteStreamError;
 use axum::{
     extract::rejection::JsonRejection,
@@ -74,6 +75,14 @@ pub enum ActiveStorageError {
     #[error("error retrieving object from S3 storage")]
     S3GetObject(#[from] SdkError<GetObjectError>),
 
+    /// Error while retrieving object head from S3
+    #[error("error retrieving object metadata from S3 storage")]
+    S3HeadObject(#[from] SdkError<HeadObjectError>),
+
+    /// HTTP 403 from object store
+    #[error("error receiving object metadata from S3 storage")]
+    Forbidden,
+
     /// Error acquiring a semaphore
     #[error("error acquiring resources")]
     SemaphoreAcquireError(#[from] AcquireError),
@@ -225,7 +234,11 @@ impl From<ActiveStorageError> for ErrorResponse {
             | ActiveStorageError::ShapeInvalid(_) => Self::bad_request(&error),
 
             // Not found
-            ActiveStorageError::UnsupportedOperation { operation: _ } => Self::not_found(&error),
+            ActiveStorageError::UnsupportedOperation { operation: _ }
+            // If we receive a forbidden from object store, return a
+            // not found to client to avoid leaking information about
+            // bucket contents.
+            | ActiveStorageError::Forbidden => Self::not_found(&error),
 
             // Internal server error
             ActiveStorageError::FromBytes { type_name: _ }
@@ -277,6 +290,31 @@ impl From<ActiveStorageError> for ErrorResponse {
                     _ => Self::internal_server_error(&error),
                 }
             }
+
+            ActiveStorageError::S3HeadObject(sdk_error) => {
+                // Tailor the response based on the specific SdkError variant.
+                match &sdk_error {
+                    // These are generic SdkError variants.
+                    // Internal server error
+                    SdkError::ConstructionFailure(_)
+                    | SdkError::DispatchFailure(_)
+                    | SdkError::ResponseError(_)
+                    | SdkError::TimeoutError(_) => Self::internal_server_error(&error),
+
+                    // This is a more specific ServiceError variant,
+                    // with HeadObjectError as the inner error.
+                    SdkError::ServiceError(head_obj_error) => {
+                        let head_obj_error = head_obj_error.err();
+                        match head_obj_error {
+                            HeadObjectError::NotFound(_) => Self::bad_request(&error),
+                            // Enum is marked as non-exhaustive
+                            _ => Self::internal_server_error(&error),
+                        }
+                    }
+                    // Enum is marked as non-exhaustive
+                    _ => Self::internal_server_error(&error),
+                }
+            }
         };
 
         // Log server errors.

src/s3_client.rs

@@ -7,8 +7,10 @@ use crate::error::ActiveStorageError;
 use crate::resource_manager::ResourceManager;
 
 use aws_credential_types::Credentials;
-use aws_sdk_s3::config::BehaviorVersion;
+use aws_sdk_s3::operation::head_object::HeadObjectError;
 use aws_sdk_s3::Client;
+use aws_sdk_s3::{config::BehaviorVersion, error::SdkError};
+use aws_smithy_runtime_api::http::Response;
 use aws_types::region::Region;
 use axum::body::Bytes;
 use hashbrown::HashMap;
@@ -142,6 +144,44 @@ impl S3Client {
         }
     }
 
+    /// Checks whether the client is authorised to download an
+    /// object from object storage.
+    ///
+    /// # Arguments
+    ///
+    /// * `bucket`: Name of the bucket
+    /// * `key`: Name of the object in the bucket
+    pub async fn is_authorised(
+        &self,
+        bucket: &str,
+        key: &str,
+    ) -> Result<bool, SdkError<HeadObjectError, Response>> {
+        let response = self
+            .client
+            .head_object()
+            .bucket(bucket)
+            .key(key)
+            .send()
+            .instrument(tracing::Span::current())
+            .await;
+
+        // Strategy here is to return true if client is authorised to download object,
+        // false if explicitly not authorised (HTTP 403) and pass any other errors or
+        // responses back to the caller.
+        match response {
+            Ok(_) => Ok(true),
+            Err(err) => match &err {
+                aws_smithy_runtime_api::client::result::SdkError::ServiceError(inner) => {
+                    match inner.raw().status().as_u16() {
+                        403 => Ok(false), // HTTP 403 == Forbidden
+                        _ => Err(err),
+                    }
+                }
+                _ => Err(err),
+            },
+        }
+    }
+
     /// Downloads an object from object storage and returns the data as Bytes
     ///
     /// # Arguments
@@ -152,7 +192,7 @@ impl S3Client {
     /// * `resource_manager`: ResourceManager object
     /// * `mem_permits`: Optional SemaphorePermit for any memory resources reserved
     pub async fn download_object<'a>(
-        self: &S3Client,
+        &self,
         bucket: &str,
         key: &str,
         range: Option<String>,