path: root/src/api/s3/encryption.rs

use std::borrow::Cow;

use aes_gcm::{
	aead::{Aead, AeadCore, KeyInit, OsRng},
	Aes256Gcm, Key, Nonce,
};
use base64::prelude::*;

use http::header::{HeaderName, HeaderValue};
use hyper::{body::Body, Request};

use garage_net::stream::{ByteStream, ByteStreamReader};
use garage_rpc::rpc_helper::OrderTag;
use garage_util::data::Hash;

use garage_model::garage::Garage;
use garage_model::s3::object_table::{ObjectVersionEncryption, ObjectVersionHeaders};

use crate::s3::error::Error;

const X_AMZ_SERVER_SIDE_ENCRYPTION_CUSTOMER_ALGORITHM: HeaderName =
	HeaderName::from_static("x-amz-server-side-encryption-customer-algorithm");
const X_AMZ_SERVER_SIDE_ENCRYPTION_CUSTOMER_KEY: HeaderName =
	HeaderName::from_static("x-amz-server-side-encryption-customer-key");
const X_AMZ_SERVER_SIDE_ENCRYPTION_CUSTOMER_KEY_MD5: HeaderName =
	HeaderName::from_static("x-amz-server-side-encryption-customer-key-MD5");

const X_AMZ_COPY_SOURCE_SERVER_SIDE_ENCRYPTION_CUSTOMER_ALGORITHM: HeaderName =
	HeaderName::from_static("x-amz-copy-source-server-side-encryption-customer-algorithm");
const X_AMZ_COPY_SOURCE_SERVER_SIDE_ENCRYPTION_CUSTOMER_KEY: HeaderName =
	HeaderName::from_static("x-amz-copy-source-server-side-encryption-customer-key");
const X_AMZ_COPY_SOURCE_SERVER_SIDE_ENCRYPTION_CUSTOMER_KEY_MD5: HeaderName =
	HeaderName::from_static("x-amz-copy-source-server-side-encryption-customer-key-MD5");

const CUSTOMER_ALGORITHM_AES256: HeaderValue = HeaderValue::from_static("AES256");

const STREAM_ENC_CHUNK_SIZE: usize = 0x1000; // 4096 bytes

#[derive(Clone, Copy)]
pub enum EncryptionParams {
	Plaintext,
	SseC {
		client_key: Key<Aes256Gcm>,
		compression_level: Option<i32>,
	},
}

impl EncryptionParams {
	pub fn new_from_req(
		garage: &Garage,
		req: &Request<impl Body>,
	) -> Result<EncryptionParams, Error> {
		let key = parse_request_headers(
			req,
			&X_AMZ_SERVER_SIDE_ENCRYPTION_CUSTOMER_ALGORITHM,
			&X_AMZ_SERVER_SIDE_ENCRYPTION_CUSTOMER_KEY,
			&X_AMZ_SERVER_SIDE_ENCRYPTION_CUSTOMER_KEY_MD5,
		)?;
		match key {
			Some(client_key) => Ok(EncryptionParams::SseC {
				client_key: parse_and_check_key(req)?,
				compression_level: garage.config.compression_level,
			}),
			None => Ok(EncryptionParams::Plaintext),
		}
	}

	pub fn check_decrypt_for_get<'a>(
		garage: &Garage,
		req: &Request<impl Body>,
		obj_enc: &'a ObjectVersionEncryption,
	) -> Result<(Self, Cow<'a, ObjectVersionHeaders>), Error> {
		let key = parse_request_headers(
			req,
			&X_AMZ_SERVER_SIDE_ENCRYPTION_CUSTOMER_ALGORITHM,
			&X_AMZ_SERVER_SIDE_ENCRYPTION_CUSTOMER_KEY,
			&X_AMZ_SERVER_SIDE_ENCRYPTION_CUSTOMER_KEY_MD5,
		)?;
		Self::check_decrypt(garage, key, obj_enc)
	}

	pub fn check_decrypt_for_copy_source<'a>(
		garage: &Garage,
		req: &Request<impl Body>,
		obj_enc: &'a ObjectVersionEncryption,
	) -> Result<(Self, Cow<'a, ObjectVersionHeaders>), Error> {
		let key = parse_request_headers(
			req,
			&X_AMZ_COPY_SOURCE_SERVER_SIDE_ENCRYPTION_CUSTOMER_ALGORITHM,
			&X_AMZ_COPY_SOURCE_SERVER_SIDE_ENCRYPTION_CUSTOMER_KEY,
			&X_AMZ_COPY_SOURCE_SERVER_SIDE_ENCRYPTION_CUSTOMER_KEY_MD5,
		)?;
		Self::check_decrypt(garage, key, obj_enc)
	}

	fn check_decrypt<'a>(
		garage: &Garage,
		key: Option<Key<Aes256Gcm>>,
		obj_enc: &'a ObjectVersionEncryption,
	) -> Result<(Self, Cow<'a, ObjectVersionHeaders>), Error> {
		match (key, obj_enc) {
			(
				Some(client_key),
				ObjectVersionEncryption::SseC {
					headers,
					compressed,
				},
			) => {
				let enc = Self::SseC {
					client_key,
					compression_level: if compressed {
						Some(garage.config.compression_level.unwrap_or(1))
					} else {
						None
					},
				};
				let plaintext = enc.decrypt_blob(&headers)?;
				let headers = ObjectVersionHeaders::decode(&plaintext)?;
				Ok((enc, headers.into()))
			}
			(None, ObjectVersionEncryption::Plaintext { headers }) => {
				Ok((Self::Plaintext, headers.into()))
			}
			(_, ObjectVersionEncryption::SseC { .. }) => {
				Err(Error::bad_request("Object is encrypted"))
			}
			(Some(_), _) => {
				// TODO: should this be an OK scenario?
				Err(Error::bad_request("Trying to decrypt a plaintext object"))
			}
		}
	}

	pub fn encrypt_headers(
		&self,
		h: ObjectVersionHeaders,
	) -> Result<ObjectVersionEncryption, Error> {
		match self {
			Self::SseC {
				client_key,
				compression_level,
			} => {
				let plaintext = h.encode()?;
				let ciphertext = self.encrypt_blob(plaintext)?;
				Ok(ObjectVersionEncryption::SseC {
					headers: ciphertext,
					compressed: compression_level.is_some(),
				})
			}
			Self::Plaintext => Ok(ObjectVersionEncryption::Plaintext { headers: h }),
		}
	}

	// ---- generic function for encrypting / decrypting blobs ----
	// prepends a randomly-generated nonce to the encrypted value

	pub fn encrypt_blob<'a>(&self, blob: &'a [u8]) -> Result<Cow<'a, [u8]>, Error> {
		match self {
			Self::SseC { client_key, .. } => {
				let cipher = Aes256Gcm::new(&client_key);
				let nonce = Aes256Gcm::generate_nonce(&mut OsRng);
				let ciphertext = cipher
					.encrypt(&nonce, &blob)
					.ok_or_internal_error("Encryption failed")?;
				Ok([nonce.to_vec(), ciphertext].concat().into())
			}
			Self::Plaintext => Ok(blob.into()),
		}
	}

	pub fn decrypt_blob<'a>(&self, blob: &'a [u8]) -> Result<Cow<'a, [u8]>, Error> {
		match self {
			Self::SseC { client_key, .. } => {
				let cipher = Aes256Gcm::new(&client_key);
				let nonce_size = Aes256Gcm::NonceSize::to_usize();
				let nonce: Nonce<Aes256Gcm::NonceSize> = blob
					.get(..nonce_size)
					.ok_or_internal_error("invalid encrypted data")?
					.try_into()
					.unwrap();
				let plaintext = cipher
					.decrypt(&nonce, &blob[nonce_size..])
					.ok_or_bad_request(
						"Invalid encryption key, could not decrypt object metadata.",
					)?;
				Ok(plaintext.into())
			}
			Self::Plaintext => Ok(blob.into()),
		}
	}

	// ----  function for encrypting / decrypting byte streams ----

	/// Get a data block from the storage node, and decrypt+decompress it
	/// if necessary. If object is plaintext, just get it without any processing.
	pub async fn get_and_decrypt_block(
		&self,
		garage: &Garage,
		hash: &Hash,
		order: Option<OrderTag>,
	) -> Result<ByteStream, Error> {
		let raw_block = garage
			.block_manager
			.rpc_get_block_streaming(hash, order)
			.await?;
		match self {
			Self::Plaintext => Ok(raw_block),
			Self::SseC {
				client_key,
				compression_level,
			} => {
				todo!()
			}
		}
	}
}

fn parse_request_headers(
	req: &Request<impl Body>,
	alg_header: &HeaderName,
	key_header: &HeaderName,
	md5_header: &HeaderName,
) -> Result<Option<Key<Aes256Gcm>>, Error> {
	match req.headers().get(alg_header) {
		Some(CUSTOMER_ALGORITHM_AES256) => {
			use md5::{Digest, Md5};

			let key_b64 = req
				.headers()
				.get(key_header)
				.ok_or_bad_request(format!("Missing {} header", key_header))?;
			let key_bytes: [u8; 32] = BASE64_STANDARD
				.decode(&key_b64)
				.ok_or_bad_request(format!("Invalid {} header", key_header))?
				.try_into()
				.ok_or_bad_request(format!("Invalid {} header", key_header))?;

			let md5_b64 = req
				.headers()
				.get(md5_header)
				.ok_or_bad_request(format!("Missing {} header", md5_header))?;
			let md5_bytes = BASE64_STANDARD
				.decode(&md5_b64)
				.ok_or_bad_request(format!("Invalid {} header", md5_header))?;

			let mut hasher = Md5::new();
			hasher.update(&key_bytes[..]);
			if hasher.finalize() != md5_bytes {
				return Err(Error::bad_request(
					"Encryption key MD5 checksum does not match",
				));
			}

			Ok(Some(key_bytes.into()))
		}
		Some(alg) => Err(Error::InvalidEncryptionAlgorithm(alg.to_string())),
		None => Ok(None),
	}
}