Files

143 lines
5.2 KiB
Python

import asyncio
import ipaddress
import json
import socket
from urllib.parse import urlparse
import httpx
from fastapi import APIRouter, Depends, HTTPException
from pydantic import BaseModel
from typing import Any, Dict
from app.api.v1.auth import get_current_user
from app.api.v1.user_config import _load_ai_configs, _get_default_providers
router = APIRouter()
class ProxyRequest(BaseModel):
url: str
headers: Dict[str, str] = {}
body: Dict[str, Any] = {}
# Ranges that are never legitimate AI endpoints: cloud metadata + this host.
_BLOCKED_NETWORKS = [
ipaddress.ip_network("169.254.0.0/16"), # link-local / cloud metadata
ipaddress.ip_network("0.0.0.0/8"),
]
# Private ranges: only reachable when the target host is one the user has
# explicitly configured as an AI provider (e.g. local Ollama/LM Studio).
_PRIVATE_NETWORKS = [
ipaddress.ip_network("10.0.0.0/8"),
ipaddress.ip_network("172.16.0.0/12"),
ipaddress.ip_network("192.168.0.0/16"),
ipaddress.ip_network("127.0.0.0/8"),
ipaddress.ip_network("::1/128"),
ipaddress.ip_network("fc00::/7"), # ULA
]
_LOOPBACK_HOSTS = {"localhost", "127.0.0.1", "::1", "0.0.0.0"}
def _configured_ai_hosts(user_id: str) -> set:
"""Hosts the user has configured as AI providers (from saved config + defaults)."""
hosts = set()
configs = _load_ai_configs()
providers = configs.get(user_id) or _get_default_providers()
for p in providers:
base = (p.get("api_base_url") or "").strip()
if not base:
continue
try:
host = urlparse(base).hostname
if host:
hosts.add(host.lower())
except Exception:
continue
return hosts
async def _resolve_host_ips(hostname: str):
"""Resolve hostname to IPs (non-blocking). Returns list of ipaddress objects."""
loop = asyncio.get_event_loop()
try:
infos = await loop.run_in_executor(None, socket.getaddrinfo, hostname, None)
ips = []
for info in infos:
try:
ips.append(ipaddress.ip_address(info[4][0]))
except ValueError:
continue
return ips
except Exception:
return []
async def _validate_target_url(url: str, user_id: str):
parsed = urlparse(url)
if parsed.scheme not in ("http", "https"):
raise HTTPException(status_code=400, detail="URL chỉ hỗ trợ giao thức http/https")
if parsed.username or parsed.password:
raise HTTPException(status_code=400, detail="URL không được chứa thông tin đăng nhập")
hostname = (parsed.hostname or "").lower()
if not hostname:
raise HTTPException(status_code=400, detail="URL không hợp lệ")
allowed_hosts = _configured_ai_hosts(user_id)
# Hostname-level fast path for loopback hosts
if hostname in _LOOPBACK_HOSTS:
if hostname in allowed_hosts:
return
raise HTTPException(status_code=403, detail="Target nội bộ không nằm trong danh sách AI provider đã cấu hình")
# Try direct IP parse (hostname may itself be an IP)
try:
ip = ipaddress.ip_address(hostname)
ips = [ip]
except ValueError:
ips = await _resolve_host_ips(hostname)
if not ips:
raise HTTPException(status_code=502, detail="Không phân giải được hostname")
for ip in ips:
if any(ip in net for net in _BLOCKED_NETWORKS):
raise HTTPException(status_code=403, detail="Target bị chặn (metadata/link-local không được phép)")
if any(ip in net for net in _PRIVATE_NETWORKS):
if hostname in allowed_hosts:
continue
raise HTTPException(status_code=403, detail="Target IP nội bộ không nằm trong danh sách AI provider đã cấu hình")
@router.post("/proxy")
async def proxy_llm(req: ProxyRequest, current_user: dict = Depends(get_current_user)):
await _validate_target_url(req.url, current_user["user_id"])
# Never forward the app's own auth token upstream.
headers = {
k: v for k, v in req.headers.items()
if k.lower() not in ("host", "origin", "referer", "x-auth-token")
}
try:
async with httpx.AsyncClient(timeout=180.0, follow_redirects=False) as client:
resp = await client.post(req.url, headers=headers, json=req.body)
raw = resp.text
try:
return resp.json()
except json.JSONDecodeError:
try:
return json.loads(raw[:raw.find('\n')])
except (json.JSONDecodeError, ValueError):
return {"content": raw}
except httpx.TimeoutException:
raise HTTPException(status_code=504, detail="AI provider timeout")
except httpx.ConnectError as e:
msg = f"Cannot connect to AI provider: {e}"
if 'localhost' in req.url or '127.0.0.1' in req.url:
msg += "\nNếu app chạy trong Docker, localhost trỏ vào container, không ra host.\nHãy thay localhost bằng host.docker.internal hoặc IP bridge Docker (172.17.0.1)."
raise HTTPException(status_code=502, detail=msg)
except HTTPException:
raise
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))