From 377f13ceee927541ab89af1bc2a30a4c8247a25d Mon Sep 17 00:00:00 2001
From: "Richard Kuo (Danswer)" <rkuo@danswer.ai>
Date: Tue, 12 Nov 2024 09:59:34 -0800
Subject: [PATCH 1/5] in progress PoC

---
 backend/danswer/auth/schemas.py               |  1 +
 backend/danswer/auth/users.py                 | 16 +++++++-
 .../app/admin/api-key/DanswerApiKeyForm.tsx   | 41 ++++++++++++-------
 web/src/lib/types.ts                          |  2 +
 4 files changed, 44 insertions(+), 16 deletions(-)

diff --git a/backend/danswer/auth/schemas.py b/backend/danswer/auth/schemas.py
index 9c81899a421..ce503afe76e 100644
--- a/backend/danswer/auth/schemas.py
+++ b/backend/danswer/auth/schemas.py
@@ -15,6 +15,7 @@ class UserRole(str, Enum):
         for all groups they are a member of
     """
 
+    LIMITED = "limited"
     BASIC = "basic"
     ADMIN = "admin"
     CURATOR = "curator"
diff --git a/backend/danswer/auth/users.py b/backend/danswer/auth/users.py
index 73f1ec18484..23805e518fb 100644
--- a/backend/danswer/auth/users.py
+++ b/backend/danswer/auth/users.py
@@ -652,10 +652,24 @@ async def current_user_with_expired_token(
     return await double_check_user(user, include_expired=True)
 
 
+async def current_limited_user(
+    user: User | None = Depends(optional_user),
+) -> User | None:
+    return double_check_user(user)
+
+
 async def current_user(
     user: User | None = Depends(optional_user),
 ) -> User | None:
-    return await double_check_user(user)
+    user = await double_check_user(user)
+    if not user:
+        return None
+
+    if user.role == UserRole.LIMITED:
+        raise BasicAuthenticationError(
+            detail="Access denied. User role is LIMITED. BASIC or higher permissions are required.",
+        )
+    return user
 
 
 async def current_curator_or_admin_user(
diff --git a/web/src/app/admin/api-key/DanswerApiKeyForm.tsx b/web/src/app/admin/api-key/DanswerApiKeyForm.tsx
index 03f15e08038..80bb84d626f 100644
--- a/web/src/app/admin/api-key/DanswerApiKeyForm.tsx
+++ b/web/src/app/admin/api-key/DanswerApiKeyForm.tsx
@@ -2,6 +2,7 @@ import { Form, Formik } from "formik";
 import { PopupSpec } from "@/components/admin/connectors/Popup";
 import {
   BooleanFormField,
+  SelectorFormField,
   TextFormField,
 } from "@/components/admin/connectors/Field";
 import { createApiKey, updateApiKey } from "./lib";
@@ -9,7 +10,7 @@ import { Modal } from "@/components/Modal";
 import { Button } from "@/components/ui/button";
 import { Separator } from "@/components/ui/separator";
 import Text from "@/components/ui/text";
-import { UserRole } from "@/lib/types";
+import { USER_ROLE_LABELS, UserRole } from "@/lib/types";
 import { APIKey } from "./types";
 
 interface DanswerApiKeyFormProps {
@@ -39,20 +40,15 @@ export const DanswerApiKeyForm = ({
         <Formik
           initialValues={{
             name: apiKey?.api_key_name || "",
-            is_admin: apiKey?.api_key_role === "admin",
+            role: apiKey?.api_key_role || UserRole.BASIC.toString(),
           }}
           onSubmit={async (values, formikHelpers) => {
             formikHelpers.setSubmitting(true);
 
-            // Map the boolean to a UserRole string
-            const role: UserRole = values.is_admin
-              ? UserRole.ADMIN
-              : UserRole.BASIC;
-
             // Prepare the payload with the UserRole
             const payload = {
               ...values,
-              role, // Assign the role directly as a UserRole type
+              role: values.role as UserRole, // Assign the role directly as a UserRole type
             };
 
             let response;
@@ -98,13 +94,28 @@ export const DanswerApiKeyForm = ({
                 autoCompleteDisabled={true}
               />
 
-              <BooleanFormField
-                small
-                removeIndent
-                alignTop
-                name="is_admin"
-                label="Is Admin?"
-                subtext="If set, this API key will have access to admin level server API's."
+              <SelectorFormField
+                // defaultValue is managed by Formik
+                label="Role:"
+                subtext="Select the role for this API key.
+                         Limited has access to simple public API's.
+                         Basic has access to regular user API's.
+                         Admin has access to admin level APIs."
+                name="role"
+                options={[
+                  {
+                    name: USER_ROLE_LABELS[UserRole.LIMITED],
+                    value: UserRole.LIMITED.toString(),
+                  },
+                  {
+                    name: USER_ROLE_LABELS[UserRole.BASIC],
+                    value: UserRole.BASIC.toString(),
+                  },
+                  {
+                    name: USER_ROLE_LABELS[UserRole.ADMIN],
+                    value: UserRole.ADMIN.toString(),
+                  },
+                ]}
               />
 
               <Button
diff --git a/web/src/lib/types.ts b/web/src/lib/types.ts
index c9f8d46b220..eff53f24d36 100644
--- a/web/src/lib/types.ts
+++ b/web/src/lib/types.ts
@@ -18,6 +18,7 @@ export enum UserStatus {
 }
 
 export enum UserRole {
+  LIMITED = "limited",
   BASIC = "basic",
   ADMIN = "admin",
   CURATOR = "curator",
@@ -25,6 +26,7 @@ export enum UserRole {
 }
 
 export const USER_ROLE_LABELS: Record<UserRole, string> = {
+  [UserRole.LIMITED]: "Limited",
   [UserRole.BASIC]: "Basic",
   [UserRole.ADMIN]: "Admin",
   [UserRole.GLOBAL_CURATOR]: "Global Curator",

From 8ed5859e156e833c7aa328f38d5bc84c78c8c4b0 Mon Sep 17 00:00:00 2001
From: "Richard Kuo (Danswer)" <rkuo@danswer.ai>
Date: Tue, 12 Nov 2024 14:39:02 -0800
Subject: [PATCH 2/5] working limited user, needs routes to be marked next

---
 backend/danswer/auth/users.py                 |  2 +-
 backend/danswer/server/auth_check.py          |  6 ++-
 .../danswer/server/features/persona/api.py    |  3 +-
 .../integration/tests/api_key/test_api_key.py | 42 +++++++++++++++++++
 4 files changed, 49 insertions(+), 4 deletions(-)
 create mode 100644 backend/tests/integration/tests/api_key/test_api_key.py

diff --git a/backend/danswer/auth/users.py b/backend/danswer/auth/users.py
index 23805e518fb..4542fbf869e 100644
--- a/backend/danswer/auth/users.py
+++ b/backend/danswer/auth/users.py
@@ -655,7 +655,7 @@ async def current_user_with_expired_token(
 async def current_limited_user(
     user: User | None = Depends(optional_user),
 ) -> User | None:
-    return double_check_user(user)
+    return await double_check_user(user)
 
 
 async def current_user(
diff --git a/backend/danswer/server/auth_check.py b/backend/danswer/server/auth_check.py
index 64550b09861..bf55d80d90e 100644
--- a/backend/danswer/server/auth_check.py
+++ b/backend/danswer/server/auth_check.py
@@ -6,6 +6,7 @@
 
 from danswer.auth.users import current_admin_user
 from danswer.auth.users import current_curator_or_admin_user
+from danswer.auth.users import current_limited_user
 from danswer.auth.users import current_user
 from danswer.auth.users import current_user_with_expired_token
 from danswer.configs.app_configs import APP_API_PREFIX
@@ -102,7 +103,8 @@ def check_router_auth(
             for dependency in route_dependant_obj.dependencies:
                 depends_fn = dependency.cache_key[0]
                 if (
-                    depends_fn == current_user
+                    depends_fn == current_limited_user
+                    or depends_fn == current_user
                     or depends_fn == current_admin_user
                     or depends_fn == current_curator_or_admin_user
                     or depends_fn == api_key_dep
@@ -118,5 +120,5 @@ def check_router_auth(
             # print(f"(\"{route.path}\", {set(route.methods)}),")
 
             raise RuntimeError(
-                f"Did not find current_user or current_admin_user dependency in route - {route}"
+                f"Did not find user dependency in private route - {route}"
             )
diff --git a/backend/danswer/server/features/persona/api.py b/backend/danswer/server/features/persona/api.py
index 1e0e84a5896..c7d3d79b11e 100644
--- a/backend/danswer/server/features/persona/api.py
+++ b/backend/danswer/server/features/persona/api.py
@@ -11,6 +11,7 @@
 
 from danswer.auth.users import current_admin_user
 from danswer.auth.users import current_curator_or_admin_user
+from danswer.auth.users import current_limited_user
 from danswer.auth.users import current_user
 from danswer.configs.constants import FileOrigin
 from danswer.configs.constants import NotificationType
@@ -272,7 +273,7 @@ def list_personas(
 @basic_router.get("/{persona_id}")
 def get_persona(
     persona_id: int,
-    user: User | None = Depends(current_user),
+    user: User | None = Depends(current_limited_user),
     db_session: Session = Depends(get_session),
 ) -> PersonaSnapshot:
     return PersonaSnapshot.from_model(
diff --git a/backend/tests/integration/tests/api_key/test_api_key.py b/backend/tests/integration/tests/api_key/test_api_key.py
new file mode 100644
index 00000000000..bd0618b962d
--- /dev/null
+++ b/backend/tests/integration/tests/api_key/test_api_key.py
@@ -0,0 +1,42 @@
+import requests
+
+from danswer.auth.schemas import UserRole
+from tests.integration.common_utils.constants import API_SERVER_URL
+from tests.integration.common_utils.managers.api_key import APIKeyManager
+from tests.integration.common_utils.managers.user import UserManager
+from tests.integration.common_utils.test_models import DATestAPIKey
+from tests.integration.common_utils.test_models import DATestUser
+
+
+def test_limited(reset: None) -> None:
+    """Verify that with a limited role key, limited endpoints are accessible and
+    others are not."""
+
+    # Creating an admin user (first user created is automatically an admin)
+    admin_user: DATestUser = UserManager.create(name="admin_user")
+
+    api_key: DATestAPIKey = APIKeyManager.create(
+        api_key_role=UserRole.LIMITED,
+        user_performing_action=admin_user,
+    )
+
+    # test limited endpoint
+    response = requests.get(
+        f"{API_SERVER_URL}/persona/0",
+        headers=api_key.headers,
+    )
+    assert response.status_code == 200
+
+    # test basic endpoints
+    response = requests.get(
+        f"{API_SERVER_URL}/input_prompt",
+        headers=api_key.headers,
+    )
+    assert response.status_code == 403
+
+    # test admin endpoints
+    response = requests.get(
+        f"{API_SERVER_URL}/admin/api-key",
+        headers=api_key.headers,
+    )
+    assert response.status_code == 403

From f261681baaab93941f8c87d61429d6e740987bc8 Mon Sep 17 00:00:00 2001
From: Richard Kuo <rkuo@rkuo.com>
Date: Tue, 12 Nov 2024 19:50:10 -0800
Subject: [PATCH 3/5] make selected endpoint available to limited user role

---
 backend/danswer/server/query_and_chat/chat_backend.py  | 5 +++--
 backend/danswer/server/query_and_chat/query_backend.py | 3 ++-
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/backend/danswer/server/query_and_chat/chat_backend.py b/backend/danswer/server/query_and_chat/chat_backend.py
index 64c667cfd16..c4728336c86 100644
--- a/backend/danswer/server/query_and_chat/chat_backend.py
+++ b/backend/danswer/server/query_and_chat/chat_backend.py
@@ -18,6 +18,7 @@
 from pydantic import BaseModel
 from sqlalchemy.orm import Session
 
+from danswer.auth.users import current_limited_user
 from danswer.auth.users import current_user
 from danswer.chat.chat_utils import create_chat_chain
 from danswer.chat.chat_utils import extract_headers
@@ -309,7 +310,7 @@ def is_connected_sync() -> bool:
 def handle_new_chat_message(
     chat_message_req: CreateChatMessageRequest,
     request: Request,
-    user: User | None = Depends(current_user),
+    user: User | None = Depends(current_limited_user),
     _: None = Depends(check_token_rate_limits),
     is_connected_func: Callable[[], bool] = Depends(is_connected),
 ) -> StreamingResponse:
@@ -391,7 +392,7 @@ def set_message_as_latest(
 @router.post("/create-chat-message-feedback")
 def create_chat_feedback(
     feedback: ChatFeedbackRequest,
-    user: User | None = Depends(current_user),
+    user: User | None = Depends(current_limited_user),
     db_session: Session = Depends(get_session),
 ) -> None:
     user_id = user.id if user else None
diff --git a/backend/danswer/server/query_and_chat/query_backend.py b/backend/danswer/server/query_and_chat/query_backend.py
index 4d6767ac2e2..3636de5080c 100644
--- a/backend/danswer/server/query_and_chat/query_backend.py
+++ b/backend/danswer/server/query_and_chat/query_backend.py
@@ -9,6 +9,7 @@
 from sqlalchemy.orm import Session
 
 from danswer.auth.users import current_curator_or_admin_user
+from danswer.auth.users import current_limited_user
 from danswer.auth.users import current_user
 from danswer.configs.constants import DocumentSource
 from danswer.configs.constants import MessageType
@@ -262,7 +263,7 @@ def stream_query_validation(
 @basic_router.post("/stream-answer-with-quote")
 def get_answer_with_quote(
     query_request: DirectQARequest,
-    user: User = Depends(current_user),
+    user: User = Depends(current_limited_user),
     _: None = Depends(check_token_rate_limits),
 ) -> StreamingResponse:
     query = query_request.messages[0].message

From 4e1bc2a4eb089188ca4657b412fba6efe0f136e8 Mon Sep 17 00:00:00 2001
From: Richard Kuo <rkuo@rkuo.com>
Date: Tue, 12 Nov 2024 20:19:06 -0800
Subject: [PATCH 4/5] xfail on test_slack_prune

---
 .../tests/integration/connector_job_tests/slack/test_prune.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/backend/tests/integration/connector_job_tests/slack/test_prune.py b/backend/tests/integration/connector_job_tests/slack/test_prune.py
index bcef148a2a0..08a8ac4fbec 100644
--- a/backend/tests/integration/connector_job_tests/slack/test_prune.py
+++ b/backend/tests/integration/connector_job_tests/slack/test_prune.py
@@ -3,6 +3,8 @@
 from datetime import timezone
 from typing import Any
 
+import pytest
+
 from danswer.connectors.models import InputType
 from danswer.db.enums import AccessType
 from danswer.server.documents.models import DocumentSource
@@ -22,7 +24,7 @@
 from tests.integration.connector_job_tests.slack.slack_api_utils import SlackManager
 
 
-# @pytest.mark.xfail(reason="flaky - see DAN-835 for example", strict=False)
+@pytest.mark.xfail(reason="flaky - see DAN-986 for details", strict=False)
 def test_slack_prune(
     reset: None,
     vespa_client: vespa_fixture,

From e5889d7a3f7103ea53bfdd819456247a7e03a5bb Mon Sep 17 00:00:00 2001
From: Richard Kuo <rkuo@rkuo.com>
Date: Tue, 12 Nov 2024 20:21:36 -0800
Subject: [PATCH 5/5] add comment to sync function

---
 backend/tests/integration/common_utils/managers/cc_pair.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/backend/tests/integration/common_utils/managers/cc_pair.py b/backend/tests/integration/common_utils/managers/cc_pair.py
index d14680d8389..b5f49fc92b2 100644
--- a/backend/tests/integration/common_utils/managers/cc_pair.py
+++ b/backend/tests/integration/common_utils/managers/cc_pair.py
@@ -326,6 +326,10 @@ def sync(
         cc_pair: DATestCCPair,
         user_performing_action: DATestUser | None = None,
     ) -> None:
+        """This function triggers a permission sync.
+        Naming / intent of this function probably could use improvement, but currently it's letting
+        409 Conflict pass through since if it's running that's what we were trying to do anyway.
+        """
         result = requests.post(
             url=f"{API_SERVER_URL}/manage/admin/cc-pair/{cc_pair.id}/sync-permissions",
             headers=user_performing_action.headers