Add S3 encryption

crates/pile-io/src/s3reader.rs

@@ -0,0 +1,88 @@
+use smartstring::{LazyCompact, SmartString};
+use std::{io::SeekFrom, sync::Arc};
+
+use crate::{AsyncReader, AsyncSeekReader};
+
+pub struct S3Reader {
+	pub client: Arc<aws_sdk_s3::Client>,
+	pub bucket: SmartString<LazyCompact>,
+	pub key: SmartString<LazyCompact>,
+	pub cursor: u64,
+	pub size: u64,
+}
+
+impl AsyncReader for S3Reader {
+	async fn read(&mut self, buf: &mut [u8]) -> Result<usize, std::io::Error> {
+		let len_left = self.size.saturating_sub(self.cursor);
+		if len_left == 0 || buf.is_empty() {
+			return Ok(0);
+		}
+
+		let start_byte = self.cursor;
+		let len_to_read = (buf.len() as u64).min(len_left);
+		let end_byte = start_byte + len_to_read - 1;
+
+		let resp = self
+			.client
+			.get_object()
+			.bucket(self.bucket.as_str())
+			.key(self.key.as_str())
+			.range(format!("bytes={start_byte}-{end_byte}"))
+			.send()
+			.await
+			.map_err(std::io::Error::other)?;
+
+		let bytes = resp
+			.body
+			.collect()
+			.await
+			.map(|x| x.into_bytes())
+			.map_err(std::io::Error::other)?;
+
+		let n = bytes.len().min(buf.len());
+		buf[..n].copy_from_slice(&bytes[..n]);
+		self.cursor += n as u64;
+		Ok(n)
+	}
+}
+
+impl AsyncSeekReader for S3Reader {
+	async fn seek(&mut self, pos: SeekFrom) -> Result<u64, std::io::Error> {
+		match pos {
+			SeekFrom::Start(x) => self.cursor = x.min(self.size),
+
+			SeekFrom::Current(x) => {
+				if x < 0 {
+					let abs = x.unsigned_abs();
+					if abs > self.cursor {
+						return Err(std::io::Error::new(
+							std::io::ErrorKind::InvalidInput,
+							"cannot seek past start",
+						));
+					}
+					self.cursor -= abs;
+				} else {
+					self.cursor += x as u64;
+				}
+			}
+
+			std::io::SeekFrom::End(x) => {
+				if x < 0 {
+					let abs = x.unsigned_abs();
+					if abs > self.size {
+						return Err(std::io::Error::new(
+							std::io::ErrorKind::InvalidInput,
+							"cannot seek past start",
+						));
+					}
+					self.cursor = self.size - abs;
+				} else {
+					self.cursor = self.size + x as u64;
+				}
+			}
+		}
+
+		self.cursor = self.cursor.min(self.size);
+		Ok(self.cursor)
+	}
+}