feat(rust/rbac-registration): use KeyRotation (#297)

* feat(rust/rbac-registration): use KeyRotation

Refactors the implementation of key rotation values associated with
IdUri and RBAC registrations.

Update all occurrences within the `rbac-registration` crate that employ
the `usize` type for rotation values (e.g.,
`get_encryption_pk_for_role_at_rotation`,
`get_singing_key_cert_or_key_for_role_at_rotation`) to utilize the newly
introduced `KeyRotation` type instead.

* apply fmt

* update comments & rename function

* fmt

---------

Co-authored-by: Alex Pozhylenkov <[email protected]>

Author: vlopes11

rust/catalyst-types/src/id_uri/key_rotation.rs

@@ -30,6 +30,24 @@ impl KeyRotation {
     pub fn is_default(self) -> bool {
         self == Self::DEFAULT
     }
+
+    /// Get the key by the rotation value from the provided keys slice, if present.
+    pub fn get_key<'a, T>(&self, keys: &'a [T]) -> Option<&'a T> {
+        keys.get(self.0 as usize)
+    }
+
+    /// Get the latest rotation value calculated from the provided keys slice
+    /// (`keys.len() - 1`).
+    ///
+    /// An empty slice will be saturated to `0` rotation value.
+    pub fn from_latest_rotation<T>(keys: &[T]) -> Self {
+        // we allow the cast here as per spec we cannot overflow `u16` with a keys count.
+        #[allow(clippy::cast_possible_truncation)]
+        let rotation = keys.len() as u16;
+        let rotation = rotation.saturating_sub(1);
+
+        Self(rotation)
+    }
 }
 
 impl Default for KeyRotation {

rust/rbac-registration/src/cardano/cip509/types/role_data_record.rs

@@ -2,6 +2,7 @@
 
 use std::collections::HashMap;
 
+use catalyst_types::id_uri::key_rotation::KeyRotation;
 use pallas::ledger::addresses::ShelleyAddress;
 
 use super::{CertOrPk, PointData, PointTxnIdx};
@@ -86,15 +87,15 @@ impl RoleDataRecord {
     /// The first signing key is at rotation 0, the second at rotation 1, and so on.
     /// Returns `None` if the given key rotation does not exist.
     #[must_use]
-    pub fn signing_key_from_rotation(&self, rotation: usize) -> Option<&CertOrPk> {
-        self.signing_keys.get(rotation).map(PointData::data)
+    pub fn signing_key_from_rotation(&self, rotation: &KeyRotation) -> Option<&CertOrPk> {
+        rotation.get_key(&self.signing_keys).map(PointData::data)
     }
 
     /// Get the encryption key data associated with this role and the given key rotation.
     /// The first encryption key is at rotation 0, the second at rotation 1, and so on.
     /// Returns `None` if the given key rotation does not exist.
     #[must_use]
-    pub fn encryption_key_from_rotation(&self, rotation: usize) -> Option<&CertOrPk> {
-        self.encryption_keys.get(rotation).map(PointData::data)
+    pub fn encryption_key_from_rotation(&self, rotation: &KeyRotation) -> Option<&CertOrPk> {
+        rotation.get_key(&self.encryption_keys).map(PointData::data)
     }
 }

rust/rbac-registration/src/registration/cardano/mod.rs

@@ -7,7 +7,10 @@ use std::{collections::HashMap, sync::Arc};
 use anyhow::bail;
 use c509_certificate::c509::C509;
 use cardano_blockchain_types::{Cip0134Uri, StakeAddress, TransactionId};
-use catalyst_types::{id_uri::IdUri, uuid::UuidV4};
+use catalyst_types::{
+    id_uri::{key_rotation::KeyRotation, IdUri},
+    uuid::UuidV4,
+};
 use ed25519_dalek::VerifyingKey;
 use pallas::ledger::addresses::Address;
 use tracing::error;
@@ -133,12 +136,12 @@ impl RegistrationChain {
     #[must_use]
     pub fn get_latest_signing_pk_for_role(
         &self, role: &RoleNumber,
-    ) -> Option<(VerifyingKey, usize)> {
+    ) -> Option<(VerifyingKey, KeyRotation)> {
         self.inner.role_data_record.get(role).and_then(|rdr| {
             rdr.signing_keys().last().and_then(|key| {
-                key.data()
-                    .extract_pk()
-                    .map(|pk| (pk, rdr.signing_keys().len().saturating_sub(1)))
+                let rotation = KeyRotation::from_latest_rotation(rdr.signing_keys());
+
+                key.data().extract_pk().map(|pk| (pk, rotation))
             })
         })
     }
@@ -148,12 +151,12 @@ impl RegistrationChain {
     #[must_use]
     pub fn get_latest_encryption_pk_for_role(
         &self, role: &RoleNumber,
-    ) -> Option<(VerifyingKey, usize)> {
+    ) -> Option<(VerifyingKey, KeyRotation)> {
         self.inner.role_data_record.get(role).and_then(|rdr| {
             rdr.encryption_keys().last().and_then(|key| {
-                key.data()
-                    .extract_pk()
-                    .map(|pk| (pk, rdr.encryption_keys().len().saturating_sub(1)))
+                let rotation = KeyRotation::from_latest_rotation(rdr.encryption_keys());
+
+                key.data().extract_pk().map(|pk| (pk, rotation))
             })
         })
     }
@@ -162,7 +165,7 @@ impl RegistrationChain {
     /// Returns the public key, `None` if not found.
     #[must_use]
     pub fn get_signing_pk_for_role_at_rotation(
-        &self, role: &RoleNumber, rotation: usize,
+        &self, role: &RoleNumber, rotation: &KeyRotation,
     ) -> Option<VerifyingKey> {
         self.inner.role_data_record.get(role).and_then(|rdr| {
             rdr.signing_key_from_rotation(rotation)
@@ -174,7 +177,7 @@ impl RegistrationChain {
     /// Returns the public key, `None` if not found.
     #[must_use]
     pub fn get_encryption_pk_for_role_at_rotation(
-        &self, role: &RoleNumber, rotation: usize,
+        &self, role: &RoleNumber, rotation: &KeyRotation,
     ) -> Option<VerifyingKey> {
         self.inner.role_data_record.get(role).and_then(|rdr| {
             rdr.encryption_key_from_rotation(rotation)
@@ -186,7 +189,7 @@ impl RegistrationChain {
     /// given rotation.
     #[must_use]
     pub fn get_singing_key_cert_or_key_for_role_at_rotation(
-        &self, role: &RoleNumber, rotation: usize,
+        &self, role: &RoleNumber, rotation: &KeyRotation,
     ) -> Option<&CertOrPk> {
         self.inner
             .role_data_record
@@ -198,7 +201,7 @@ impl RegistrationChain {
     /// with given rotation.
     #[must_use]
     pub fn get_encryption_key_cert_or_key_for_role_at_rotation(
-        &self, role: &RoleNumber, rotation: usize,
+        &self, role: &RoleNumber, rotation: &KeyRotation,
     ) -> Option<&CertOrPk> {
         self.inner
             .role_data_record
@@ -487,12 +490,15 @@ mod test {
         let (_k, r) = update
             .get_latest_signing_pk_for_role(&RoleNumber::ROLE_0)
             .unwrap();
-        assert_eq!(r, 1);
+        assert_eq!(r, KeyRotation::from(1));
         assert!(update
-            .get_signing_pk_for_role_at_rotation(&RoleNumber::ROLE_0, 2)
+            .get_signing_pk_for_role_at_rotation(&RoleNumber::ROLE_0, &KeyRotation::from(2))
             .is_none());
         assert!(update
-            .get_singing_key_cert_or_key_for_role_at_rotation(&RoleNumber::ROLE_0, 0)
+            .get_singing_key_cert_or_key_for_role_at_rotation(
+                &RoleNumber::ROLE_0,
+                &KeyRotation::from(0)
+            )
             .is_some());
     }
 }