From f1fbd2e7f40d283e22bfcc26e2e355e8aa27d548 Mon Sep 17 00:00:00 2001
From: Ashish Kolhe <35160958+oasisk@users.noreply.github.com>
Date: Wed, 24 Jan 2024 18:13:23 +0530
Subject: [PATCH] fix: panic in users db (#2556)

---
 src/common/meta/user.rs      |  5 +++++
 src/handler/http/auth/jwt.rs | 20 ++++++++++++++------
 src/service/db/user.rs       | 26 ++++++++++++++++++++++++++
 3 files changed, 45 insertions(+), 6 deletions(-)

diff --git a/src/common/meta/user.rs b/src/common/meta/user.rs
index ebb42ecef6..38cbf7b64e 100644
--- a/src/common/meta/user.rs
+++ b/src/common/meta/user.rs
@@ -194,6 +194,9 @@ pub enum UserRole {
     Member,
     #[serde(rename = "root")]
     Root,
+    #[cfg(feature = "enterprise")]
+    #[serde(rename = "viewer")]
+    Viewer,
 }
 
 impl fmt::Display for UserRole {
@@ -202,6 +205,8 @@ impl fmt::Display for UserRole {
             UserRole::Admin => write!(f, "admin"),
             UserRole::Member => write!(f, "member"),
             UserRole::Root => write!(f, "root"),
+            #[cfg(feature = "enterprise")]
+            UserRole::Viewer => write!(f, "viewer"),
         }
     }
 }
diff --git a/src/handler/http/auth/jwt.rs b/src/handler/http/auth/jwt.rs
index 01be6f9575..3039b5166d 100644
--- a/src/handler/http/auth/jwt.rs
+++ b/src/handler/http/auth/jwt.rs
@@ -58,14 +58,22 @@ pub async fn process_token(
     let user_email = res.0.user_email.to_owned();
     let mut source_orgs: Vec<UserOrg> = vec![];
     let mut tuples_to_add = HashMap::new();
-    for group in groups {
-        let role_org = parse_dn(group.as_str().unwrap()).unwrap();
-
+    if groups.is_empty() {
         source_orgs.push(UserOrg {
-            role: role_org.role,
-            name: role_org.org,
+            role: crate::common::meta::user::UserRole::Viewer,
+            name: O2_CONFIG.dex.default_org.clone(),
             ..UserOrg::default()
         });
+    } else {
+        for group in groups {
+            let role_org = parse_dn(group.as_str().unwrap()).unwrap();
+
+            source_orgs.push(UserOrg {
+                role: role_org.role,
+                name: role_org.org,
+                ..UserOrg::default()
+            });
+        }
     }
 
     // Check if the user exists in the database
@@ -287,7 +295,7 @@ fn parse_dn(dn: &str) -> Option<RoleOrg> {
     let role = if role.contains("admin") {
         crate::common::meta::user::UserRole::Admin
     } else {
-        crate::common::meta::user::UserRole::Member
+        crate::common::meta::user::UserRole::Viewer
     };
     if org.is_empty() {
         org = &O2_CONFIG.dex.default_org;
diff --git a/src/service/db/user.rs b/src/service/db/user.rs
index 05e22853f3..0783f65572 100644
--- a/src/service/db/user.rs
+++ b/src/service/db/user.rs
@@ -162,6 +162,7 @@ pub async fn watch() -> Result<(), anyhow::Error> {
                 let item_key = ev.key.strip_prefix(key).unwrap();
                 let item_value: DBUser = json::from_slice(&ev.value.unwrap()).unwrap();
                 let users = item_value.get_all_users();
+                #[cfg(not(feature = "enterprise"))]
                 for mut user in users {
                     if user.role.eq(&UserRole::Root) {
                         ROOT_USER.insert("root".to_string(), user.clone());
@@ -170,6 +171,14 @@ pub async fn watch() -> Result<(), anyhow::Error> {
                     };
                     USERS.insert(format!("{}/{}", user.org, item_key), user);
                 }
+
+                #[cfg(feature = "enterprise")]
+                for user in users {
+                    if user.role.eq(&UserRole::Root) {
+                        ROOT_USER.insert("root".to_string(), user.clone());
+                    }
+                    USERS.insert(format!("{}/{}", user.org, item_key), user);
+                }
                 // Invalidate the entire RUM-TOKEN-CACHE
                 USERS_RUM_TOKEN.clear();
             }
@@ -198,6 +207,7 @@ pub async fn cache() -> Result<(), anyhow::Error> {
         // let item_key = item_key.strip_prefix(key).unwrap();
         let json_val: DBUser = json::from_slice(&item_value).unwrap();
         let users = json_val.get_all_users();
+        #[cfg(not(feature = "enterprise"))]
         for mut user in users {
             if user.role.eq(&UserRole::Root) {
                 ROOT_USER.insert("root".to_string(), user.clone());
@@ -211,6 +221,19 @@ pub async fn cache() -> Result<(), anyhow::Error> {
                     .insert(format!("{}/{}", user.org, rum_token), user);
             }
         }
+
+        #[cfg(feature = "enterprise")]
+        for user in users {
+            if user.role.eq(&UserRole::Root) {
+                ROOT_USER.insert("root".to_string(), user.clone());
+            }
+            USERS.insert(format!("{}/{}", user.org, user.email), user.clone());
+            if let Some(rum_token) = &user.rum_token {
+                USERS_RUM_TOKEN
+                    .clone()
+                    .insert(format!("{}/{}", user.org, rum_token), user);
+            }
+        }
     }
     log::info!("Users Cached");
     Ok(())
@@ -222,6 +245,9 @@ pub async fn root_user_exists() -> bool {
     let mut ret = db.list_values(key).await.unwrap_or_default();
     ret.retain(|item| {
         let user: DBUser = json::from_slice(item).unwrap();
+        if user.organizations.is_empty() {
+            return false;
+        }
         user.organizations
             .first()
             .as_ref()