From 8fa764bfbe75aa62092a964d45da754298d34afe Mon Sep 17 00:00:00 2001
From: kingbri <bdashore3@proton.me>
Date: Thu, 21 Dec 2023 23:40:16 -0500
Subject: [PATCH] Auth: Add option to disable authentication

This creates a massive security hole, but it's gated behind a flag
for users who only use localhost.

A warning will pop up when users disable authentication.

Signed-off-by: kingbri <bdashore3@proton.me>
---
 auth.py           | 23 ++++++++++++++++++++++-
 config_sample.yml |  5 +++++
 main.py           |  9 +++++----
 3 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/auth.py b/auth.py
index 27f97cd..80611e6 100644
--- a/auth.py
+++ b/auth.py
@@ -24,9 +24,22 @@ class AuthKeys(BaseModel):
             return False
 
 auth_keys: Optional[AuthKeys] = None
+disable_auth: bool = False
 
-def load_auth_keys():
+def load_auth_keys(disable_from_config: bool):
     global auth_keys
+    global disable_auth
+
+    disable_auth = disable_from_config
+    if disable_from_config:
+        print(
+            "!! Warning: Disabling authentication makes your instance vulnerable.",
+            "Set the \"disable_auth\" flag to False in config.yml if you want to share this",
+            "instance with others."
+        )
+
+        return
+
     try:
         with open("api_tokens.yml", "r", encoding = 'utf8') as auth_file:
             auth_keys_dict = yaml.safe_load(auth_file)
@@ -48,6 +61,10 @@ def load_auth_keys():
     )
 
 def check_api_key(x_api_key: str = Header(None), authorization: str = Header(None)):
+    # Allow request if auth is disabled
+    if disable_auth:
+        return
+
     if x_api_key:
         if auth_keys.verify_key(x_api_key, "api_key"):
             return x_api_key
@@ -66,6 +83,10 @@ def check_api_key(x_api_key: str = Header(None), authorization: str = Header(Non
         raise HTTPException(401, "Please provide an API key")
 
 def check_admin_key(x_admin_key: str = Header(None), authorization: str = Header(None)):
+    # Allow request if auth is disabled
+    if disable_auth:
+        return
+
     if x_admin_key:
         if auth_keys.verify_key(x_admin_key, "admin_key"):
             return x_admin_key
diff --git a/config_sample.yml b/config_sample.yml
index 1a1038d..435e91d 100644
--- a/config_sample.yml
+++ b/config_sample.yml
@@ -13,6 +13,11 @@ network:
   # The port to host on (default: 5000)
   port: 5000
 
+  # Disable HTTP token authenticaion with requests
+  # WARNING: This will make your instance vulnerable!
+  # Turn on this option if you are ONLY connecting from localhost
+  disable_auth: False
+
 # Options for logging
 logging:
   # Enable prompt logging (default: False)
diff --git a/main.py b/main.py
index 61a8f51..1228404 100644
--- a/main.py
+++ b/main.py
@@ -380,9 +380,6 @@ async def generate_chat_completion(request: Request, data: ChatCompletionRequest
         return response
 
 if __name__ == "__main__":
-    # Initialize auth keys
-    load_auth_keys()
-
     # Load from YAML config. Possibly add a config -> kwargs conversion function
     try:
         with open('config.yml', 'r', encoding = "utf8") as config_file:
@@ -395,6 +392,11 @@ if __name__ == "__main__":
         )
         config = {}
 
+    network_config = unwrap(config.get("network"), {})
+
+    # Initialize auth keys
+    load_auth_keys(unwrap(network_config.get("disable_auth"), False))
+
     # Override the generation log options if given
     log_config = unwrap(config.get("logging"), {})
     if log_config:
@@ -426,7 +428,6 @@ if __name__ == "__main__":
         lora_dir = pathlib.Path(unwrap(lora_config.get("lora_dir"), "loras"))
         model_container.load_loras(lora_dir.resolve(), **lora_config)
 
-    network_config = unwrap(config.get("network"), {})
     uvicorn.run(
         app,
         host=network_config.get("host", "127.0.0.1"),