tdpeuter/open-webui
tdpeuter/open-webui
1
Fork 0
forked from open-webui/open-webui
Code Issues Pull requests 1 Projects 1 Releases 2 Packages 1 Activity Actions

feat: basic RBAC support

Browse source
This commit is contained in:
Timothy J. Baek 2023-11-18 21:41:43 -08:00
parent 921eef03b3
commit 8547b7807d
13 changed files with 266 additions and 44 deletions
Download patch file Download diff file Expand all files Collapse all files

23
backend/apps/ollama/main.py
View file

@ -8,7 +8,7 @@ import json
from apps.web.models.users import Users
from constants import ERROR_MESSAGES
from utils import extract_token_from_auth_header
from utils.utils import extract_token_from_auth_header
from config import OLLAMA_API_BASE_URL, OLLAMA_WEBUI_AUTH
app = Flask(__name__)
@ -25,24 +25,37 @@ TARGET_SERVER_URL = OLLAMA_API_BASE_URL
def proxy(path):
# Combine the base URL of the target server with the requested path
target_url = f"{TARGET_SERVER_URL}/{path}"
print(target_url)
print(path)
# Get data from the original request
data = request.get_data()
headers = dict(request.headers)
# Basic RBAC support
if OLLAMA_WEBUI_AUTH:
if "Authorization" in headers:
token = extract_token_from_auth_header(headers["Authorization"])
user = Users.get_user_by_token(token)
if user:
print(user)
pass
# Only user and admin roles can access
if user.role in ["user", "admin"]:
if path in ["pull", "delete", "push", "copy", "create"]:
# Only admin role can perform actions above
if user.role == "admin":
pass
else:
return (
jsonify({"detail": ERROR_MESSAGES.ACCESS_PROHIBITED}),
401,
)
else:
pass
else:
return jsonify({"detail": ERROR_MESSAGES.ACCESS_PROHIBITED}), 401
else:
return jsonify({"detail": ERROR_MESSAGES.UNAUTHORIZED}), 401
else:
return jsonify({"detail": ERROR_MESSAGES.UNAUTHORIZED}), 401
else:
pass