Fixed contract/admin proof validation for migrated contracts

Author: vlad

cosmwasm/enclaves/shared/contract-engine/src/contract_operations.rs

@@ -1,3 +1,4 @@
+use enclave_utils::{Keychain, KEY_MANAGER};
 use serde::{Deserialize, Serialize};
 
 #[cfg(feature = "random")]
@@ -19,7 +20,8 @@ use crate::cosmwasm_config::ContractOperation;
 use crate::contract_validation::verify_block_info;
 
 use crate::contract_validation::{
-    generate_admin_proof, generate_contract_key_proof, ReplyParams, ValidatedMessage,
+    generate_admin_proof, generate_contract_key_proof, validate_admin_proof, ReplyParams,
+    ValidatedMessage,
 };
 use crate::external::results::{
     HandleSuccess, InitSuccess, MigrateSuccess, QuerySuccess, UpdateAdminSuccess,
@@ -240,7 +242,11 @@ pub fn init(
 
     // todo: can move the key to somewhere in the output message if we want
 
-    let admin_proof = generate_admin_proof(&canonical_admin_address.0 .0, &og_contract_key);
+    let admin_proof = generate_admin_proof(
+        &Keychain::generate_admin_proof_secret(KEY_MANAGER.get_consensus_seed().unwrap().last()),
+        &canonical_admin_address.0 .0,
+        &og_contract_key,
+    );
 
     Ok(InitSuccess {
         output,
@@ -328,12 +334,13 @@ pub fn migrate(
     ) {
         debug!("Found hardcoded admin for migrate");
     } else {
-        let sender_admin_proof =
-            generate_admin_proof(&canonical_sender_address.0 .0, &og_contract_key);
-
-        if admin_proof != sender_admin_proof {
+        if let Err(e) = validate_admin_proof(
+            &canonical_sender_address.0 .0,
+            &og_contract_key,
+            admin_proof,
+        ) {
             error!("Failed to validate sender as current admin for migrate");
-            return Err(EnclaveError::ValidationFailure);
+            return Err(e);
         }
         debug!("Validated migrate proof successfully");
     }
@@ -442,6 +449,9 @@ pub fn migrate(
     // todo: can move the key to somewhere in the output message if we want
 
     let new_contract_key_proof = generate_contract_key_proof(
+        &Keychain::generate_contract_key_proof_secret(
+            KEY_MANAGER.get_consensus_seed().unwrap().last(),
+        ),
         &canonical_contract_address.0 .0,
         &contract_code.hash(),
         &og_contract_key,
@@ -495,12 +505,15 @@ pub fn update_admin(
 
     let og_contract_key = base_env.get_og_contract_key()?;
 
-    let sender_admin_proof = generate_admin_proof(&canonical_sender_address.0 .0, &og_contract_key);
-
-    if sender_admin_proof != current_admin_proof {
+    if let Err(e) = validate_admin_proof(
+        &canonical_sender_address.0 .0,
+        &og_contract_key,
+        current_admin_proof,
+    ) {
         error!("Failed to validate sender as current admin for update_admin");
-        return Err(EnclaveError::ValidationFailure);
+        return Err(e);
     }
+
     debug!("Validated update_admin proof successfully");
 
     let parsed_sig_info: SigInfo = extract_sig_info(sig_info)?;
@@ -522,7 +535,11 @@ pub fn update_admin(
         Some(&canonical_new_admin_address),
     )?;
 
-    let new_admin_proof = generate_admin_proof(&canonical_new_admin_address.0 .0, &og_contract_key);
+    let new_admin_proof = generate_admin_proof(
+        &Keychain::generate_admin_proof_secret(KEY_MANAGER.get_consensus_seed().unwrap().last()),
+        &canonical_new_admin_address.0 .0,
+        &og_contract_key,
+    );
 
     debug!("update_admin success: {:?}", new_admin_proof);
 

cosmwasm/enclaves/shared/contract-engine/src/contract_validation.rs

@@ -13,7 +13,7 @@ use enclave_cosmos_types::types::{
 use enclave_crypto::traits::VerifyingKey;
 use enclave_crypto::{sha_256, AESKey, Hmac, Kdf, HASH_SIZE};
 use enclave_ffi_types::EnclaveError;
-use enclave_utils::KEY_MANAGER;
+use enclave_utils::{Keychain, KEY_MANAGER};
 use protobuf::Message;
 
 use crate::hardcoded_admins::is_code_hash_allowed;
@@ -308,19 +308,13 @@ pub fn validate_contract_key(
 
         let sent_contract_key_proof = base_env.get_current_contract_key_proof()?;
 
-        let contract_key_proof = generate_contract_key_proof(
+        validate_contract_key_proof(
             &canonical_contract_address.0 .0,
             &contract_code.hash(),
             &og_contract_key,
-            &current_contract_key, // this is already validated
-        );
-
-        if sent_contract_key_proof != contract_key_proof {
-            error!("Failed to validate contract key proof for a migrated contract");
-            return Err(EnclaveError::ValidationFailure);
-        }
-
-        Ok(())
+            &current_contract_key,
+            &sent_contract_key_proof,
+        )
     } else {
         trace!("Contract still has original code, validating contract_key");
 
@@ -335,17 +329,42 @@ pub fn validate_contract_key(
     }
 }
 
-pub fn generate_admin_proof(admin: &[u8], contract_key: &[u8]) -> [u8; enclave_crypto::HASH_SIZE] {
+pub fn generate_admin_proof(
+    proof_secret: &AESKey,
+    admin: &[u8],
+    contract_key: &[u8],
+) -> [u8; enclave_crypto::HASH_SIZE] {
     let mut data_to_sign = vec![];
     data_to_sign.extend_from_slice(admin);
     data_to_sign.extend_from_slice(contract_key);
 
-    let admin_proof_secret = KEY_MANAGER.get_admin_proof_secret().unwrap();
+    proof_secret.sign_sha_256(data_to_sign.as_slice())
+}
+
+pub fn validate_admin_proof(
+    admin: &[u8],
+    contract_key: &[u8],
+    proof: &[u8],
+) -> Result<(), EnclaveError> {
+    let seeds = KEY_MANAGER.get_consensus_seed().unwrap();
+    for i_seed in 1..seeds.arr.len() {
+        let expected_proof = generate_admin_proof(
+            &Keychain::generate_admin_proof_secret(&seeds.arr[i_seed]),
+            admin,
+            contract_key,
+        );
+
+        if proof == &expected_proof {
+            return Ok(());
+        }
+    }
 
-    admin_proof_secret.sign_sha_256(data_to_sign.as_slice())
+    error!("Failed to validate admin key proof for a migrated contract");
+    Err(EnclaveError::ValidationFailure)
 }
 
 pub fn generate_contract_key_proof(
+    proof_secret: &AESKey,
     contract_address: &[u8],
     code_hash: &[u8],
     og_contract_key: &[u8],
@@ -357,9 +376,33 @@ pub fn generate_contract_key_proof(
     data_to_sign.extend_from_slice(og_contract_key);
     data_to_sign.extend_from_slice(new_contract_key);
 
-    let contract_key_proof_secret = KEY_MANAGER.get_contract_key_proof_secret().unwrap();
+    proof_secret.sign_sha_256(data_to_sign.as_slice())
+}
+
+pub fn validate_contract_key_proof(
+    contract_address: &[u8],
+    code_hash: &[u8],
+    og_contract_key: &[u8],
+    new_contract_key: &[u8],
+    proof: &[u8; enclave_crypto::HASH_SIZE],
+) -> Result<(), EnclaveError> {
+    let seeds = KEY_MANAGER.get_consensus_seed().unwrap();
+    for i_seed in 1..seeds.arr.len() {
+        let expected_proof = generate_contract_key_proof(
+            &Keychain::generate_contract_key_proof_secret(&seeds.arr[i_seed]),
+            contract_address,
+            code_hash,
+            og_contract_key,
+            new_contract_key,
+        );
+
+        if proof == &expected_proof {
+            return Ok(());
+        }
+    }
 
-    contract_key_proof_secret.sign_sha_256(data_to_sign.as_slice())
+    error!("Failed to validate contract key proof for a migrated contract");
+    Err(EnclaveError::ValidationFailure)
 }
 
 pub struct ValidatedMessage {

cosmwasm/enclaves/shared/utils/src/key_manager.rs

@@ -70,8 +70,6 @@ pub struct Keychain {
     pub random_encryption_key: Option<AESKey>,
     pub initial_randomness_seed: Option<AESKey>,
     registration_key: Option<KeyPair>,
-    admin_proof_secret: Option<AESKey>,
-    contract_key_proof_secret: Option<AESKey>,
     pub extra_data: SgxMutex<KeychainMutableData>,
 }
 
@@ -283,8 +281,6 @@ impl Keychain {
             initial_randomness_seed: None,
             //#[cfg(feature = "random")]
             random_encryption_key: None,
-            admin_proof_secret: None,
-            contract_key_proof_secret: None,
             extra_data: SgxMutex::new(KeychainMutableData {
                 height: 0,
                 validator_set_serialized: Vec::new(),
@@ -433,20 +429,6 @@ impl Keychain {
         })
     }
 
-    pub fn get_admin_proof_secret(&self) -> Result<AESKey, CryptoError> {
-        self.admin_proof_secret.ok_or_else(|| {
-            error!("Error accessing admin_proof_secret (does not exist, or was not initialized)");
-            CryptoError::ParsingError
-        })
-    }
-
-    pub fn get_contract_key_proof_secret(&self) -> Result<AESKey, CryptoError> {
-        self.contract_key_proof_secret.ok_or_else(|| {
-            error!("Error accessing contract_key_proof_secret (does not exist, or was not initialized)");
-            CryptoError::ParsingError
-        })
-    }
-
     pub fn delete_consensus_seed(&mut self) {
         self.consensus_seed.arr.clear();
     }
@@ -471,6 +453,14 @@ impl Keychain {
         seed.derive_key_from_this(&RANDOMNESS_ENCRYPTION_KEY_SECRET_DERIVE_ORDER.to_be_bytes())
     }
 
+    pub fn generate_contract_key_proof_secret(seed: &Seed) -> AESKey {
+        seed.derive_key_from_this(&CONTRACT_KEY_PROOF_SECRET_DERIVE_ORDER.to_be_bytes())
+    }
+
+    pub fn generate_admin_proof_secret(seed: &Seed) -> AESKey {
+        seed.derive_key_from_this(&ADMIN_PROOF_SECRET_DERIVE_ORDER.to_be_bytes())
+    }
+
     pub fn generate_consensus_master_keys(&mut self) -> Result<(), EnclaveError> {
         if !self.is_consensus_seed_set() {
             trace!("Seed not initialized, skipping derivation of enclave keys");
@@ -512,18 +502,6 @@ impl Keychain {
                 Some(Self::generate_randomness_seed(self.consensus_seed.last()));
         }
 
-        self.admin_proof_secret = Some(
-            self.consensus_seed
-                .last()
-                .derive_key_from_this(&ADMIN_PROOF_SECRET_DERIVE_ORDER.to_be_bytes()),
-        );
-
-        self.contract_key_proof_secret = Some(
-            self.consensus_seed
-                .last()
-                .derive_key_from_this(&CONTRACT_KEY_PROOF_SECRET_DERIVE_ORDER.to_be_bytes()),
-        );
-
         Ok(())
     }
 }