模型服务访问控制实现

模型服务访问控制实现踩坑记录

最近在为公司AI平台搭建模型服务访问控制时，踩了不少坑，分享一下避免大家重蹈覆辙。

问题背景

我们使用FastAPI + Uvicorn部署模型服务，需要实现基于JWT的访问控制。最初方案是直接在FastAPI路由上添加依赖注入，但发现存在安全漏洞。

错误实践1：直接在路由函数中验证

from fastapi import Depends, HTTPException
from jose import jwt

async def verify_token(token: str = Depends(OAuth2PasswordBearer(tokenUrl="token"))):
    # 这样做很危险，因为没有正确处理异常
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        return payload
    except Exception:
        raise HTTPException(status_code=401, detail="Invalid token")

@app.get("/predict")
async def predict(data: dict, token: dict = Depends(verify_token)):
    # 业务逻辑
    pass

正确方案：使用中间件+依赖注入

from fastapi import FastAPI, Request, HTTPException
from fastapi.middleware.trustedhost import TrustedHostMiddleware

app = FastAPI()
app.add_middleware(TrustedHostMiddleware, allowed_hosts=["localhost", "127.0.0.1"])

# 自定义中间件
class AuthMiddleware:
    def __init__(self, app):
        self.app = app
    
    async def __call__(self, scope, receive, send):
        # 在这里进行token验证
        request = Request(scope, receive)
        auth_header = request.headers.get("authorization")
        if not auth_header or not auth_header.startswith("Bearer "):
            raise HTTPException(status_code=401, detail="Missing token")
        
        # 验证token逻辑
        token = auth_header.split(" ")[1]
        try:
            payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        except Exception:
            raise HTTPException(status_code=401, detail="Invalid token")
        
        scope["user"] = payload
        await self.app(scope, receive, send)

app.add_middleware(AuthMiddleware)

关键踩坑点

1. Token过期处理：不要只捕获Exception，要区分Token过期和无效token
2. 中间件优先级：确保在路由之前执行验证逻辑
3. 日志记录：添加详细的访问日志便于排查问题
4. 性能考虑：避免在每次请求都重新解析JWT，建议使用缓存

最终推荐方案

使用FastAPI官方推荐的Depends + Security依赖注入方式，配合数据库token黑名单机制，既安全又高效。