Browse Source

enh: configurable api key endpoint restrictions

Timothy Jaeryang Baek 4 months ago
parent
commit
1e974439d9

+ 12 - 0
backend/open_webui/config.py

@@ -272,6 +272,18 @@ ENABLE_API_KEY = PersistentConfig(
     os.environ.get("ENABLE_API_KEY", "True").lower() == "true",
 )
 
+ENABLE_API_KEY_ENDPOINT_RESTRICTIONS = PersistentConfig(
+    "ENABLE_API_KEY_ENDPOINT_RESTRICTIONS",
+    "auth.api_key.endpoint_restrictions",
+    os.environ.get("ENABLE_API_KEY_ENDPOINT_RESTRICTIONS", "False").lower() == "true",
+)
+
+API_KEY_ALLOWED_ENDPOINTS = PersistentConfig(
+    "API_KEY_ALLOWED_ENDPOINTS",
+    "auth.api_key.allowed_endpoints",
+    os.environ.get("API_KEY_ALLOWED_ENDPOINTS", ""),
+)
+
 
 JWT_EXPIRES_IN = PersistentConfig(
     "JWT_EXPIRES_IN", "auth.jwt_expiry", os.environ.get("JWT_EXPIRES_IN", "-1")

+ 7 - 0
backend/open_webui/main.py

@@ -199,6 +199,8 @@ from open_webui.config import (
     ENABLE_SIGNUP,
     ENABLE_LOGIN_FORM,
     ENABLE_API_KEY,
+    ENABLE_API_KEY_ENDPOINT_RESTRICTIONS,
+    API_KEY_ALLOWED_ENDPOINTS,
     ENABLE_CHANNELS,
     ENABLE_COMMUNITY_SHARING,
     ENABLE_MESSAGE_RATING,
@@ -392,7 +394,12 @@ app.state.OPENAI_MODELS = {}
 app.state.config.WEBUI_URL = WEBUI_URL
 app.state.config.ENABLE_SIGNUP = ENABLE_SIGNUP
 app.state.config.ENABLE_LOGIN_FORM = ENABLE_LOGIN_FORM
+
 app.state.config.ENABLE_API_KEY = ENABLE_API_KEY
+app.state.config.ENABLE_API_KEY_ENDPOINT_RESTRICTIONS = (
+    ENABLE_API_KEY_ENDPOINT_RESTRICTIONS
+)
+app.state.config.API_KEY_ALLOWED_ENDPOINTS = API_KEY_ALLOWED_ENDPOINTS
 
 app.state.config.JWT_EXPIRES_IN = JWT_EXPIRES_IN
 

+ 14 - 0
backend/open_webui/routers/auths.py

@@ -616,6 +616,8 @@ async def get_admin_config(request: Request, user=Depends(get_admin_user)):
         "WEBUI_URL": request.app.state.config.WEBUI_URL,
         "ENABLE_SIGNUP": request.app.state.config.ENABLE_SIGNUP,
         "ENABLE_API_KEY": request.app.state.config.ENABLE_API_KEY,
+        "ENABLE_API_KEY_ENDPOINT_RESTRICTIONS": request.app.state.config.ENABLE_API_KEY_ENDPOINT_RESTRICTIONS,
+        "API_KEY_ALLOWED_ENDPOINTS": request.app.state.config.API_KEY_ALLOWED_ENDPOINTS,
         "ENABLE_CHANNELS": request.app.state.config.ENABLE_CHANNELS,
         "DEFAULT_USER_ROLE": request.app.state.config.DEFAULT_USER_ROLE,
         "JWT_EXPIRES_IN": request.app.state.config.JWT_EXPIRES_IN,
@@ -629,6 +631,8 @@ class AdminConfig(BaseModel):
     WEBUI_URL: str
     ENABLE_SIGNUP: bool
     ENABLE_API_KEY: bool
+    ENABLE_API_KEY_ENDPOINT_RESTRICTIONS: bool
+    API_KEY_ALLOWED_ENDPOINTS: str
     ENABLE_CHANNELS: bool
     DEFAULT_USER_ROLE: str
     JWT_EXPIRES_IN: str
@@ -643,7 +647,15 @@ async def update_admin_config(
     request.app.state.config.SHOW_ADMIN_DETAILS = form_data.SHOW_ADMIN_DETAILS
     request.app.state.config.WEBUI_URL = form_data.WEBUI_URL
     request.app.state.config.ENABLE_SIGNUP = form_data.ENABLE_SIGNUP
+
     request.app.state.config.ENABLE_API_KEY = form_data.ENABLE_API_KEY
+    request.app.state.config.ENABLE_API_KEY_ENDPOINT_RESTRICTIONS = (
+        form_data.ENABLE_API_KEY_ENDPOINT_RESTRICTIONS
+    )
+    request.app.state.config.API_KEY_ALLOWED_ENDPOINTS = (
+        form_data.API_KEY_ALLOWED_ENDPOINTS
+    )
+
     request.app.state.config.ENABLE_CHANNELS = form_data.ENABLE_CHANNELS
 
     if form_data.DEFAULT_USER_ROLE in ["pending", "user", "admin"]:
@@ -665,6 +677,8 @@ async def update_admin_config(
         "WEBUI_URL": request.app.state.config.WEBUI_URL,
         "ENABLE_SIGNUP": request.app.state.config.ENABLE_SIGNUP,
         "ENABLE_API_KEY": request.app.state.config.ENABLE_API_KEY,
+        "ENABLE_API_KEY_ENDPOINT_RESTRICTIONS": request.app.state.config.ENABLE_API_KEY_ENDPOINT_RESTRICTIONS,
+        "API_KEY_ALLOWED_ENDPOINTS": request.app.state.config.API_KEY_ALLOWED_ENDPOINTS,
         "ENABLE_CHANNELS": request.app.state.config.ENABLE_CHANNELS,
         "DEFAULT_USER_ROLE": request.app.state.config.DEFAULT_USER_ROLE,
         "JWT_EXPIRES_IN": request.app.state.config.JWT_EXPIRES_IN,

+ 7 - 5
backend/open_webui/utils/auth.py

@@ -96,11 +96,13 @@ def get_current_user(
                 status.HTTP_403_FORBIDDEN, detail=ERROR_MESSAGES.API_KEY_NOT_ALLOWED
             )
 
-        allowed_paths = ["/api/models", "/api/chat/completions"]
-        if request.url.path not in allowed_paths:
-            raise HTTPException(
-                status.HTTP_403_FORBIDDEN, detail=ERROR_MESSAGES.API_KEY_NOT_ALLOWED
-            )
+        if request.app.state.ENABLE_API_KEY_ENDPOINT_RESTRICTIONS:
+            allowed_paths = str(request.app.state.API_KEY_ALLOWED_ENDPOINTS).split(",")
+
+            if request.url.path not in allowed_paths:
+                raise HTTPException(
+                    status.HTTP_403_FORBIDDEN, detail=ERROR_MESSAGES.API_KEY_NOT_ALLOWED
+                )
 
         return get_current_user_by_api_key(token)
 

+ 27 - 2
src/lib/components/admin/Settings/General.svelte

@@ -112,12 +112,37 @@
 					</div>
 				</div>
 
-				<div class=" flex w-full justify-between pr-2">
-					<div class=" self-center text-xs font-medium">{$i18n.t('Enable API Key Auth')}</div>
+				<div class=" flex w-full justify-between pr-2 my-3">
+					<div class=" self-center text-xs font-medium">{$i18n.t('Enable API Key')}</div>
 
 					<Switch bind:state={adminConfig.ENABLE_API_KEY} />
 				</div>
 
+				{#if adminConfig?.ENABLE_API_KEY}
+					<div class=" flex w-full justify-between pr-2 my-3">
+						<div class=" self-center text-xs font-medium">
+							{$i18n.t('API Key Endpoint Restrictions')}
+						</div>
+
+						<Switch bind:state={adminConfig.ENABLE_API_KEY_ENDPOINT_RESTRICTIONS} />
+					</div>
+
+					{#if adminConfig?.ENABLE_API_KEY_ENDPOINT_RESTRICTIONS}
+						<div class=" flex w-full flex-col pr-2">
+							<div class=" text-xs font-medium">
+								{$i18n.t('Allowed Endpoints')}
+							</div>
+
+							<input
+								class="w-full mt-1 rounded-lg text-sm dark:text-gray-300 bg-transparent outline-none"
+								type="text"
+								placeholder={`e.g.) /api/v1/messages, /api/v1/channels`}
+								bind:value={adminConfig.API_KEY_ALLOWED_ENDPOINTS}
+							/>
+						</div>
+					{/if}
+				{/if}
+
 				<hr class=" border-gray-50 dark:border-gray-850 my-2" />
 
 				<div class="my-3 flex w-full items-center justify-between pr-2">