Skip to content

Commit

Permalink
add default role api (#517)
Browse files Browse the repository at this point in the history
This PR adds an API to set default role to be 
assigned to any new OIDC user. This is needed 
for identity providers like Google where there is 
no group membership for a given user.

Fixes #513
  • Loading branch information
trueleo authored Sep 27, 2023
1 parent 1bbdf89 commit 1a1c340
Show file tree
Hide file tree
Showing 7 changed files with 65 additions and 20 deletions.
6 changes: 6 additions & 0 deletions server/src/handlers/http.rs
Original file line number Diff line number Diff line change
Expand Up @@ -261,8 +261,14 @@ pub fn configure_routes(
.authorize(Action::QueryLLM),
),
);

let role_api = web::scope("/role")
.service(resource("").route(web::get().to(role::list).authorize(Action::ListRole)))
.service(
resource("/default")
.route(web::put().to(role::put_default).authorize(Action::PutRole))
.route(web::get().to(role::get_default).authorize(Action::GetRole)),
)
.service(
resource("/{name}")
.route(web::put().to(role::put).authorize(Action::PutRole))
Expand Down
29 changes: 17 additions & 12 deletions server/src/handlers/http/oidc.rs
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ use crate::{
oidc::{Claims, DiscoveredClient},
option::CONFIG,
rbac::{
map::SessionKey,
map::{SessionKey, DEFAULT_ROLE},
user::{User, UserType},
Users,
},
Expand Down Expand Up @@ -263,22 +263,27 @@ async fn put_user(
group: Option<HashSet<String>>,
) -> Result<User, ObjectStorageError> {
let mut metadata = get_metadata().await?;
let user = match metadata
let group = group.unwrap_or_else(|| {
DEFAULT_ROLE
.lock()
.unwrap()
.clone()
.map(|role| HashSet::from([role]))
.unwrap_or_default()
});

let user = metadata
.users
.iter()
.find(|user| user.username() == username)
{
Some(user) => user.clone(),
None => {
let mut user = User::new_oauth(username.to_owned());
if let Some(group) = group {
user.roles = group
}
.cloned()
.unwrap_or_else(|| {
let user = User::new_oauth(username.to_owned(), group);
metadata.users.push(user.clone());
put_metadata(&metadata).await?;
user
}
};
});

put_metadata(&metadata).await?;
Users.put_user(user.clone());
Ok(user)
}
Expand Down
27 changes: 26 additions & 1 deletion server/src/handlers/http/role.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,10 @@ use http::StatusCode;

use crate::{
option::CONFIG,
rbac::{map::mut_roles, role::model::DefaultPrivilege},
rbac::{
map::{mut_roles, DEFAULT_ROLE},
role::model::DefaultPrivilege,
},
storage::{self, ObjectStorageError, StorageMetadata},
};

Expand Down Expand Up @@ -71,6 +74,28 @@ pub async fn delete(name: web::Path<String>) -> Result<impl Responder, RoleError
Ok(HttpResponse::Ok().finish())
}

// Handler for PUT /api/v1/role/default
// Delete existing role
pub async fn put_default(name: web::Json<String>) -> Result<impl Responder, RoleError> {
let name = name.into_inner();
let mut metadata = get_metadata().await?;
metadata.default_role = Some(name.clone());
*DEFAULT_ROLE.lock().unwrap() = Some(name);
put_metadata(&metadata).await?;
Ok(HttpResponse::Ok().finish())
}

// Handler for GET /api/v1/role/default
// Delete existing role
pub async fn get_default() -> Result<impl Responder, RoleError> {
let res = match DEFAULT_ROLE.lock().unwrap().clone() {
Some(role) => serde_json::Value::String(role),
None => serde_json::Value::Null,
};

Ok(web::Json(res))
}

async fn get_metadata() -> Result<crate::storage::StorageMetadata, ObjectStorageError> {
let metadata = CONFIG
.storage()
Expand Down
2 changes: 1 addition & 1 deletion server/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,7 @@ async fn main() -> anyhow::Result<()> {
migration::run_metadata_migration(&CONFIG).await?;
let metadata = storage::resolve_parseable_metadata().await?;
banner::print(&CONFIG, &metadata).await;
rbac::map::init(metadata.users.clone(), metadata.roles.clone());
rbac::map::init(&metadata);
metadata.set_global();
let prometheus = metrics::build_metrics_handler();
CONFIG.storage().register_store_metrics(&prometheus);
Expand Down
14 changes: 10 additions & 4 deletions server/src/rbac/map.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,22 +16,23 @@
*
*/

use crate::option::CONFIG;
use crate::rbac::user::User;
use std::collections::HashMap;
use crate::{option::CONFIG, storage::StorageMetadata};
use std::{collections::HashMap, sync::Mutex};

use super::{
role::{model::DefaultPrivilege, Action, Permission, RoleBuilder},
user,
};
use chrono::{DateTime, Utc};
use once_cell::sync::OnceCell;
use once_cell::sync::{Lazy, OnceCell};
use std::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};

pub type Roles = HashMap<String, Vec<DefaultPrivilege>>;

pub static USERS: OnceCell<RwLock<Users>> = OnceCell::new();
pub static ROLES: OnceCell<RwLock<Roles>> = OnceCell::new();
pub static DEFAULT_ROLE: Lazy<Mutex<Option<String>>> = Lazy::new(|| Mutex::new(None));
pub static SESSIONS: OnceCell<RwLock<Sessions>> = OnceCell::new();

pub fn users() -> RwLockReadGuard<'static, Users> {
Expand Down Expand Up @@ -86,7 +87,12 @@ pub fn mut_sessions() -> RwLockWriteGuard<'static, Sessions> {
// the user_map is initialized from the config file and has a list of all users
// the auth_map is initialized with admin user only and then gets lazily populated
// as users authenticate
pub fn init(users: Vec<User>, mut roles: Roles) {
pub fn init(metadata: &StorageMetadata) {
let users = metadata.users.clone();
let mut roles = metadata.roles.clone();

*DEFAULT_ROLE.lock().unwrap() = metadata.default_role.clone();

let admin_privilege = DefaultPrivilege::Admin;
let admin_permissions = RoleBuilder::from(&admin_privilege).build();
roles.insert("admin".to_string(), vec![admin_privilege]);
Expand Down
4 changes: 2 additions & 2 deletions server/src/rbac/user.rs
Original file line number Diff line number Diff line change
Expand Up @@ -57,10 +57,10 @@ impl User {
)
}

pub fn new_oauth(username: String) -> Self {
pub fn new_oauth(username: String, roles: HashSet<String>) -> Self {
Self {
ty: UserType::OAuth(OAuth { userid: username }),
roles: HashSet::new(),
roles,
}
}

Expand Down
3 changes: 3 additions & 0 deletions server/src/storage/store_metadata.rs
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,8 @@ pub struct StorageMetadata {
pub streams: Vec<String>,
#[serde(default)]
pub roles: HashMap<String, Vec<DefaultPrivilege>>,
#[serde(default)]
pub default_role: Option<String>,
}

impl StorageMetadata {
Expand All @@ -70,6 +72,7 @@ impl StorageMetadata {
users: Vec::new(),
streams: Vec::new(),
roles: HashMap::default(),
default_role: None,
}
}

Expand Down

0 comments on commit 1a1c340

Please sign in to comment.