From c450d3224da73b39c1c7b1352fea16388414c2b2 Mon Sep 17 00:00:00 2001
From: Alex Chi Z <chi@neon.tech>
Date: Wed, 16 Apr 2025 14:19:43 -0400
Subject: [PATCH] add remote_key crate

Signed-off-by: Alex Chi Z <chi@neon.tech>
---
 Cargo.lock                                    | 12 ++++++++++
 Cargo.toml                                    |  1 +
 libs/remote_keys/Cargo.toml                   | 15 +++++++++++++
 .../src/kms.rs => remote_keys/src/lib.rs}     | 22 +++++++------------
 libs/remote_storage/Cargo.toml                |  1 -
 libs/remote_storage/src/lib.rs                |  1 -
 6 files changed, 36 insertions(+), 16 deletions(-)
 create mode 100644 libs/remote_keys/Cargo.toml
 rename libs/{remote_storage/src/kms.rs => remote_keys/src/lib.rs} (58%)

diff --git a/Cargo.lock b/Cargo.lock
index b5d173d6d9..0a67d07ac3 100644
--- a/Cargo.lock
+++ b/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"
diff --git a/Cargo.toml b/Cargo.toml
index d957fa9070..2d597a8454 100644
--- a/Cargo.toml
+++ b/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",
diff --git a/libs/remote_keys/Cargo.toml b/libs/remote_keys/Cargo.toml
new file mode 100644
index 0000000000..62df3306fd
--- /dev/null
+++ b/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
diff --git a/libs/remote_storage/src/kms.rs b/libs/remote_keys/src/lib.rs
similarity index 58%
rename from libs/remote_storage/src/kms.rs
rename to libs/remote_keys/src/lib.rs
index a7ef375f05..6bf795e507 100644
--- a/libs/remote_storage/src/kms.rs
+++ b/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);
     }
 }
diff --git a/libs/remote_storage/Cargo.toml b/libs/remote_storage/Cargo.toml
index b4150cbf87..ef744bb48c 100644
--- a/libs/remote_storage/Cargo.toml
+++ b/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
diff --git a/libs/remote_storage/src/lib.rs b/libs/remote_storage/src/lib.rs
index 671227c4c8..2e18feaa19 100644
--- a/libs/remote_storage/src/lib.rs
+++ b/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;