From d00963a4be47cc234da4ad5f4496aedd4c4b5323 Mon Sep 17 00:00:00 2001
From: Filip Kojro <filiopeq@gmail.com>
Date: Mon, 5 Oct 2026 00:19:55 +0200
Subject: [PATCH] adding some security

---
 README.md         |  4 +++-
 llm-reviewer.yaml |  3 +++
 main.py           | 46 ++++++++++++++++++++++++++++++++++++++++------
 3 files changed, 46 insertions(+), 7 deletions(-)

diff --git a/README.md b/README.md
index 82089a9..6bff8a1 100644
--- a/README.md
+++ b/README.md
@@ -3,9 +3,11 @@
 running server
 ```bash
 docker build -t llm-reviewer .
-docker run -d --rm -p 8080:8080 --name llm-reviewer llm-reviewer
+docker run -e LLM_REVIEW_ALLOWED_HOST='https://your-gitea-server-host' -d --rm -p 8080:8080 --name llm-reviewer llm-reviewer
 ```
 
+where `LLM_REVIEW_ALLOWED_HOST='https://your-gitea-server-host'` can be multiple comma separated host names
+
 in gitea repository fill:
 - secrets:
   ```
diff --git a/llm-reviewer.yaml b/llm-reviewer.yaml
index 91aed4d..7232905 100644
--- a/llm-reviewer.yaml
+++ b/llm-reviewer.yaml
@@ -14,6 +14,9 @@ jobs:
     if: github.event.pull_request.draft == false
     runs-on: ubuntu-latest
     steps:
+      - name: Test connection to openrouter
+        run: curl -f -X GET "${{ vars.LLM_REVIEW_HOST }}/health"
+
       - name: Trigger AI Reviewer
         run: |
           curl -f -X POST "${{ vars.LLM_REVIEW_HOST }}" \
diff --git a/main.py b/main.py
index 59d34bb..eea5645 100644
--- a/main.py
+++ b/main.py
@@ -1,8 +1,26 @@
+import os
+
 import requests
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException, status
 from openai import OpenAI
 from pydantic import BaseModel
 
+allowed_host_env = os.getenv("LLM_REVIEW_ALLOWED_HOST")
+allowed_hosts = []
+allowed_hosts_not_empty = True
+
+if allowed_host_env is None:
+    allowed_hosts_not_empty = False
+else:
+    allowed_hosts = allowed_host_env.split(',')
+    if len(allowed_hosts) == 0:
+        allowed_hosts_not_empty = False
+
+for (i, host) in enumerate(allowed_hosts):
+    allowed_hosts[i] = host.rstrip('/')
+
+print(f'allowed host: {allowed_hosts}')
+
 app = FastAPI()
 
 class ReviewRequest(BaseModel):
@@ -13,15 +31,30 @@ class ReviewRequest(BaseModel):
     repo: str
     pr_number: int
 
+@app.get("/health")
+def test():
+    try:
+        r = requests.get("https://openrouter.ai/api/v1/models")
+        r.raise_for_status()
+        return {"status": "ok"}
+    except requests.exceptions.RequestException as e:
+        raise HTTPException(
+            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
+            detail=f"OpenRouter health check failed: {e}",
+        )
+
 @app.post("/")
-async def review(
+def review(
     request_body: ReviewRequest,
 ):
-    print('starting bg task with:')
-    print(f'request: {request_body}')
-
     gitea_host = request_body.gitea_host.rstrip("/")
 
+    if allowed_hosts_not_empty and gitea_host not in allowed_hosts:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Request source is not allowed host",
+        )
+
     headers = {"Authorization": f"token {request_body.gitea_token}"}
 
     diff_res = requests.get(f"{gitea_host}/api/v1/repos/{request_body.repo}/pulls/{request_body.pr_number}.diff", headers=headers)
@@ -35,7 +68,8 @@ async def review(
     prompt = (
             "You are an automated senior code reviewer. Review the following git diff.\n"
             "Focus on: critical bugs, security vulnerabilities, edge cases, and code smells.\n"
-            "Keep it short and clear, don't give me long solutions."
+            "Keep it short and clear, don't give me long solutions.\n"
+            "Dont point out good things only problems, errors, etc.\n"
             "Format your answer cleanly in Markdown with actionable feedback.\n\n"
             f"```diff\n{diff_text}\n```"
     )