ACL: don't panic if acl isn't initialized

karim-en · karim-en · commit 35635bc6e6e4 · 2023-02-28T02:21:24.000Z
diff --git a/near-plugins-derive/src/access_controllable.rs b/near-plugins-derive/src/access_controllable.rs
@@ -108,17 +108,18 @@ pub fn access_controllable(attrs: TokenStream, item: TokenStream) -> TokenStream
                 })
             }
 
-            fn acl_get(&self) -> #acl_type {
-                self.acl_get_storage().unwrap_or_else(|| ::near_sdk::env::panic_str("ACL: storage isn't initialized"))
+            fn acl_get_or_init(&self) -> #acl_type {
+                self.acl_get_storage().unwrap_or_else(|| self.acl_init_storage_unchecked())
             }
 
-            fn acl_init_storage_unchecked(&self) {
+            fn acl_init_storage_unchecked(&self) -> #acl_type {
                 let base_prefix = <#ident as AccessControllable>::acl_storage_prefix();
                 let acl_storage: #acl_type = Default::default();
                 near_sdk::env::storage_write(
                     &__acl_storage_prefix(base_prefix, __AclStorageKey::AclStorage),
                     &acl_storage.try_to_vec().unwrap(),
                 );
+                acl_storage
             }
         }
 
@@ -496,6 +497,39 @@ pub fn access_controllable(attrs: TokenStream, item: TokenStream) -> TokenStream
             }
         }
 
+        fn get_default_permissioned_accounts() -> #cratename::access_controllable::PermissionedAccounts {
+            let roles = <#role_type>::acl_role_variants();
+            let mut map = ::std::collections::HashMap::new();
+            for role in roles {
+                let role: #role_type = ::std::convert::TryFrom::try_from(role)
+                    .unwrap_or_else(|_| ::near_sdk::env::panic_str(#ERR_PARSE_ROLE));
+
+                map.insert(
+                    role.into(),
+                    #cratename::access_controllable::PermissionedAccountsPerRole {
+                        admins: Default::default(),
+                        grantees: Default::default(),
+                    }
+                );
+            }
+
+            #cratename::access_controllable::PermissionedAccounts {
+                super_admins: Default::default(),
+                roles: map,
+            }
+        }
+
+        macro_rules! return_if_none {
+            ($res:expr, $default_value:expr) => {
+                match $res {
+                    Some(val) => val,
+                    None => {
+                        return $default_value;
+                    }
+                }
+            };
+        }
+
         // Note that `#[near-bindgen]` exposes non-public functions in trait
         // implementations. This is [documented] behavior. Therefore some
         // functions are made `#[private]` despite _not_ being public.
@@ -507,63 +541,58 @@ pub fn access_controllable(attrs: TokenStream, item: TokenStream) -> TokenStream
                 (#storage_prefix).as_bytes()
             }
 
-            fn acl_init_storage(&self) {
-                ::near_sdk::require!(self.acl_get_storage().is_none(), "ACL: storage was already initialized");
-                self.acl_init_storage_unchecked();
-            }
-
             #[private]
             fn acl_init_super_admin(&mut self, account_id: ::near_sdk::AccountId) -> bool {
-                self.acl_get().init_super_admin(&account_id)
+                self.acl_get_or_init().init_super_admin(&account_id)
             }
 
             fn acl_role_variants(&self) -> Vec<&'static str> {
                 <#role_type>::acl_role_variants()
             }
 
             fn acl_is_super_admin(&self, account_id: ::near_sdk::AccountId) -> bool {
-                self.acl_get().is_super_admin(&account_id)
+                return_if_none!(self.acl_get_storage(), false).is_super_admin(&account_id)
             }
 
             fn acl_add_admin(&mut self, role: String, account_id: ::near_sdk::AccountId) -> Option<bool> {
                 let role: #role_type = ::std::convert::TryFrom::try_from(role.as_str()).unwrap_or_else(|_| ::near_sdk::env::panic_str(#ERR_PARSE_ROLE));
-                self.acl_get().add_admin(role, &account_id)
+                self.acl_get_or_init().add_admin(role, &account_id)
             }
 
             fn acl_is_admin(&self, role: String, account_id: ::near_sdk::AccountId) -> bool {
                 let role: #role_type = ::std::convert::TryFrom::try_from(role.as_str()).unwrap_or_else(|_| ::near_sdk::env::panic_str(#ERR_PARSE_ROLE));
-                self.acl_get().is_admin(role, &account_id)
+                return_if_none!(self.acl_get_storage(), false).is_admin(role, &account_id)
             }
 
             fn acl_revoke_admin(&mut self, role: String, account_id: ::near_sdk::AccountId) -> Option<bool> {
                 let role: #role_type = ::std::convert::TryFrom::try_from(role.as_str()).unwrap_or_else(|_| ::near_sdk::env::panic_str(#ERR_PARSE_ROLE));
-                self.acl_get().revoke_admin(role, &account_id)
+                self.acl_get_or_init().revoke_admin(role, &account_id)
             }
 
             fn acl_renounce_admin(&mut self, role: String) -> bool {
                 let role: #role_type = ::std::convert::TryFrom::try_from(role.as_str()).unwrap_or_else(|_| ::near_sdk::env::panic_str(#ERR_PARSE_ROLE));
-                self.acl_get().renounce_admin(role)
+                self.acl_get_or_init().renounce_admin(role)
             }
 
             fn acl_revoke_role(&mut self, role: String, account_id: ::near_sdk::AccountId) -> Option<bool> {
                 let role: #role_type = ::std::convert::TryFrom::try_from(role.as_str()).unwrap_or_else(|_| ::near_sdk::env::panic_str(#ERR_PARSE_ROLE));
-                self.acl_get().revoke_role(role, &account_id)
+                self.acl_get_or_init().revoke_role(role, &account_id)
             }
 
             fn acl_renounce_role(&mut self, role: String) -> bool {
                 let role: #role_type = ::std::convert::TryFrom::try_from(role.as_str()).unwrap_or_else(|_| ::near_sdk::env::panic_str(#ERR_PARSE_ROLE));
-                self.acl_get().renounce_role(role)
+                self.acl_get_or_init().renounce_role(role)
             }
 
             fn acl_grant_role(&mut self, role: String, account_id: ::near_sdk::AccountId) -> Option<bool> {
                 let role: #role_type = ::std::convert::TryFrom::try_from(role.as_str()).unwrap_or_else(|_| ::near_sdk::env::panic_str(#ERR_PARSE_ROLE));
-                self.acl_get().grant_role(role, &account_id)
+                self.acl_get_or_init().grant_role(role, &account_id)
             }
 
 
             fn acl_has_role(&self, role: String, account_id: ::near_sdk::AccountId) -> bool {
                 let role: #role_type = ::std::convert::TryFrom::try_from(role.as_str()).unwrap_or_else(|_| ::near_sdk::env::panic_str(#ERR_PARSE_ROLE));
-                self.acl_get().has_role(role, &account_id)
+                return_if_none!(self.acl_get_storage(), false).has_role(role, &account_id)
             }
 
             fn acl_has_any_role(&self, roles: Vec<String>, account_id: ::near_sdk::AccountId) -> bool {
@@ -573,33 +602,33 @@ pub fn access_controllable(attrs: TokenStream, item: TokenStream) -> TokenStream
                         ::std::convert::TryFrom::try_from(role.as_str()).unwrap_or_else(|_| ::near_sdk::env::panic_str(#ERR_PARSE_ROLE))
                     })
                     .collect();
-                self.acl_get().has_any_role(roles, &account_id)
+                return_if_none!(self.acl_get_storage(), false).has_any_role(roles, &account_id)
             }
 
             fn acl_get_super_admins(&self, skip: u64, limit: u64) -> Vec<::near_sdk::AccountId> {
                 let permission = <#bitflags_type>::from_bits(
                     <#role_type>::acl_super_admin_permission()
                 )
                 .unwrap_or_else(|| ::near_sdk::env::panic_str(#ERR_PARSE_BITFLAG));
-                self.acl_get().get_bearers(permission, skip, limit)
+                return_if_none!(self.acl_get_storage(), vec![]).get_bearers(permission, skip, limit)
             }
 
             fn acl_get_admins(&self, role: String, skip: u64, limit: u64) -> Vec<::near_sdk::AccountId> {
                 let role: #role_type = ::std::convert::TryFrom::try_from(role.as_str()).unwrap_or_else(|_| ::near_sdk::env::panic_str(#ERR_PARSE_ROLE));
                 let permission = <#bitflags_type>::from_bits(role.acl_admin_permission())
                     .unwrap_or_else(|| ::near_sdk::env::panic_str(#ERR_PARSE_BITFLAG));
-                self.acl_get().get_bearers(permission, skip, limit)
+                return_if_none!(self.acl_get_storage(), vec![]).get_bearers(permission, skip, limit)
             }
 
             fn acl_get_grantees(&self, role: String, skip: u64, limit: u64) -> Vec<::near_sdk::AccountId> {
                 let role: #role_type = ::std::convert::TryFrom::try_from(role.as_str()).unwrap_or_else(|_| ::near_sdk::env::panic_str(#ERR_PARSE_ROLE));
                 let permission = <#bitflags_type>::from_bits(role.acl_permission())
                     .unwrap_or_else(|| ::near_sdk::env::panic_str(#ERR_PARSE_BITFLAG));
-                self.acl_get().get_bearers(permission, skip, limit)
+                return_if_none!(self.acl_get_storage(), vec![]).get_bearers(permission, skip, limit)
             }
 
             fn acl_get_permissioned_accounts(&self) -> #cratename::access_controllable::PermissionedAccounts {
-                self.acl_get().get_permissioned_accounts()
+                return_if_none!(self.acl_get_storage(), get_default_permissioned_accounts()).get_permissioned_accounts()
             }
         }
     };
diff --git a/near-plugins-derive/tests/contracts/access_controllable/src/lib.rs b/near-plugins-derive/tests/contracts/access_controllable/src/lib.rs
@@ -38,7 +38,6 @@ impl Counter {
     #[init]
     pub fn new(admins: HashMap<String, AccountId>, grantees: HashMap<String, AccountId>) -> Self {
         let mut contract = Self { counter: 0 };
-        contract.acl_init_storage();
 
         if admins.len() > 0 || grantees.len() > 0 {
             // First we make the contract itself super admin to allow it adding admin and grantees.
@@ -65,11 +64,11 @@ impl Counter {
             // granting roles, for example:
             //
             // ```
-            // contract.acl_get().add_admin_unchecked(role, account_id);
-            // contract.acl_get().grant_role_unchecked(role, account_id);
+            // contract.acl_get_or_init().add_admin_unchecked(role, account_id);
+            // contract.acl_get_or_init().grant_role_unchecked(role, account_id);
             // ```
             //
-            // **Attention**: for security reasons, `acl_get().*_unchecked` methods should only be called
+            // **Attention**: for security reasons, `acl_get_or_init().*_unchecked` methods should only be called
             // from within methods with attribute `#[init]` or `#[private]`.
         }
 
@@ -133,7 +132,7 @@ impl Counter {
             self.acl_is_super_admin(env::predecessor_account_id()),
             "Only super admins are allowed to add other super admins."
         );
-        self.acl_get().add_super_admin_unchecked(&account_id)
+        self.acl_get_or_init().add_super_admin_unchecked(&account_id)
     }
 }
 
@@ -142,32 +141,32 @@ impl Counter {
 impl Counter {
     #[private]
     pub fn acl_add_super_admin_unchecked(&mut self, account_id: AccountId) -> bool {
-        self.acl_get().add_super_admin_unchecked(&account_id)
+        self.acl_get_or_init().add_super_admin_unchecked(&account_id)
     }
 
     #[private]
     pub fn acl_revoke_super_admin_unchecked(&mut self, account_id: AccountId) -> bool {
-        self.acl_get().revoke_super_admin_unchecked(&account_id)
+        self.acl_get_or_init().revoke_super_admin_unchecked(&account_id)
     }
 
     #[private]
     pub fn acl_revoke_role_unchecked(&mut self, role: Role, account_id: AccountId) -> bool {
-        self.acl_get()
+        self.acl_get_or_init()
             .revoke_role_unchecked(role.into(), &account_id)
     }
 
     #[private]
     pub fn acl_add_admin_unchecked(&mut self, role: Role, account_id: AccountId) -> bool {
-        self.acl_get().add_admin_unchecked(role, &account_id)
+        self.acl_get_or_init().add_admin_unchecked(role, &account_id)
     }
 
     #[private]
     pub fn acl_revoke_admin_unchecked(&mut self, role: Role, account_id: AccountId) -> bool {
-        self.acl_get().revoke_admin_unchecked(role, &account_id)
+        self.acl_get_or_init().revoke_admin_unchecked(role, &account_id)
     }
 
     #[private]
     pub fn acl_grant_role_unchecked(&mut self, role: Role, account_id: AccountId) -> bool {
-        self.acl_get().grant_role_unchecked(role, &account_id)
+        self.acl_get_or_init().grant_role_unchecked(role, &account_id)
     }
 }
diff --git a/near-plugins-derive/tests/contracts/pausable/src/lib.rs b/near-plugins-derive/tests/contracts/pausable/src/lib.rs
@@ -41,7 +41,6 @@ impl Counter {
         let mut contract = Self {
             counter: 0,
         };
-        contract.acl_init_storage();
 
         // Make the contract itself super admin. This allows us to grant any role in the
         // constructor.
diff --git a/near-plugins/src/access_controllable.rs b/near-plugins/src/access_controllable.rs
@@ -29,9 +29,6 @@ pub trait AccessControllable {
     /// ```
     fn acl_storage_prefix() -> &'static [u8];
 
-    /// Initialize the access control storage.
-    fn acl_init_storage(&self);
-
     /// Returns the names of all variants of the enum that represents roles.
     ///
     /// In the default implementation provided by this crate, this enum is defined by contract
