Skip to content

Commit

Permalink
Refactor contract_key validation & fix it in query
Browse files Browse the repository at this point in the history 
Query was not validating contract_key for migrated contracts, and didn't use the correct contract_key for state encryption
  • Loading branch information
@assafmo
assafmo committed Jun 27, 2023
1 parent 3d79d50 commit 585b429
Show file tree
Hide file tree
Showing 2 changed files with 74 additions and 55 deletions.
66 changes: 13 additions & 53 deletions cosmwasm/enclaves/shared/contract-engine/src/contract_operations.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ use crate::cosmwasm_config::ContractOperation;
#[cfg(feature = "light-client-validation")]
use crate::contract_validation::verify_block_info;

use crate::contract_validation::{ReplyParams, ValidatedMessage};
use crate::contract_validation::{generate_contract_key_proof, ReplyParams, ValidatedMessage};
use crate::external::results::{
HandleSuccess, InitSuccess, MigrateSuccess, QuerySuccess, UpdateAdminSuccess,
};
Expand Down Expand Up @@ -69,25 +69,6 @@ fn generate_admin_proof(admin: &[u8], contract_key: &[u8]) -> [u8; enclave_crypt
sha_256(&data_to_hash)
}

fn generate_contract_key_proof(
contract_address: &[u8],
code_hash: &[u8],
prev_contract_key: &[u8],
new_contract_key: &[u8],
) -> [u8; enclave_crypto::HASH_SIZE] {
let mut data_to_hash = vec![];
data_to_hash.extend_from_slice(contract_address);
data_to_hash.extend_from_slice(code_hash);
data_to_hash.extend_from_slice(prev_contract_key);
data_to_hash.extend_from_slice(new_contract_key);

let contract_key_proof_secret = KEY_MANAGER.get_contract_key_proof_secret().unwrap();

data_to_hash.extend_from_slice(contract_key_proof_secret.get());

sha_256(&data_to_hash)
}

#[cfg_attr(feature = "cargo-clippy", allow(clippy::too_many_arguments))]
pub fn init(
context: Ctx, // need to pass this to read_db & write_db
Expand Down Expand Up @@ -530,33 +511,8 @@ pub fn handle(

let canonical_contract_address = to_canonical(contract_address)?;

// contract_key is unique for each contract
// it's used in state encryption to prevent the same
// encryption keys from being used for different contracts
let mut contract_key = base_env.get_current_contract_key()?;
validate_contract_key(&contract_key, &canonical_contract_address, &contract_code)?;

if base_env.was_migrated() {
println!("Contract was migrated, validating proof");

// was_migrated checks that these won't fail
let og_contract_key = base_env.get_original_contract_key()?;
let sent_contract_key_proof = base_env.get_contract_key_proof()?;

let contract_key_proof = generate_contract_key_proof(
&canonical_contract_address.0 .0,
&contract_code.hash(),
&og_contract_key,
&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);
}

contract_key = og_contract_key; // used in engine for state encryption
}
let og_contract_key =
validate_contract_key(&base_env, &canonical_contract_address, &contract_code)?;

let parsed_sig_info: SigInfo = extract_sig_info(sig_info)?;

Expand Down Expand Up @@ -619,7 +575,7 @@ pub fn handle(
context,
gas_limit,
&contract_code,
&contract_key,
&og_contract_key,
ContractOperation::Handle,
query_depth,
secret_msg.nonce,
Expand Down Expand Up @@ -655,7 +611,12 @@ pub fn handle(
}

#[cfg(feature = "random")]
set_random_in_env(block_height, &contract_key, &mut engine, &mut versioned_env);
set_random_in_env(
block_height,
&og_contract_key,
&mut engine,
&mut versioned_env,
);

versioned_env.set_contract_hash(&contract_hash);

Expand Down Expand Up @@ -750,9 +711,8 @@ pub fn query(

let canonical_contract_address = to_canonical(contract_address)?;

let contract_key = base_env.get_current_contract_key()?;

validate_contract_key(&contract_key, &canonical_contract_address, &contract_code)?;
let og_contract_key =
validate_contract_key(&base_env, &canonical_contract_address, &contract_code)?;

let secret_msg = SecretMessage::from_slice(msg)?;
let decrypted_msg = secret_msg.decrypt()?;
Expand All @@ -764,7 +724,7 @@ pub fn query(
context,
gas_limit,
&contract_code,
&contract_key,
&og_contract_key,
ContractOperation::Query,
query_depth,
secret_msg.nonce,
Expand Down
63 changes: 61 additions & 2 deletions cosmwasm/enclaves/shared/contract-engine/src/contract_validation.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@ use cw_types_v1::ibc::IbcPacketReceiveMsg;
use cw_types_v1::results::REPLY_ENCRYPTION_MAGIC_BYTES;
use log::*;

#[cfg(feature = "light-client-validation")]
use cw_types_generic::BaseEnv;

use cw_types_v010::types::{CanonicalAddr, Coin, HumanAddr};
Expand Down Expand Up @@ -154,7 +153,7 @@ pub fn generate_contract_id(
authentication_key.sign_sha_256(&input_data)
}

pub fn validate_contract_key(
pub fn validate_current_contract_key(
contract_key: &[u8; CONTRACT_KEY_LENGTH],
contract_address: &CanonicalAddr,
contract_code: &ContractCode,
Expand Down Expand Up @@ -193,6 +192,66 @@ pub fn validate_contract_key(
}
}

/// validate_contract_key validates the contract key against the contract address and code hash. If the contract was previously migrated, it also validates the contract key proof against the original contract key.
pub fn validate_contract_key(
base_env: &BaseEnv,
canonical_contract_address: &CanonicalAddr,
contract_code: &ContractCode,
) -> Result<[u8; CONTRACT_KEY_LENGTH], EnclaveError> {
let current_contract_key: [u8; CONTRACT_KEY_LENGTH] = base_env.get_current_contract_key()?;

validate_current_contract_key(
&current_contract_key,
canonical_contract_address,
contract_code,
)?;

if base_env.was_migrated() {
println!("Contract was migrated, validating proof");

let og_contract_key: [u8; CONTRACT_KEY_LENGTH] = base_env.get_original_contract_key()?;
let sent_contract_key_proof = base_env.get_contract_key_proof()?;

let contract_key_proof = generate_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);
}

// og_contract_key is always used in engine for state encryption
Ok(og_contract_key)
} else {
// og_contract_key is always used in engine for state encryption
// contract wasn't previously migrated, thus og_contract_key == current_contract_key
Ok(current_contract_key)
}
}

pub fn generate_contract_key_proof(
contract_address: &[u8],
code_hash: &[u8],
prev_contract_key: &[u8],
new_contract_key: &[u8],
) -> [u8; enclave_crypto::HASH_SIZE] {
let mut data_to_hash = vec![];
data_to_hash.extend_from_slice(contract_address);
data_to_hash.extend_from_slice(code_hash);
data_to_hash.extend_from_slice(prev_contract_key);
data_to_hash.extend_from_slice(new_contract_key);

let contract_key_proof_secret = KEY_MANAGER.get_contract_key_proof_secret().unwrap();

data_to_hash.extend_from_slice(contract_key_proof_secret.get());

sha_256(&data_to_hash)
}

pub struct ValidatedMessage {
pub validated_msg: Vec<u8>,
pub reply_params: Option<Vec<ReplyParams>>,
Expand Down

0 comments on commit 585b429

Please sign in to comment.