feat: DKG TSS Phase Three and Phase Four (#450)

pallets/dkg/src/functions.rs

@@ -73,6 +73,8 @@ impl<T: Config> Pallet<T> {
 		match data {
 			JobResult::DKGPhaseOne(info) => Self::verify_generated_dkg_key(info),
 			JobResult::DKGPhaseTwo(info) => Self::verify_dkg_signature(info),
+			JobResult::DKGPhaseThree(_) => Ok(()),
+			JobResult::DKGPhaseFour(info) => Self::verify_dkg_key_rotation(info),
 			_ => Err(Error::<T>::InvalidJobType.into()), // this should never happen
 		}
 	}
@@ -293,6 +295,100 @@ impl<T: Config> Pallet<T> {
 		Ok(())
 	}
 
+	/// Verifies a DKG Key Rotation.
+	///
+	/// The verification process depends on the key type specified in the DKG result.
+	/// It dispatches the verification to the appropriate function for the specified key type (ECDSA
+	/// or Schnorr).
+	///
+	/// # Arguments
+	///
+	/// * `data` - The DKG result containing current key, new key and signature.
+	///
+	/// # Returns
+	///
+	/// Returns a `DispatchResult` indicating whether the key rotation verification was successful
+	/// or encountered an error.
+	fn verify_dkg_key_rotation(data: DKGTSSKeyRotationResult) -> DispatchResult {
+		match data.signature_type {
+			DigitalSignatureType::Ecdsa => Self::verify_dkg_key_rotation_ecdsa(data),
+			DigitalSignatureType::SchnorrSr25519 => Self::verify_dkg_key_rotation_schnorr(data),
+			_ => Err(Error::<T>::InvalidSignature.into()), // unimplemented
+		}
+	}
+
+	/// Verifies the Key Rotation signature result by recovering the ECDSA public key from the
+	/// provided new key and signature.
+	///
+	/// This function checks whether the recovered public key matches the expected current key,
+	/// ensuring the validity of the signature.
+	///
+	/// # Arguments
+	///
+	/// * `data` - The Key Rotation result containing the new key and ECDSA signature.
+	fn verify_dkg_key_rotation_ecdsa(data: DKGTSSKeyRotationResult) -> DispatchResult {
+		// Recover the ECDSA public key from the provided data and signature
+		let recovered_key = Self::recover_ecdsa_pub_key(&data.new_key, &data.signature)
+			.map_err(|_| Error::<T>::InvalidSignature)?;
+
+		// Extract the expected key from the provided signing key
+		let expected_key: Vec<_> = data.key.iter().skip(1).cloned().collect();
+		// The recovered key is 64 bytes uncompressed. The first 32 bytes represent the compressed
+		// portion of the key.
+		let signer = &recovered_key[..32];
+
+		// Ensure that the recovered key matches the expected signing key
+		ensure!(expected_key == signer, Error::<T>::SigningKeyMismatch);
+
+		Self::deposit_event(Event::KeyRotated {
+			from_job_id: data.phase_one_id,
+			to_job_id: data.new_phase_one_id,
+			signature: data.signature,
+		});
+		Ok(())
+	}
+
+	/// Verifies the Key Rotation signature result by recovering the Schnorr public key from the
+	/// provided new key and signature.
+	///
+	/// This function checks whether the recovered public key matches the expected current key,
+	/// ensuring the validity of the signature.
+	///
+	/// # Arguments
+	///
+	/// * `data` - The Key Rotation result containing the new key and Schnorr signature.
+	fn verify_dkg_key_rotation_schnorr(data: DKGTSSKeyRotationResult) -> DispatchResult {
+		// Convert the signature from bytes to sr25519::Signature
+		let signature: sr25519::Signature = data
+			.signature
+			.as_slice()
+			.try_into()
+			.map_err(|_| Error::<T>::CannotRetreiveSigner)?;
+
+		// Encode the message data and compute its keccak256 hash
+		let msg = data.new_key;
+		let hash = keccak_256(&msg);
+
+		// Verify the Schnorr signature using sr25519_verify
+		if !sr25519_verify(
+			&signature,
+			&hash,
+			&sr25519::Public(
+				Self::to_slice_32(&data.key)
+					.unwrap_or_else(|| panic!("Failed to convert input to sr25519 public key")),
+			),
+		) {
+			return Err(Error::<T>::InvalidSignature.into())
+		}
+
+		Self::deposit_event(Event::KeyRotated {
+			from_job_id: data.phase_one_id,
+			to_job_id: data.new_phase_one_id,
+			signature: data.signature,
+		});
+		Ok(())
+	}
+
 	/// Recovers the ECDSA public key from a given message and signature.
 	///
 	/// # Arguments

pallets/dkg/src/lib.rs

@@ -46,6 +46,8 @@ pub mod pallet {
 		traits::{Get, ReservableCurrency},
 	};
 	use frame_system::pallet_prelude::*;
+	use sp_std::prelude::*;
+	use tangle_primitives::jobs::JobId;
 
 	/// Configure the pallet by specifying the parameters and types on which it depends.
 	#[pallet::config]
@@ -75,6 +77,8 @@ pub mod pallet {
 	pub enum Event<T: Config> {
 		/// Fee has been updated to the new value
 		FeeUpdated(FeeInfoOf<T>),
+		/// A DKG has been rotated.
+		KeyRotated { from_job_id: JobId, to_job_id: JobId, signature: Vec<u8> },
 	}
 
 	// Errors inform users that something went wrong.

pallets/dkg/src/tests.rs

@@ -13,13 +13,14 @@
 //
 // You should have received a copy of the GNU General Public License
 // along with Tangle.  If not, see <http://www.gnu.org/licenses/>.
-use crate::{mock::*, types::FeeInfo, Error, FeeInfo as FeeInfoStorage};
+use crate::{mock::*, types::FeeInfo, Error, Event, FeeInfo as FeeInfoStorage};
 use frame_support::{assert_noop, assert_ok};
 use parity_scale_codec::Encode;
 use sp_core::{crypto::ByteArray, ecdsa, keccak_256, sr25519};
 use sp_io::crypto::{ecdsa_generate, ecdsa_sign_prehashed, sr25519_generate, sr25519_sign};
 use tangle_primitives::jobs::{
-	DKGTSSKeySubmissionResult, DKGTSSSignatureResult, DigitalSignatureType, JobResult,
+	DKGTSSKeyRotationResult, DKGTSSKeySubmissionResult, DKGTSSSignatureResult,
+	DigitalSignatureType, JobResult,
 };
 
 fn mock_pub_key_ecdsa() -> ecdsa::Public {
@@ -298,3 +299,49 @@ fn dkg_signature_verifcation_works_schnorr() {
 		assert_ok!(DKG::verify(JobResult::DKGPhaseTwo(job_to_verify)));
 	});
 }
+
+#[test]
+fn dkg_key_rotation_works() {
+	new_test_ext().execute_with(|| {
+		// setup key/signature
+		let curr_key = mock_pub_key_ecdsa();
+		let new_key = mock_pub_key_ecdsa();
+		let invalid_key = mock_pub_key_ecdsa();
+		let signature = mock_signature_ecdsa(invalid_key, new_key);
+
+		let job_to_verify = DKGTSSKeyRotationResult {
+			signature_type: DigitalSignatureType::Ecdsa,
+			signature,
+			key: curr_key.to_raw_vec(),
+			new_key: new_key.to_raw_vec(),
+			phase_one_id: 1,
+			new_phase_one_id: 2,
+		};
+
+		// should fail for invalid keys
+		assert_noop!(
+			DKG::verify(JobResult::DKGPhaseFour(job_to_verify)),
+			Error::<Runtime>::SigningKeyMismatch
+		);
+
+		let signature = mock_signature_ecdsa(curr_key, new_key);
+
+		let job_to_verify = DKGTSSKeyRotationResult {
+			signature_type: DigitalSignatureType::Ecdsa,
+			signature: signature.clone(),
+			key: curr_key.to_raw_vec(),
+			new_key: new_key.to_raw_vec(),
+			phase_one_id: 1,
+			new_phase_one_id: 2,
+		};
+		// should work with correct params
+		assert_ok!(DKG::verify(JobResult::DKGPhaseFour(job_to_verify)));
+		// should emit KeyRotated event
+		assert!(System::events().iter().any(|r| r.event ==
+			RuntimeEvent::DKG(Event::KeyRotated {
+				from_job_id: 1,
+				to_job_id: 2,
+				signature: signature.clone()
+			})));
+	});
+}

pallets/jobs/src/functions.rs

@@ -2,8 +2,9 @@ use super::*;
 use sp_runtime::traits::Zero;
 use tangle_primitives::{
 	jobs::{
-		DKGTSSPhaseOneJobType, DKGTSSSignatureResult, JobType, JobWithResult, ZkSaaSCircuitResult,
-		ZkSaaSPhaseOneJobType, ZkSaaSProofResult,
+		DKGTSSKeyRefreshResult, DKGTSSKeyRotationResult, DKGTSSPhaseOneJobType,
+		DKGTSSSignatureResult, JobType, JobWithResult, ZkSaaSCircuitResult, ZkSaaSPhaseOneJobType,
+		ZkSaaSProofResult,
 	},
 	roles::RoleType,
 };
@@ -379,6 +380,127 @@ impl<T: Config> Pallet<T> {
 		Ok(result)
 	}
 
+	pub fn verify_dkg_key_refresh_job_result(
+		role_type: RoleType,
+		job_info: &JobInfoOf<T>,
+		info: DKGTSSKeyRefreshResult,
+	) -> Result<PhaseResultOf<T>, DispatchError> {
+		// 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
+			.job_type
+			.clone()
+			.get_participants()
+			.ok_or(Error::<T>::InvalidJobParams)?;
+		let mut participant_keys: Vec<Vec<u8>> = Default::default();
+
+		for participant in participants.clone() {
+			let key = T::RolesHandler::get_validator_role_key(participant);
+			ensure!(key.is_some(), Error::<T>::ValidatorRoleKeyNotFound);
+			participant_keys.push(key.expect("checked above"));
+		}
+
+		let job_result = JobResult::DKGPhaseThree(DKGTSSKeyRefreshResult {
+			signature_type: info.signature_type.clone(),
+		});
+
+		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)?;
+		T::MPCHandler::verify(JobWithResult {
+			job_type: job_info.job_type.clone(),
+			phase_one_job_type: Some(phase_one_job_info.job_type),
+			result: job_result.clone(),
+		})?;
+
+		let result = PhaseResult {
+			owner: job_info.owner.clone(),
+			job_type: job_info.job_type.clone(),
+			ttl: job_info.ttl,
+			permitted_caller: job_info.job_type.clone().get_permitted_caller(),
+			result: job_result,
+		};
+		Ok(result)
+	}
+
+	pub fn verify_dkg_key_rotation_job_result(
+		role_type: RoleType,
+		job_info: &JobInfoOf<T>,
+		info: DKGTSSKeyRotationResult,
+	) -> Result<PhaseResultOf<T>, DispatchError> {
+		// 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
+			.job_type
+			.clone()
+			.get_participants()
+			.ok_or(Error::<T>::InvalidJobParams)?;
+		let mut participant_keys: Vec<Vec<u8>> = Default::default();
+
+		for participant in participants.clone() {
+			let key = T::RolesHandler::get_validator_role_key(participant);
+			ensure!(key.is_some(), Error::<T>::ValidatorRoleKeyNotFound);
+			participant_keys.push(key.expect("checked above"));
+		}
+
+		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 {
+			JobResult::DKGPhaseOne(info) => info.key.clone(),
+			_ => return Err(Error::<T>::InvalidJobPhase.into()),
+		};
+
+		let new_phase_one_job_id = match job_info.job_type {
+			JobType::DKGTSSPhaseFour(ref info) => info.new_phase_one_id,
+			_ => return Err(Error::<T>::InvalidJobPhase.into()),
+		};
+
+		let new_phase_one_job_info =
+			KnownResults::<T>::get(job_info.job_type.get_role_type(), new_phase_one_job_id)
+				.ok_or(Error::<T>::JobNotFound)?;
+
+		let new_key = match new_phase_one_job_info.result {
+			JobResult::DKGPhaseOne(info) => info.key.clone(),
+			_ => return Err(Error::<T>::InvalidJobPhase.into()),
+		};
+		let job_result = JobResult::DKGPhaseFour(DKGTSSKeyRotationResult {
+			phase_one_id: job_info
+				.job_type
+				.get_phase_one_id()
+				.ok_or(Error::<T>::InvalidJobPhase)?,
+			new_phase_one_id: new_phase_one_job_id,
+			new_key,
+			key: curr_key,
+			signature: info.signature.clone(),
+			signature_type: info.signature_type.clone(),
+		});
+
+		T::MPCHandler::verify(JobWithResult {
+			job_type: job_info.job_type.clone(),
+			phase_one_job_type: Some(phase_one_job_info.job_type),
+			result: job_result.clone(),
+		})?;
+
+		let result = PhaseResult {
+			owner: job_info.owner.clone(),
+			job_type: job_info.job_type.clone(),
+			ttl: job_info.ttl,
+			permitted_caller: job_info.job_type.clone().get_permitted_caller(),
+			result: job_result,
+		};
+		Ok(result)
+	}
+
 	pub fn verify_zksaas_circuit_job_result(
 		role_type: RoleType,
 		job_id: JobId,

pallets/jobs/src/lib.rs

@@ -331,7 +331,10 @@ pub mod module {
 					.clone()
 					.get_participants()
 					.ok_or(Error::<T>::InvalidJobParams)?,
-				JobResult::DKGPhaseTwo(_) | JobResult::ZkSaaSPhaseTwo(_) => {
+				JobResult::DKGPhaseTwo(_) |
+				JobResult::DKGPhaseThree(_) |
+				JobResult::DKGPhaseFour(_) |
+				JobResult::ZkSaaSPhaseTwo(_) => {
 					let existing_result_id = job_info
 						.job_type
 						.clone()
@@ -357,6 +360,16 @@ pub mod module {
 					let result = Self::verify_dkg_signature_job_result(role_type, &job_info, info)?;
 					KnownResults::<T>::insert(role_type, job_id, result);
 				},
+				JobResult::DKGPhaseThree(info) => {
+					let result =
+						Self::verify_dkg_key_refresh_job_result(role_type, &job_info, info)?;
+					KnownResults::<T>::insert(role_type, job_id, result);
+				},
+				JobResult::DKGPhaseFour(info) => {
+					let result =
+						Self::verify_dkg_key_rotation_job_result(role_type, &job_info, info)?;
+					KnownResults::<T>::insert(role_type, job_id, result);
+				},
 				JobResult::ZkSaaSPhaseOne(info) => {
 					let result =
 						Self::verify_zksaas_circuit_job_result(role_type, job_id, &job_info, info)?;

pallets/jobs/src/mock.rs

@@ -118,10 +118,12 @@ impl JobToFee<AccountId, BlockNumber> for MockJobToFeeHandler {
 
 	fn job_to_fee(job: &JobSubmission<AccountId, BlockNumber>) -> Balance {
 		match job.job_type {
-			JobType::DKGTSSPhaseOne(_) => MockDKGPallet::job_to_fee(job),
-			JobType::DKGTSSPhaseTwo(_) => MockDKGPallet::job_to_fee(job),
-			JobType::ZkSaaSPhaseOne(_) => MockZkSaasPallet::job_to_fee(job),
-			JobType::ZkSaaSPhaseTwo(_) => MockZkSaasPallet::job_to_fee(job),
+			JobType::DKGTSSPhaseOne(_) |
+			JobType::DKGTSSPhaseTwo(_) |
+			JobType::DKGTSSPhaseThree(_) |
+			JobType::DKGTSSPhaseFour(_) => MockDKGPallet::job_to_fee(job),
+			JobType::ZkSaaSPhaseOne(_) | JobType::ZkSaaSPhaseTwo(_) =>
+				MockZkSaasPallet::job_to_fee(job),
 		}
 	}
 }

precompiles/jobs/src/mock.rs

@@ -181,10 +181,12 @@ impl JobToFee<AccountId, BlockNumber> for MockJobToFeeHandler {
 
 	fn job_to_fee(job: &JobSubmission<AccountId, BlockNumber>) -> Balance {
 		match job.job_type {
-			JobType::DKGTSSPhaseOne(_) => MockDKGPallet::job_to_fee(job),
-			JobType::DKGTSSPhaseTwo(_) => MockDKGPallet::job_to_fee(job),
-			JobType::ZkSaaSPhaseOne(_) => MockZkSaasPallet::job_to_fee(job),
-			JobType::ZkSaaSPhaseTwo(_) => MockZkSaasPallet::job_to_fee(job),
+			JobType::DKGTSSPhaseOne(_) |
+			JobType::DKGTSSPhaseTwo(_) |
+			JobType::DKGTSSPhaseThree(_) |
+			JobType::DKGTSSPhaseFour(_) => MockDKGPallet::job_to_fee(job),
+			JobType::ZkSaaSPhaseOne(_) | JobType::ZkSaaSPhaseTwo(_) =>
+				MockZkSaasPallet::job_to_fee(job),
 		}
 	}
 }