feat: implement internal API authentication mechanism

Author: dartpain

application/.env_sample

@@ -1,6 +1,7 @@
 API_KEY=your_api_key
 EMBEDDINGS_KEY=your_api_key
 API_URL=http://localhost:7091
+INTERNAL_KEY=your_internal_key
 FLASK_APP=application/app.py
 FLASK_DEBUG=true
 

application/api/internal/routes.py

@@ -1,7 +1,7 @@
 import os
 import datetime
 import json
-from flask import Blueprint, request, send_from_directory
+from flask import Blueprint, request, send_from_directory, jsonify
 from werkzeug.utils import secure_filename
 from bson.objectid import ObjectId
 import logging
@@ -24,6 +24,16 @@
 internal = Blueprint("internal", __name__)
 
 
+@internal.before_request
+def verify_internal_key():
+    """Verify INTERNAL_KEY for all internal endpoint requests."""
+    if settings.INTERNAL_KEY:
+        internal_key = request.headers.get("X-Internal-Key")
+        if not internal_key or internal_key != settings.INTERNAL_KEY:
+            logger.warning(f"Unauthorized internal API access attempt from {request.remote_addr}")
+            return jsonify({"error": "Unauthorized", "message": "Invalid or missing internal key"}), 401
+
+
 @internal.route("/api/download", methods=["get"])
 def download_file():
     user = secure_filename(request.args.get("user"))

application/core/settings.py

@@ -62,6 +62,7 @@ class Settings(BaseSettings):
     CACHE_REDIS_URL: str = "redis://localhost:6379/2"
 
     API_URL: str = "http://localhost:7091"  # backend url for celery worker
+    INTERNAL_KEY: Optional[str] = None  # internal api key for worker-to-backend auth
 
     API_KEY: Optional[str] = None  # LLM api key (used by LLM_PROVIDER)
 

application/worker.py

@@ -109,6 +109,10 @@ def download_file(url, params, dest_path):
 def upload_index(full_path, file_data):
     files = None
     try:
+        headers = {}
+        if settings.INTERNAL_KEY:
+            headers["X-Internal-Key"] = settings.INTERNAL_KEY
+
         if settings.VECTOR_STORE == "faiss":
             faiss_path = full_path + "/index.faiss"
             pkl_path = full_path + "/index.pkl"
@@ -129,10 +133,13 @@ def upload_index(full_path, file_data):
                 urljoin(settings.API_URL, "/api/upload_index"),
                 files=files,
                 data=file_data,
+                headers=headers,
             )
         else:
             response = requests.post(
-                urljoin(settings.API_URL, "/api/upload_index"), data=file_data
+                urljoin(settings.API_URL, "/api/upload_index"),
+                data=file_data,
+                headers=headers,
             )
         response.raise_for_status()
     except (requests.RequestException, FileNotFoundError) as e:

deployment/docker-compose.yaml

@@ -26,6 +26,7 @@ services:
       - MONGO_URI=mongodb://mongo:27017/docsgpt
       - CACHE_REDIS_URL=redis://redis:6379/2
       - OPENAI_BASE_URL=$OPENAI_BASE_URL
+      - INTERNAL_KEY=$INTERNAL_KEY
     ports:
       - "7091:7091"
     volumes:
@@ -50,6 +51,7 @@ services:
       - MONGO_URI=mongodb://mongo:27017/docsgpt
       - API_URL=http://backend:7091
       - CACHE_REDIS_URL=redis://redis:6379/2
+      - INTERNAL_KEY=$INTERNAL_KEY
     volumes:
       - ../application/indexes:/app/indexes
       - ../application/inputs:/app/inputs