Merge pull request #119 from Dstack-TEE/sign-pubkey

kms: Sign env encrypt pubkey

Author: kvinwang

kms/rpc/proto/kms_rpc.proto

@@ -15,6 +15,7 @@ message AppId {
 
 message PublicKeyResponse {
   bytes public_key = 1;
+  bytes signature = 2;
 }
 
 message AppKeyResponse {

kms/src/crypto.rs

@@ -1,13 +1,13 @@
 use anyhow::{Context, Result};
-use k256::ecdsa::{RecoveryId, Signature, SigningKey};
+use k256::ecdsa::SigningKey;
 use sha3::{Digest, Keccak256};
 
 use ra_tls::kdf;
 
 pub(crate) fn derive_k256_key(
     parent_key: &SigningKey,
     app_id: &[u8],
-) -> Result<(SigningKey, Signature, RecoveryId)> {
+) -> Result<(SigningKey, Vec<u8>)> {
     let context_data = [app_id, b"app-key"];
     let derived_key_bytes: [u8; 32] =
         kdf::derive_ecdsa_key(&parent_key.to_bytes(), &context_data, 32)?
@@ -16,10 +16,25 @@ pub(crate) fn derive_k256_key(
             .context("Invalid derived key len")?;
     let derived_signing_key = SigningKey::from_bytes(&derived_key_bytes.into())?;
     let pubkey = derived_signing_key.verifying_key();
-    let digest = Keccak256::new_with_prefix(
-        [b"dstack-kms-issued:", app_id, &pubkey.to_sec1_bytes()].concat(),
-    );
-    let (signature, recid) = parent_key.sign_digest_recoverable(digest)?;
 
-    Ok((derived_signing_key, signature, recid))
+    let signature = sign_message(
+        parent_key,
+        b"dstack-kms-issued",
+        app_id,
+        &pubkey.to_sec1_bytes(),
+    )?;
+    Ok((derived_signing_key, signature))
+}
+
+pub(crate) fn sign_message(
+    key: &SigningKey,
+    prefix: &[u8],
+    appid: &[u8],
+    message: &[u8],
+) -> Result<Vec<u8>> {
+    let digest = Keccak256::new_with_prefix([prefix, b":", appid, message].concat());
+    let (signature, recid) = key.sign_digest_recoverable(digest)?;
+    let mut signature_bytes = signature.to_vec();
+    signature_bytes.push(recid.to_byte());
+    Ok(signature_bytes)
 }

kms/src/main_service.rs

@@ -17,7 +17,10 @@ use ra_tls::{
 use scale::Decode;
 use upgrade_authority::BootInfo;
 
-use crate::{config::KmsConfig, crypto::derive_k256_key};
+use crate::{
+    config::KmsConfig,
+    crypto::{derive_k256_key, sign_message},
+};
 
 mod upgrade_authority;
 
@@ -184,11 +187,8 @@ impl KmsRpc for RpcHandler {
         };
 
         let (k256_key, k256_signature) = {
-            let (k256_app_key, signature, recid) = derive_k256_key(&self.state.k256_key, &app_id)
+            let (k256_app_key, signature) = derive_k256_key(&self.state.k256_key, &app_id)
                 .context("Failed to derive app ecdsa key")?;
-
-            let mut signature = signature.to_vec();
-            signature.push(recid.to_byte());
             (k256_app_key.to_bytes().to_vec(), signature)
         };
 
@@ -210,8 +210,19 @@ impl KmsRpc for RpcHandler {
         .context("Failed to derive env encrypt key")?;
         let secret = x25519_dalek::StaticSecret::from(secret);
         let pubkey = x25519_dalek::PublicKey::from(&secret);
+
+        let public_key = pubkey.to_bytes().to_vec();
+        let signature = sign_message(
+            &self.state.k256_key,
+            b"dstack-env-encrypt-pubkey",
+            &request.app_id,
+            &public_key,
+        )
+        .context("Failed to sign the public key")?;
+
         Ok(PublicKeyResponse {
-            public_key: pubkey.to_bytes().to_vec(),
+            public_key,
+            signature,
         })
     }
 

teepod/rpc/proto/teepod_rpc.proto

@@ -114,6 +114,7 @@ message AppId {
 
 message PublicKeyResponse {
   bytes public_key = 1;
+  bytes signature = 2;
 }
 
 message GetInfoResponse {

teepod/src/main_service.rs

@@ -217,6 +217,7 @@ impl TeepodRpc for RpcHandler {
             .await?;
         Ok(PublicKeyResponse {
             public_key: response.public_key,
+            signature: response.signature,
         })
     }