refactor(axelar-gateway): use contractstorage

Author: nbayindirli

contracts/stellar-axelar-gateway/src/auth.rs

@@ -4,7 +4,7 @@ use stellar_axelar_std::events::Event;
 
 use crate::error::ContractError;
 use crate::event::SignersRotatedEvent;
-use crate::storage_types::DataKey;
+use crate::storage;
 use crate::types::{Proof, ProofSignature, ProofSigner, WeightedSigner, WeightedSigners};
 
 pub fn initialize_auth(
@@ -14,20 +14,10 @@ pub fn initialize_auth(
     previous_signer_retention: u64,
     initial_signers: Vec<WeightedSigners>,
 ) -> Result<(), ContractError> {
-    env.storage().instance().set(&DataKey::Epoch, &0_u64);
-
-    env.storage().instance().set(
-        &DataKey::PreviousSignerRetention,
-        &previous_signer_retention,
-    );
-
-    env.storage()
-        .instance()
-        .set(&DataKey::DomainSeparator, &domain_separator);
-
-    env.storage()
-        .instance()
-        .set(&DataKey::MinimumRotationDelay, &minimum_rotation_delay);
+    storage::set_epoch(&env, &0_u64);
+    storage::set_previous_signer_retention(&env, &previous_signer_retention);
+    storage::set_domain_separator(&env, &domain_separator);
+    storage::set_minimum_rotation_delay(&env, &minimum_rotation_delay);
 
     ensure!(!initial_signers.is_empty(), ContractError::EmptySigners);
 
@@ -38,27 +28,6 @@ pub fn initialize_auth(
     Ok(())
 }
 
-pub fn domain_separator(env: &Env) -> BytesN<32> {
-    env.storage()
-        .instance()
-        .get(&DataKey::DomainSeparator)
-        .expect("domain_separator not found")
-}
-
-pub fn minimum_rotation_delay(env: &Env) -> u64 {
-    env.storage()
-        .instance()
-        .get(&DataKey::MinimumRotationDelay)
-        .expect("minimum_rotation_delay not found")
-}
-
-pub fn previous_signers_retention(env: &Env) -> u64 {
-    env.storage()
-        .instance()
-        .get(&DataKey::PreviousSignerRetention)
-        .expect("previous_signers_retention not found")
-}
-
 pub fn validate_proof(
     env: &Env,
     data_hash: &BytesN<32>,
@@ -68,14 +37,15 @@ pub fn validate_proof(
 
     let signers_hash = signers_set.hash(env);
 
-    let signers_epoch = epoch_by_signers_hash(env, signers_hash.clone())?;
+    let signers_epoch = storage::try_epoch_by_signers_hash(env, signers_hash.clone())
+        .ok_or(ContractError::InvalidSignersHash)?;
 
-    let current_epoch = epoch(env);
+    let current_epoch = storage::epoch(env);
 
     let is_latest_signers: bool = signers_epoch == current_epoch;
 
     ensure!(
-        current_epoch - signers_epoch <= previous_signers_retention(env),
+        current_epoch - signers_epoch <= storage::previous_signer_retention(env),
         ContractError::OutdatedSigners
     );
 
@@ -100,23 +70,20 @@ pub fn rotate_signers(
 
     let new_signers_hash = new_signers.hash(env);
 
-    let new_epoch: u64 = epoch(env) + 1;
+    let new_epoch: u64 = storage::epoch(env) + 1;
 
-    env.storage().instance().set(&DataKey::Epoch, &new_epoch);
+    storage::set_epoch(env, &new_epoch);
 
-    env.storage()
-        .persistent()
-        .set(&DataKey::SignersHashByEpoch(new_epoch), &new_signers_hash);
+    storage::set_signers_hash_by_epoch(env, new_epoch, &new_signers_hash);
 
     ensure!(
-        epoch_by_signers_hash(env, new_signers_hash.clone()).is_err(),
+        storage::try_epoch_by_signers_hash(env, new_signers_hash.clone())
+            .ok_or(ContractError::InvalidSignersHash)
+            .is_err(),
         ContractError::DuplicateSigners
     );
 
-    env.storage().persistent().set(
-        &DataKey::EpochBySignersHash(new_signers_hash.clone()),
-        &new_epoch,
-    );
+    storage::set_epoch_by_signers_hash(env, new_signers_hash.clone(), &new_epoch);
 
     SignersRotatedEvent {
         epoch: new_epoch,
@@ -128,54 +95,26 @@ pub fn rotate_signers(
     Ok(())
 }
 
-pub fn epoch(env: &Env) -> u64 {
-    env.storage()
-        .instance()
-        .get(&DataKey::Epoch)
-        .expect("epoch not found")
-}
-
-pub fn epoch_by_signers_hash(env: &Env, signers_hash: BytesN<32>) -> Result<u64, ContractError> {
-    env.storage()
-        .persistent()
-        .get(&DataKey::EpochBySignersHash(signers_hash))
-        .ok_or(ContractError::InvalidSignersHash)
-}
-
-pub fn signers_hash_by_epoch(env: &Env, epoch: u64) -> Result<BytesN<32>, ContractError> {
-    env.storage()
-        .persistent()
-        .get(&DataKey::SignersHashByEpoch(epoch))
-        .ok_or(ContractError::InvalidEpoch)
-}
-
 fn message_hash_to_sign(env: &Env, signers_hash: BytesN<32>, data_hash: &BytesN<32>) -> BytesN<32> {
-    let mut msg: Bytes = domain_separator(env).into();
+    let mut msg: Bytes = storage::domain_separator(env).into();
     msg.extend_from_array(&signers_hash.to_array());
     msg.extend_from_array(&data_hash.to_array());
 
     env.crypto().keccak256(&msg).into()
 }
 
 fn update_rotation_timestamp(env: &Env, enforce_rotation_delay: bool) -> Result<(), ContractError> {
-    let last_rotation_timestamp: u64 = env
-        .storage()
-        .instance()
-        .get(&DataKey::LastRotationTimestamp)
-        .unwrap_or(0);
-
     let current_timestamp = env.ledger().timestamp();
 
     if enforce_rotation_delay {
         ensure!(
-            current_timestamp - last_rotation_timestamp >= minimum_rotation_delay(env),
+            current_timestamp - storage::last_rotation_timestamp(env)
+                >= storage::minimum_rotation_delay(env),
             ContractError::InsufficientRotationDelay
         );
     }
 
-    env.storage()
-        .instance()
-        .set(&DataKey::LastRotationTimestamp, &current_timestamp);
+    storage::set_last_rotation_timestamp(env, &current_timestamp);
 
     Ok(())
 }

contracts/stellar-axelar-gateway/src/contract.rs

@@ -6,13 +6,13 @@ use stellar_axelar_std::{
     ensure, interfaces, when_not_paused, Operatable, Ownable, Pausable, Upgradable,
 };
 
-use crate::auth;
 use crate::error::ContractError;
 use crate::event::{ContractCalledEvent, MessageApprovedEvent, MessageExecutedEvent};
 use crate::interface::AxelarGatewayInterface;
 use crate::messaging_interface::AxelarGatewayMessagingInterface;
-use crate::storage_types::{DataKey, MessageApprovalKey, MessageApprovalValue};
+use crate::storage::{MessageApprovalKey, MessageApprovalValue};
 use crate::types::{CommandType, Message, Proof, WeightedSigners};
+use crate::{auth, storage};
 
 #[contract]
 #[derive(Operatable, Ownable, Pausable, Upgradable)]
@@ -110,7 +110,8 @@ impl AxelarGatewayMessagingInterface for AxelarGateway {
             source_chain: source_chain.clone(),
             message_id: message_id.clone(),
         };
-        let message_approval = Self::message_approval_by_key(&env, key.clone());
+        let message_approval = storage::try_message_approval(&env, key.clone())
+            .unwrap_or(MessageApprovalValue::NotApproved);
         let message = Message {
             source_chain,
             message_id,
@@ -120,10 +121,7 @@ impl AxelarGatewayMessagingInterface for AxelarGateway {
         };
 
         if message_approval == Self::message_approval_hash(&env, message.clone()) {
-            env.storage().persistent().set(
-                &DataKey::MessageApproval(key),
-                &MessageApprovalValue::Executed,
-            );
+            storage::set_message_approval(&env, key, &MessageApprovalValue::Executed);
 
             MessageExecutedEvent { message }.emit(&env);
 
@@ -137,15 +135,15 @@ impl AxelarGatewayMessagingInterface for AxelarGateway {
 #[contractimpl]
 impl AxelarGatewayInterface for AxelarGateway {
     fn domain_separator(env: &Env) -> BytesN<32> {
-        auth::domain_separator(env)
+        storage::domain_separator(env)
     }
 
     fn minimum_rotation_delay(env: &Env) -> u64 {
-        auth::minimum_rotation_delay(env)
+        storage::minimum_rotation_delay(env)
     }
 
     fn previous_signers_retention(env: &Env) -> u64 {
-        auth::previous_signers_retention(env)
+        storage::previous_signer_retention(env)
     }
 
     #[when_not_paused]
@@ -170,13 +168,15 @@ impl AxelarGatewayInterface for AxelarGateway {
             };
 
             // Prevent replay if message is already approved/executed
-            let message_approval = Self::message_approval_by_key(env, key.clone());
+            let message_approval = storage::try_message_approval(env, key.clone())
+                .unwrap_or(MessageApprovalValue::NotApproved);
             if message_approval != MessageApprovalValue::NotApproved {
                 continue;
             }
 
-            env.storage().persistent().set(
-                &DataKey::MessageApproval(key),
+            storage::set_message_approval(
+                env,
+                key,
                 &Self::message_approval_hash(env, message.clone()),
             );
 
@@ -214,15 +214,16 @@ impl AxelarGatewayInterface for AxelarGateway {
     }
 
     fn epoch(env: &Env) -> u64 {
-        auth::epoch(env)
+        storage::epoch(env)
     }
 
     fn epoch_by_signers_hash(env: &Env, signers_hash: BytesN<32>) -> Result<u64, ContractError> {
-        auth::epoch_by_signers_hash(env, signers_hash)
+        storage::try_epoch_by_signers_hash(env, signers_hash)
+            .ok_or(ContractError::InvalidSignersHash)
     }
 
     fn signers_hash_by_epoch(env: &Env, epoch: u64) -> Result<BytesN<32>, ContractError> {
-        auth::signers_hash_by_epoch(env, epoch)
+        storage::try_signers_hash_by_epoch(env, epoch).ok_or(ContractError::InvalidEpoch)
     }
 
     fn validate_proof(
@@ -246,15 +247,7 @@ impl AxelarGateway {
             message_id,
         };
 
-        Self::message_approval_by_key(env, key)
-    }
-
-    /// Get the message approval value by key, defaulting to `MessageNotApproved`
-    fn message_approval_by_key(env: &Env, key: MessageApprovalKey) -> MessageApprovalValue {
-        env.storage()
-            .persistent()
-            .get(&DataKey::MessageApproval(key))
-            .unwrap_or(MessageApprovalValue::NotApproved)
+        storage::try_message_approval(env, key).unwrap_or(MessageApprovalValue::NotApproved)
     }
 
     fn message_approval_hash(env: &Env, message: Message) -> MessageApprovalValue {

contracts/stellar-axelar-gateway/src/lib.rs

@@ -28,7 +28,7 @@ cfg_if::cfg_if! {
     } else {
         mod auth;
         pub mod event;
-        mod storage_types;
+        mod storage;
         mod contract;
 
         pub use contract::{AxelarGateway, AxelarGatewayClient};

contracts/stellar-axelar-gateway/src/storage.rs

@@ -0,0 +1,57 @@
+use soroban_sdk::{contracttype, BytesN, String};
+use stellar_axelar_std::contractstorage;
+
+#[contracttype]
+#[derive(Clone, Debug)]
+pub struct MessageApprovalKey {
+    pub source_chain: String,
+    pub message_id: String,
+}
+
+#[contracttype]
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub enum MessageApprovalValue {
+    NotApproved,
+    Approved(BytesN<32>),
+    Executed,
+}
+
+#[contractstorage]
+#[derive(Clone, Debug)]
+pub enum DataKey {
+    /// Gateway
+    #[persistent]
+    #[value(MessageApprovalValue)]
+    MessageApproval {
+        message_approval_key: MessageApprovalKey,
+    },
+
+    /// Auth Module
+    #[instance]
+    #[value(u64)]
+    PreviousSignerRetention,
+
+    #[instance]
+    #[value(BytesN<32>)]
+    DomainSeparator,
+
+    #[instance]
+    #[value(u64)]
+    MinimumRotationDelay,
+
+    #[instance]
+    #[value(u64)]
+    Epoch,
+
+    #[instance]
+    #[value(u64)]
+    LastRotationTimestamp,
+
+    #[persistent]
+    #[value(BytesN<32>)]
+    SignersHashByEpoch { epoch: u64 },
+
+    #[persistent]
+    #[value(u64)]
+    EpochBySignersHash { signers_hash: BytesN<32> },
+}