From a1d25078ccd6c83346c397ddd6caaadb16d0ef1f Mon Sep 17 00:00:00 2001
From: Alex Auvolat <alex@adnab.me>
Date: Mon, 26 Feb 2024 15:27:25 +0100
Subject: [sse-c] add test for block encryption and decryption

---
 src/api/s3/encryption.rs | 59 +++++++++++++++++++++++++++++++++++++++++++-----
 src/api/s3/list.rs       |  8 ++++---
 2 files changed, 58 insertions(+), 9 deletions(-)

(limited to 'src')

diff --git a/src/api/s3/encryption.rs b/src/api/s3/encryption.rs
index f13333b2..17974df1 100644
--- a/src/api/s3/encryption.rs
+++ b/src/api/s3/encryption.rs
@@ -48,7 +48,6 @@ const X_AMZ_COPY_SOURCE_SERVER_SIDE_ENCRYPTION_CUSTOMER_KEY_MD5: HeaderName =
 
 const CUSTOMER_ALGORITHM_AES256: HeaderValue = HeaderValue::from_static("AES256");
 
-type StreamNonce = aes_gcm::aead::stream::Nonce<Aes256Gcm, StreamLE31<Aes256Gcm>>;
 type StreamNonceSize = aes_gcm::aead::stream::NonceSize<Aes256Gcm, StreamLE31<Aes256Gcm>>;
 
 const STREAM_ENC_PLAIN_CHUNK_SIZE: usize = 0x1000; // 4096 bytes
@@ -224,20 +223,24 @@ impl EncryptionParams {
 			.block_manager
 			.rpc_get_block_streaming(hash, order)
 			.await?;
+		Ok(self.decrypt_block_stream(raw_block))
+	}
+
+	pub fn decrypt_block_stream(&self, stream: ByteStream) -> ByteStream {
 		match self {
-			Self::Plaintext => Ok(raw_block),
+			Self::Plaintext => stream,
 			Self::SseC {
 				client_key,
 				compression_level,
 			} => {
-				let plaintext = DecryptStream::new(raw_block, *client_key);
+				let plaintext = DecryptStream::new(stream, *client_key);
 				if compression_level.is_some() {
 					let reader = stream_asyncread(Box::pin(plaintext));
 					let reader = BufReader::new(reader);
 					let reader = async_compression::tokio::bufread::ZstdDecoder::new(reader);
-					Ok(Box::pin(tokio_util::io::ReaderStream::new(reader)))
+					Box::pin(tokio_util::io::ReaderStream::new(reader))
 				} else {
-					Ok(Box::pin(plaintext))
+					Box::pin(plaintext)
 				}
 			}
 		}
@@ -245,7 +248,7 @@ impl EncryptionParams {
 
 	/// Encrypt a data block if encryption is set, for use before
 	/// putting the data blocks into storage
-	pub fn encrpyt_block(&self, block: Bytes) -> Result<Bytes, Error> {
+	pub fn encrypt_block(&self, block: Bytes) -> Result<Bytes, Error> {
 		match self {
 			Self::Plaintext => Ok(block),
 			Self::SseC {
@@ -417,3 +420,47 @@ impl Stream for DecryptStream {
 		}
 	}
 }
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+
+	use futures::prelude::*;
+
+	use garage_net::stream::read_stream_to_end;
+
+	fn stream() -> ByteStream {
+		Box::pin(
+			futures::stream::iter(16usize..1024)
+				.map(|i| Ok(Bytes::from(vec![(i % 256) as u8; (i * 37) % 1024]))),
+		)
+	}
+
+	async fn test_block_enc(compression_level: Option<i32>) {
+		let enc = EncryptionParams::SseC {
+			client_key: Aes256Gcm::generate_key(&mut OsRng),
+			compression_level,
+		};
+
+		let block_plain = read_stream_to_end(stream()).await.unwrap().into_bytes();
+
+		let block_enc = enc.encrypt_block(block_plain.clone()).unwrap();
+
+		let block_dec =
+			enc.decrypt_block_stream(Box::pin(futures::stream::once(async { Ok(block_enc) })));
+		let block_dec = read_stream_to_end(block_dec).await.unwrap().into_bytes();
+
+		assert_eq!(block_plain, block_dec);
+		assert!(block_dec.len() > 128000);
+	}
+
+	#[tokio::test]
+	async fn test_encrypt_block() {
+		test_block_enc(None).await
+	}
+
+	#[tokio::test]
+	async fn test_encrypt_block_compressed() {
+		test_block_enc(Some(1)).await
+	}
+}
diff --git a/src/api/s3/list.rs b/src/api/s3/list.rs
index b832a4f4..1c875d06 100644
--- a/src/api/s3/list.rs
+++ b/src/api/s3/list.rs
@@ -944,9 +944,11 @@ mod tests {
 			timestamp: TS,
 			state: ObjectVersionState::Uploading {
 				multipart: true,
-				headers: ObjectVersionHeaders {
-					content_type: "text/plain".to_string(),
-					other: BTreeMap::<String, String>::new(),
+				encryption: ObjectVersionEncryption::Plaintext {
+					headers: ObjectVersionHeaders {
+						content_type: "text/plain".to_string(),
+						other: BTreeMap::<String, String>::new(),
+					},
 				},
 			},
 		}
-- 
cgit v1.2.3