2-step ownership transfer

amm/contracts/stable_pool/lib.rs

@@ -27,7 +27,10 @@ pub mod stable_pool {
         {vec, vec::Vec},
     };
     use psp22::{PSP22Data, PSP22Error, PSP22Event, PSP22Metadata, PSP22};
-    use traits::{MathError, StablePool, StablePoolError, StablePoolView};
+    use traits::{
+        MathError, Ownable2Step, Ownable2StepData, Ownable2StepResult, StablePool, StablePoolError,
+        StablePoolView,
+    };
 
     #[ink(event)]
     pub struct AddLiquidity {
@@ -97,8 +100,12 @@ pub mod stable_pool {
     }
 
     #[ink(event)]
-    pub struct OwnerChanged {
-        #[ink(topic)]
+    pub struct TransferOwnershipInitiated {
+        pub new_owner: AccountId,
+    }
+
+    #[ink(event)]
+    pub struct TransferOwnershipAccepted {
         pub new_owner: AccountId,
     }
 
@@ -141,7 +148,7 @@ pub mod stable_pool {
 
     #[ink(storage)]
     pub struct StablePoolContract {
-        owner: AccountId,
+        ownable: Ownable2StepData,
         pool: StablePoolData,
         psp22: PSP22Data,
     }
@@ -192,7 +199,7 @@ pub mod stable_pool {
                 })
                 .collect();
             Ok(Self {
-                owner,
+                ownable: Ownable2StepData::new(owner),
                 pool: StablePoolData {
                     tokens,
                     reserves: vec![0; token_count],
@@ -330,14 +337,6 @@ pub mod stable_pool {
                 .collect()
         }
 
-        fn ensure_owner(&self) -> Result<(), StablePoolError> {
-            ensure!(
-                self.env().caller() == self.owner,
-                StablePoolError::OnlyOwner
-            );
-            Ok(())
-        }
-
         fn token_id(&self, token: AccountId) -> Result<usize, StablePoolError> {
             self.pool
                 .tokens
@@ -783,14 +782,6 @@ pub mod stable_pool {
             self._swap_exact_in(token_in, token_out, None, min_token_out_amount, to)
         }
 
-        #[ink(message)]
-        fn set_owner(&mut self, new_owner: AccountId) -> Result<(), StablePoolError> {
-            self.ensure_owner()?;
-            self.owner = new_owner;
-            self.env().emit_event(OwnerChanged { new_owner });
-            Ok(())
-        }
-
         #[ink(message)]
         fn set_fee_receiver(
             &mut self,
@@ -1063,4 +1054,39 @@ pub mod stable_pool {
             TOKEN_TARGET_DECIMALS
         }
     }
+
+    impl Ownable2Step for StablePoolContract {
+        #[ink(message)]
+        fn get_owner(&self) -> Ownable2StepResult<AccountId> {
+            self.ownable.get_owner()
+        }
+
+        #[ink(message)]
+        fn get_pending_owner(&self) -> Ownable2StepResult<AccountId> {
+            self.ownable.get_pending_owner()
+        }
+
+        #[ink(message)]
+        fn transfer_ownership(&mut self, new_owner: AccountId) -> Ownable2StepResult<()> {
+            self.ownable
+                .transfer_ownership(self.env().caller(), new_owner)?;
+            self.env()
+                .emit_event(TransferOwnershipInitiated { new_owner });
+            Ok(())
+        }
+
+        #[ink(message)]
+        fn accept_ownership(&mut self) -> Ownable2StepResult<()> {
+            let new_owner = self.env().caller();
+            self.ownable.accept_ownership(new_owner)?;
+            self.env()
+                .emit_event(TransferOwnershipAccepted { new_owner });
+            Ok(())
+        }
+
+        #[ink(message)]
+        fn ensure_owner(&self) -> Ownable2StepResult<()> {
+            self.ownable.ensure_owner(self.env().caller())
+        }
+    }
 }

amm/traits/lib.rs

@@ -1,10 +1,12 @@
 #![cfg_attr(not(feature = "std"), no_std, no_main)]
 
+mod ownable2step;
 mod rate_provider;
 mod stable_pool;
 
 pub type Balance = <ink::env::DefaultEnvironment as ink::env::Environment>::Balance;
 
 pub use amm_helpers::math::MathError;
+pub use ownable2step::{Ownable2Step, Ownable2StepData, Ownable2StepError, Ownable2StepResult};
 pub use rate_provider::RateProvider;
 pub use stable_pool::{StablePool, StablePoolError, StablePoolView};

amm/traits/ownable2step.rs

@@ -0,0 +1,106 @@
+use ink::{prelude::string::String, primitives::AccountId};
+use scale::{Decode, Encode};
+
+/// Implement this trait to enable two-step ownership trasfer process in your contract.
+///
+/// The process looks like this:
+/// * current owner (Alice) calls `self.transfer_ownership(bob)`,
+/// * the contract still has the owner: Alice and a pending owner: bob,
+/// * when Bob claims the ownership by calling `self.accept_ownership()` he becomes the new owner and pending owner is removed.
+///
+/// The methods are all wrapper in `Ownable2StepResult` to make it possible to use them in settings where the `Data` is e.g. behid `Lazy`.
+#[ink::trait_definition]
+pub trait Ownable2Step {
+    /// Returns the address of the current owner.
+    #[ink(message)]
+    fn get_owner(&self) -> Ownable2StepResult<AccountId>;
+
+    /// Returns the address of the pending owner.
+    #[ink(message)]
+    fn get_pending_owner(&self) -> Ownable2StepResult<AccountId>;
+
+    /// Starts the ownership transfer of the contract to a new account. Replaces the pending transfer if there is one.
+    /// Can only be called by the current owner.
+    #[ink(message)]
+    fn transfer_ownership(&mut self, new_owner: AccountId) -> Ownable2StepResult<()>;
+
+    /// The new owner accepts the ownership transfer.
+    #[ink(message)]
+    fn accept_ownership(&mut self) -> Ownable2StepResult<()>;
+
+    /// Return error if called by any account other than the owner.
+    #[ink(message)]
+    fn ensure_owner(&self) -> Ownable2StepResult<()>;
+}
+
+#[derive(Debug, PartialEq, Eq, Encode, Decode)]
+#[cfg_attr(feature = "std", derive(scale_info::TypeInfo))]
+pub enum Ownable2StepError {
+    /// The caller didn't have the permissions to call a given method
+    CallerNotOwner(AccountId),
+    /// The caller tried to accept ownership but caller in not the pending owner
+    CallerNotPendingOwner(AccountId),
+    /// The caller tried to accept ownership but the process hasn't been started
+    NoPendingOwner,
+    /// Useful in cases, when the `Data` struct is not accessed directly but inside of `Lazy` or a `Mapping`, means that we failed to access the `Data` struct itself.
+    Custom(String),
+}
+
+pub type Ownable2StepResult<T> = Result<T, Ownable2StepError>;
+
+#[derive(Debug)]
+#[ink::storage_item]
+pub struct Ownable2StepData {
+    owner: AccountId,
+    pending_owner: Option<AccountId>,
+}
+
+impl Ownable2StepData {
+    pub fn new(owner: AccountId) -> Self {
+        Self {
+            owner,
+            pending_owner: None,
+        }
+    }
+
+    pub fn transfer_ownership(
+        &mut self,
+        caller: AccountId,
+        new_owner: AccountId,
+    ) -> Ownable2StepResult<()> {
+        self.ensure_owner(caller)?;
+        self.pending_owner = Some(new_owner);
+        Ok(())
+    }
+
+    pub fn accept_ownership(&mut self, caller: AccountId) -> Ownable2StepResult<()> {
+        let pending_owner = self
+            .pending_owner
+            .ok_or(Ownable2StepError::NoPendingOwner)?;
+
+        if caller != pending_owner {
+            return Err(Ownable2StepError::CallerNotPendingOwner(caller));
+        }
+
+        self.owner = pending_owner;
+        self.pending_owner = None;
+
+        Ok(())
+    }
+
+    pub fn get_owner(&self) -> Ownable2StepResult<AccountId> {
+        Ok(self.owner)
+    }
+
+    pub fn get_pending_owner(&self) -> Ownable2StepResult<AccountId> {
+        self.pending_owner.ok_or(Ownable2StepError::NoPendingOwner)
+    }
+
+    pub fn ensure_owner(&self, caller: AccountId) -> Ownable2StepResult<()> {
+        if caller != self.owner {
+            Err(Ownable2StepError::CallerNotOwner(caller))
+        } else {
+            Ok(())
+        }
+    }
+}

amm/traits/stable_pool.rs

@@ -3,7 +3,7 @@ use ink::primitives::AccountId;
 use ink::LangError;
 use psp22::PSP22Error;
 
-use crate::MathError;
+use crate::{MathError, Ownable2StepError};
 
 #[ink::trait_definition]
 pub trait StablePoolView {
@@ -179,9 +179,6 @@ pub trait StablePool {
 
     // --- OWNER RESTRICTED FUNCTIONS --- //
 
-    #[ink(message)]
-    fn set_owner(&mut self, new_owner: AccountId) -> Result<(), StablePoolError>;
-
     #[ink(message)]
     fn set_fee_receiver(&mut self, fee_receiver: Option<AccountId>) -> Result<(), StablePoolError>;
 
@@ -202,6 +199,7 @@ pub trait StablePool {
 #[derive(Debug, PartialEq, Eq, scale::Encode, scale::Decode)]
 #[cfg_attr(feature = "std", derive(scale_info::TypeInfo))]
 pub enum StablePoolError {
+    Ownable2StepError(Ownable2StepError),
     MathError(MathError),
     PSP22Error(PSP22Error),
     LangError(LangError),
@@ -217,7 +215,6 @@ pub enum StablePoolError {
     IncorrectTokenCount,
     TooLargeTokenDecimal,
     InvalidFee,
-    OnlyOwner,
 }
 
 impl From<PSP22Error> for StablePoolError {
@@ -237,3 +234,9 @@ impl From<MathError> for StablePoolError {
         StablePoolError::MathError(error)
     }
 }
+
+impl From<Ownable2StepError> for StablePoolError {
+    fn from(error: Ownable2StepError) -> Self {
+        StablePoolError::Ownable2StepError(error)
+    }
+}