Skip to content

Commit 0f8cb4d

Browse files
fix: harden agent invoke execution path and logger init order
- Add authentication to all agent invoke endpoints using CALL_SERVER_TOKEN - Fix blocking sync calls in async handlers with asyncio.run_in_executor - Remove agent enumeration from 404 error responses (security) - Improve error handling and async/sync compatibility Fixes security vulnerabilities identified in code review: - P1: Unauthenticated agent invocation endpoints - P1: Blocking sync calls in async event loop - P1: Information disclosure via 404 responses
1 parent 13188fc commit 0f8cb4d

1 file changed

Lines changed: 54 additions & 10 deletions

File tree

src/praisonai/praisonai/api/agent_invoke.py

Lines changed: 54 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
import logging
1111

1212
try:
13-
from fastapi import APIRouter, HTTPException, Depends
13+
from fastapi import APIRouter, HTTPException, Depends, Header, Request
1414
from pydantic import BaseModel, Field
1515
FASTAPI_AVAILABLE = True
1616
except ImportError:
@@ -19,10 +19,48 @@
1919
HTTPException = None
2020
BaseModel = object
2121
Field = lambda *args, **kwargs: None
22+
Depends = lambda x: x
23+
Header = lambda *args, **kwargs: None
24+
Request = object
2225
FASTAPI_AVAILABLE = False
2326

2427
logger = logging.getLogger(__name__)
2528

29+
# Authentication
30+
import os
31+
CALL_SERVER_TOKEN = os.getenv('CALL_SERVER_TOKEN')
32+
33+
async def verify_token(
34+
request: Request,
35+
authorization: Optional[str] = Header(None)
36+
) -> None:
37+
"""Verify API token for authentication."""
38+
if not FASTAPI_AVAILABLE or not CALL_SERVER_TOKEN:
39+
return # No authentication if FastAPI unavailable or no token set
40+
41+
token = None
42+
43+
# Check Authorization header first (Bearer or Basic)
44+
if authorization:
45+
if authorization.startswith("Bearer "):
46+
token = authorization.split(" ")[1]
47+
elif authorization.startswith("Basic "):
48+
try:
49+
import base64
50+
decoded = base64.b64decode(authorization[6:]).decode("utf-8")
51+
if ":" in decoded:
52+
token = decoded.split(":", 1)[1] # Use password as token
53+
else:
54+
token = decoded
55+
except Exception:
56+
pass
57+
58+
# Check query param as fallback
59+
if not token:
60+
token = request.query_params.get("token")
61+
62+
if token != CALL_SERVER_TOKEN:
63+
raise HTTPException(status_code=401, detail="Unauthorized")
2664

2765
# Request/Response Models
2866
if FASTAPI_AVAILABLE:
@@ -114,7 +152,8 @@ def _supports_sync_start(agent: Any) -> bool:
114152
@router.post("/agents/{agent_id}/invoke")
115153
async def invoke_agent(
116154
agent_id: str,
117-
request: AgentInvokeRequest
155+
request: AgentInvokeRequest,
156+
_: None = Depends(verify_token)
118157
) -> Union[AgentInvokeResponse, ErrorResponse]:
119158
"""
120159
Invoke a PraisonAI agent with a message.
@@ -150,7 +189,7 @@ async def invoke_agent(
150189
logger.error(f"Agent not found: {agent_id}")
151190
raise HTTPException(
152191
status_code=404,
153-
detail=f"Agent '{agent_id}' not found. Available agents: {list_registered_agents()}"
192+
detail=f"Agent '{agent_id}' not found"
154193
)
155194

156195
try:
@@ -168,8 +207,10 @@ async def invoke_agent(
168207
# Async agent
169208
result = await agent.astart(request.message)
170209
elif _supports_sync_start(agent):
171-
# Sync agent (use start method)
172-
result = agent.start(request.message)
210+
# Sync agent - run in thread pool to avoid blocking the event loop
211+
import asyncio
212+
loop = asyncio.get_event_loop()
213+
result = await loop.run_in_executor(None, agent.start, request.message)
173214
else:
174215
raise AttributeError(f"Agent {agent_id} must provide start() or async astart()")
175216

@@ -194,7 +235,7 @@ async def invoke_agent(
194235
)
195236

196237
@router.get("/agents")
197-
async def list_agents() -> Dict[str, Any]:
238+
async def list_agents(_: None = Depends(verify_token)) -> Dict[str, Any]:
198239
"""
199240
List all registered agents.
200241
@@ -209,7 +250,7 @@ async def list_agents() -> Dict[str, Any]:
209250
}
210251

211252
@router.post("/agents/{agent_id}/register")
212-
async def register_agent_endpoint(agent_id: str) -> Dict[str, Any]:
253+
async def register_agent_endpoint(agent_id: str, _: None = Depends(verify_token)) -> Dict[str, Any]:
213254
"""
214255
Register an agent for API access.
215256
@@ -225,7 +266,7 @@ async def register_agent_endpoint(agent_id: str) -> Dict[str, Any]:
225266
}
226267

227268
@router.delete("/agents/{agent_id}")
228-
async def unregister_agent_endpoint(agent_id: str) -> Dict[str, Any]:
269+
async def unregister_agent_endpoint(agent_id: str, _: None = Depends(verify_token)) -> Dict[str, Any]:
229270
"""
230271
Unregister an agent from API access.
231272
"""
@@ -242,7 +283,7 @@ async def unregister_agent_endpoint(agent_id: str) -> Dict[str, Any]:
242283
)
243284

244285
@router.get("/agents/{agent_id}")
245-
async def get_agent_info(agent_id: str) -> Dict[str, Any]:
286+
async def get_agent_info(agent_id: str, _: None = Depends(verify_token)) -> Dict[str, Any]:
246287
"""
247288
Get information about a registered agent.
248289
"""
@@ -305,7 +346,10 @@ async def invoke_agent_standalone(
305346
if _supports_async_start(agent):
306347
result = await agent.astart(message)
307348
elif _supports_sync_start(agent):
308-
result = agent.start(message)
349+
# Sync agent - run in thread pool to avoid blocking the event loop
350+
import asyncio
351+
loop = asyncio.get_event_loop()
352+
result = await loop.run_in_executor(None, agent.start, message)
309353
else:
310354
raise AttributeError(f"Agent {agent_id} must provide start() or async astart()")
311355

0 commit comments

Comments
 (0)