From 7c2b2325f1619f96c2473d5cdbdb732e7e8aca55 Mon Sep 17 00:00:00 2001 From: Alex Chi Z Date: Wed, 16 Apr 2025 18:28:34 -0400 Subject: [PATCH] consolidate encryption_key into download opts Signed-off-by: Alex Chi Z --- libs/remote_storage/src/azure_blob.rs | 11 ----- libs/remote_storage/src/lib.rs | 48 +++++-------------- libs/remote_storage/src/local_fs.rs | 11 ----- libs/remote_storage/src/s3_bucket.rs | 15 +----- libs/remote_storage/src/simulate_failures.rs | 11 ----- libs/remote_storage/tests/test_real_s3.rs | 7 ++- .../import_pgdata/importbucket_client.rs | 3 +- 7 files changed, 22 insertions(+), 84 deletions(-) diff --git a/libs/remote_storage/src/azure_blob.rs b/libs/remote_storage/src/azure_blob.rs index fbb4e1d52b..f8958e8fc5 100644 --- a/libs/remote_storage/src/azure_blob.rs +++ b/libs/remote_storage/src/azure_blob.rs @@ -550,17 +550,6 @@ impl RemoteStorage for AzureBlobStorage { self.download_for_builder(builder, timeout, cancel).await } - #[allow(unused_variables)] - async fn download_with_encryption( - &self, - from: &RemotePath, - opts: &DownloadOpts, - encryption_key: Option<&[u8]>, - cancel: &CancellationToken, - ) -> Result { - unimplemented!() - } - #[allow(unused_variables)] async fn upload_with_encryption( &self, diff --git a/libs/remote_storage/src/lib.rs b/libs/remote_storage/src/lib.rs index 2e18feaa19..2d6a11e81f 100644 --- a/libs/remote_storage/src/lib.rs +++ b/libs/remote_storage/src/lib.rs @@ -190,6 +190,8 @@ pub struct DownloadOpts { /// timeouts: for something like an index/manifest/heatmap, we should time out faster than /// for layer files pub kind: DownloadKind, + /// The encryption key to use for the download. + pub encryption_key: Option>, } pub enum DownloadKind { @@ -204,6 +206,7 @@ impl Default for DownloadOpts { byte_start: Bound::Unbounded, byte_end: Bound::Unbounded, kind: DownloadKind::Large, + encryption_key: None, } } } @@ -241,6 +244,15 @@ impl DownloadOpts { None => format!("bytes={start}-"), }) } + + pub fn with_encryption_key(mut self, encryption_key: Option>) -> Self { + self.encryption_key = encryption_key.map(|k| k.as_ref().to_vec()); + self + } + + pub fn encryption_key(&self) -> Option<&[u8]> { + self.encryption_key.as_deref() + } } /// Storage (potentially remote) API to manage its state. @@ -331,15 +343,6 @@ pub trait RemoteStorage: Send + Sync + 'static { cancel: &CancellationToken, ) -> Result; - /// Same as download, but with encryption if the backend supports it (e.g. SSE-C on AWS). - async fn download_with_encryption( - &self, - from: &RemotePath, - opts: &DownloadOpts, - encryption_key: Option<&[u8]>, - cancel: &CancellationToken, - ) -> Result; - /// Same as upload, but with remote encryption if the backend supports it (e.g. SSE-C on AWS). async fn upload_with_encryption( &self, @@ -638,33 +641,6 @@ impl GenericRemoteStorage> { } } - pub async fn download_with_encryption( - &self, - from: &RemotePath, - opts: &DownloadOpts, - encryption_key: Option<&[u8]>, - cancel: &CancellationToken, - ) -> Result { - match self { - Self::LocalFs(s) => { - s.download_with_encryption(from, opts, encryption_key, cancel) - .await - } - Self::AwsS3(s) => { - s.download_with_encryption(from, opts, encryption_key, cancel) - .await - } - Self::AzureBlob(s) => { - s.download_with_encryption(from, opts, encryption_key, cancel) - .await - } - Self::Unreliable(s) => { - s.download_with_encryption(from, opts, encryption_key, cancel) - .await - } - } - } - pub async fn upload_with_encryption( &self, from: impl Stream> + Send + Sync + 'static, diff --git a/libs/remote_storage/src/local_fs.rs b/libs/remote_storage/src/local_fs.rs index 97a1834d8b..b091fde41a 100644 --- a/libs/remote_storage/src/local_fs.rs +++ b/libs/remote_storage/src/local_fs.rs @@ -560,17 +560,6 @@ impl RemoteStorage for LocalFs { } } - #[allow(unused_variables)] - async fn download_with_encryption( - &self, - from: &RemotePath, - opts: &DownloadOpts, - encryption_key: Option<&[u8]>, - cancel: &CancellationToken, - ) -> Result { - unimplemented!() - } - #[allow(unused_variables)] async fn upload_with_encryption( &self, diff --git a/libs/remote_storage/src/s3_bucket.rs b/libs/remote_storage/src/s3_bucket.rs index ce2bdee024..4a85b46584 100644 --- a/libs/remote_storage/src/s3_bucket.rs +++ b/libs/remote_storage/src/s3_bucket.rs @@ -819,11 +819,10 @@ impl RemoteStorage for S3Bucket { Ok(()) } - async fn download_with_encryption( + async fn download( &self, from: &RemotePath, opts: &DownloadOpts, - encryption_key: Option<&[u8]>, cancel: &CancellationToken, ) -> Result { // if prefix is not none then download file `prefix/from` @@ -834,23 +833,13 @@ impl RemoteStorage for S3Bucket { key: self.relative_path_to_s3_object(from), etag: opts.etag.as_ref().map(|e| e.to_string()), range: opts.byte_range_header(), - sse_c_key: encryption_key.map(|k| k.to_vec()), + sse_c_key: opts.encryption_key.clone(), }, cancel, ) .await } - async fn download( - &self, - from: &RemotePath, - opts: &DownloadOpts, - cancel: &CancellationToken, - ) -> Result { - self.download_with_encryption(from, opts, None, cancel) - .await - } - async fn delete_objects( &self, paths: &[RemotePath], diff --git a/libs/remote_storage/src/simulate_failures.rs b/libs/remote_storage/src/simulate_failures.rs index 2ac394201c..6431e9aa64 100644 --- a/libs/remote_storage/src/simulate_failures.rs +++ b/libs/remote_storage/src/simulate_failures.rs @@ -178,17 +178,6 @@ impl RemoteStorage for UnreliableWrapper { self.inner.download(from, opts, cancel).await } - #[allow(unused_variables)] - async fn download_with_encryption( - &self, - from: &RemotePath, - opts: &DownloadOpts, - encryption_key: Option<&[u8]>, - cancel: &CancellationToken, - ) -> Result { - unimplemented!() - } - #[allow(unused_variables)] async fn upload_with_encryption( &self, diff --git a/libs/remote_storage/tests/test_real_s3.rs b/libs/remote_storage/tests/test_real_s3.rs index b85fb82014..939d6384b6 100644 --- a/libs/remote_storage/tests/test_real_s3.rs +++ b/libs/remote_storage/tests/test_real_s3.rs @@ -593,7 +593,11 @@ async fn encryption_works(ctx: &mut MaybeEnabledStorage) { { let download = ctx .client - .download_with_encryption(&path, &DownloadOpts::default(), Some(&key), &cancel) + .download( + &path, + &DownloadOpts::default().with_encryption_key(Some(&key)), + &cancel, + ) .await .expect("should succeed"); let vec = download_to_vec(download).await.expect("should succeed"); @@ -601,6 +605,7 @@ async fn encryption_works(ctx: &mut MaybeEnabledStorage) { } { + // Download without encryption key should fail let download = ctx .client .download(&path, &DownloadOpts::default(), &cancel) diff --git a/pageserver/src/tenant/timeline/import_pgdata/importbucket_client.rs b/pageserver/src/tenant/timeline/import_pgdata/importbucket_client.rs index a17a10d56b..b28056e1e1 100644 --- a/pageserver/src/tenant/timeline/import_pgdata/importbucket_client.rs +++ b/pageserver/src/tenant/timeline/import_pgdata/importbucket_client.rs @@ -244,7 +244,8 @@ impl RemoteStorageWrapper { kind: DownloadKind::Large, etag: None, byte_start: Bound::Included(start_inclusive), - byte_end: Bound::Excluded(end_exclusive) + byte_end: Bound::Excluded(end_exclusive), + encryption_key: None, }, &self.cancel) .await?;