From e68da34c1317185a14928a34fa96956b7cbb48bd Mon Sep 17 00:00:00 2001
From: Alex <a@tushynski.me>
Date: Thu, 4 Dec 2025 15:52:45 +0000
Subject: [PATCH] feat: implement internal API authentication mechanism

---
 application/.env_sample            |  1 +
 application/api/internal/routes.py | 12 +++++++++++-
 application/core/settings.py       |  1 +
 application/worker.py              |  9 ++++++++-
 deployment/docker-compose.yaml     |  2 ++
 5 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/application/.env_sample b/application/.env_sample
index 8ab24d2a..c08b2c1d 100644
--- a/application/.env_sample
+++ b/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
 
diff --git a/application/api/internal/routes.py b/application/api/internal/routes.py
index f0c2950e..f812a37f 100755
--- a/application/api/internal/routes.py
+++ b/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 @@ current_dir = os.path.dirname(
 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"))
diff --git a/application/core/settings.py b/application/core/settings.py
index cecbb333..12759d7f 100644
--- a/application/core/settings.py
+++ b/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)
 
diff --git a/application/worker.py b/application/worker.py
index 5c5dc367..cbc0dfdc 100755
--- a/application/worker.py
+++ b/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:
diff --git a/deployment/docker-compose.yaml b/deployment/docker-compose.yaml
index 2eef387b..7ca110aa 100644
--- a/deployment/docker-compose.yaml
+++ b/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