fix: Wait for ThresholdSignature success before switching to NewKeysA…

…ctivated (#4534)

* track tss request id before completing rotation

* add unit test to check the correct flow is followed when tss is required

state-chain/cf-integration-tests/src/authorities.rs

@@ -113,6 +113,9 @@ fn authority_rotates_with_correct_sequence() {
 				RotationPhase::ActivatingKeys(..)
 			));
 
+			// Wait for an extra block to allow TSS to complete, we switch to RotationComplete once
+			// that's done
+			testnet.move_forward_blocks(1);
 			assert_eq!(
 				AllVaults::status(),
 				AsyncResult::Ready(KeyRotationStatusOuter::RotationComplete),

state-chain/cf-integration-tests/src/broadcasting.rs

@@ -42,7 +42,7 @@ fn bitcoin_broadcast_delay_works() {
 			)
 			.unwrap();
 
-			let broadcast_id =
+			let (broadcast_id, _) =
 				<BitcoinBroadcaster as Broadcaster<Bitcoin>>::threshold_sign_and_broadcast(
 					bitcoin_call,
 				);

state-chain/pallets/cf-broadcast/src/benchmarking.rs

@@ -128,7 +128,7 @@ mod benchmarks {
 
 	#[benchmark]
 	fn start_next_broadcast_attempt() {
-		let broadcast_id = Pallet::<T, I>::threshold_sign_and_broadcast(
+		let (broadcast_id, _) = Pallet::<T, I>::threshold_sign_and_broadcast(
 			BenchmarkValue::benchmark_value(),
 			None,
 			|_| None,

state-chain/pallets/cf-broadcast/src/lib.rs

@@ -664,7 +664,7 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 		maybe_failed_callback_generator: impl FnOnce(
 			BroadcastId,
 		) -> Option<<T as Config<I>>::BroadcastCallable>,
-	) -> BroadcastId {
+	) -> (BroadcastId, ThresholdSignatureRequestId) {
 		let broadcast_id = Self::next_broadcast_id();
 
 		PendingBroadcasts::<T, I>::append(broadcast_id);
@@ -676,9 +676,7 @@ impl<T: Config<I>, I: 'static> Pallet<T, I> {
 			RequestFailureCallbacks::<T, I>::insert(broadcast_id, callback);
 		}
 
-		let _threshold_signature_id = Self::threshold_sign(api_call, broadcast_id, true);
-
-		broadcast_id
+		(broadcast_id, Self::threshold_sign(api_call, broadcast_id, true))
 	}
 
 	/// Signs a API call, use `Call::on_signature_ready` as the callback, and returns the signature
@@ -922,7 +920,9 @@ impl<T: Config<I>, I: 'static> Broadcaster<T::TargetChain> for Pallet<T, I> {
 	type ApiCall = T::ApiCall;
 	type Callback = <T as Config<I>>::BroadcastCallable;
 
-	fn threshold_sign_and_broadcast(api_call: Self::ApiCall) -> BroadcastId {
+	fn threshold_sign_and_broadcast(
+		api_call: Self::ApiCall,
+	) -> (BroadcastId, ThresholdSignatureRequestId) {
 		Self::threshold_sign_and_broadcast(api_call, None, |_| None)
 	}
 
@@ -931,7 +931,7 @@ impl<T: Config<I>, I: 'static> Broadcaster<T::TargetChain> for Pallet<T, I> {
 		success_callback: Option<Self::Callback>,
 		failed_callback_generator: impl FnOnce(BroadcastId) -> Option<Self::Callback>,
 	) -> BroadcastId {
-		Self::threshold_sign_and_broadcast(api_call, success_callback, failed_callback_generator)
+		Self::threshold_sign_and_broadcast(api_call, success_callback, failed_callback_generator).0
 	}
 
 	fn threshold_sign(api_call: Self::ApiCall) -> (BroadcastId, ThresholdSignatureRequestId) {
@@ -949,8 +949,11 @@ impl<T: Config<I>, I: 'static> Broadcaster<T::TargetChain> for Pallet<T, I> {
 		Self::clean_up_broadcast_storage(broadcast_id);
 	}
 
-	fn threshold_sign_and_broadcast_rotation_tx(api_call: Self::ApiCall) -> BroadcastId {
-		let broadcast_id = <Self as Broadcaster<_>>::threshold_sign_and_broadcast(api_call);
+	fn threshold_sign_and_broadcast_rotation_tx(
+		api_call: Self::ApiCall,
+	) -> (BroadcastId, ThresholdSignatureRequestId) {
+		let (broadcast_id, request_id) =
+			<Self as Broadcaster<_>>::threshold_sign_and_broadcast(api_call);
 
 		if let Some(earliest_pending_broadcast_id) = PendingBroadcasts::<T, I>::get()
 			.first()
@@ -962,6 +965,7 @@ impl<T: Config<I>, I: 'static> Broadcaster<T::TargetChain> for Pallet<T, I> {
 					}
 				}
 		}
-		broadcast_id
+
+		(broadcast_id, request_id)
 	}
 }

state-chain/pallets/cf-broadcast/src/tests.rs

@@ -435,7 +435,7 @@ fn threshold_sign_and_broadcast_with_callback() {
 			tx_out_id: MOCK_TRANSACTION_OUT_ID,
 		};
 
-		let broadcast_id =
+		let (broadcast_id, _) =
 			Broadcaster::threshold_sign_and_broadcast(api_call.clone(), Some(MockCallback), |_| {
 				None
 			});
@@ -544,7 +544,7 @@ fn callback_is_called_upon_broadcast_failure() {
 			sig: Default::default(),
 			tx_out_id: MOCK_TRANSACTION_OUT_ID,
 		};
-		let broadcast_id =
+		let (broadcast_id, _) =
 			Broadcaster::threshold_sign_and_broadcast(api_call.clone(), None, |_| {
 				Some(MockCallback)
 			});
@@ -883,7 +883,7 @@ fn initiate_and_sign_broadcast(
 	api_call: &MockApiCall<MockEthereumChainCrypto>,
 	tx_type: TxType,
 ) -> BroadcastId {
-	let broadcast_id = match tx_type {
+	let (broadcast_id, _) = match tx_type {
 		TxType::Normal => <Broadcaster as BroadcasterTrait<
 			<Test as Config<Instance1>>::TargetChain,
 		>>::threshold_sign_and_broadcast((*api_call).clone()),

state-chain/pallets/cf-emissions/src/mock.rs

@@ -169,9 +169,11 @@ impl Broadcaster<MockEthereum> for MockBroadcast {
 	type ApiCall = MockUpdateFlipSupply;
 	type Callback = MockCallback;
 
-	fn threshold_sign_and_broadcast(api_call: Self::ApiCall) -> BroadcastId {
+	fn threshold_sign_and_broadcast(
+		api_call: Self::ApiCall,
+	) -> (BroadcastId, ThresholdSignatureRequestId) {
 		Self::call(api_call);
-		1
+		(1, 1)
 	}
 
 	fn threshold_sign_and_broadcast_with_callback(
@@ -182,7 +184,9 @@ impl Broadcaster<MockEthereum> for MockBroadcast {
 		unimplemented!()
 	}
 
-	fn threshold_sign_and_broadcast_rotation_tx(_api_call: Self::ApiCall) -> BroadcastId {
+	fn threshold_sign_and_broadcast_rotation_tx(
+		_api_call: Self::ApiCall,
+	) -> (BroadcastId, ThresholdSignatureRequestId) {
 		unimplemented!()
 	}
 

state-chain/pallets/cf-environment/src/mock.rs

@@ -101,7 +101,9 @@ impl Broadcaster<Polkadot> for MockPolkadotBroadcaster {
 	type ApiCall = MockCreatePolkadotVault;
 	type Callback = MockCallback;
 
-	fn threshold_sign_and_broadcast(_api_call: Self::ApiCall) -> BroadcastId {
+	fn threshold_sign_and_broadcast(
+		_api_call: Self::ApiCall,
+	) -> (BroadcastId, ThresholdSignatureRequestId) {
 		unimplemented!()
 	}
 
@@ -113,7 +115,9 @@ impl Broadcaster<Polkadot> for MockPolkadotBroadcaster {
 		unimplemented!()
 	}
 
-	fn threshold_sign_and_broadcast_rotation_tx(_api_call: Self::ApiCall) -> BroadcastId {
+	fn threshold_sign_and_broadcast_rotation_tx(
+		_api_call: Self::ApiCall,
+	) -> (BroadcastId, ThresholdSignatureRequestId) {
 		unimplemented!()
 	}
 

state-chain/pallets/cf-funding/src/lib.rs

@@ -483,7 +483,7 @@ pub mod pallet {
 				Self::deposit_event(Event::RedemptionRequested {
 					account_id,
 					amount: redeem_amount,
-					broadcast_id: T::Broadcaster::threshold_sign_and_broadcast(call),
+					broadcast_id: T::Broadcaster::threshold_sign_and_broadcast(call).0,
 					expiry_time: contract_expiry,
 				});
 			} else {

state-chain/pallets/cf-funding/src/mock.rs

@@ -156,11 +156,13 @@ impl Broadcaster<Ethereum> for MockBroadcaster {
 	type ApiCall = MockRegisterRedemption;
 	type Callback = MockCallback;
 
-	fn threshold_sign_and_broadcast(api_call: Self::ApiCall) -> BroadcastId {
+	fn threshold_sign_and_broadcast(
+		api_call: Self::ApiCall,
+	) -> (BroadcastId, ThresholdSignatureRequestId) {
 		REDEMPTION_BROADCAST_REQUESTS.with(|cell| {
 			cell.borrow_mut().push(api_call.amount);
 		});
-		0
+		(0, 0)
 	}
 
 	fn threshold_sign_and_broadcast_with_callback(
@@ -171,7 +173,9 @@ impl Broadcaster<Ethereum> for MockBroadcaster {
 		unimplemented!()
 	}
 
-	fn threshold_sign_and_broadcast_rotation_tx(_api_call: Self::ApiCall) -> BroadcastId {
+	fn threshold_sign_and_broadcast_rotation_tx(
+		_api_call: Self::ApiCall,
+	) -> (BroadcastId, ThresholdSignatureRequestId) {
 		unimplemented!()
 	}
 

state-chain/pallets/cf-ingress-egress/src/lib.rs

@@ -608,7 +608,8 @@ pub mod pallet {
 			if let Ok(egress_transaction) =
 				<T::ChainApiCall as ConsolidateCall<T::TargetChain>>::consolidate_utxos()
 			{
-				let broadcast_id = T::Broadcaster::threshold_sign_and_broadcast(egress_transaction);
+				let (broadcast_id, _) =
+					T::Broadcaster::threshold_sign_and_broadcast(egress_transaction);
 				Self::deposit_event(Event::<T, I>::UtxoConsolidation { broadcast_id });
 				Self::deposit_event(Event::<T, I>::BatchBroadcastRequested {
 					broadcast_id,

state-chain/pallets/cf-threshold-signature/src/key_rotator.rs

@@ -147,14 +147,26 @@ impl<T: Config<I>, I: 'static> KeyRotator for Pallet<T, I> {
 						),
 					},
 				KeyRotationStatusVariant::AwaitingActivation => {
-					let new_public_key = match PendingKeyRotation::<T, I>::get() {
-						Some(KeyRotationStatus::AwaitingActivation { new_public_key }) =>
+					let (new_public_key, maybe_request_id) = match PendingKeyRotation::<T, I>::get()
+					{
+						Some(KeyRotationStatus::AwaitingActivation {
+							request_id,
 							new_public_key,
+						}) => (new_public_key, request_id),
 						_ => unreachable!(
 							"Unreachable because we are in the branch for the AwaitingActivation variant."
 						),
 					};
 
+					if let Some(request_id) = maybe_request_id {
+						// After the ceremony completes, it is consumed and Void is left
+						// behind. At this point we are sure the ceremony existed and succeded, we
+						// can activate the key
+						if Signature::<T, I>::get(request_id) == AsyncResult::Void {
+							T::VaultActivator::activate_key();
+						};
+					};
+
 					let status = T::VaultActivator::status()
 						.replace_inner(KeyRotationStatusOuter::RotationComplete);
 					if status.is_ready() {
@@ -188,17 +200,34 @@ impl<T: Config<I>, I: 'static> KeyRotator for Pallet<T, I> {
 			PendingKeyRotation::<T, I>::get()
 		{
 			let maybe_active_epoch_key = Self::active_epoch_key();
-			T::VaultActivator::activate(
+
+			match T::VaultActivator::start_key_activation(
 				new_public_key,
 				maybe_active_epoch_key.map(|EpochKey { key, .. }| key),
-			);
-
-			if maybe_active_epoch_key.is_some() {
-				Self::activate_new_key(new_public_key);
-			} else {
-				PendingKeyRotation::<T, I>::put(KeyRotationStatus::<T, I>::AwaitingActivation {
-					new_public_key,
-				});
+			) {
+				// If a request_id was returned we need to wait for the signing request to complete
+				// before ending the rotation succesfully
+				Some(request_id) => {
+					PendingKeyRotation::<T, I>::put(
+						KeyRotationStatus::<T, I>::AwaitingActivation {
+							request_id: Some(request_id),
+							new_public_key,
+						},
+					);
+				},
+				// if none was returned no ceremony is required and we can already complete the
+				// rotation/wait for governance extrinsic to activate the vault
+				None =>
+					if maybe_active_epoch_key.is_some() {
+						Self::activate_new_key(new_public_key);
+					} else {
+						PendingKeyRotation::<T, I>::put(
+							KeyRotationStatus::<T, I>::AwaitingActivation {
+								request_id: None,
+								new_public_key,
+							},
+						);
+					},
 			}
 		} else {
 			log::error!("Vault activation called during wrong state.");

state-chain/pallets/cf-threshold-signature/src/lib.rs

@@ -141,6 +141,7 @@ pub enum KeyRotationStatus<T: Config<I>, I: 'static = ()> {
 	},
 	/// We are waiting for the key to be updated on the contract, and witnessed by the network.
 	AwaitingActivation {
+		request_id: Option<RequestId>,
 		new_public_key: AggKeyFor<T, I>,
 	},
 	/// The key has been successfully updated on the external chain, and/or funds rotated to new

state-chain/pallets/cf-threshold-signature/src/mock.rs

@@ -8,7 +8,10 @@ use cf_chains::{
 	mocks::{MockAggKey, MockEthereumChainCrypto, MockThresholdSignature},
 	ChainCrypto,
 };
-use cf_primitives::{AuthorityCount, CeremonyId, FlipBalance, FLIPPERINOS_PER_FLIP, GENESIS_EPOCH};
+use cf_primitives::{
+	AuthorityCount, CeremonyId, FlipBalance, ThresholdSignatureRequestId, FLIPPERINOS_PER_FLIP,
+	GENESIS_EPOCH,
+};
 use cf_traits::{
 	impl_mock_chainflip, impl_mock_runtime_safe_mode,
 	mocks::{cfe_interface_mock::MockCfeInterface, signer_nomination::MockNominator},
@@ -188,7 +191,14 @@ impl pallet_cf_threshold_signature::Config<Instance1> for Test {
 pub struct MockVaultActivator;
 impl VaultActivator<MockEthereumChainCrypto> for MockVaultActivator {
 	type ValidatorId = <Test as Chainflip>::ValidatorId;
-	fn activate(_new_key: MockAggKey, _maybe_old_key: Option<MockAggKey>) {}
+	fn start_key_activation(
+		_new_key: MockAggKey,
+		_maybe_old_key: Option<MockAggKey>,
+	) -> Option<ThresholdSignatureRequestId> {
+		VAULT_ACTIVATION_STATUS.with(|value| *(value.borrow_mut()) = AsyncResult::Pending);
+		let ceremony_id = current_ceremony_id();
+		Some(ceremony_id as u32)
+	}
 
 	fn status() -> AsyncResult<()> {
 		VAULT_ACTIVATION_STATUS.with(|value| *value.borrow())
@@ -198,6 +208,10 @@ impl VaultActivator<MockEthereumChainCrypto> for MockVaultActivator {
 	fn set_status(outcome: AsyncResult<()>) {
 		VAULT_ACTIVATION_STATUS.with(|value| *(value.borrow_mut()) = outcome)
 	}
+
+	fn activate_key() {
+		VAULT_ACTIVATION_STATUS.with(|value| *(value.borrow_mut()) = AsyncResult::Ready(()))
+	}
 }
 
 impl MockVaultActivator {

state-chain/pallets/cf-threshold-signature/src/tests.rs

@@ -18,7 +18,7 @@ use cf_traits::{
 		signer_nomination::MockNominator,
 	},
 	AccountRoleRegistry, AsyncResult, Chainflip, EpochInfo, EpochKey, KeyProvider,
-	KeyRotationStatusOuter, KeyRotator, SetSafeMode,
+	KeyRotationStatusOuter, KeyRotator, SetSafeMode, VaultActivator,
 };
 pub use frame_support::traits::Get;
 
@@ -1632,6 +1632,7 @@ mod key_rotation {
 
 			BlockHeightProvider::<MockEthereum>::set_block_height(HANDOVER_ACTIVATION_BLOCK);
 			EthereumThresholdSigner::activate_keys();
+			EthereumThresholdSigner::status();
 
 			assert!(matches!(
 				PendingKeyRotation::<Test, _>::get().unwrap(),
@@ -1674,6 +1675,7 @@ mod key_rotation {
 			run_cfes_on_sc_events(&cfes);
 
 			EthereumThresholdSigner::activate_keys();
+			EthereumThresholdSigner::status();
 		});
 
 		final_checks(ext);
@@ -1742,6 +1744,37 @@ mod key_rotation {
 			do_full_key_rotation();
 		});
 	}
+
+	#[test]
+	fn wait_for_activating_key_tss_before_completing_rotation() {
+		let ext = setup(Ok(NEW_AGG_PUB_KEY_POST_HANDOVER)).execute_with(|| {
+			// KeyHandoverComplete
+			PendingKeyRotation::<Test, _>::put(KeyRotationStatus::KeyHandoverComplete {
+				new_public_key: NEW_AGG_PUB_KEY_POST_HANDOVER,
+			});
+			// Start ActivatingKeys
+			EthereumThresholdSigner::activate_keys();
+			assert_eq!(
+				PendingKeyRotation::<Test, _>::get().unwrap(),
+				KeyRotationStatus::AwaitingActivation {
+					request_id: Some(4),
+					new_public_key: NEW_AGG_PUB_KEY_POST_HANDOVER
+				}
+			);
+
+			// Vault activation started and it is now pending
+			assert_eq!(MockVaultActivator::status(), AsyncResult::Pending);
+
+			// Request is complete
+			assert_eq!(EthereumThresholdSigner::signature(4), AsyncResult::Void);
+
+			// Proceed to complete activation
+			EthereumThresholdSigner::status();
+
+			assert_eq!(MockVaultActivator::status(), AsyncResult::Ready(()));
+		});
+		final_checks(ext);
+	}
 }
 
 #[test]