Fix staking pallet benchmarking

entropyxyz · Dec 16, 2024 · 546b034 · 546b034
1 parent ee23433
commit 546b034
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 7 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/pallets/staking/Cargo.toml b/pallets/staking/Cargo.toml
@@ -34,6 +34,9 @@ pallet-parameters={ version="0.3.0", path="../parameters", default-features=fals
 entropy-shared={ version="0.3.0", path="../../crates/shared", features=[
   "wasm-no-std",
 ], default-features=false }
+tdx-quote={ git="https://github.com/entropyxyz/tdx-quote.git", branch="peg/cert-chain-parse", features=[
+  "mock",
+], optional=true }
 
 [dev-dependencies]
 frame-election-provider-support={ version="29.0.0", default-features=false }
@@ -48,7 +51,7 @@ rand_core="0.6.4"
 
 [features]
 default=['std']
-runtime-benchmarks=['frame-benchmarking']
+runtime-benchmarks=['frame-benchmarking', "tdx-quote"]
 std=[
   "sp-npos-elections/std",
   "sp-staking/std",

diff --git a/pallets/staking/src/benchmarking.rs b/pallets/staking/src/benchmarking.rs
@@ -16,7 +16,6 @@
 //! Benchmarking setup for pallet-propgation
 #![allow(unused_imports)]
 use super::*;
-use crate::pck::{signing_key_from_seed, MOCK_PCK_DERIVED_FROM_NULL_ARRAY};
 #[allow(unused_imports)]
 use crate::Pallet as Staking;
 use entropy_shared::{AttestationHandler, QuoteContext, MAX_SIGNERS};
@@ -33,8 +32,14 @@ use pallet_staking::{
     Event as FrameStakingEvent, MaxNominationsOf, MaxValidatorsCount, Nominations,
     Pallet as FrameStaking, RewardDestination, ValidatorPrefs,
 };
+use rand::{rngs::StdRng, SeedableRng};
 use sp_std::{vec, vec::Vec};
+use tdx_quote::SigningKey;
 
+const MOCK_PCK_DERIVED_FROM_NULL_ARRAY: [u8; 33] = [
+    3, 237, 193, 27, 177, 204, 234, 67, 54, 141, 157, 13, 62, 87, 113, 224, 4, 121, 206, 251, 190,
+    151, 134, 87, 68, 46, 37, 163, 127, 97, 252, 174, 108,
+];
 const NULL_ARR: [u8; 32] = [0; 32];
 const SEED: u32 = 0;
 
@@ -96,8 +101,10 @@ fn prepare_attestation_for_validate<T: Config>(
             nonce,
             QuoteContext::Validate,
         );
-
-        tdx_quote::Quote::mock(attestation_key.clone(), pck, input_data.0).as_bytes().to_vec()
+        let pck_encoded = tdx_quote::encode_verifying_key(pck.verifying_key()).unwrap();
+        tdx_quote::Quote::mock(attestation_key.clone(), pck, input_data.0, pck_encoded.to_vec())
+            .as_bytes()
+            .to_vec()
     };
 
     let joining_server_info = JoiningServerInfo {
@@ -115,6 +122,11 @@ fn prepare_attestation_for_validate<T: Config>(
     (quote, joining_server_info)
 }
 
+fn signing_key_from_seed(input: [u8; 32]) -> SigningKey {
+    let mut pck_seeder = StdRng::from_seed(input);
+    SigningKey::random(&mut pck_seeder)
+}
+
 fn prep_bond_and_validate<T: Config>(
     validate_also: bool,
     caller: T::AccountId,
@@ -230,15 +242,12 @@ benchmarks! {
         endpoint.clone().to_vec(),
     );
 
-    let pck_certificate_chain = joining_server_info.pck_certificate_chain;
-
     let signers = vec![validator_id_signers.clone(); s as usize];
     Signers::<T>::put(signers.clone());
   }:  _(
         RawOrigin::Signed(_bonder.clone()),
         new_threshold.clone(),
         x25519_public_key.clone(),
-        pck_certificate_chain,
         quote
     )
   verify {