Add: Authenticate

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")

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