Skip to content

Commit

Permalink
fix: TSS Phase Three and Phase Four Job Submission Validation (#451)
Browse files Browse the repository at this point in the history 
* fix: TSS Phase Three and Phase Four Job Submission Validation

* add tests
  • Loading branch information
@shekohex
shekohex authored Jan 29, 2024
1 parent 0385b12 commit 420c8eb
Show file tree
Hide file tree
Showing 2 changed files with 239 additions and 26 deletions.
62 changes: 40 additions & 22 deletions pallets/jobs/src/functions.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -385,21 +385,33 @@ impl<T: Config> Pallet<T> {
job_info: &JobInfoOf<T>,
info: DKGTSSKeyRefreshResult,
) -> Result<PhaseResultOf<T>, DispatchError> {
let now = <frame_system::Pallet<T>>::block_number();
// sanity check, does job and result type match
ensure!(role_type.is_dkg_tss(), Error::<T>::ResultNotExpectedType);

// ensure the participants are the expected participants from job
let participants = job_info
let existing_result_id = job_info
.job_type
.clone()
.get_participants()
.ok_or(Error::<T>::InvalidJobParams)?;
let mut participant_keys: Vec<Vec<u8>> = Default::default();
.get_phase_one_id()
.ok_or(Error::<T>::InvalidJobPhase)?;
// Ensure the result exists
let phase_one_result =
KnownResults::<T>::get(job_info.job_type.get_role_type(), existing_result_id)
.ok_or(Error::<T>::PreviousResultNotFound)?;

for participant in participants.clone() {
// Validate existing result
ensure!(phase_one_result.ttl >= now, Error::<T>::ResultExpired);

// ensure the participants are the expected participants from job
let mut participant_keys: Vec<sp_core::ecdsa::Public> = Default::default();

let participants = phase_one_result.participants().ok_or(Error::<T>::InvalidJobPhase)?;
for participant in participants {
let key = T::RolesHandler::get_validator_role_key(participant);
ensure!(key.is_some(), Error::<T>::ValidatorRoleKeyNotFound);
participant_keys.push(key.expect("checked above"));
let pub_key = sp_core::ecdsa::Public::from_slice(&key.expect("checked above")[0..33])
.map_err(|_| Error::<T>::InvalidValidator)?;
participant_keys.push(pub_key);
}

let job_result = JobResult::DKGPhaseThree(DKGTSSKeyRefreshResult {
Expand Down Expand Up @@ -432,30 +444,36 @@ impl<T: Config> Pallet<T> {
job_info: &JobInfoOf<T>,
info: DKGTSSKeyRotationResult,
) -> Result<PhaseResultOf<T>, DispatchError> {
let now = <frame_system::Pallet<T>>::block_number();
// sanity check, does job and result type match
ensure!(role_type.is_dkg_tss(), Error::<T>::ResultNotExpectedType);

// ensure the participants are the expected participants from job
let participants = job_info
let existing_result_id = job_info
.job_type
.clone()
.get_participants()
.ok_or(Error::<T>::InvalidJobParams)?;
let mut participant_keys: Vec<Vec<u8>> = Default::default();
.get_phase_one_id()
.ok_or(Error::<T>::InvalidJobPhase)?;
// Ensure the result exists
let phase_one_result =
KnownResults::<T>::get(job_info.job_type.get_role_type(), existing_result_id)
.ok_or(Error::<T>::PreviousResultNotFound)?;

for participant in participants.clone() {
// Validate existing result
ensure!(phase_one_result.ttl >= now, Error::<T>::ResultExpired);

// ensure the participants are the expected participants from job
let mut participant_keys: Vec<sp_core::ecdsa::Public> = Default::default();

let participants = phase_one_result.participants().ok_or(Error::<T>::InvalidJobPhase)?;
for participant in participants {
let key = T::RolesHandler::get_validator_role_key(participant);
ensure!(key.is_some(), Error::<T>::ValidatorRoleKeyNotFound);
participant_keys.push(key.expect("checked above"));
let pub_key = sp_core::ecdsa::Public::from_slice(&key.expect("checked above")[0..33])
.map_err(|_| Error::<T>::InvalidValidator)?;
participant_keys.push(pub_key);
}

let phase_one_job_info = KnownResults::<T>::get(
job_info.job_type.get_role_type(),
job_info.job_type.get_phase_one_id().ok_or(Error::<T>::InvalidJobPhase)?,
)
.ok_or(Error::<T>::JobNotFound)?;

let curr_key = match phase_one_job_info.result {
let curr_key = match phase_one_result.result {
JobResult::DKGPhaseOne(info) => info.key.clone(),
_ => return Err(Error::<T>::InvalidJobPhase.into()),
};
Expand Down Expand Up @@ -487,7 +505,7 @@ impl<T: Config> Pallet<T> {

T::MPCHandler::verify(JobWithResult {
job_type: job_info.job_type.clone(),
phase_one_job_type: Some(phase_one_job_info.job_type),
phase_one_job_type: Some(phase_one_result.job_type),
result: job_result.clone(),
})?;

Expand Down
203 changes: 199 additions & 4 deletions pallets/jobs/src/tests.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,10 +21,11 @@ use mock::*;
use pallet_roles::profile::{IndependentRestakeProfile, Profile, Record, SharedRestakeProfile};
use tangle_primitives::{
jobs::{
DKGTSSPhaseOneJobType, DKGTSSPhaseTwoJobType, DKGTSSSignatureResult, DigitalSignatureType,
Groth16ProveRequest, Groth16System, HyperData, JobSubmission, JobType, RpcResponseJobsData,
ZkSaaSCircuitResult, ZkSaaSPhaseOneJobType, ZkSaaSPhaseTwoJobType, ZkSaaSPhaseTwoRequest,
ZkSaaSSystem,
DKGTSSKeyRefreshResult, DKGTSSKeyRotationResult, DKGTSSPhaseFourJobType,
DKGTSSPhaseOneJobType, DKGTSSPhaseThreeJobType, DKGTSSPhaseTwoJobType,
DKGTSSSignatureResult, DigitalSignatureType, Groth16ProveRequest, Groth16System, HyperData,
JobSubmission, JobType, RpcResponseJobsData, ZkSaaSCircuitResult, ZkSaaSPhaseOneJobType,
ZkSaaSPhaseTwoJobType, ZkSaaSPhaseTwoRequest, ZkSaaSSystem,
},
roles::{RoleType, ThresholdSignatureRoleType, ZeroKnowledgeRoleType},
};
Expand Down Expand Up @@ -240,6 +241,200 @@ fn jobs_submission_e2e_works_for_dkg() {
});
}

#[test]
fn jobs_submission_e2e_for_dkg_refresh() {
new_test_ext(vec![ALICE, BOB, CHARLIE, DAVE, EVE]).execute_with(|| {
System::set_block_number(1);

let threshold_signature_role_type = ThresholdSignatureRoleType::ZengoGG20Secp256k1;
// all validators sign up in roles pallet
let profile = shared_profile();
for validator in [ALICE, BOB, CHARLIE, DAVE, EVE] {
assert_ok!(Roles::create_profile(
RuntimeOrigin::signed(mock_pub_key(validator)),
profile.clone()
));
}

Balances::make_free_balance_be(&mock_pub_key(TEN), 100);

let submission = JobSubmission {
expiry: 10,
ttl: 200,
job_type: JobType::DKGTSSPhaseOne(DKGTSSPhaseOneJobType {
participants: [ALICE, BOB, CHARLIE, DAVE, EVE]
.iter()
.map(|x| mock_pub_key(*x))
.collect(),
threshold: 3,
permitted_caller: Some(mock_pub_key(TEN)),
role_type: threshold_signature_role_type,
}),
};
assert_ok!(Jobs::submit_job(RuntimeOrigin::signed(mock_pub_key(TEN)), submission));

assert_eq!(Balances::free_balance(mock_pub_key(TEN)), 100 - 5);

// submit a solution for this job
assert_ok!(Jobs::submit_job_result(
RuntimeOrigin::signed(mock_pub_key(TEN)),
RoleType::Tss(ThresholdSignatureRoleType::ZengoGG20Secp256k1),
0,
JobResult::DKGPhaseOne(DKGTSSKeySubmissionResult {
signatures: vec![],
threshold: 3,
participants: vec![],
key: vec![],
signature_type: DigitalSignatureType::Ecdsa
})
));

// ---- use phase one solution in phase 3 key refresh -------

let submission = JobSubmission {
expiry: 10,
ttl: 0,
job_type: JobType::DKGTSSPhaseThree(DKGTSSPhaseThreeJobType {
phase_one_id: 0,
role_type: threshold_signature_role_type,
}),
};
assert_ok!(Jobs::submit_job(RuntimeOrigin::signed(mock_pub_key(TEN)), submission));

assert_eq!(Balances::free_balance(mock_pub_key(TEN)), 100 - 25);

// submit a solution for this job
assert_ok!(Jobs::submit_job_result(
RuntimeOrigin::signed(mock_pub_key(TEN)),
RoleType::Tss(threshold_signature_role_type),
1,
JobResult::DKGPhaseThree(DKGTSSKeyRefreshResult {
signature_type: DigitalSignatureType::Ecdsa
})
));

// ensure the job reward is distributed correctly
for validator in [ALICE, BOB, CHARLIE, DAVE, EVE].iter().map(|x| mock_pub_key(*x)) {
assert_eq!(ValidatorRewards::<Runtime>::get(validator), Some(5));
}
});
}

#[test]
fn jobs_submission_e2e_for_dkg_rotation() {
new_test_ext(vec![ALICE, BOB, CHARLIE, DAVE, EVE]).execute_with(|| {
System::set_block_number(1);

let threshold_signature_role_type = ThresholdSignatureRoleType::ZengoGG20Secp256k1;
// all validators sign up in roles pallet
let profile = shared_profile();
for validator in [ALICE, BOB, CHARLIE, DAVE, EVE] {
assert_ok!(Roles::create_profile(
RuntimeOrigin::signed(mock_pub_key(validator)),
profile.clone()
));
}

Balances::make_free_balance_be(&mock_pub_key(TEN), 100);

let submission = JobSubmission {
expiry: 10,
ttl: 200,
job_type: JobType::DKGTSSPhaseOne(DKGTSSPhaseOneJobType {
participants: [ALICE, BOB, CHARLIE, DAVE, EVE]
.iter()
.map(|x| mock_pub_key(*x))
.collect(),
threshold: 3,
permitted_caller: Some(mock_pub_key(TEN)),
role_type: threshold_signature_role_type,
}),
};
assert_ok!(Jobs::submit_job(RuntimeOrigin::signed(mock_pub_key(TEN)), submission));

assert_eq!(Balances::free_balance(mock_pub_key(TEN)), 100 - 5);

let submission = JobSubmission {
expiry: 10,
ttl: 200,
job_type: JobType::DKGTSSPhaseOne(DKGTSSPhaseOneJobType {
participants: [ALICE, BOB, CHARLIE, DAVE, EVE]
.iter()
.map(|x| mock_pub_key(*x))
.collect(),
threshold: 3,
permitted_caller: Some(mock_pub_key(TEN)),
role_type: threshold_signature_role_type,
}),
};
assert_ok!(Jobs::submit_job(RuntimeOrigin::signed(mock_pub_key(TEN)), submission));

assert_eq!(Balances::free_balance(mock_pub_key(TEN)), 100 - 10);
// submit a solution for this job
assert_ok!(Jobs::submit_job_result(
RuntimeOrigin::signed(mock_pub_key(TEN)),
RoleType::Tss(ThresholdSignatureRoleType::ZengoGG20Secp256k1),
0,
JobResult::DKGPhaseOne(DKGTSSKeySubmissionResult {
signatures: vec![],
threshold: 3,
participants: vec![],
key: vec![],
signature_type: DigitalSignatureType::Ecdsa
})
));

// submit a solution for this job
assert_ok!(Jobs::submit_job_result(
RuntimeOrigin::signed(mock_pub_key(TEN)),
RoleType::Tss(ThresholdSignatureRoleType::ZengoGG20Secp256k1),
1,
JobResult::DKGPhaseOne(DKGTSSKeySubmissionResult {
signatures: vec![],
threshold: 3,
participants: vec![],
key: vec![],
signature_type: DigitalSignatureType::Ecdsa
})
));

// ---- use phase one solution in phase 4 key rotation -------

let submission = JobSubmission {
expiry: 10,
ttl: 0,
job_type: JobType::DKGTSSPhaseFour(DKGTSSPhaseFourJobType {
phase_one_id: 0,
new_phase_one_id: 1,
role_type: threshold_signature_role_type,
}),
};
assert_ok!(Jobs::submit_job(RuntimeOrigin::signed(mock_pub_key(TEN)), submission));

assert_eq!(Balances::free_balance(mock_pub_key(TEN)), 100 - 30);

// submit a solution for this job
assert_ok!(Jobs::submit_job_result(
RuntimeOrigin::signed(mock_pub_key(TEN)),
RoleType::Tss(threshold_signature_role_type),
2,
JobResult::DKGPhaseFour(DKGTSSKeyRotationResult {
key: vec![],
new_key: vec![],
signature: vec![],
phase_one_id: 0,
new_phase_one_id: 1,
signature_type: DigitalSignatureType::Ecdsa
})
));

// ensure the job reward is distributed correctly
for validator in [ALICE, BOB, CHARLIE, DAVE, EVE].iter().map(|x| mock_pub_key(*x)) {
assert_eq!(ValidatorRewards::<Runtime>::get(validator), Some(6));
}
});
}

#[test]
fn jobs_rpc_tests() {
new_test_ext(vec![ALICE, BOB, CHARLIE, DAVE, EVE]).execute_with(|| {
Expand Down

0 comments on commit 420c8eb

Please sign in to comment.