From 3bc3f71418afc2d66e1c6564eaf89e668948cda0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Arpad=20M=C3=BCller?= <arpad@neon.tech>
Date: Fri, 28 Mar 2025 15:22:13 +0100
Subject: [PATCH] Initial remote_keys crate

---
 Cargo.lock                       | 11 ++++++++
 Cargo.toml                       |  1 +
 libs/remote_keys/Cargo.toml      | 12 ++++++++
 libs/remote_keys/src/aws_keys.rs | 47 ++++++++++++++++++++++++++++++++
 libs/remote_keys/src/lib.rs      |  6 ++++
 5 files changed, 77 insertions(+)
 create mode 100644 libs/remote_keys/Cargo.toml
 create mode 100644 libs/remote_keys/src/aws_keys.rs
 create mode 100644 libs/remote_keys/src/lib.rs

diff --git a/Cargo.lock b/Cargo.lock
index 5d2cdcea27..c95aa483b7 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -5495,6 +5495,17 @@ 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",
+ "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..3f8b60fdca
--- /dev/null
+++ b/libs/remote_keys/Cargo.toml
@@ -0,0 +1,12 @@
+[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
\ No newline at end of file
diff --git a/libs/remote_keys/src/aws_keys.rs b/libs/remote_keys/src/aws_keys.rs
new file mode 100644
index 0000000000..9a4d2f864f
--- /dev/null
+++ b/libs/remote_keys/src/aws_keys.rs
@@ -0,0 +1,47 @@
+use aws_config::BehaviorVersion;
+
+use crate::KeyId;
+
+pub struct AwsRemoteKeyClient {
+    client: aws_sdk_kms::Client,
+}
+
+impl AwsRemoteKeyClient {
+    pub async fn new() -> Self {
+        let sdk_config = aws_config::defaults(BehaviorVersion::v2024_03_28())
+            .retry_config(
+                aws_config::retry::RetryConfig::standard()
+                    .with_max_attempts(5) // Retry up to 5 times
+                    .with_initial_backoff(std::time::Duration::from_millis(200)) // Start with 200ms delay
+                    .with_max_backoff(std::time::Duration::from_secs(5)), // Cap at 5 seconds
+            )
+            .load()
+            .await;
+        let client = aws_sdk_kms::Client::new(&sdk_config);
+        Self { client }
+    }
+
+    pub async fn decrypt(&self, key_id: &KeyId, ciphertext: impl Into<Vec<u8>>) -> Vec<u8> {
+        let output = self
+            .client
+            .decrypt()
+            .key_id(&key_id.0)
+            .ciphertext_blob(aws_smithy_types::Blob::new(ciphertext.into()))
+            .send()
+            .await
+            .expect("decrypt");
+        output.plaintext.expect("plaintext").into_inner()
+    }
+
+    pub async fn encrypt(&self, key_id: &KeyId, ciphertext: impl Into<Vec<u8>>) -> Vec<u8> {
+        let output = self
+            .client
+            .encrypt()
+            .key_id(&key_id.0)
+            .plaintext(aws_smithy_types::Blob::new(ciphertext.into()))
+            .send()
+            .await
+            .expect("decrypt");
+        output.ciphertext_blob.expect("ciphertext").into_inner()
+    }
+}
diff --git a/libs/remote_keys/src/lib.rs b/libs/remote_keys/src/lib.rs
new file mode 100644
index 0000000000..d1c1042913
--- /dev/null
+++ b/libs/remote_keys/src/lib.rs
@@ -0,0 +1,6 @@
+mod aws_keys;
+pub use aws_keys::AwsRemoteKeyClient;
+
+/// A string uniquely identifying a key
+#[derive(Debug, PartialEq, Eq)]
+pub struct KeyId(pub String);
\ No newline at end of file