Repository: marlinprotocol/oyster-monorepo
Pull request #117: Update KMS root and derive server to use contract validation
Author: login: vg-27
Created: 2025-03-06T09:28:47Z | Merged: 2025-03-19T08:25:29Z

- In root-server, replace validation by enclave public key with hash of PCRs and user data
- Update derive server to fetch seed from KMS root-server-contract, optionally.
- Update oyster-cvm cli to fill contract address and KMS contract root server details (endpoint, PCRs, user data)


Diff:
diff --git a/cli/oyster-cvm/Cargo.lock b/cli/oyster-cvm/Cargo.lock
index 623838546..8cb959fc9 100644
--- a/cli/oyster-cvm/Cargo.lock
+++ b/cli/oyster-cvm/Cargo.lock
@@ -2914,6 +2914,7 @@ dependencies = [
  "clap",
  "futures-util",
  "hex",
+ "lazy_static",
  "libsodium-sys-stable",
  "openssl",
  "oyster-sdk",
diff --git a/cli/oyster-cvm/Cargo.toml b/cli/oyster-cvm/Cargo.toml
index 81ff010a9..b2947bbbc 100644
--- a/cli/oyster-cvm/Cargo.toml
+++ b/cli/oyster-cvm/Cargo.toml
@@ -22,3 +22,4 @@ base64 = "0.22.1"
 ureq = "3.0.6"
 prettytable-rs = "0.10"
 chrono = "0.4.40"
+lazy_static = "1.5.0"
diff --git a/cli/oyster-cvm/src/args/init_params.rs b/cli/oyster-cvm/src/args/init_params.rs
index d601caab6..1cd3b1df9 100644
--- a/cli/oyster-cvm/src/args/init_params.rs
+++ b/cli/oyster-cvm/src/args/init_params.rs
@@ -1,12 +1,18 @@
 use std::{
+    collections::HashMap,
     fs,
     path::{Component, PathBuf},
 };
 
-use alloy::signers::k256::sha2::{Digest, Sha256};
+use alloy::primitives::Address;
+use alloy::{
+    hex::FromHex,
+    signers::k256::sha2::{Digest, Sha256},
+};
 use anyhow::{anyhow, bail, Context, Result};
 use base64::{prelude::BASE64_STANDARD, Engine};
 use clap::Args;
+use lazy_static::lazy_static;
 use libsodium_sys::{crypto_box_SEALBYTES, crypto_box_seal, sodium_init};
 use serde::Serialize;
 use tracing::info;
@@ -36,6 +42,14 @@ pub struct InitParamsArgs {
     #[command(flatten)]
     pub pcrs: PcrArgs,
 
+    /// Encalve verifier contract address
+    #[arg(long, requires = "chain_id")]
+    pub contract_address: Option<String>,
+
+    /// Chain ID for KMS contract root server
+    #[arg(long, requires = "contract_address")]
+    pub chain_id: Option<u64>,
+
     /// Docker compose file defining services to run,
     /// set as first init param
     #[arg(long)]
@@ -50,7 +64,10 @@ impl InitParamsArgs {
         }
 
         // check if there are any init params
-        if self.init_params.is_none() && self.docker_compose.is_none() {
+        if self.init_params.is_none()
+            && self.docker_compose.is_none()
+            && self.contract_address.is_none()
+        {
             return Ok(None);
         };
 
@@ -58,6 +75,22 @@ impl InitParamsArgs {
             .docker_compose
             .map(|x| vec![format!("docker-compose.yml:1:0:file:{x}")])
             .unwrap_or_default();
+
+        if let Some(address) = self.contract_address {
+            if let Err(_) = Address::from_hex(&address) {
+                bail!("invalid contract address");
+            }
+            init_params.push(format!("contract-address:1:0:utf8:{}", address));
+
+            let Some(root_server_str) = KMS_ROOT_SERVERS.get(&self.chain_id.unwrap()) else {
+                bail!("unknown chain id");
+            };
+            init_params.push(format!(
+                "root-server-config.json:1:0:utf8:{}",
+                root_server_str
+            ));
+        }
+
         init_params.append(&mut self.init_params.unwrap_or_default());
 
         // encoding has to be done
@@ -148,7 +181,7 @@ impl InitParamsArgs {
             })?;
 
         // fetch key
-        let pk = fetch_encryption_key(
+        let pk = fetch_encryption_key_with_pcr(
             self.kms_endpoint
                 .as_ref()
                 .unwrap_or(&"http://image-v2.kms.box:1101".into()),
@@ -244,7 +277,7 @@ struct InitParamsList {
     params: Vec<InitParam>,
 }
 
-fn fetch_encryption_key(
+fn fetch_encryption_key_with_pcr(
     endpoint: &str,
     pcr0: &str,
     pcr1: &str,
@@ -266,3 +299,23 @@ fn fetch_encryption_key(
         .try_into()
         .context("failed to parse reponse")
 }
+
+// TODO fill store KMS details
+lazy_static! {
+    static ref KMS_ROOT_SERVERS: HashMap<u64, &'static str> = {
+        let mut root_servers = HashMap::new();
+        root_servers.insert(
+            31337,
+            r#"
+                {
+                    "kms_endpoint": "image-v2.kms.box:1101",
+                    "pcr0": "cb7ebc13d527e9cf9cc271b0d816c72a5bfa685ae56118ce4986fa82b8b9aac8b851206539a6c4600ad77566aa15bd0e",
+                    "pcr1": "3b26340a10ac3a5494139fed12bc30028017b72dabfb78a38763bd21ea67bbfe03214e7ce628c2952cfa2ff478f370ba",
+                    "pcr2": "0d128997bfc8ef24a2aa1ecec60c61e48eda9b439f8662d41eb38d8f0ba0401367778a5f8804127c591d824d5c3ed617",
+                    "user_data": "544d4b69000100009293c4308e0ebc7f830be9a963f343358c4875d67c36eb1dc0c3250cb9438a43f0399ed4fc7e0b80ccb99b106136e06bbaecb90bc460843d33b632beb3537a7d1c42531ad1b43bce0015ba7e7fb7ed82fed9a0d200ac11dc646ca22b743daec9baa39d689c350a3e70ce85979e856e9cac3c31f45febe28510360beb78495932a6e06c2a5fd6caa0adc9dff0fef20d9ff4916199bcbcc450cf36031706fbe58b82fcf808e76f801924ec66a3ea5bcaec0fb33619cf5df6314b43eca506a25392b1e042126a63744f35bd39403c6d73e81a4efaf0e0251563ee3eacc2ce419ae5a1c139d5d321761a9292c430868c3d012a5d524f0939e4ee4d60b738b4c44448ec286a5361e15ffbf2641e2df25363a204a738231e5f1a9621999741da01b87b22636f6e646974696f6e223a7b22636861696e223a312c22636f6e646974696f6e54797065223a22636f6e7472616374222c22636f6e747261637441646472657373223a22307843374430383443326536424341633030374433424146433431353438356133383038303343306265222c2266756e6374696f6e416269223a7b22696e70757473223a5b7b22696e7465726e616c54797065223a2261646472657373222c226e616d65223a22222c2274797065223a2261646472657373227d5d2c226e616d65223a2269735665726966696564222c226f757470757473223a5b7b22696e7465726e616c54797065223a22626f6f6c222c226e616d65223a22222c2274797065223a22626f6f6c227d5d2c2273746174654d75746162696c697479223a2276696577222c2274797065223a2266756e6374696f6e227d2c226d6574686f64223a2269735665726966696564222c22706172616d6574657273223a5b223a7573657241646472657373225d2c2272657475726e56616c756554657374223a7b22636f6d70617261746f72223a223d3d222c2276616c7565223a747275657d7d2c2276657273696f6e223a22312e302e30227dc441cd0ad8db7d03074a07437a39025a8deea6a72c433dedc45375b95d28257dc0300ae70507050f75c5d5c31c50b167bdde4c06ac789b22641defe91de9442fd74d1b"
+                }
+            "#,
+        );
+        root_servers
+    };
+}
diff --git a/enclaves/blue/setup.sh b/enclaves/blue/setup.sh
index ee01ea97c..732e96372 100644
--- a/enclaves/blue/setup.sh
+++ b/enclaves/blue/setup.sh
@@ -92,12 +92,16 @@ sleep 2
 
 # start derive server
 /app/supervisord ctl -c /etc/supervisord.conf start derive-server
-
 sleep 10
 
 # process init params into their constituent files
 /app/init-params-decoder
 
+# start derive server contract if contract address and root server config are present
+if [ -f /init-params/contract-address ] && [ -f /init-params/root-server-config.json ]; then
+    /app/supervisord ctl -c /etc/supervisord.conf start derive-server-contract
+fi
+
 # Start the Docker daemon
 /app/supervisord ctl -c /etc/supervisord.conf start docker
 
diff --git a/enclaves/blue/supervisord.conf b/enclaves/blue/supervisord.conf
index a4a453285..dc74ebaba 100644
--- a/enclaves/blue/supervisord.conf
+++ b/enclaves/blue/supervisord.conf
@@ -66,6 +66,15 @@ stdout_logfile_maxbytes=0
 stderr_logfile=/dev/stdout
 stderr_logfile_maxbytes=0
 
+[program:derive-server-contract]
+command=/app/kms-derive-server --listen-addr 127.0.0.1:1101 --attestation-endpoint http://127.0.0.1:1300/attestation/raw --secret-path /app/id.sec --root-server-config /init-params/root-server-config.json --contract-address-file /init-params/contract-address
+autostart=false
+autorestart=true
+stdout_logfile=/dev/stdout
+stdout_logfile_maxbytes=0
+stderr_logfile=/dev/stdout
+stderr_logfile_maxbytes=0
+
 # docker
 [program:docker]
 command=/bin/dockerd --iptables=false
diff --git a/kms/derive-server/Cargo.lock b/kms/derive-server/Cargo.lock
index d129e1a20..26352aaff 100644
--- a/kms/derive-server/Cargo.lock
+++ b/kms/derive-server/Cargo.lock
@@ -1342,6 +1342,7 @@ dependencies = [
  "oyster-sdk",
  "reqwest",
  "serde",
+ "serde_json",
  "tokio",
  "tracing",
  "tracing-subscriber",
@@ -2094,9 +2095,9 @@ dependencies = [
 
 [[package]]
 name = "serde_json"
-version = "1.0.137"
+version = "1.0.140"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "930cfb6e6abf99298aaad7d29abbef7a9999a9a8806a40088f55f0dcec03146b"
+checksum = "20068b6e96dc6c9bd23e01df8827e6c7e1f2fddd43c21810382803c136b99373"
 dependencies = [
  "itoa",
  "memchr",
diff --git a/kms/derive-server/Cargo.toml b/kms/derive-server/Cargo.toml
index e389ab267..3af69bc93 100644
--- a/kms/derive-server/Cargo.toml
+++ b/kms/derive-server/Cargo.toml
@@ -16,6 +16,7 @@ kms-derive-utils = { version = "0.1.0", path = "../derive-utils" }
 oyster-sdk = "0.15.0"
 reqwest = "0.12.12"
 serde = { version = "1.0.217", features = ["derive"] }
+serde_json = "1.0.140"
 tokio = { version = "1.43.0", features = ["full"] }
 tracing = "0.1.41"
 tracing-subscriber = { version = "0.3.19", features = ["env-filter"] }
diff --git a/kms/derive-server/src/fetch.rs b/kms/derive-server/src/fetch.rs
index 0c982c9e3..2e30d03dc 100644
--- a/kms/derive-server/src/fetch.rs
+++ b/kms/derive-server/src/fetch.rs
@@ -12,6 +12,7 @@ pub async fn fetch_seed(
     auth_store: AuthStore,
     secret: [u8; 32],
     kms_enpdoint: String,
+    contract_address: Option<String>,
 ) -> Result<[u8; 64]> {
     // oh how I wish I could use reqwest
 
@@ -39,9 +40,13 @@ pub async fn fetch_seed(
         }
     });
 
+    let uri = match contract_address {
+        Some(contract_address) => format!("/derive?address={}", contract_address),
+        None => "/derive".to_string(),
+    };
     let request = Request::builder()
         .method("GET")
-        .uri("/derive")
+        .uri(uri)
         .body("".to_string())
         .context("failed to build request")?;
     let response = request_sender
diff --git a/kms/derive-server/src/main.rs b/kms/derive-server/src/main.rs
index a33ceac6b..b7461cc5d 100644
--- a/kms/derive-server/src/main.rs
+++ b/kms/derive-server/src/main.rs
@@ -3,7 +3,10 @@ use axum::{routing::get, Router};
 use clap::Parser;
 use fetch::fetch_seed;
 use scallop::{AuthStore, Auther};
-use tokio::{fs::read, net::TcpListener};
+use tokio::{
+    fs::{read, read_to_string},
+    net::TcpListener,
+};
 use tracing::info;
 use tracing_subscriber::EnvFilter;
 
@@ -15,8 +18,12 @@ mod scallop;
 #[command(version, about, long_about = None)]
 struct Args {
     /// KMS endpoint
-    #[arg(long)]
-    kms_endpoint: String,
+    #[arg(
+        long,
+        conflicts_with = "root_server_config",
+        required_unless_present = "root_server_config"
+    )]
+    kms_endpoint: Option<String>,
 
     /// Listening address
     #[arg(long, default_value = "127.0.0.1:1100")]
@@ -31,20 +38,44 @@ struct Args {
     secret_path: String,
 
     /// PCR0 of the root server
-    #[arg(long)]
-    pcr0: String,
+    #[arg(
+        long,
+        conflicts_with = "root_server_config",
+        required_unless_present = "root_server_config"
+    )]
+    pcr0: Option<String>,
 
     /// PCR1 of the root server
-    #[arg(long)]
-    pcr1: String,
+    #[arg(
+        long,
+        conflicts_with = "root_server_config",
+        required_unless_present = "root_server_config"
+    )]
+    pcr1: Option<String>,
 
     /// PCR2 of the root server
-    #[arg(long)]
-    pcr2: String,
+    #[arg(
+        long,
+        conflicts_with = "root_server_config",
+        required_unless_present = "root_server_config"
+    )]
+    pcr2: Option<String>,
 
     /// user data of the root server
+    #[arg(
+        long,
+        conflicts_with = "root_server_config",
+        required_unless_present = "root_server_config"
+    )]
+    user_data: Option<String>,
+
+    /// file containing enclave verification contract address in hexadecimal
     #[arg(long)]
-    user_data: String,
+    contract_address_file: Option<String>,
+
+    /// JSON config file containing the root server's details
+    #[arg(long, required_unless_present_all = ["kms_endpoint", "pcr0", "pcr1", "pcr2", "user_data"])]
+    root_server_config: Option<String>,
 }
 
 #[derive(Clone)]
@@ -68,27 +99,80 @@ async fn main() -> Result<()> {
         url: args.attestation_endpoint,
     };
 
-    let pcr0: [u8; 48] = hex::decode(args.pcr0)
-        .context("failed to decode pcr0")?
-        .try_into()
-        .map_err(|_| anyhow!("incorrect pcr0 size"))?;
-    let pcr1: [u8; 48] = hex::decode(args.pcr1)
-        .context("failed to decode pcr1")?
-        .try_into()
-        .map_err(|_| anyhow!("incorrect pcr1 size"))?;
-    let pcr2: [u8; 48] = hex::decode(args.pcr2)
-        .context("failed to decode pcr2")?
-        .try_into()
-        .map_err(|_| anyhow!("incorrect pcr2 size"))?;
-    let user_data = hex::decode(args.user_data)
+    let pcr0: [u8; 48];
+    let pcr1: [u8; 48];
+    let pcr2: [u8; 48];
+    let user_data: Box<[u8]>;
+    let kms_endpoint: String;
+
+    if let Some(filename) = args.root_server_config {
+        let config_content = read_to_string(filename)
+            .await
+            .context("failed to read root server config file")?;
+        let config: serde_json::Value = serde_json::from_str(&config_content)
+            .context("failed to parse root server config file")?;
+
+        kms_endpoint = config["kms_endpoint"]
+            .as_str()
+            .context("missing kms_endpoint in config")?
+            .to_string();
+        pcr0 = hex::decode(config["pcr0"].as_str().context("missing pcr0 in config")?)
+            .context("failed to decode pcr0")?
+            .try_into()
+            .map_err(|_| anyhow!("incorrect pcr0 size"))?;
+        pcr1 = hex::decode(config["pcr1"].as_str().context("missing pcr1 in config")?)
+            .context("failed to decode pcr1")?
+            .try_into()
+            .map_err(|_| anyhow!("incorrect pcr1 size"))?;
+        pcr2 = hex::decode(config["pcr2"].as_str().context("missing pcr2 in config")?)
+            .context("failed to decode pcr2")?
+            .try_into()
+            .map_err(|_| anyhow!("incorrect pcr2 size"))?;
+        user_data = hex::decode(
+            config["user_data"]
+                .as_str()
+                .context("missing user_data in config")?,
+        )
         .context("failed to decode user data")?
         .into_boxed_slice();
+    } else {
+        kms_endpoint = args.kms_endpoint.unwrap();
+        pcr0 = hex::decode(args.pcr0.unwrap())
+            .context("failed to decode pcr0")?
+            .try_into()
+            .map_err(|_| anyhow!("incorrect pcr0 size"))?;
+        pcr1 = hex::decode(args.pcr1.unwrap())
+            .context("failed to decode pcr1")?
+            .try_into()
+            .map_err(|_| anyhow!("incorrect pcr1 size"))?;
+        pcr2 = hex::decode(args.pcr2.unwrap())
+            .context("failed to decode pcr2")?
+            .try_into()
+            .map_err(|_| anyhow!("incorrect pcr2 size"))?;
+        user_data = hex::decode(args.user_data.unwrap())
+            .context("failed to decode user data")?
+            .into_boxed_slice();
+    }
 
     let auth_store = AuthStore {
         state: ([pcr0, pcr1, pcr2], user_data),
     };
 
-    let seed = fetch_seed(auther, auth_store, secret, args.kms_endpoint)
+    let contract_address = if let Some(contract_address_file) = args.contract_address_file {
+        let address = read(contract_address_file)
+            .await
+            .context("failed to read contract address file")?;
+        Some(
+            String::from_utf8(address)
+                .context("failed to parse contract address")?
+                .trim()
+                .to_string(),
+        )
+    } else {
+        None
+    };
+
+    let seed = fetch_seed(auther, auth_store, secret, kms_endpoint, contract_address)
         .await
         .context("failed to fetch seed")?;
 
diff --git a/kms/root-server-contract/src/derive.rs b/kms/root-server-contract/src/derive.rs
index 27d96c176..9a74c78f0 100644
--- a/kms/root-server-contract/src/derive.rs
+++ b/kms/root-server-contract/src/derive.rs
@@ -1,4 +1,9 @@
-use alloy::{providers::ProviderBuilder, sol};
+use alloy::{
+    primitives::B256,
+    providers::ProviderBuilder,
+    signers::k256::sha2::{Digest, Sha256},
+    sol,
+};
 use axum::{
     extract::{ConnectInfo, Query, State},
     http::StatusCode,
@@ -13,7 +18,7 @@ use crate::{scallop::AuthStoreState, AppState};
 sol!(
     #[sol(rpc)]
     interface IKMSVerifiable {
-        function oysterKMSVerify(bytes32 _key) external returns (bool);
+        function oysterKMSVerify(bytes32 _imageId) external returns (bool);
     }
 );
 
@@ -37,12 +42,17 @@ pub async fn derive(
     let provider = ProviderBuilder::new().on_http(rpc);
     let contract = IKMSVerifiableInstance::new(address, provider);
 
-    // SAFETY: transport should always have key associated, safe to unwrap
-    let Ok(res) = contract
-        .oysterKMSVerify(scallop_state.0.unwrap().into())
-        .call()
-        .await
-    else {
+    // SAFETY: transport should always have PCRs & User Data associated, safe to unwrap
+    let (pcrs, user_data) = scallop_state.0.unwrap();
+    let mut hasher = Sha256::new();
+    hasher.update(pcrs[0]);
+    hasher.update(pcrs[1]);
+    hasher.update(pcrs[2]);
+    hasher.update((user_data.len() as u16).to_be_bytes());
+    hasher.update(user_data);
+    let image_id: [u8; 32] = hasher.finalize().into();
+
+    let Ok(res) = contract.oysterKMSVerify(B256::from(image_id)).call().await else {
         return (StatusCode::INTERNAL_SERVER_ERROR, [0; 64]);
     };
 
diff --git a/kms/root-server-contract/src/scallop.rs b/kms/root-server-contract/src/scallop.rs
index 9c9d377ff..b51ec6125 100644
--- a/kms/root-server-contract/src/scallop.rs
+++ b/kms/root-server-contract/src/scallop.rs
@@ -12,7 +12,7 @@ use oyster::{
 #[derive(Clone, Default)]
 pub struct AuthStore {}
 
-pub type AuthStoreState = Key;
+pub type AuthStoreState = ([[u8; 48]; 3], Box<[u8]>);
 
 impl ScallopAuthStore for AuthStore {
     type State = AuthStoreState;
@@ -44,7 +44,7 @@ impl ScallopAuthStore for AuthStore {
             return None;
         }
 
-        return Some(key);
+        return Some((decoded.pcrs, decoded.user_data));
     }
 }
 


Reviews:
user: login: roshanrags
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-19T06:20:17Z
user: login: vg-27
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-19T07:02:17Z
user: login: roshanrags
state: APPROVED
author_association: MEMBER
submitted_at: 2025-03-19T08:01:00Z

Review comments:
diff_hunk: @@ -37,12 +42,17 @@ pub async fn derive(
     let provider = ProviderBuilder::new().on_http(rpc);
     let contract = IKMSVerifiableInstance::new(address, provider);
 
+    let (pcrs, user_data) = scallop_state.0.unwrap();
+    let mut hasher = Sha256::new();
+    hasher.update(pcrs[0]);
+    hasher.update(pcrs[1]);
+    hasher.update(pcrs[2]);
+    hasher.update((user_data.len() as u16).to_be_bytes());
+    hasher.update(user_data);
+    let image_id: [u8; 32] = hasher.finalize().into();
+
     // SAFETY: transport should always have key associated, safe to unwrap
path: kms/root-server-contract/src/derive.rs
user: login: roshanrags
body: move the comment to the unwrap location
created_at: 2025-03-19T05:54:08Z
author_association: MEMBER
diff_hunk: @@ -60,6 +77,21 @@ impl InitParamsArgs {
             .unwrap_or_default();
         init_params.append(&mut self.init_params.unwrap_or_default());
 
+        if let Some(address) = self.contract_address {
+            if let Err(_) = Address::from_hex(&address) {
+                bail!("invalid contract address");
+            }
+            init_params.push(format!("contract-address:1:0:utf8:{}", address));
path: cli/oyster-cvm/src/args/init_params.rs
user: login: roshanrags
body: might want to put this right after the compose file instead
created_at: 2025-03-19T06:05:37Z
author_association: MEMBER
diff_hunk: @@ -266,3 +298,23 @@ fn fetch_encryption_key(
         .try_into()
         .context("failed to parse reponse")
 }
+
+// TODO fill store KMS details
+lazy_static! {
+    static ref KMS_ROOT_SERVERS: HashMap<u64, &'static str> = {
+        let mut root_servers = HashMap::new();
+        root_servers.insert(
+            31337,
path: cli/oyster-cvm/src/args/init_params.rs
user: login: roshanrags
body: hardhat chain id?
created_at: 2025-03-19T06:18:33Z
author_association: MEMBER
diff_hunk: @@ -266,3 +298,23 @@ fn fetch_encryption_key(
         .try_into()
         .context("failed to parse reponse")
 }
+
+// TODO fill store KMS details
+lazy_static! {
+    static ref KMS_ROOT_SERVERS: HashMap<u64, &'static str> = {
+        let mut root_servers = HashMap::new();
+        root_servers.insert(
+            31337,
path: cli/oyster-cvm/src/args/init_params.rs
user: login: vg-27
body: put there temporarily, will fill it with actual values after the KMS deployment.
created_at: 2025-03-19T07:02:17Z
author_association: CONTRIBUTOR

PR comments:
user: login: roshanrags
created_at: 2025-03-18T12:59:37Z
body: Do we want to include it in the blue image itself or make a new base image
author_association: MEMBER
user: login: vg-27
created_at: 2025-03-18T13:11:27Z
body: > Do we want to include it in the blue image itself or make a new base image

We are starting contract based derive-server conditionally, only if certain input parameters are present. This will save us managing a separate image.
author_association: CONTRIBUTOR