[sse-c] fix encryption.rs and add block encryption

3 files changed, 91 insertions, 36 deletions

diff --git a/src/api/s3/encryption.rs b/src/api/s3/encryption.rs
index 1515d903..29b26a37 100644
--- a/src/api/s3/encryption.rs
+++ b/src/api/s3/encryption.rs
@@ -1,8 +1,12 @@
 use std::borrow::Cow;
+use std::convert::TryInto;
 use std::pin::Pin;
 
 use aes_gcm::{
-	aead::{stream, Aead, AeadCore, KeyInit, OsRng},
+	aead::stream::{DecryptorLE31, EncryptorLE31, StreamLE31},
+	aead::{Aead, AeadCore, KeyInit, OsRng},
+	aes::cipher::crypto_common::rand_core::RngCore,
+	aes::cipher::typenum::Unsigned,
 	Aes256Gcm, Key, Nonce,
 };
 use base64::prelude::*;
@@ -19,11 +23,13 @@ use garage_net::bytes_buf::BytesBuf;
 use garage_net::stream::{stream_asyncread, ByteStream};
 use garage_rpc::rpc_helper::OrderTag;
 use garage_util::data::Hash;
+use garage_util::error::Error as GarageError;
 use garage_util::migrate::Migrate;
 
 use garage_model::garage::Garage;
 use garage_model::s3::object_table::{ObjectVersionEncryption, ObjectVersionHeaders};
 
+use crate::common_error::*;
 use crate::s3::error::Error;
 
 const X_AMZ_SERVER_SIDE_ENCRYPTION_CUSTOMER_ALGORITHM: HeaderName =
@@ -42,7 +48,12 @@ const X_AMZ_COPY_SOURCE_SERVER_SIDE_ENCRYPTION_CUSTOMER_KEY_MD5: HeaderName =
 
 const CUSTOMER_ALGORITHM_AES256: HeaderValue = HeaderValue::from_static("AES256");
 
-const STREAM_ENC_CHUNK_SIZE: usize = 0x1000; // 4096 bytes
+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
+const STREAM_ENC_CYPER_CHUNK_SIZE: usize =
+	STREAM_ENC_CYPER_CHUNK_SIZE + <Aes256Gcm as AeadCore>::TagSize::to_usize();
 
 #[derive(Clone, Copy)]
 pub enum EncryptionParams {
@@ -66,7 +77,7 @@ impl EncryptionParams {
 		)?;
 		match key {
 			Some(client_key) => Ok(EncryptionParams::SseC {
-				client_key: parse_and_check_key(req)?,
+				client_key,
 				compression_level: garage.config.compression_level,
 			}),
 			None => Ok(EncryptionParams::Plaintext),
@@ -123,11 +134,12 @@ impl EncryptionParams {
 					},
 				};
 				let plaintext = enc.decrypt_blob(&headers)?;
-				let headers = ObjectVersionHeaders::decode(&plaintext)?;
-				Ok((enc, headers.into()))
+				let headers = ObjectVersionHeaders::decode(&plaintext)
+					.ok_or_internal_error("Could not decode encrypted headers")?;
+				Ok((enc, Cow::Borrowed(&headers)))
 			}
 			(None, ObjectVersionEncryption::Plaintext { headers }) => {
-				Ok((Self::Plaintext, headers.into()))
+				Ok((Self::Plaintext, Cow::Borrowed(headers)))
 			}
 			(_, ObjectVersionEncryption::SseC { .. }) => {
 				Err(Error::bad_request("Object is encrypted"))
@@ -148,10 +160,10 @@ impl EncryptionParams {
 				client_key,
 				compression_level,
 			} => {
-				let plaintext = h.encode()?;
-				let ciphertext = self.encrypt_blob(plaintext)?;
+				let plaintext = h.encode().map_err(GarageError::from)?;
+				let ciphertext = self.encrypt_blob(&plaintext)?;
 				Ok(ObjectVersionEncryption::SseC {
-					headers: ciphertext,
+					headers: ciphertext.into_owned(),
 					compressed: compression_level.is_some(),
 				})
 			}
@@ -168,7 +180,7 @@ impl EncryptionParams {
 				let cipher = Aes256Gcm::new(&client_key);
 				let nonce = Aes256Gcm::generate_nonce(&mut OsRng);
 				let ciphertext = cipher
-					.encrypt(&nonce, &blob)
+					.encrypt(&nonce, blob)
 					.ok_or_internal_error("Encryption failed")?;
 				Ok([nonce.to_vec(), ciphertext].concat().into())
 			}
@@ -180,14 +192,13 @@ impl EncryptionParams {
 		match self {
 			Self::SseC { client_key, .. } => {
 				let cipher = Aes256Gcm::new(&client_key);
-				let nonce_size = Aes256Gcm::NonceSize::to_usize();
-				let nonce: Nonce<Aes256Gcm> = blob
-					.get(..nonce_size)
-					.ok_or_internal_error("invalid encrypted data")?
-					.try_into()
-					.unwrap();
+				let nonce_size = <Aes256Gcm as AeadCore>::NonceSize::to_usize();
+				let nonce = Nonce::from_slice(
+					blob.get(..nonce_size)
+						.ok_or_internal_error("invalid encrypted data")?,
+				);
 				let plaintext = cipher
-					.decrypt(&nonce, &blob[nonce_size..])
+					.decrypt(nonce, &blob[nonce_size..])
 					.ok_or_bad_request(
 						"Invalid encryption key, could not decrypt object metadata.",
 					)?;
@@ -229,6 +240,47 @@ 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> {
+		match self {
+			Self::Plaintext => Ok(block),
+			Self::SseC {
+				client_key,
+				compression_level,
+			} => {
+				let block = if let Some(level) = compression_level {
+					Cow::Owned(
+						garage_block::zstd_encode(block.as_ref(), *level)
+							.ok_or_internal_error("failed to compress data block")?,
+					)
+				} else {
+					Cow::Borrowed(block.as_ref())
+				};
+
+				let mut ret = Vec::with_capacity(block.len() + 32 + block.len() / 64);
+
+				let mut nonce: Nonce<StreamNonceSize> = Default::default();
+				OsRng.fill_bytes(&mut nonce);
+				ret.extend_from_slice(nonce.as_slice());
+
+				let mut cipher = EncryptorLE31::<Aes256Gcm>::new(&client_key, &nonce);
+				for chunk in block.chunks(STREAM_ENC_PLAIN_CHUNK_SIZE) {
+					// TODO: use encrypt_last for last chunk
+					let chunk_enc = cipher
+						.encrypt_next(chunk)
+						.ok_or_internal_error("failed to encrypt chunk")?;
+					if chunk.len() == STREAM_ENC_PLAIN_CHUNK_SIZE {
+						assert_eq!(chunk_enc.len(), STREAM_ENC_CYPER_CHUNK_SIZE);
+					}
+					ret.extend_from_slice(&chunk_enc);
+				}
+
+				Ok(ret.into())
+			}
+		}
+	}
 }
 
 fn parse_request_headers(
@@ -238,7 +290,7 @@ fn parse_request_headers(
 	md5_header: &HeaderName,
 ) -> Result<Option<Key<Aes256Gcm>>, Error> {
 	match req.headers().get(alg_header) {
-		Some(CUSTOMER_ALGORITHM_AES256) => {
+		Some(alg) if *alg == CUSTOMER_ALGORITHM_AES256 => {
 			use md5::{Digest, Md5};
 
 			let key_b64 = req
@@ -249,6 +301,7 @@ fn parse_request_headers(
 				.decode(&key_b64)
 				.ok_or_bad_request(format!("Invalid {} header", key_header))?
 				.try_into()
+				.ok()
 				.ok_or_bad_request(format!("Invalid {} header", key_header))?;
 
 			let md5_b64 = req
@@ -261,7 +314,7 @@ fn parse_request_headers(
 
 			let mut hasher = Md5::new();
 			hasher.update(&key_bytes[..]);
-			if hasher.finalize() != md5_bytes {
+			if hasher.finalize().as_slice() != md5_bytes.as_slice() {
 				return Err(Error::bad_request(
 					"Encryption key MD5 checksum does not match",
 				));
@@ -269,7 +322,9 @@ fn parse_request_headers(
 
 			Ok(Some(key_bytes.into()))
 		}
-		Some(alg) => Err(Error::InvalidEncryptionAlgorithm(alg.to_string())),
+		Some(alg) => Err(Error::InvalidEncryptionAlgorithm(
+			alg.to_str().unwrap_or("??").to_string(),
+		)),
 		None => Ok(None),
 	}
 }
@@ -282,7 +337,7 @@ struct DecryptStream {
 	stream: ByteStream,
 	buf: BytesBuf,
 	key: Key<Aes256Gcm>,
-	cipher: Option<stream::DecryptorLE31<Aes256Gcm>>,
+	cipher: Option<DecryptorLE31<Aes256Gcm>>,
 }
 
 impl DecryptStream {
@@ -308,13 +363,11 @@ impl Stream for DecryptStream {
 		let mut this = self.project();
 
 		while this.cipher.is_none() {
-			if this.buf.len() >= Aes256Gcm::NonceSize::as_usize() {
-				let nonce = this
-					.buf
-					.take_exact(Aes256Gcm::NonceSize::as_usize())
-					.unwrap();
-				let nonce: Aes256Gcm::Nonce = nonce.try_into().unwrap();
-				*this.cipher = Some(stream::DecryptorLE31::new(&self.key, &nonce));
+			let nonce_size = StreamNonceSize::to_usize();
+			if this.buf.len() >= nonce_size {
+				let nonce = this.buf.take_exact(nonce_size).unwrap();
+				let nonce = Nonce::from_slice(nonce.as_ref());
+				*this.cipher = Some(DecryptorLE31::new(&self.key, nonce));
 				break;
 			}
 
@@ -334,8 +387,7 @@ impl Stream for DecryptStream {
 			}
 		}
 
-		let chunk_size = STREAM_ENC_CHUNK_SIZE + Aes256Gcm::TagSize::as_usize();
-		while this.buf.len() < chunk_size {
+		while this.buf.len() < STREAM_ENC_CYPER_CHUNK_SIZE {
 			match futures::ready!(this.stream.as_mut().poll_next(cx)) {
 				Some(Ok(bytes)) => {
 					this.buf.extend(bytes);
@@ -351,14 +403,15 @@ impl Stream for DecryptStream {
 			return Poll::Ready(None);
 		}
 
-		let chunk = this.buf.take_max(chunk_size);
-		let res = this.cipher.as_ref().unwrap().decrypt_next(&chunk);
+		let chunk = this.buf.take_max(STREAM_ENC_CYPER_CHUNK_SIZE);
+		// TODO: use decrypt_last for last chunk
+		let res = this.cipher.as_ref().unwrap().decrypt_next(chunk.as_ref());
 		match res {
-			Ok(bytes) => Poll::Ready(Some(Ok(bytes))),
-			Err(_) => Poll::Ready(SomeErr(std::io::Error::new(
+			Ok(bytes) => Poll::Ready(Some(Ok(bytes.into()))),
+			Err(_) => Poll::Ready(Some(Err(std::io::Error::new(
 				std::io::ErrorKind::Other,
 				"Decryption failed",
-			))),
+			)))),
 		}
 	}
 }
diff --git a/src/block/block.rs b/src/block/block.rs
index 504d11f8..bd95680e 100644
--- a/src/block/block.rs
+++ b/src/block/block.rs
@@ -96,7 +96,7 @@ impl DataBlock {
 	}
 }
 
-fn zstd_encode<R: std::io::Read>(mut source: R, level: i32) -> std::io::Result<Vec<u8>> {
+pub fn zstd_encode<R: std::io::Read>(mut source: R, level: i32) -> std::io::Result<Vec<u8>> {
 	let mut result = Vec::<u8>::new();
 	let mut encoder = Encoder::new(&mut result, level)?;
 	encoder.include_checksum(true)?;
diff --git a/src/block/lib.rs b/src/block/lib.rs
index c9ff2845..6c4711ef 100644
--- a/src/block/lib.rs
+++ b/src/block/lib.rs
@@ -9,3 +9,5 @@ mod block;
 mod layout;
 mod metrics;
 mod rc;
+
+pub use block::zstd_encode;