From 06d30287bb11a7c5b1cfbcefe5cb6d67f9e1f9f3 Mon Sep 17 00:00:00 2001
From: peg <ameba23@systemli.org>
Date: Mon, 28 Oct 2024 10:12:10 +0100
Subject: [PATCH] Fix staking pallet benchmarks in always generate a mock quote
 before validate

---
 pallets/staking/src/benchmarking.rs | 94 +++++++++++++++--------------
 1 file changed, 50 insertions(+), 44 deletions(-)

diff --git a/pallets/staking/src/benchmarking.rs b/pallets/staking/src/benchmarking.rs
index 498d0ad1f..f02d58ebe 100644
--- a/pallets/staking/src/benchmarking.rs
+++ b/pallets/staking/src/benchmarking.rs
@@ -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,
@@ -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(),
@@ -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();
@@ -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>(