opencatd-open/middleware/auth_llm.go

package middleware

import (
	"net/http"
	"opencatd-open/internal/store"
	"strings"

	"github.com/gin-gonic/gin"
	"gorm.io/gorm"
)

func AuthLLM(db *gorm.DB) gin.HandlerFunc {
	return func(c *gin.Context) {
		authToken := c.GetHeader("Authorization")
		if authToken == "" {
			c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{
				"error": map[string]interface{}{
					"message": "未提供认证信息",
					"type":    "invalid_request_error",
				},
			})
			return
		}

		// Extract API key from Bearer token
		if len(authToken) > 7 {
			authToken = authToken[7:]
		}

		// Find API key by prefix
		var apiKey store.APIKey
		if err := db.Where("key_prefix = ? AND status = ?", authToken[:8], store.KeyStatusActive).First(&apiKey).Error; err != nil {
			c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{
				"error": map[string]interface{}{
					"message": "无效的API密钥",
					"type":    "invalid_request_error",
				},
			})
			return
		}

		// Verify full key hash
		keyHash := store.HashAPIKey(authToken)
		if apiKey.KeyHash != keyHash {
			c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{
				"error": map[string]interface{}{
					"message": "无效的API密钥",
					"type":    "invalid_request_error",
				},
			})
			return
		}

		c.Set("api_key", &apiKey)
		c.Set("user_id", apiKey.UserID)
		c.Next()
	}
}

// extractAPIKey extracts the API key from the Authorization header
func extractAPIKey(c *gin.Context) string {
	auth := c.GetHeader("Authorization")
	if strings.HasPrefix(auth, "Bearer ") {
		return auth[7:]
	}
	return auth
}