Skip to content

Commit

Permalink
fix test
Browse files Browse the repository at this point in the history
  • Loading branch information
@JesseAbram
JesseAbram committed Jan 13, 2025
1 parent 1c6fa39 commit 8b85ae6
Show file tree
Hide file tree
Showing 4 changed files with 108 additions and 73 deletions.
64 changes: 45 additions & 19 deletions crates/client/src/client.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,16 +15,6 @@

//! Simple client for Entropy.
//! Used in integration tests and for the test-cli
pub use crate::{
chain_api::{get_api, get_rpc},
errors::{ClientError, SubstrateError},
};
pub use entropy_protocol::{sign_and_encrypt::EncryptedSignedMessage, KeyParams};
pub use entropy_shared::{HashingAlgorithm, QuoteContext};
use parity_scale_codec::Decode;
use rand::Rng;
pub use synedrion::KeyShare;

use crate::{
chain_api::{
entropy::{
Expand All @@ -36,6 +26,7 @@ use crate::{
pallet_registry::pallet::{ProgramInstance, RegisteredInfo},
pallet_staking::{RewardDestination, ValidatorPrefs},
pallet_staking_extension::pallet::JoiningServerInfo,
sp_arithmetic::per_things::Perbill,
},
},
EntropyConfig,
Expand All @@ -50,6 +41,16 @@ use crate::{
},
Hasher,
};
pub use crate::{
chain_api::{get_api, get_rpc},
errors::{ClientError, SubstrateError},
};
pub use entropy_protocol::{sign_and_encrypt::EncryptedSignedMessage, KeyParams};
pub use entropy_shared::{HashingAlgorithm, QuoteContext};
use parity_scale_codec::Decode;
use rand::Rng;
use std::str::FromStr;
pub use synedrion::KeyShare;

use base64::prelude::{Engine, BASE64_STANDARD};
use entropy_protocol::RecoverableSignature;
Expand Down Expand Up @@ -572,24 +573,49 @@ pub async fn get_quote_and_declare_validate(
api: &OnlineClient<EntropyConfig>,
rpc: &LegacyRpcMethods<EntropyConfig>,
signer: sr25519::Pair,
prefs: ValidatorPrefs,
joining_server_info: JoiningServerInfo<SubxtAccountId32>,
comission: u32,
blocked: bool,
tss_account: String,
x25519_public_key: String,
endpoint: String,
) -> Result<ValidatorCandidateAccepted, ClientError> {
let quote =
get_tdx_quote(std::str::from_utf8(&joining_server_info.endpoint)?, QuoteContext::Validate)
.await?;
declare_validate(api, rpc, signer, prefs, joining_server_info, quote).await
let quote = get_tdx_quote(&endpoint, QuoteContext::Validate).await?;
declare_validate(
api,
rpc,
signer,
comission,
blocked,
tss_account,
x25519_public_key,
endpoint,
quote,
)
.await
}
pub async fn declare_validate(
api: &OnlineClient<EntropyConfig>,
rpc: &LegacyRpcMethods<EntropyConfig>,
signer: sr25519::Pair,
prefs: ValidatorPrefs,
joining_server_info: JoiningServerInfo<SubxtAccountId32>,
comission: u32,
blocked: bool,
tss_account: String,
x25519_public_key: String,
endpoint: String,
quote: Vec<u8>,
) -> Result<ValidatorCandidateAccepted, ClientError> {
let tss_account = SubxtAccountId32::from_str(&tss_account)
.map_err(|e| ClientError::FromSs58(e.to_string()))?;
let x25519_public_key = hex::decode(x25519_public_key)?
.try_into()
.map_err(|_| ClientError::Conversion("Error converting x25519_public_key"))?;
let joining_server_info =
JoiningServerInfo { tss_account, x25519_public_key, endpoint: endpoint.into() };

let validator_prefs = ValidatorPrefs { commission: Perbill(comission), blocked };

let validate_request =
entropy::tx().staking_extension().validate(prefs, joining_server_info, quote);
entropy::tx().staking_extension().validate(validator_prefs, joining_server_info, quote);
let in_block = submit_transaction_with_pair(api, rpc, &signer, &validate_request, None).await?;
let result_event =
in_block.find_first::<ValidatorCandidateAccepted>()?.ok_or(SubstrateError::NoEvent)?;
Expand Down
6 changes: 6 additions & 0 deletions crates/client/src/errors.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -127,4 +127,10 @@ pub enum ClientError {
QuoteGet(String),
#[error("Unable to get info for TSS server from chain")]
NoServerInfo,
#[error("From Hex Error: {0}")]
FromHex(#[from] hex::FromHexError),
#[error("From Ss58 Error: {0}")]
FromSs58(String),
#[error("Vec<u8> Conversion Error: {0}")]
Conversion(&'static str),
}
96 changes: 55 additions & 41 deletions crates/client/src/tests.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,9 +6,7 @@ use crate::{
runtime_types::{
bounded_collections::bounded_vec::BoundedVec,
pallet_registry::pallet::ProgramInstance,
pallet_staking::ValidatorPrefs,
pallet_staking_extension::pallet::{JoiningServerInfo, ServerInfo},
sp_arithmetic::per_things::Perbill,
pallet_staking_extension::pallet::ServerInfo,
},
staking::events as staking_events,
staking_extension::events,
Expand Down Expand Up @@ -340,13 +338,8 @@ async fn test_set_session_key_and_declare_validate() {
let tss_public_key = tss_signer_pair.signer().public();
let endpoint = "test".to_string();

let joining_server_info = JoiningServerInfo {
tss_account,
x25519_public_key: *x25519_public_key.as_bytes(),
endpoint: endpoint.into(),
};

let validator_prefs = ValidatorPrefs { commission: Perbill(0), blocked: false };
let commission = 0;
let blocked = false;

// We need to give our new TSS account some funds before it can request an attestation.
let dest = tss_signer_pair.account_id().clone().into();
Expand All @@ -362,35 +355,56 @@ async fn test_set_session_key_and_declare_validate() {
.await
.unwrap();

// When we request an attestation we get a nonce back that we must use when generating our quote.
let nonce = request_attestation(&api, &rpc, tss_signer_pair.signer()).await.unwrap();
let nonce: [u8; 32] = nonce.try_into().unwrap();

// Our runtime is using the mock `PckCertChainVerifier`, which means that the expected
// "certificate" basically is just our TSS account ID. This account needs to match the one
// used to sign the following `quote`.
let mut pck_seeder = StdRng::from_seed(tss_public_key.0.clone());
let pck = tdx_quote::SigningKey::random(&mut pck_seeder);
let encoded_pck = encode_verifying_key(&pck.verifying_key()).unwrap().to_vec();

let quote = {
let input_data = entropy_shared::QuoteInputData::new(
tss_public_key,
*x25519_public_key.as_bytes(),
nonce,
QuoteContext::Validate,
);

let signing_key = tdx_quote::SigningKey::random(&mut OsRng);

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

let result_declare_validate =
declare_validate(&api, &rpc, one.into(), validator_prefs, joining_server_info, quote).await;
dbg!(&result_declare_validate);
// TODO: better assert_eq statment
assert!(result_declare_validate.is_ok());
// When we request an attestation we get a nonce back that we must use when generating our quote.
let nonce = request_attestation(&api, &rpc, tss_signer_pair.signer()).await.unwrap();
let nonce: [u8; 32] = nonce.try_into().unwrap();

// Our runtime is using the mock `PckCertChainVerifier`, which means that the expected
// "certificate" basically is just our TSS account ID. This account needs to match the one
// used to sign the following `quote`.
let mut pck_seeder = StdRng::from_seed(tss_public_key.0.clone());
let pck = tdx_quote::SigningKey::random(&mut pck_seeder);
let encoded_pck = encode_verifying_key(&pck.verifying_key()).unwrap().to_vec();

let quote = {
let input_data = entropy_shared::QuoteInputData::new(
tss_public_key,
*x25519_public_key.as_bytes(),
nonce,
QuoteContext::Validate,
);

let signing_key = tdx_quote::SigningKey::random(&mut OsRng);

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

let result_declare_validate = declare_validate(
&api,
&rpc,
one.into(),
commission,
blocked,
tss_account.to_string(),
hex::encode(*x25519_public_key.as_bytes()),
endpoint.clone(),
quote,
)
.await
.unwrap();

assert_eq!(
format!("{:?}", result_declare_validate),
format!(
"{:?}",
events::ValidatorCandidateAccepted(
AccountId32(one.pair().public().0),
AccountId32(one.pair().public().0),
AccountId32(tss_public_key.0),
endpoint.as_bytes().to_vec()
)
)
);
}
15 changes: 2 additions & 13 deletions crates/test-cli/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -615,20 +615,9 @@ pub async fn run_command(
let signer = handle_mnemonic(mnemonic_option)?;
cli.log(format!("Account being used for session keys: {}", signer.public()));

// TODO move this into declare validate function for better testing
let tss_account = SubxtAccountId32::from_str(&tss_account)?;
let x25519_public_key = hex::decode(x25519_public_key)?
.try_into()
.map_err(|_| anyhow!("X25519 pub key needs to be 32 bytes"))?;
let joining_server_info =
JoiningServerInfo { tss_account, x25519_public_key, endpoint };

let validator_prefs = ValidatorPrefs {
commission: Perbill(comission),
blocked,
};
let result_event =
get_quote_and_declare_validate(&api, &rpc, signer , joining_server_info, validator_prefs).await?;
get_quote_and_declare_validate(api, rpc, signer, comission, blocked, tss_account, x25519_public_key, endpoint).await

cli.log(format!("Event result: {:?}", result_event));

if cli.json {
Expand Down

0 comments on commit 8b85ae6

Please sign in to comment.