From 159daf6bc03087a864bc4b3ef5f432e720d0163e Mon Sep 17 00:00:00 2001 From: orbisai0security Date: Wed, 6 May 2026 10:47:31 +0000 Subject: [PATCH] fix: V-008 security vulnerability Automated security fix generated by Orbis Security AI --- .../ColossalChat/start_code_verifier.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/applications/ColossalChat/start_code_verifier.py b/applications/ColossalChat/start_code_verifier.py index d1924f610698..d6b81fe54092 100644 --- a/applications/ColossalChat/start_code_verifier.py +++ b/applications/ColossalChat/start_code_verifier.py @@ -1,11 +1,25 @@ +import os +import secrets from typing import List, Optional from coati.distributed.reward.code_reward.utils import check_correctness # Assuming utils.py is in the same directory -from fastapi import FastAPI, HTTPException +from fastapi import Depends, FastAPI, HTTPException, Security +from fastapi.security import APIKeyHeader from pydantic import BaseModel app = FastAPI() +API_KEY = os.environ.get("CODE_VERIFIER_API_KEY", "") +API_KEY_HEADER = APIKeyHeader(name="X-API-Key", auto_error=True) + + +def verify_api_key(api_key: str = Security(API_KEY_HEADER)): + if not API_KEY: + raise HTTPException(status_code=500, detail="Server API key not configured") + if not secrets.compare_digest(api_key, API_KEY): + raise HTTPException(status_code=403, detail="Invalid API key") + return api_key + class CheckCorrectnessRequest(BaseModel): in_outs: Optional[dict] @@ -21,7 +35,7 @@ class CheckCorrectnessResponse(BaseModel): @app.post("/check_correctness", response_model=CheckCorrectnessResponse) -def check_correctness_api(request: CheckCorrectnessRequest): +def check_correctness_api(request: CheckCorrectnessRequest, _: str = Depends(verify_api_key)): try: result, metadata = check_correctness( in_outs=request.in_outs,