Add S3 encryption

crates/pile-io/src/chachawriter_async.rs

@@ -0,0 +1,95 @@
+use std::io::SeekFrom;
+
+use tokio::io::{AsyncSeek, AsyncSeekExt, AsyncWrite, AsyncWriteExt};
+
+use crate::{ChaChaHeader, ChaChaReaderConfig};
+
+pub struct ChaChaWriterAsync<W: AsyncWrite + AsyncSeek + Unpin + Send> {
+	inner: W,
+	config: ChaChaReaderConfig,
+	encryption_key: [u8; 32],
+	buffer: Vec<u8>,
+	plaintext_bytes_written: u64,
+}
+
+impl<W: AsyncWrite + AsyncSeek + Unpin + Send> ChaChaWriterAsync<W> {
+	pub async fn new(mut inner: W, encryption_key: [u8; 32]) -> Result<Self, std::io::Error> {
+		let config = ChaChaReaderConfig::default();
+		let header_bytes = serialize_header(ChaChaHeader {
+			chunk_size: config.chunk_size,
+			nonce_size: config.nonce_size,
+			tag_size: config.tag_size,
+			plaintext_size: 0,
+		})?;
+		inner.write_all(&header_bytes).await?;
+
+		Ok(Self {
+			inner,
+			config,
+			encryption_key,
+			buffer: Vec::new(),
+			plaintext_bytes_written: 0,
+		})
+	}
+
+	pub async fn write(&mut self, buf: &[u8]) -> Result<(), std::io::Error> {
+		self.buffer.extend_from_slice(buf);
+		self.plaintext_bytes_written += buf.len() as u64;
+
+		let chunk_size = self.config.chunk_size as usize;
+		while self.buffer.len() >= chunk_size {
+			let encrypted = encrypt_chunk(&self.encryption_key, &self.buffer[..chunk_size])?;
+			self.inner.write_all(&encrypted).await?;
+			self.buffer.drain(..chunk_size);
+		}
+
+		Ok(())
+	}
+
+	/// Encrypt and write any buffered plaintext, patch the header with the
+	/// final `plaintext_size`, then return the inner writer.
+	pub async fn finish(mut self) -> Result<W, std::io::Error> {
+		if !self.buffer.is_empty() {
+			let encrypted = encrypt_chunk(&self.encryption_key, &self.buffer)?;
+			self.inner.write_all(&encrypted).await?;
+		}
+
+		self.inner.seek(SeekFrom::Start(0)).await?;
+		let header_bytes = serialize_header(ChaChaHeader {
+			chunk_size: self.config.chunk_size,
+			nonce_size: self.config.nonce_size,
+			tag_size: self.config.tag_size,
+			plaintext_size: self.plaintext_bytes_written,
+		})?;
+		self.inner.write_all(&header_bytes).await?;
+
+		Ok(self.inner)
+	}
+}
+
+fn encrypt_chunk(key: &[u8; 32], plaintext: &[u8]) -> Result<Vec<u8>, std::io::Error> {
+	use chacha20poly1305::{
+		XChaCha20Poly1305,
+		aead::{Aead, AeadCore, KeyInit, OsRng},
+	};
+
+	let nonce = XChaCha20Poly1305::generate_nonce(&mut OsRng);
+	let cipher = XChaCha20Poly1305::new(chacha20poly1305::Key::from_slice(key));
+	let ciphertext = cipher
+		.encrypt(&nonce, plaintext)
+		.map_err(|_| std::io::Error::other("encryption failed"))?;
+
+	let mut output = Vec::with_capacity(nonce.len() + ciphertext.len());
+	output.extend_from_slice(&nonce);
+	output.extend_from_slice(&ciphertext);
+	Ok(output)
+}
+
+fn serialize_header(header: ChaChaHeader) -> Result<Vec<u8>, std::io::Error> {
+	use binrw::BinWriterExt;
+	use std::io::Cursor;
+
+	let mut buf = Cursor::new(Vec::new());
+	buf.write_le(&header).map_err(std::io::Error::other)?;
+	Ok(buf.into_inner())
+}