staking: enshrine consensus key indexing requirement in unit-test (#4148

)

crates/core/component/stake/src/component/action_handler/validator_definition.rs

@@ -88,7 +88,7 @@ impl ActionHandler for validator::Definition {
         // Check if the consensus key is known, and if so, that it is by the
         // validator that declares it in this definition.
         if let Some(ck_owner) = state
-            .get_validator_by_consensus_key(&new_validator.consensus_key)
+            .get_validator_definition_by_consensus_key(&new_validator.consensus_key)
             .await?
         {
             // If we detect that the new definition tries to squat someone else's

crates/core/component/stake/src/component/stake.rs

@@ -30,6 +30,9 @@ use crate::component::validator_handler::{
     ValidatorDataRead, ValidatorManager, ValidatorUptimeTracker,
 };
 
+#[cfg(test)]
+mod tests;
+
 pub struct Staking {}
 
 impl Staking {}
@@ -321,11 +324,16 @@ pub trait StateWriteExt: StateWrite {
         )
     }
 
-    async fn register_consensus_key(
-        &mut self,
-        identity_key: &IdentityKey,
-        consensus_key: &PublicKey,
-    ) {
+    /// Register a [consensus key][`PublicKey`] in the state, via two verifiable indices:
+    /// 1. CometBFT address -> [`PublicKey`]
+    /// 2. [`PublicKey`] -> [`IdentityKey`]
+    ///
+    /// # Important note
+    /// We do not delete obsolete entries on purpose. This is so that
+    /// the staking component can do evidence attribution even if a byzantine validator
+    /// has changed the consensus key that was used at the time of the misbehavior.
+    #[instrument(skip_all)]
+    fn register_consensus_key(&mut self, identity_key: &IdentityKey, consensus_key: &PublicKey) {
         let address = self::address::validator_address(consensus_key);
         tracing::debug!(?identity_key, ?consensus_key, hash = ?hex::encode(address), "registering consensus key");
         self.put(

crates/core/component/stake/src/component/stake/tests.rs

@@ -0,0 +1,104 @@
+use anyhow::ensure;
+use cnidarium::{StateDelta, TempStorage};
+use decaf377_rdsa::SigningKey;
+use rand_core::OsRng;
+use tendermint::PublicKey;
+
+use crate::{
+    component::{stake::address::validator_address, validator_handler::ValidatorDataRead},
+    IdentityKey, StateWriteExt,
+};
+
+#[tokio::test]
+/// Test that we do not delete rotated consensus keys from the [consensus key -> identity key] index.A
+/// This is important to maintain because we want to be able to resolve byzantine evidence to a validator's
+/// persistent identity even if they have rotated their consensus keys.
+async fn test_persistent_identity_by_ck() -> anyhow::Result<()> {
+    let _ = tracing_subscriber::fmt::try_init();
+    let storage = TempStorage::new().await?;
+    let mut state = StateDelta::new(storage.latest_snapshot());
+
+    let rng = OsRng;
+    let persistent_identity = IdentityKey(SigningKey::new(rng).into());
+
+    let old_ck_raw = ed25519_consensus::SigningKey::new(rng)
+        .verification_key()
+        .to_bytes();
+    let new_ck_raw = ed25519_consensus::SigningKey::new(rng)
+        .verification_key()
+        .to_bytes();
+
+    let old_ck = PublicKey::from_raw_ed25519(&old_ck_raw).expect("valid vk");
+    let new_ck = PublicKey::from_raw_ed25519(&new_ck_raw).expect("valid vk");
+    anyhow::ensure!(
+        old_ck.to_hex() != new_ck.to_hex(),
+        "the keys must encode to different hex strings for the test to be useful"
+    );
+
+    let old_address = validator_address(&old_ck);
+    let new_address = validator_address(&new_ck);
+
+    state.register_consensus_key(&persistent_identity, &old_ck);
+
+    let retrieved_ck = state
+        .lookup_consensus_key_by_comet_address(&old_address)
+        .await
+        .expect("key is registered");
+
+    ensure!(
+        retrieved_ck == old_ck,
+        "the retrieved consensus key must match the initial ck"
+    );
+
+    let retrieved_id = state
+        .lookup_identity_key_by_consensus_key(&retrieved_ck)
+        .await
+        .expect("key is found");
+    ensure!(
+        retrieved_id == persistent_identity,
+        "the retrieved identity must match its persistent identity"
+    );
+
+    state.register_consensus_key(&persistent_identity, &new_ck);
+    // We want to do a basic check that we can reach for the updated values
+    // but CRUCIALLY, we want to make sure that we can associate an identity to
+    // the old consenus key.
+    let retrieved_ck = state
+        .lookup_consensus_key_by_comet_address(&new_address)
+        .await
+        .expect("key is registered");
+    ensure!(
+        retrieved_ck == new_ck,
+        "we must be able to find the updated ck"
+    );
+
+    let retrieved_id = state
+        .lookup_identity_key_by_consensus_key(&retrieved_ck)
+        .await
+        .expect("key is found");
+    ensure!(
+        retrieved_id == persistent_identity,
+        "the retrieved id must match the persistent identity, even after an update"
+    );
+
+    // CRUCIAL PART: can we find the persistent identity from a rotated comet address/consensus key?
+    let culprit_ck = state
+        .lookup_consensus_key_by_comet_address(&old_address)
+        .await
+        .expect("key must be found!");
+    ensure!(
+        culprit_ck == old_ck,
+        "the old address must be associated with the old ck"
+    );
+
+    let culprit_id = state
+        .lookup_identity_key_by_consensus_key(&culprit_ck)
+        .await
+        .expect("consensus key -> identity index is persistent across validator updates");
+    ensure!(
+        culprit_id == persistent_identity,
+        "the retrieved identity must match the persistent identity that we setup"
+    );
+
+    Ok(())
+}

crates/core/component/stake/src/component/validator_handler/validator_manager.rs

@@ -462,8 +462,7 @@ pub trait ValidatorManager: StateWrite {
         // Then, we create a mapping from the validator's consensus key to its
         // identity key, so we can look up the validator by its consensus key, and
         // vice-versa.
-        self.register_consensus_key(&validator_identity, &validator.consensus_key)
-            .await;
+        self.register_consensus_key(&validator_identity, &validator.consensus_key);
         // We register the validator's delegation token in the token registry...
         self.register_denom(&DelegationToken::from(&validator_identity).denom())
             .await;
@@ -589,8 +588,7 @@ pub trait ValidatorManager: StateWrite {
 
         // Update the consensus key lookup, in case the validator rotated their
         // consensus key.
-        self.register_consensus_key(&validator.identity_key, &validator.consensus_key)
-            .await;
+        self.register_consensus_key(&validator.identity_key, &validator.consensus_key);
 
         self.put(state_key::validators::definitions::by_id(id), validator);
 
@@ -642,7 +640,7 @@ pub trait ValidatorManager: StateWrite {
     /// Returns an error if the validator is not found in the JMT.
     async fn process_evidence(&mut self, evidence: &Misbehavior) -> Result<()> {
         let validator = self
-            .get_validator_by_cometbft_address(&evidence.validator.address)
+            .get_validator_definition_by_cometbft_address(&evidence.validator.address)
             .await?
             .ok_or_else(|| {
                 anyhow::anyhow!(

crates/core/component/stake/src/component/validator_handler/validator_store.rs

@@ -128,28 +128,38 @@ pub trait ValidatorDataRead: StateRead {
         self.nonverifiable_get(key.as_bytes())
     }
 
+    async fn lookup_identity_key_by_consensus_key(&self, ck: &PublicKey) -> Option<IdentityKey> {
+        self.get(&state_key::validators::lookup_by::consensus_key(ck))
+            .await
+            .expect("no deserialization error")
+    }
+
+    async fn lookup_consensus_key_by_comet_address(&self, address: &[u8; 20]) -> Option<PublicKey> {
+        self.get(&state_key::validators::lookup_by::cometbft_address(address))
+            .await
+            .expect("no deserialization error")
+    }
+
     // Tendermint validators are referenced to us by their Tendermint consensus key,
     // but we reference them by their Penumbra identity key.
-    async fn get_validator_by_consensus_key(&self, ck: &PublicKey) -> Result<Option<Validator>> {
-        if let Some(identity_key) = self
-            .get(&state_key::validators::lookup_by::consensus_key(ck))
-            .await?
-        {
+    async fn get_validator_definition_by_consensus_key(
+        &self,
+        ck: &PublicKey,
+    ) -> Result<Option<Validator>> {
+        if let Some(identity_key) = self.lookup_identity_key_by_consensus_key(ck).await {
             self.get_validator_definition(&identity_key).await
         } else {
             return Ok(None);
         }
     }
 
-    async fn get_validator_by_cometbft_address(
+    async fn get_validator_definition_by_cometbft_address(
         &self,
         address: &[u8; 20],
     ) -> Result<Option<Validator>> {
-        if let Some(consensus_key) = self
-            .get(&state_key::validators::lookup_by::cometbft_address(address))
-            .await?
-        {
-            self.get_validator_by_consensus_key(&consensus_key).await
+        if let Some(consensus_key) = self.lookup_consensus_key_by_comet_address(address).await {
+            self.get_validator_definition_by_consensus_key(&consensus_key)
+                .await
         } else {
             return Ok(None);
         }