From 47838f9880c12f46070d4a0a4470b2daa930d2a4 Mon Sep 17 00:00:00 2001
From: Nguyen Van Toan <toannv@myserver.local>
Date: Sat, 18 Jan 2025 23:17:27 +0700
Subject: [PATCH] Add: Authenticate

---
 auth.py   | 23 +++++++++++++++++++++++
 server.py |  4 ++++
 2 files changed, 27 insertions(+)
 create mode 100644 auth.py

diff --git a/auth.py b/auth.py
new file mode 100644
index 000000000..afe0cd078
--- /dev/null
+++ b/auth.py
@@ -0,0 +1,23 @@
+import os
+from aiohttp import web
+
+COMFYUI_API_KEY = 'OKZOO_COMFYUI_API_KEY'
+CLIENT_API_KEY = 'X-OKZOO-API-KEY'
+
+def validate_api_key(header_api_key):
+    api_key = os.getenv(COMFYUI_API_KEY)
+    
+    if(header_api_key == api_key):
+        return True
+    return False
+
+def get_api_key_from_client_request(request: web.Request):
+    return request.headers.get(CLIENT_API_KEY)
+
+def validate_request(request: web.Request):
+    api_key = get_api_key_from_client_request(request)
+
+    if api_key is None:
+        raise web.HTTPUnauthorized(text="Unauthorized access: Missing API key")
+    if validate_api_key(api_key) == False:
+        raise web.HTTPUnauthorized(text="Unauthorized access: Invalid API key")
\ No newline at end of file
diff --git a/server.py b/server.py
index bae898ef5..491022cb7 100644
--- a/server.py
+++ b/server.py
@@ -34,6 +34,7 @@ from app.model_manager import ModelFileManager
 from app.custom_node_manager import CustomNodeManager
 from typing import Optional
 from api_server.routes.internal.internal_routes import InternalRoutes
+from auth import validate_request
 
 class BinaryEventTypes:
     PREVIEW_IMAGE = 1
@@ -55,6 +56,8 @@ async def cache_control(request: web.Request, handler):
 def create_cors_middleware(allowed_origin: str):
     @web.middleware
     async def cors_middleware(request: web.Request, handler):
+        validate_request(request)
+
         if request.method == "OPTIONS":
             # Pre-flight request. Reply successfully:
             response = web.Response()
@@ -98,6 +101,7 @@ def is_loopback(host):
 def create_origin_only_middleware():
     @web.middleware
     async def origin_only_middleware(request: web.Request, handler):
+        validate_request(request)
         #this code is used to prevent the case where a random website can queue comfy workflows by making a POST to 127.0.0.1 which browsers don't prevent for some dumb reason.
         #in that case the Host and Origin hostnames won't match
         #I know the proper fix would be to add a cookie but this should take care of the problem in the meantime