Use context, not block number in quote input

entropyxyz · Nov 21, 2024 · b0eff14 · b0eff14
1 parent bdc0df8
commit b0eff14
Show file tree

Hide file tree

Showing 6 changed files with 43 additions and 32 deletions.
diff --git a/crates/shared/src/types.rs b/crates/shared/src/types.rs
@@ -123,17 +123,28 @@ impl QuoteInputData {
         tss_account_id: T,
         x25519_public_key: X25519PublicKey,
         nonce: [u8; 32],
-        block_number: u32,
+        context: QuoteContext,
     ) -> Self {
         let mut hasher = Blake2b512::new();
         hasher.update(tss_account_id.encode());
         hasher.update(x25519_public_key);
         hasher.update(nonce);
-        hasher.update(block_number.to_be_bytes());
+        hasher.update(context.encode());
         Self(hasher.finalize().into())
     }
 }
 
+/// An indicator as to the context in which a quote is intended to be used
+#[derive(Clone, Encode, Decode, Debug, Eq, PartialEq)]
+pub enum QuoteContext {
+    /// To be used in the `validate` extrinsic
+    Validate,
+    /// To be used in the `change_endpoint` extrinsic
+    ChangeEndpoint,
+    /// To be used in the `change_threshold_accounts` extrinsic
+    ChangeThresholdAccounts,
+}
+
 /// A trait for types which can handle attestation requests.
 #[cfg(not(feature = "wasm"))]
 pub trait AttestationHandler<AccountId> {
@@ -143,6 +154,7 @@ pub trait AttestationHandler<AccountId> {
         x25519_public_key: X25519PublicKey,
         provisioning_certification_key: BoundedVecEncodedVerifyingKey,
         quote: Vec<u8>,
+        context: QuoteContext,
     ) -> Result<(), sp_runtime::DispatchError>;
 
     /// Indicate to the attestation handler that a quote is desired.
@@ -160,6 +172,7 @@ impl<AccountId> AttestationHandler<AccountId> for () {
         _x25519_public_key: X25519PublicKey,
         _provisioning_certification_key: BoundedVecEncodedVerifyingKey,
         _quote: Vec<u8>,
+        _context: QuoteContext,
     ) -> Result<(), sp_runtime::DispatchError> {
         Ok(())
     }

diff --git a/crates/threshold-signature-server/src/attestation/api.rs b/crates/threshold-signature-server/src/attestation/api.rs
@@ -26,7 +26,7 @@ use crate::{
 use axum::{body::Bytes, extract::State, http::StatusCode};
 use entropy_client::user::request_attestation;
 use entropy_kvdb::kv_manager::KvManager;
-use entropy_shared::OcwMessageAttestationRequest;
+use entropy_shared::{OcwMessageAttestationRequest, QuoteContext};
 use parity_scale_codec::Decode;
 use sp_core::Pair;
 use subxt::tx::PairSigner;
@@ -68,8 +68,8 @@ pub async fn attest(
             .ok_or_else(|| AttestationErr::Unexpected)?
     };
 
-    // We add 1 to the block number as this will be processed in the next block
-    let quote = create_quote(block_number + 1, nonce, &signer, &x25519_secret).await?;
+    let context = QuoteContext::Validate; // TODO
+    let quote = create_quote(nonce, &signer, &x25519_secret, context).await?;
 
     // Submit the quote
     let attest_tx = entropy::tx().attestation().attest(quote.clone());
@@ -93,23 +93,19 @@ pub async fn get_attest(
     // Request attestation to get nonce
     let nonce = request_attestation(&api, &rpc, signer.signer()).await?;
 
-    // We also need the current block number as input
-    let block_number =
-        rpc.chain_get_header(None).await?.ok_or_else(|| AttestationErr::BlockNumber)?.number;
-
-    // We add 1 to the block number as this will be processed in the next block
-    let quote = create_quote(block_number + 1, nonce, &signer, &x25519_secret).await?;
+    let context = QuoteContext::Validate; // TODO
+    let quote = create_quote(nonce, &signer, &x25519_secret, context).await?;
 
     Ok((StatusCode::OK, quote))
 }
 
 /// Create a mock quote for testing on non-TDX hardware
 #[cfg(not(feature = "production"))]
 pub async fn create_quote(
-    block_number: u32,
     nonce: [u8; 32],
     signer: &PairSigner<EntropyConfig, sp_core::sr25519::Pair>,
     x25519_secret: &StaticSecret,
+    context: QuoteContext,
 ) -> Result<Vec<u8>, AttestationErr> {
     use rand::{rngs::StdRng, SeedableRng};
     use rand_core::OsRng;
@@ -124,7 +120,7 @@ pub async fn create_quote(
         signer.signer().public(),
         *public_key.as_bytes(),
         nonce,
-        block_number,
+        context,
     );
 
     // This is generated deterministically from TSS account id
@@ -167,18 +163,18 @@ pub async fn validate_new_attestation(
 /// Create a TDX quote in production
 #[cfg(feature = "production")]
 pub async fn create_quote(
-    block_number: u32,
     nonce: [u8; 32],
     signer: &PairSigner<EntropyConfig, sp_core::sr25519::Pair>,
     x25519_secret: &StaticSecret,
+    context: QuoteContext,
 ) -> Result<Vec<u8>, AttestationErr> {
     let public_key = x25519_dalek::PublicKey::from(x25519_secret);
 
     let input_data = entropy_shared::QuoteInputData::new(
         signer.signer().public(),
         *public_key.as_bytes(),
         nonce,
-        block_number,
+        context,
     );
 
     Ok(configfs_tsm::create_quote(input_data.0)?)

diff --git a/pallets/attestation/src/lib.rs b/pallets/attestation/src/lib.rs
@@ -47,7 +47,7 @@ mod tests;
 
 #[frame_support::pallet]
 pub mod pallet {
-    use entropy_shared::{AttestationHandler, QuoteInputData};
+    use entropy_shared::{AttestationHandler, QuoteContext, QuoteInputData};
     use frame_support::pallet_prelude::*;
     use frame_support::traits::Randomness;
     use frame_system::pallet_prelude::*;
@@ -205,6 +205,7 @@ pub mod pallet {
             x25519_public_key: entropy_shared::X25519PublicKey,
             provisioning_certification_key: entropy_shared::BoundedVecEncodedVerifyingKey,
             quote: Vec<u8>,
+            context: QuoteContext,
         ) -> Result<(), DispatchError> {
             // Check that we were expecting a quote from this validator by getting the associated
             // nonce from PendingAttestations.
@@ -214,15 +215,9 @@ pub mod pallet {
             // Parse the quote (which internally verifies the attestation key signature)
             let quote = Quote::from_bytes(&quote).map_err(|_| Error::<T>::BadQuote)?;
 
-            // Get current block number
-            let block_number: u32 = {
-                let block_number = <frame_system::Pallet<T>>::block_number();
-                BlockNumberFor::<T>::try_into(block_number).unwrap_or_default()
-            };
-
             // Check report input data matches the nonce, TSS details and block number
             let expected_input_data =
-                QuoteInputData::new(attestee, x25519_public_key, nonce, block_number);
+                QuoteInputData::new(attestee, x25519_public_key, nonce, context);
             ensure!(
                 quote.report_input_data() == expected_input_data.0,
                 Error::<T>::IncorrectInputData

diff --git a/pallets/attestation/src/tests.rs b/pallets/attestation/src/tests.rs
@@ -14,7 +14,7 @@
 // along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
 use crate::mock::*;
-use entropy_shared::{AttestationHandler, QuoteInputData};
+use entropy_shared::{AttestationHandler, QuoteContext, QuoteInputData};
 use frame_support::{assert_noop, assert_ok};
 use rand_core::OsRng;
 
@@ -32,13 +32,12 @@ fn verify_quote_works() {
         let pck_encoded = tdx_quote::encode_verifying_key(pck.verifying_key()).unwrap();
 
         let x25519_public_key = [0; 32];
-        let block_number = 0;
 
         let input_data = QuoteInputData::new(
             ATTESTEE, // TSS Account ID
             x25519_public_key,
             nonce,
-            block_number,
+            QuoteContext::Validate,
         );
 
         let quote = tdx_quote::Quote::mock(attestation_key.clone(), pck, input_data.0);
@@ -47,6 +46,7 @@ fn verify_quote_works() {
             x25519_public_key,
             sp_runtime::BoundedVec::try_from(pck_encoded.to_vec()).unwrap(),
             quote.as_bytes().to_vec(),
+            QuoteContext::Validate,
         ));
     })
 }
@@ -63,13 +63,12 @@ fn verify_quote_fails_with_mismatched_input_data() {
         let pck_encoded = tdx_quote::encode_verifying_key(pck.verifying_key()).unwrap();
 
         let x25519_public_key = [0; 32];
-        let block_number = 0;
 
         let input_data = QuoteInputData::new(
             ATTESTEE, // TSS Account ID
             x25519_public_key,
             nonce,
-            block_number,
+            QuoteContext::Validate,
         );
 
         let quote = tdx_quote::Quote::mock(attestation_key.clone(), pck, input_data.0);
@@ -83,6 +82,7 @@ fn verify_quote_fails_with_mismatched_input_data() {
                 x25519_public_key,
                 sp_runtime::BoundedVec::try_from(pck_encoded.to_vec()).unwrap(),
                 quote.as_bytes().to_vec(),
+                QuoteContext::Validate,
             ),
             crate::Error::<Test>::UnexpectedAttestation,
         );
@@ -96,6 +96,7 @@ fn verify_quote_fails_with_mismatched_input_data() {
                 mismatched_x25519_public_key,
                 sp_runtime::BoundedVec::try_from(pck_encoded.to_vec()).unwrap(),
                 quote.as_bytes().to_vec(),
+                QuoteContext::Validate,
             ),
             crate::Error::<Test>::IncorrectInputData,
         );
@@ -116,6 +117,7 @@ fn verify_quote_fails_with_mismatched_input_data() {
                 x25519_public_key,
                 sp_runtime::BoundedVec::try_from(mismatched_pck_encoded.to_vec()).unwrap(),
                 quote.as_bytes().to_vec(),
+                QuoteContext::Validate,
             ),
             crate::Error::<Test>::PckVerification
         );

diff --git a/pallets/staking/src/lib.rs b/pallets/staking/src/lib.rs
@@ -60,7 +60,7 @@ use sp_staking::SessionIndex;
 #[frame_support::pallet]
 pub mod pallet {
     use entropy_shared::{
-        ValidatorInfo, X25519PublicKey, MAX_SIGNERS, TEST_RESHARE_BLOCK_NUMBER,
+        QuoteContext, ValidatorInfo, X25519PublicKey, MAX_SIGNERS, TEST_RESHARE_BLOCK_NUMBER,
         VERIFICATION_KEY_LENGTH,
     };
     use frame_support::{
@@ -417,7 +417,8 @@ pub mod pallet {
                             &server_info.tss_account.clone(),
                             server_info.x25519_public_key,
                             server_info.provisioning_certification_key.clone(),
-                            quote
+                            quote,
+                            QuoteContext::ChangeEndpoint,
                         )
                         .is_ok(),
                         Error::<T>::FailedAttestationCheck
@@ -492,7 +493,8 @@ pub mod pallet {
                                 &tss_account.clone(),
                                 x25519_public_key,
                                 provisioning_certification_key.clone(),
-                                quote
+                                quote,
+                                QuoteContext::ChangeThresholdAccounts,
                             )
                             .is_ok(),
                             Error::<T>::FailedAttestationCheck
@@ -645,7 +647,8 @@ pub mod pallet {
                     &server_info.tss_account.clone(),
                     server_info.x25519_public_key,
                     server_info.provisioning_certification_key.clone(),
-                    quote
+                    quote,
+                    QuoteContext::Validate,
                 )
                 .is_ok(),
                 Error::<T>::FailedAttestationCheck

diff --git a/pallets/staking/src/mock.rs b/pallets/staking/src/mock.rs
@@ -16,6 +16,7 @@
 use core::convert::{TryFrom, TryInto};
 use std::cell::RefCell;
 
+use entropy_shared::QuoteContext;
 use frame_election_provider_support::{
     bounds::{ElectionBounds, ElectionBoundsBuilder},
     onchain, SequentialPhragmen, VoteWeight,
@@ -400,6 +401,7 @@ impl entropy_shared::AttestationHandler<AccountId> for MockAttestationHandler {
         _x25519_public_key: entropy_shared::X25519PublicKey,
         _provisioning_certification_key: entropy_shared::BoundedVecEncodedVerifyingKey,
         quote: Vec<u8>,
+        _context: QuoteContext,
     ) -> Result<(), sp_runtime::DispatchError> {
         let quote: Result<[u8; 32], _> = quote.try_into();
         match quote {