add remote_key crate

Signed-off-by: Alex Chi Z <chi@neon.tech>

Cargo.lock

@@ -5495,6 +5495,18 @@ version = "1.9.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "c707298afce11da2efef2f600116fa93ffa7a032b5d7b628aa17711ec81383ca"
 
+[[package]]
+name = "remote_keys"
+version = "0.1.0"
+dependencies = [
+ "anyhow",
+ "aws-config",
+ "aws-sdk-kms",
+ "aws-smithy-types",
+ "rand 0.8.5",
+ "utils",
+]
+
 [[package]]
 name = "remote_storage"
 version = "0.1.0"

Cargo.toml

@@ -30,6 +30,7 @@ members = [
     "libs/tenant_size_model",
     "libs/metrics",
     "libs/postgres_connection",
+    "libs/remote_keys",
     "libs/remote_storage",
     "libs/tracing-utils",
     "libs/postgres_ffi/wal_craft",

libs/remote_keys/Cargo.toml

@@ -0,0 +1,15 @@
+[package]
+name = "remote_keys"
+version = "0.1.0"
+edition = "2024"
+license.workspace = true
+
+[dependencies]
+anyhow.workspace = true
+aws-smithy-types.workspace = true
+aws-sdk-kms.workspace = true
+aws-config.workspace = true
+utils.workspace = true
+
+[dev-dependencies]
+rand.workspace = true

libs/remote_keys/src/lib.rs

@@ -8,25 +8,17 @@ pub struct NaiveKms {
     account_id: String,
 }
 
-pub struct KeyPair {
-    pub wrapped: Vec<u8>,
-    pub plain: Vec<u8>,
-}
-
 impl NaiveKms {
     pub fn new(account_id: String) -> Self {
         Self { account_id }
     }
 
-    /// Generate a new key in the KMS. Returns the plain and wrapped keys.
-    pub fn generate_key(&self) -> anyhow::Result<KeyPair> {
-        // TODO: use a proper library for generating the key (i.e., ring).
-        let plain = rand::random::<[u8; 32]>().to_vec();
-        let wrapped = [self.account_id.as_bytes(), "-wrapped-".as_bytes(), &plain].concat();
-        Ok(KeyPair { wrapped, plain })
+    pub fn encrypt(&self, plain: &[u8]) -> anyhow::Result<Vec<u8>> {
+        let wrapped = [self.account_id.as_bytes(), "-wrapped-".as_bytes(), plain].concat();
+        Ok(wrapped)
     }
 
-    pub fn unwrap_key(&self, wrapped: &[u8]) -> anyhow::Result<Vec<u8>> {
+    pub fn decrypt(&self, wrapped: &[u8]) -> anyhow::Result<Vec<u8>> {
         let Some(wrapped) = wrapped.strip_prefix(self.account_id.as_bytes()) else {
             return Err(anyhow::anyhow!("invalid key"));
         };
@@ -44,7 +36,9 @@ mod tests {
     #[test]
     fn test_generate_key() {
         let kms = NaiveKms::new("test-tenant".to_string());
-        let key_pair = kms.generate_key().unwrap();
-        assert_eq!(kms.unwrap_key(&key_pair.wrapped).unwrap(), key_pair.plain);
+        let data = rand::random::<[u8; 32]>().to_vec();
+        let encrypted = kms.encrypt(&data).unwrap();
+        let decrypted = kms.decrypt(&encrypted).unwrap();
+        assert_eq!(data, decrypted);
     }
 }

libs/remote_storage/Cargo.toml

@@ -19,7 +19,6 @@ camino = { workspace = true, features = ["serde1"] }
 humantime-serde.workspace = true
 hyper = { workspace = true, features = ["client"] }
 futures.workspace = true
-rand.workspace = true
 reqwest.workspace = true
 serde.workspace = true
 serde_json.workspace = true

libs/remote_storage/src/lib.rs

@@ -12,7 +12,6 @@
 mod azure_blob;
 mod config;
 mod error;
-mod kms;
 mod local_fs;
 mod metrics;
 mod s3_bucket;