Skip to content

Commit

Permalink
Fix staking pallet benchmarks in always generate a mock quote before …
Browse files Browse the repository at this point in the history
…validate
  • Loading branch information
ameba23 committed Oct 28, 2024
1 parent 27d64eb commit 06d3028
Showing 1 changed file with 50 additions and 44 deletions.
94 changes: 50 additions & 44 deletions pallets/staking/src/benchmarking.rs
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,46 @@ pub fn create_validators<T: Config>(
validators
}

/// Sets up a mock quote and requests an attestation in preparation for calling the `validate`
/// extrinsic
fn prepare_attestation_for_validate<T: Config>(
threshold: T::AccountId,
x25519_public_key: [u8; 32],
endpoint: Vec<u8>,
block_number: u32,
) -> (Vec<u8>, JoiningServerInfo<T::AccountId>) {
let nonce = NULL_ARR;
let quote = {
let pck = signing_key_from_seed(NULL_ARR);
/// This is a randomly generated secret p256 ECDSA key - for mocking attestation
const ATTESTATION_KEY: [u8; 32] = [
167, 184, 203, 130, 240, 249, 191, 129, 206, 9, 200, 29, 99, 197, 64, 81, 135, 166, 59,
73, 31, 27, 206, 207, 69, 248, 56, 195, 64, 92, 109, 46,
];

let attestation_key = tdx_quote::SigningKey::from_bytes(&ATTESTATION_KEY.into()).unwrap();

let input_data =
entropy_shared::QuoteInputData::new(&threshold, x25519_public_key, nonce, block_number);

tdx_quote::Quote::mock(attestation_key.clone(), pck, input_data.0).as_bytes().to_vec()
};

let joining_server_info = JoiningServerInfo {
tss_account: threshold.clone(),
x25519_public_key,
endpoint,
// Since we are using the mock PckCertChainVerifier, this needs to be the same seed for
// generating the PCK as we used to sign the quote above
pck_certificate_chain: vec![NULL_ARR.to_vec()],
};

// We need to tell the attestation handler that we want a quote. This will let the system to
// know to expect one back when we call `validate()`.
T::AttestationHandler::request_quote(&threshold, nonce);
(quote, joining_server_info)
}

fn prep_bond_and_validate<T: Config>(
validate_also: bool,
caller: T::AccountId,
Expand All @@ -92,17 +132,14 @@ fn prep_bond_and_validate<T: Config>(
));

if validate_also {
let joining_server_info = JoiningServerInfo {
tss_account: threshold,
let block_number = 0;
let endpoint = vec![20, 20];
let (quote, joining_server_info) = prepare_attestation_for_validate::<T>(
threshold,
x25519_public_key,
endpoint: vec![20, 20],
pck_certificate_chain: vec![[0u8; 32].to_vec()],
};

// Note: This isn't a valid quote, but for testing benches this will pass.
//
// For actually running benches a valid quote will be required in the future.
let quote = [0; 32].to_vec();
endpoint,
block_number,
);

assert_ok!(<Staking<T>>::validate(
RawOrigin::Signed(bonder.clone()).into(),
Expand Down Expand Up @@ -280,7 +317,6 @@ benchmarks! {
.or(Err(Error::<T>::InvalidValidatorId))
.unwrap();

let block_number = 1;
let nonce = NULL_ARR;
let x25519_public_key = NULL_ARR;
let endpoint = b"http://localhost:3001".to_vec();
Expand All @@ -294,39 +330,9 @@ benchmarks! {
x25519_public_key.clone()
);

let quote = {
let pck = signing_key_from_seed(NULL_ARR);
/// This is a randomly generated secret p256 ECDSA key - for mocking attestation
const ATTESTATION_KEY: [u8; 32] = [
167, 184, 203, 130, 240, 249, 191, 129, 206, 9, 200, 29, 99, 197, 64, 81, 135, 166, 59, 73, 31,
27, 206, 207, 69, 248, 56, 195, 64, 92, 109, 46,
];

let attestation_key = tdx_quote::SigningKey::from_bytes(&ATTESTATION_KEY.into()).unwrap();

let input_data = entropy_shared::QuoteInputData::new(
&threshold_account,
x25519_public_key,
nonce,
block_number,
);

tdx_quote::Quote::mock(attestation_key.clone(), pck, input_data.0).as_bytes().to_vec()
};

let joining_server_info = JoiningServerInfo {
tss_account: threshold_account.clone(),
x25519_public_key,
endpoint: endpoint.clone(),
// Since we are using the mock PckCertChainVerifier, this needs to be the same seed for
// generating the PCK as we used to sign the quote above
pck_certificate_chain: vec![NULL_ARR.to_vec()],
};

// We need to tell the attestation handler that we want a quote. This will let the system to
// know to expect one back when we call `validate()`.
T::AttestationHandler::request_quote(&threshold_account, nonce);

let block_number = 1;
let (quote, joining_server_info) =
prepare_attestation_for_validate::<T>(threshold_account.clone(), x25519_public_key, endpoint.clone(), block_number);
}: _(RawOrigin::Signed(bonder.clone()), ValidatorPrefs::default(), joining_server_info, quote)
verify {
assert_last_event::<T>(
Expand Down

0 comments on commit 06d3028

Please sign in to comment.