gmsol_store/states/
roles.rs

1use anchor_lang::prelude::*;
2
3use gmsol_utils::bitmaps::Bitmap;
4
5use crate::CoreError;
6
7use super::InitSpace;
8
9pub use gmsol_utils::role::{RoleKey, MAX_ROLE_NAME_LEN};
10
11/// Max number of roles.
12pub const MAX_ROLES: usize = 32;
13
14/// Max number of members.
15pub const MAX_MEMBERS: usize = 64;
16
17type RoleBitmap = Bitmap<MAX_ROLES>;
18type RoleBitmapValue = u32;
19
20/// Role Metadata.
21#[zero_copy]
22#[derive(Default)]
23#[cfg_attr(feature = "debug", derive(Debug))]
24pub struct RoleMetadata {
25    name: [u8; MAX_ROLE_NAME_LEN],
26    enabled: u8,
27    index: u8,
28}
29
30impl InitSpace for RoleMetadata {
31    const INIT_SPACE: usize = 32 + 2;
32}
33
34#[cfg(test)]
35const_assert_eq!(
36    std::mem::size_of::<RoleMetadata>(),
37    RoleMetadata::INIT_SPACE
38);
39
40impl RoleMetadata {
41    /// A `u8` value indicates that this role is enabled.
42    pub const ROLE_ENABLED: u8 = u8::MAX;
43
44    fn name_to_bytes(name: &str) -> Result<[u8; MAX_ROLE_NAME_LEN]> {
45        crate::utils::fixed_str::fixed_str_to_bytes(name)
46    }
47
48    fn bytes_to_name(bytes: &[u8; 32]) -> Result<&str> {
49        crate::utils::fixed_str::bytes_to_fixed_str(bytes)
50    }
51
52    /// Create a new role metadata.
53    pub fn new(name: &str, index: u8) -> Result<Self> {
54        Ok(Self {
55            name: Self::name_to_bytes(name)?,
56            enabled: Self::ROLE_ENABLED,
57            index,
58        })
59    }
60
61    /// Get the name of this role.
62    pub fn name(&self) -> Result<&str> {
63        Self::bytes_to_name(&self.name)
64    }
65
66    /// Enable this role.
67    ///
68    /// # Errors
69    /// Returns Error if this role is already enabled.
70    pub fn enable(&mut self) -> Result<()> {
71        require!(!self.is_enabled(), CoreError::PreconditionsAreNotMet);
72        self.set_enable();
73        Ok(())
74    }
75
76    /// Disable this role.
77    ///
78    /// # Errors
79    /// Returns Error if this role is already disabled.
80    pub fn disable(&mut self) -> Result<()> {
81        require!(self.is_enabled(), CoreError::PreconditionsAreNotMet);
82        self.set_disable();
83        Ok(())
84    }
85
86    /// Enable this role.
87    fn set_enable(&mut self) {
88        self.enabled = Self::ROLE_ENABLED;
89    }
90
91    /// Disable this role.
92    fn set_disable(&mut self) {
93        self.enabled = 0;
94    }
95
96    /// Is enbaled.
97    pub fn is_enabled(&self) -> bool {
98        self.enabled == Self::ROLE_ENABLED
99    }
100}
101
102gmsol_utils::fixed_map!(RoleMap, RoleMetadata, MAX_ROLES, 0);
103
104gmsol_utils::fixed_map!(
105    Members,
106    Pubkey,
107    crate::utils::pubkey::to_bytes,
108    u32,
109    MAX_MEMBERS,
110    0
111);
112
113/// Roles Store.
114#[zero_copy]
115#[cfg_attr(feature = "debug", derive(Debug))]
116pub struct RoleStore {
117    roles: RoleMap,
118    members: Members,
119}
120
121impl InitSpace for RoleStore {
122    const INIT_SPACE: usize = std::mem::size_of::<RoleStore>();
123}
124
125impl RoleStore {
126    /// Enable a role.
127    pub fn enable_role(&mut self, role: &str) -> Result<()> {
128        match self.roles.get_mut(role) {
129            Some(metadata) => {
130                require_eq!(metadata.name()?, role, CoreError::InvalidArgument);
131                metadata.enable()?;
132            }
133            None => {
134                let index = self
135                    .roles
136                    .len()
137                    .try_into()
138                    .map_err(|_| error!(CoreError::ExceedMaxLengthLimit))?;
139                self.roles
140                    .insert_with_options(role, RoleMetadata::new(role, index)?, true)?;
141            }
142        }
143        Ok(())
144    }
145
146    /// Disable a role.
147    pub fn disable_role(&mut self, role: &str) -> Result<()> {
148        if let Some(metadata) = self.roles.get_mut(role) {
149            require_eq!(metadata.name()?, role, CoreError::InvalidArgument);
150            metadata.disable()?;
151        }
152        Ok(())
153    }
154
155    /// Get the index of a role.
156    pub fn role_index(&self, role: &str) -> Result<Option<u8>> {
157        if let Some(metadata) = self.roles.get(role) {
158            require_eq!(metadata.name()?, role, CoreError::InvalidArgument);
159            Ok(Some(metadata.index))
160        } else {
161            Ok(None)
162        }
163    }
164
165    /// Get the index of a enabled role.
166    pub fn enabled_role_index(&self, role: &str) -> Result<Option<u8>> {
167        if let Some(metadata) = self.roles.get(role) {
168            require_eq!(metadata.name()?, role, CoreError::InvalidArgument);
169            require!(metadata.is_enabled(), CoreError::PreconditionsAreNotMet);
170            Ok(Some(metadata.index))
171        } else {
172            Ok(None)
173        }
174    }
175
176    /// Check if the given role is granted to the pubkey.
177    pub fn has_role(&self, authority: &Pubkey, role: &str) -> Result<bool> {
178        let Some(value) = self.members.get(authority) else {
179            return err!(CoreError::PermissionDenied);
180        };
181        let Some(index) = self.enabled_role_index(role)? else {
182            return err!(CoreError::NotFound);
183        };
184        let bitmap = RoleBitmap::from_value(*value);
185        Ok(bitmap.get(index as usize))
186    }
187
188    /// Grant a role to the pubkey.
189    ///
190    /// # Errors
191    /// - The `role` must be enabled.
192    /// - The `authority` must not already have the role.
193    pub fn grant(&mut self, authority: &Pubkey, role: &str) -> Result<()> {
194        let Some(index) = self.enabled_role_index(role)? else {
195            return err!(CoreError::NotFound);
196        };
197        let index = index as usize;
198        match self.members.get_mut(authority) {
199            Some(value) => {
200                let mut bitmap = RoleBitmap::from_value(*value);
201                require!(!bitmap.get(index), CoreError::PreconditionsAreNotMet);
202                bitmap.set(index, true);
203                *value = bitmap.into_value();
204            }
205            None => {
206                let mut bitmap = RoleBitmap::new();
207                bitmap.set(index, true);
208                self.members
209                    .insert_with_options(authority, bitmap.into_value(), true)?;
210            }
211        }
212        Ok(())
213    }
214
215    /// Revoke a role from the pubkey.
216    ///
217    /// # Errors
218    /// - The `authority` must have the role.
219    pub fn revoke(&mut self, authority: &Pubkey, role: &str) -> Result<()> {
220        // The `role` does not have to be enabled.
221        // This is useful when we want to modify the role configuration before enabling it.
222        let Some(index) = self.role_index(role)? else {
223            return err!(CoreError::NotFound);
224        };
225        let Some(value) = self.members.get_mut(authority) else {
226            return err!(CoreError::PermissionDenied);
227        };
228        let mut bitmap = RoleBitmap::from_value(*value);
229        let index = index as usize;
230        require!(bitmap.get(index), CoreError::PreconditionsAreNotMet);
231        bitmap.set(index, false);
232        *value = bitmap.into_value();
233
234        // Remove the membership if the authority no longer has a role.
235        if bitmap.is_empty() {
236            self.members.remove(authority);
237        }
238
239        Ok(())
240    }
241
242    /// Get the number of roles.
243    pub fn num_roles(&self) -> usize {
244        self.roles.len()
245    }
246
247    /// Get the number of members.
248    pub fn num_members(&self) -> usize {
249        self.members.len()
250    }
251
252    /// Get role value for the user.
253    pub fn role_value(&self, user: &Pubkey) -> Option<RoleBitmapValue> {
254        self.members.get(user).copied()
255    }
256
257    /// Get all members.
258    pub fn members(&self) -> impl Iterator<Item = Pubkey> + '_ {
259        self.members
260            .entries()
261            .map(|(key, _)| Pubkey::new_from_array(*key))
262    }
263
264    /// Get all roles.
265    pub fn roles(&self) -> impl Iterator<Item = Result<&str>> + '_ {
266        self.roles.entries().map(|(_, value)| value.name())
267    }
268}
269
270#[cfg(test)]
271mod tests {
272    use bytemuck::Zeroable;
273
274    use super::*;
275
276    #[test]
277    fn grant_and_revoke_roles() {
278        let mut store = RoleStore::zeroed();
279        let authority = Pubkey::new_unique();
280
281        assert!(store.grant(&authority, RoleKey::GT_CONTROLLER).is_err());
282        assert!(store.has_role(&authority, RoleKey::GT_CONTROLLER).is_err());
283
284        store.enable_role(RoleKey::GT_CONTROLLER).unwrap();
285        store.enable_role(RoleKey::MARKET_KEEPER).unwrap();
286
287        store.grant(&authority, RoleKey::GT_CONTROLLER).unwrap();
288        assert_eq!(store.has_role(&authority, RoleKey::GT_CONTROLLER), Ok(true));
289        store.grant(&authority, RoleKey::MARKET_KEEPER).unwrap();
290        assert_eq!(store.has_role(&authority, RoleKey::MARKET_KEEPER), Ok(true));
291        assert_eq!(store.has_role(&authority, RoleKey::GT_CONTROLLER), Ok(true));
292
293        store.revoke(&authority, RoleKey::GT_CONTROLLER).unwrap();
294        assert_eq!(store.has_role(&authority, RoleKey::MARKET_KEEPER), Ok(true));
295        assert_eq!(
296            store.has_role(&authority, RoleKey::GT_CONTROLLER),
297            Ok(false)
298        );
299
300        // This is the last role of the `authority`.
301        // So the membership will be removed after revoking the role.
302        store.revoke(&authority, RoleKey::MARKET_KEEPER).unwrap();
303        assert!(store.has_role(&authority, RoleKey::MARKET_KEEPER).is_err());
304        assert!(store.has_role(&authority, RoleKey::GT_CONTROLLER).is_err());
305
306        store.disable_role(RoleKey::MARKET_KEEPER).unwrap();
307        assert!(store.grant(&authority, RoleKey::MARKET_KEEPER).is_err());
308        assert!(store.has_role(&authority, RoleKey::MARKET_KEEPER).is_err());
309        store.enable_role(RoleKey::MARKET_KEEPER).unwrap();
310        store.grant(&authority, RoleKey::MARKET_KEEPER).unwrap();
311        assert_eq!(store.has_role(&authority, RoleKey::MARKET_KEEPER), Ok(true));
312    }
313
314    #[test]
315    fn enable_and_disable_role() {
316        let mut store = RoleStore::zeroed();
317        let authority = Pubkey::new_unique();
318
319        store.enable_role(RoleKey::GT_CONTROLLER).unwrap();
320        store.grant(&authority, RoleKey::GT_CONTROLLER).unwrap();
321        assert_eq!(store.has_role(&authority, RoleKey::GT_CONTROLLER), Ok(true));
322        store.disable_role(RoleKey::GT_CONTROLLER).unwrap();
323        assert!(store.has_role(&authority, RoleKey::GT_CONTROLLER).is_err());
324        store.enable_role(RoleKey::GT_CONTROLLER).unwrap();
325        assert_eq!(store.has_role(&authority, RoleKey::GT_CONTROLLER), Ok(true));
326    }
327}
gmsol_store/states/roles.rs

gmsol_store/states/
roles.rs
