reface to openteam

middleware/auth.go

@@ -0,0 +1,55 @@
+package middleware
+
+import (
+	"fmt"
+	"net/http"
+	"opencatd-open/internal/auth"
+	"opencatd-open/internal/consts"
+	"opencatd-open/internal/dto"
+	"opencatd-open/internal/model"
+	"opencatd-open/pkg/store"
+
+	"github.com/gin-gonic/gin"
+)
+
+func Auth(c *gin.Context) {
+	authToken := c.GetHeader("Authorization")
+	if authToken == "" {
+		dto.Fail(c, http.StatusUnauthorized, "未提供认证信息")
+		return
+	}
+	authToken = authToken[7:]
+	claim, err := auth.ValidateToken(authToken, consts.SecretKey)
+	if err != nil {
+		c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{
+			"code":  http.StatusUnauthorized,
+			"error": "无效的认证信息",
+		})
+		return
+	}
+	var user model.User
+	if err := store.GetDB().Model(&model.User{ID: int64(claim.UserID)}).First(&user).Error; err != nil {
+		c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{
+			"code":  http.StatusUnauthorized,
+			"error": "无效的认证信息",
+		})
+		return
+	}
+	c.Set("user", &user)
+	c.Set("user_id", claim.UserID)
+	c.Set("user_role", user.Role)
+	c.Next()
+}
+func CheckRole(role consts.UserRole) func(c *gin.Context) {
+	fmt.Println("CheckRoleMiddleware")
+	return func(c *gin.Context) {
+		userRole := c.GetInt("user_role") // 操作者
+		fmt.Println("userRole", userRole)
+		// if userRole < int(role) {
+		// 	dto.Fail(c, http.StatusForbidden, "permission denied")
+		// 	return
+		// }
+
+		c.Next()
+	}
+}

middleware/auth_team.go

@@ -0,0 +1,103 @@
+package middleware
+
+import (
+	"net/http"
+	"opencatd-open/internal/dto"
+	"opencatd-open/internal/model"
+
+	"github.com/gin-gonic/gin"
+	"gorm.io/gorm"
+)
+
+func AuthTeam(db *gorm.DB) gin.HandlerFunc {
+	return func(c *gin.Context) {
+
+		auth_token := c.GetHeader("Authorization")
+		if len(auth_token) < 7 || auth_token[:7] != "Bearer " {
+			c.JSON(http.StatusUnauthorized, gin.H{"error": "Unauthorized"})
+			c.Abort()
+			return
+		}
+		auth_token = auth_token[7:]
+		token := model.Token{}
+		if err := db.Preload("Users").First(&token, "token = ?", auth_token).Error; err != nil {
+			dto.WrapErrorAsOpenAI(c, http.StatusUnauthorized, "invalid_api_key")
+			c.Abort()
+			return
+		}
+
+		if token.User == nil {
+			c.JSON(http.StatusUnauthorized, gin.H{"error": "Unauthorized"})
+			c.Abort()
+			return
+		}
+
+		if !*token.User.Active || !*token.Active {
+			dto.WrapErrorAsOpenAI(c, http.StatusForbidden, "User or API key is not active")
+			c.Abort()
+			return
+		}
+
+		if token.Name != "default" {
+			dto.WrapErrorAsOpenAI(c, http.StatusForbidden, "Only default api key accessible")
+			c.Abort()
+			return
+		}
+
+		c.Set("user", token.User)
+		c.Set("authed", true)
+		// 可以在这里对 token 进行验证并检查权限
+		c.Next()
+
+	}
+}
+
+func AuthLLM(db *gorm.DB) gin.HandlerFunc {
+	return func(c *gin.Context) {
+
+		auth_token := c.GetHeader("Authorization")
+		if len(auth_token) < 7 || auth_token[:7] != "Bearer " {
+			c.JSON(http.StatusUnauthorized, gin.H{"error": "Unauthorized"})
+			c.Abort()
+			return
+		}
+		auth_token = auth_token[7:]
+		token := model.Token{}
+		if err := db.Preload("Users").First(&token, "token = ?", auth_token).Error; err != nil {
+			dto.WrapErrorAsOpenAI(c, http.StatusUnauthorized, "invalid_api_key")
+			c.Abort()
+			return
+		}
+
+		if token.User == nil {
+			c.JSON(http.StatusUnauthorized, gin.H{"error": "Unauthorized"})
+			c.Abort()
+			return
+		}
+
+		if !*token.User.Active || !*token.Active {
+			dto.WrapErrorAsOpenAI(c, http.StatusForbidden, "User or API key is not active")
+			c.Abort()
+			return
+		}
+
+		if !*token.User.UnlimitedQuota && *token.User.Quota <= 0 {
+			dto.WrapErrorAsOpenAI(c, http.StatusForbidden, "quota_exceeded")
+			c.Abort()
+			return
+		}
+
+		if !*token.UnlimitedQuota && *token.Quota <= 0 {
+			dto.WrapErrorAsOpenAI(c, http.StatusForbidden, "quota_exceeded")
+			c.Abort()
+			return
+		}
+
+		c.Set("user", token.User)
+		c.Set("authed", true)
+		// 可以在这里对 token 进行验证并检查权限
+
+		c.Next()
+
+	}
+}

middleware/cors.go

@@ -0,0 +1,15 @@
+package middleware
+
+import (
+	"github.com/gin-contrib/cors"
+	"github.com/gin-gonic/gin"
+)
+
+func CORS() gin.HandlerFunc {
+	config := cors.DefaultConfig()
+	config.AllowAllOrigins = true
+	config.AllowCredentials = true
+	config.AllowMethods = []string{"GET", "POST", "PUT", "DELETE", "OPTIONS"}
+	config.AllowHeaders = []string{"*"}
+	return cors.New(config)
+}

middleware/ratelimit.go

@@ -0,0 +1,53 @@
+package middleware
+
+import (
+	"net/http"
+	"sync"
+
+	"github.com/gin-gonic/gin"
+	"golang.org/x/time/rate"
+)
+
+type IPRateLimiter struct {
+	ips map[string]*rate.Limiter
+	mu  *sync.RWMutex
+	r   rate.Limit
+	b   int
+}
+
+func NewIPRateLimiter(r rate.Limit, b int) *IPRateLimiter {
+	return &IPRateLimiter{
+		ips: make(map[string]*rate.Limiter),
+		mu:  &sync.RWMutex{},
+		r:   r,
+		b:   b,
+	}
+}
+
+func (i *IPRateLimiter) GetLimiter(ip string) *rate.Limiter {
+	i.mu.Lock()
+	defer i.mu.Unlock()
+
+	limiter, exists := i.ips[ip]
+	if !exists {
+		limiter = rate.NewLimiter(i.r, i.b)
+		i.ips[ip] = limiter
+	}
+
+	return limiter
+}
+
+func RateLimit(limiter *IPRateLimiter) gin.HandlerFunc {
+	return func(c *gin.Context) {
+		ip := c.ClientIP()
+		if !limiter.GetLimiter(ip).Allow() {
+			c.JSON(http.StatusTooManyRequests, gin.H{
+				"code":    429,
+				"message": "too many requests",
+			})
+			c.Abort()
+			return
+		}
+		c.Next()
+	}
+}