add default role api (#517)

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

Author: trueleo

server/src/handlers/http.rs

@@ -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))

server/src/handlers/http/oidc.rs

@@ -34,7 +34,7 @@ use crate::{
     oidc::{Claims, DiscoveredClient},
     option::CONFIG,
     rbac::{
-        map::SessionKey,
+        map::{SessionKey, DEFAULT_ROLE},
         user::{User, UserType},
         Users,
     },
@@ -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)
 }

server/src/handlers/http/role.rs

@@ -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},
 };
 
@@ -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()

server/src/main.rs

@@ -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);

server/src/rbac/map.rs

@@ -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> {
@@ -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]);

server/src/rbac/user.rs

@@ -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,
         }
     }
 

server/src/storage/store_metadata.rs

@@ -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 {
@@ -70,6 +72,7 @@ impl StorageMetadata {
             users: Vec::new(),
             streams: Vec::new(),
             roles: HashMap::default(),
+            default_role: None,
         }
     }