Use enum for lock state representation.

Cardinal-Cryptography · Aug 25, 2023 · 9e3a56e · 9e3a56e
1 parent 6028508
commit 9e3a56e
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 8 deletions.
diff --git a/farm/contracts/lib.rs b/farm/contracts/lib.rs
@@ -68,12 +68,9 @@ mod farm {
         /// Farm state.
         state: Lazy<State, ManualKey<0x4641524d>>,
         /// Flag to prevent reentrancy attacks.
-        reentrancy_guard: u8,
+        reentrancy_guard: ReentrancyLock,
     }
 
-    const REENTRANCY_GUARD_LOCKED: u8 = 1u8;
-    const REENTRANCY_GUARD_FREE: u8 = 0u8;
-
     const SCALING_FACTOR: u128 = 10_u128.pow(18);
 
     impl Farm {
@@ -84,7 +81,7 @@ mod farm {
                 owner: Self::env().caller(),
                 is_stopped: true,
                 state: Lazy::new(),
-                reentrancy_guard: REENTRANCY_GUARD_FREE,
+                reentrancy_guard: ReentrancyLock::Unlocked,
             }
         }
 
@@ -511,16 +508,16 @@ mod farm {
 
     impl ReentrancyGuardT for Farm {
         fn lock(&mut self) -> Result<(), ReentrancyGuardError> {
-            if self.reentrancy_guard == REENTRANCY_GUARD_LOCKED {
+            if self.reentrancy_guard == ReentrancyLock::Locked {
                 return Err(ReentrancyGuardError::ReentrancyError)
             }
-            self.reentrancy_guard = REENTRANCY_GUARD_LOCKED;
+            self.reentrancy_guard = ReentrancyLock::Locked;
             Ok(())
         }
 
         fn unlock(&mut self) {
             // It's safe to "unlock" already unlocked guard.
-            self.reentrancy_guard = REENTRANCY_GUARD_FREE;
+            self.reentrancy_guard = ReentrancyLock::Unlocked;
         }
     }
 

diff --git a/farm/contracts/reentrancy_guard.rs b/farm/contracts/reentrancy_guard.rs
@@ -27,3 +27,15 @@ pub trait ReentrancyGuardT {
 pub enum ReentrancyGuardError {
     ReentrancyError,
 }
+
+/// Simple enum representing reentrancy lock state.
+#[derive(scale::Encode, scale::Decode, Debug, Eq, PartialEq)]
+#[cfg_attr(
+    feature = "std",
+    derive(scale_info::TypeInfo, ink::storage::traits::StorageLayout)
+)]
+#[repr(u8)]
+pub enum ReentrancyLock {
+    Locked,
+    Unlocked,
+}