Geliebte
a3bbf33358
- Gin + config.yaml + SQLite 后端骨架，JWT 登录 + 管理员初始化
- web/ 原生单页：拖拽多图上传 → Agnes 免费识别 → 结果复制
- HTTP API：/api/v1/vision/analyze 免鉴权 + 令牌桶限流，支持程序直调
- cmd/vision-cli 独立 CLI 工具（复用 service 层，不依赖 Web 服务）
- build.bat / build.sh 构建脚本，移除 mcp-server
- 统一响应格式 {code, message, data}，CORS/JWT/限流中间件
Co-Authored-By: Claude <noreply@anthropic.com>
115 lines
2.9 KiB
Go
115 lines
2.9 KiB
Go
package middleware
|
|
|
|
import (
|
|
"math"
|
|
"net/http"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/golang-jwt/jwt/v5"
|
|
|
|
"vision-tool/internal/model"
|
|
)
|
|
|
|
// CORS 跨域中间件（开发环境放开，生产环境应限制为具体域名）。
|
|
func CORS() gin.HandlerFunc {
|
|
return func(c *gin.Context) {
|
|
c.Header("Access-Control-Allow-Origin", "*")
|
|
c.Header("Access-Control-Allow-Methods", "GET, POST, PUT, DELETE, OPTIONS")
|
|
c.Header("Access-Control-Allow-Headers", "Content-Type, Authorization")
|
|
if c.Request.Method == http.MethodOptions {
|
|
c.AbortWithStatus(http.StatusNoContent)
|
|
return
|
|
}
|
|
c.Next()
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// JWT 鉴权
|
|
// ============================================================================
|
|
|
|
// Claims JWT 载荷。
|
|
type Claims struct {
|
|
UserID uint `json:"uid"`
|
|
Username string `json:"username"`
|
|
Role string `json:"role"`
|
|
jwt.RegisteredClaims
|
|
}
|
|
|
|
// JWTAuth 校验 Authorization: Bearer <token>，通过后将用户信息写入 context。
|
|
func JWTAuth(secret string) gin.HandlerFunc {
|
|
return func(c *gin.Context) {
|
|
header := c.GetHeader("Authorization")
|
|
tokenStr, ok := strings.CutPrefix(header, "Bearer ")
|
|
if !ok || tokenStr == "" {
|
|
c.AbortWithStatusJSON(http.StatusUnauthorized, model.Err(model.CodeUnauthorized, "未登录"))
|
|
return
|
|
}
|
|
|
|
claims := &Claims{}
|
|
token, err := jwt.ParseWithClaims(tokenStr, claims, func(t *jwt.Token) (interface{}, error) {
|
|
return []byte(secret), nil
|
|
})
|
|
if err != nil || !token.Valid {
|
|
c.AbortWithStatusJSON(http.StatusUnauthorized, model.Err(model.CodeUnauthorized, "登录已过期，请重新登录"))
|
|
return
|
|
}
|
|
|
|
c.Set("user_id", claims.UserID)
|
|
c.Set("username", claims.Username)
|
|
c.Set("role", claims.Role)
|
|
c.Next()
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// 内存令牌桶限流（全局限流，防匿名刷接口）
|
|
// ============================================================================
|
|
|
|
// RateLimiter 基于时间的令牌桶。
|
|
type RateLimiter struct {
|
|
mu sync.Mutex
|
|
rate float64 // 每秒补充令牌数
|
|
burst float64 // 桶容量
|
|
tokens float64
|
|
lastTime time.Time
|
|
}
|
|
|
|
func NewRateLimiter(rps, burst int) *RateLimiter {
|
|
return &RateLimiter{
|
|
rate: float64(rps),
|
|
burst: float64(burst),
|
|
tokens: float64(burst),
|
|
lastTime: time.Now(),
|
|
}
|
|
}
|
|
|
|
func (rl *RateLimiter) Allow() bool {
|
|
rl.mu.Lock()
|
|
defer rl.mu.Unlock()
|
|
|
|
now := time.Now()
|
|
rl.tokens = math.Min(rl.burst, rl.tokens+now.Sub(rl.lastTime).Seconds()*rl.rate)
|
|
rl.lastTime = now
|
|
|
|
if rl.tokens >= 1 {
|
|
rl.tokens--
|
|
return true
|
|
}
|
|
return false
|
|
}
|
|
|
|
// Middleware 返回限流 gin 中间件。
|
|
func (rl *RateLimiter) Middleware() gin.HandlerFunc {
|
|
return func(c *gin.Context) {
|
|
if !rl.Allow() {
|
|
c.AbortWithStatusJSON(http.StatusTooManyRequests, model.Err(429, "请求过于频繁，请稍后再试"))
|
|
return
|
|
}
|
|
c.Next()
|
|
}
|
|
}
|