consolidate encryption_key into download opts

Signed-off-by: Alex Chi Z <chi@neon.tech>
This commit is contained in:
Alex Chi Z
2025-04-16 18:28:34 -04:00
parent 7d9f677a22
commit 7c2b2325f1
7 changed files with 22 additions and 84 deletions

View File

@@ -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<Download, DownloadError> {
unimplemented!()
}
#[allow(unused_variables)]
async fn upload_with_encryption(
&self,

View File

@@ -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<Vec<u8>>,
}
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<impl AsRef<[u8]>>) -> 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<Download, DownloadError>;
/// 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<Download, DownloadError>;
/// 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<Other: RemoteStorage> GenericRemoteStorage<Arc<Other>> {
}
}
pub async fn download_with_encryption(
&self,
from: &RemotePath,
opts: &DownloadOpts,
encryption_key: Option<&[u8]>,
cancel: &CancellationToken,
) -> Result<Download, DownloadError> {
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<Item = std::io::Result<Bytes>> + Send + Sync + 'static,

View File

@@ -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<Download, DownloadError> {
unimplemented!()
}
#[allow(unused_variables)]
async fn upload_with_encryption(
&self,

View File

@@ -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<Download, DownloadError> {
// 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<Download, DownloadError> {
self.download_with_encryption(from, opts, None, cancel)
.await
}
async fn delete_objects(
&self,
paths: &[RemotePath],

View File

@@ -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<Download, DownloadError> {
unimplemented!()
}
#[allow(unused_variables)]
async fn upload_with_encryption(
&self,

View File

@@ -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)