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

sea_orm/rbac/
context.rs
