sea_orm/rbac/
context.rs

1use super::{
2    AccessType, RbacError, RbacUserId,
3    entity::{
4        permission::{self, ActiveModel as Permission, PermissionId},
5        resource::{self, ActiveModel as Resource, ResourceId},
6        role::{self, ActiveModel as Role, RoleId},
7        role_hierarchy::{self, ActiveModel as RoleHierarchy},
8        role_permission::{self, ActiveModel as RolePermission},
9        user_override::{self, ActiveModel as UserOverride},
10        user_role::{self, ActiveModel as UserRole},
11    },
12};
13use crate::{
14    AccessMode, EntityTrait, IsolationLevel, Set, TransactionSession, TransactionTrait,
15    error::DbErr, sea_query::OnConflict,
16};
17use std::collections::HashMap;
18
19/// Helper class for manipulation of RBAC tables
20#[derive(Debug)]
21pub struct RbacContext {
22    tables: HashMap<String, ResourceId>,
23    permissions: HashMap<String, PermissionId>,
24    roles: HashMap<String, RoleId>,
25}
26
27#[derive(Debug)]
28pub struct RbacAddRoleHierarchy {
29    pub super_role: &'static str,
30    pub role: &'static str,
31}
32
33#[derive(Debug)]
34pub struct RbacAddUserOverride {
35    pub user_id: i64,
36    pub table: &'static str,
37    pub action: &'static str,
38    pub grant: bool,
39}
40
41impl RbacContext {
42    /// Load context from database connection
43    pub async fn load<C: TransactionTrait>(db: &C) -> Result<Self, DbErr> {
44        // ensure snapshot is consistent across all tables
45        let txn = &db
46            .begin_with_config(
47                Some(IsolationLevel::ReadCommitted),
48                Some(AccessMode::ReadOnly),
49            )
50            .await?;
51
52        let tables = resource::Entity::find()
53            .all(txn)
54            .await?
55            .into_iter()
56            .map(|t| (t.table, t.id))
57            .collect();
58
59        let permissions = permission::Entity::find()
60            .all(txn)
61            .await?
62            .into_iter()
63            .map(|p| (p.action, p.id))
64            .collect();
65
66        let roles = role::Entity::find()
67            .all(txn)
68            .await?
69            .into_iter()
70            .map(|r| (r.role, r.id))
71            .collect();
72
73        Ok(Self {
74            tables,
75            permissions,
76            roles,
77        })
78    }
79
80    /// Add multiple tables as resources
81    pub async fn add_tables<C: TransactionTrait>(
82        &mut self,
83        db: &C,
84        tables: &[&'static str],
85    ) -> Result<(), DbErr> {
86        let txn = db.begin().await?;
87
88        for table_name in tables {
89            if let Some(table_id) = resource::Entity::insert(Resource {
90                table: Set(table_name.to_string()),
91                ..Default::default()
92            })
93            .on_conflict_do_nothing()
94            .exec(&txn)
95            .await?
96            .last_insert_id()?
97            {
98                self.tables.insert(table_name.to_string(), table_id);
99            }
100        }
101
102        txn.commit().await
103    }
104
105    /// Add CRUD actions
106    pub async fn add_crud_permissions<C: TransactionTrait>(&mut self, db: &C) -> Result<(), DbErr> {
107        let txn = db.begin().await?;
108
109        for action in [
110            AccessType::Select,
111            AccessType::Insert,
112            AccessType::Update,
113            AccessType::Delete,
114        ] {
115            if let Some(permission_id) = permission::Entity::insert(Permission {
116                action: Set(action.as_str().to_owned()),
117                ..Default::default()
118            })
119            .on_conflict_do_nothing()
120            .exec(&txn)
121            .await?
122            .last_insert_id()?
123            {
124                self.permissions
125                    .insert(action.as_str().to_owned(), permission_id);
126            }
127        }
128
129        txn.commit().await
130    }
131
132    /// Add multiple roles
133    pub async fn add_roles<C: TransactionTrait>(
134        &mut self,
135        db: &C,
136        roles: &[&'static str],
137    ) -> Result<(), DbErr> {
138        let txn = db.begin().await?;
139
140        for role in roles {
141            if let Some(role_id) = role::Entity::insert(Role {
142                role: Set(role.to_string()),
143                ..Default::default()
144            })
145            .on_conflict_do_nothing()
146            .exec(&txn)
147            .await?
148            .last_insert_id()?
149            {
150                self.roles.insert(role.to_string(), role_id);
151            }
152        }
153
154        txn.commit().await
155    }
156
157    pub fn get_role(&self, role: &'static str) -> Result<&RoleId, DbErr> {
158        self.roles
159            .get(role)
160            .ok_or_else(|| DbErr::RbacError(RbacError::RoleNotFound(role.to_string()).to_string()))
161    }
162
163    /// Add permissions to roles. Will take cartesian product of tables and actions.
164    pub async fn add_role_permissions<C: TransactionTrait>(
165        &mut self,
166        db: &C,
167        role: &'static str,
168        actions: &[&'static str],
169        tables: &[&'static str],
170    ) -> Result<(), DbErr> {
171        self.update_role_permissions(db, role, actions, tables, true)
172            .await
173    }
174
175    /// Remove permissions from roles. Will take cartesian product of tables and actions.
176    pub async fn remove_role_permissions<C: TransactionTrait>(
177        &mut self,
178        db: &C,
179        role: &'static str,
180        actions: &[&'static str],
181        tables: &[&'static str],
182    ) -> Result<(), DbErr> {
183        self.update_role_permissions(db, role, actions, tables, false)
184            .await
185    }
186
187    async fn update_role_permissions<C: TransactionTrait>(
188        &mut self,
189        db: &C,
190        role: &'static str,
191        actions: &[&'static str],
192        tables: &[&'static str],
193        grant: bool,
194    ) -> Result<(), DbErr> {
195        let txn = db.begin().await?;
196
197        for table_name in tables {
198            for action in actions {
199                let model = RolePermission {
200                    role_id: Set(*self.roles.get(role).ok_or_else(|| {
201                        DbErr::RbacError(RbacError::RoleNotFound(role.to_string()).to_string())
202                    })?),
203                    permission_id: Set(*self.permissions.get(*action).ok_or_else(|| {
204                        DbErr::RbacError(
205                            RbacError::PermissionNotFound(action.to_string()).to_string(),
206                        )
207                    })?),
208                    resource_id: Set(*self.tables.get(*table_name).ok_or_else(|| {
209                        DbErr::RbacError(
210                            RbacError::ResourceNotFound(table_name.to_string()).to_string(),
211                        )
212                    })?),
213                };
214                if grant {
215                    role_permission::Entity::insert(model)
216                        .on_conflict_do_nothing()
217                        .exec(&txn)
218                        .await?;
219                } else {
220                    role_permission::Entity::delete(model).exec(&txn).await?;
221                }
222            }
223        }
224
225        txn.commit().await
226    }
227
228    pub async fn add_user_override<C: TransactionTrait>(
229        &mut self,
230        db: &C,
231        rows: &[RbacAddUserOverride],
232    ) -> Result<(), DbErr> {
233        let txn = db.begin().await?;
234
235        for RbacAddUserOverride {
236            user_id,
237            table,
238            action,
239            grant,
240        } in rows
241        {
242            user_override::Entity::insert(UserOverride {
243                user_id: Set(RbacUserId(*user_id)),
244                permission_id: Set(*self.permissions.get(*action).ok_or_else(|| {
245                    DbErr::RbacError(RbacError::PermissionNotFound(action.to_string()).to_string())
246                })?),
247                resource_id: Set(*self.tables.get(*table).ok_or_else(|| {
248                    DbErr::RbacError(RbacError::ResourceNotFound(table.to_string()).to_string())
249                })?),
250                grant: Set(*grant),
251            })
252            .on_conflict(
253                OnConflict::columns([
254                    user_override::Column::UserId,
255                    user_override::Column::PermissionId,
256                    user_override::Column::ResourceId,
257                ])
258                .update_column(user_override::Column::Grant)
259                .to_owned(),
260            )
261            .try_insert()
262            .exec(&txn)
263            .await?;
264        }
265
266        txn.commit().await
267    }
268
269    pub async fn add_role_hierarchy<C: TransactionTrait>(
270        &mut self,
271        db: &C,
272        rows: &[RbacAddRoleHierarchy],
273    ) -> Result<(), DbErr> {
274        let txn = db.begin().await?;
275
276        for RbacAddRoleHierarchy { super_role, role } in rows {
277            role_hierarchy::Entity::insert(RoleHierarchy {
278                super_role_id: Set(*self.roles.get(*super_role).ok_or_else(|| {
279                    DbErr::RbacError(RbacError::RoleNotFound(super_role.to_string()).to_string())
280                })?),
281                role_id: Set(*self.roles.get(*role).ok_or_else(|| {
282                    DbErr::RbacError(RbacError::RoleNotFound(role.to_string()).to_string())
283                })?),
284            })
285            .on_conflict_do_nothing()
286            .exec(&txn)
287            .await?;
288        }
289
290        txn.commit().await
291    }
292
293    /// Assign role to users. Note that each user can only have 1 role,
294    /// so this assignment replaces current role.
295    /// `rows: (UserId, role)`
296    pub async fn assign_user_role<C: TransactionTrait>(
297        &mut self,
298        db: &C,
299        rows: &[(i64, &'static str)],
300    ) -> Result<(), DbErr> {
301        let txn = db.begin().await?;
302
303        for (user_id, role) in rows {
304            user_role::Entity::insert(UserRole {
305                user_id: Set(RbacUserId(*user_id)),
306                role_id: Set(*self.roles.get(*role).ok_or_else(|| {
307                    DbErr::RbacError(RbacError::RoleNotFound(role.to_string()).to_string())
308                })?),
309            })
310            .on_conflict(
311                OnConflict::column(user_role::Column::UserId)
312                    .update_column(user_role::Column::RoleId)
313                    .to_owned(),
314            )
315            .try_insert()
316            .exec(&txn)
317            .await?;
318        }
319
320        txn.commit().await
321    }
322}
sea_orm/rbac/context.rs

sea_orm/rbac/
context.rs
