Initial commit: 网络验证平台

backend/internal/middleware/usage_limit.go

@@ -0,0 +1,156 @@
+package middleware
+
+import (
+	"fmt"
+	"time"
+	"verification-platform-backend/internal/database"
+	"verification-platform-backend/internal/model"
+	"verification-platform-backend/pkg/response"
+
+	"github.com/gin-gonic/gin"
+)
+
+func CheckApiLimit() gin.HandlerFunc {
+	return func(c *gin.Context) {
+		userID, exists := c.Get("user_id")
+		if !exists {
+			response.Error(c, 401, "未授权")
+			c.Abort()
+			return
+		}
+
+		var user model.User
+		if err := database.DB.Preload("CurrentPackage").First(&user, userID).Error; err != nil {
+			response.Error(c, 500, "获取用户信息失败")
+			c.Abort()
+			return
+		}
+
+		if user.CurrentPackageID == nil {
+			response.Error(c, 403, "您还没有购买套餐，请先购买套餐")
+			c.Abort()
+			return
+		}
+
+		var permission model.PackagePermission
+		if err := database.DB.Where("package_id = ?", user.CurrentPackageID).First(&permission).Error; err != nil {
+			response.Error(c, 500, "获取套餐权限失败")
+			c.Abort()
+			return
+		}
+
+		now := time.Now()
+		if user.ApiCallsResetAt == nil || now.Sub(*user.ApiCallsResetAt) >= 24*time.Hour {
+			user.ApiCallsUsed = 0
+			user.ApiCallsResetAt = &now
+			database.DB.Save(&user)
+		}
+
+		if user.ApiCallsUsed >= permission.MaxApiCalls {
+			response.Error(c, 403, fmt.Sprintf("API调用次数已达上限（%d次/天），请升级套餐", permission.MaxApiCalls))
+			c.Abort()
+			return
+		}
+
+		c.Next()
+
+		user.ApiCallsUsed++
+		database.DB.Save(&user)
+	}
+}
+
+func CheckStorageLimit(fileSize int64) gin.HandlerFunc {
+	return func(c *gin.Context) {
+		userID, exists := c.Get("user_id")
+		if !exists {
+			response.Error(c, 401, "未授权")
+			c.Abort()
+			return
+		}
+
+		var user model.User
+		if err := database.DB.Preload("CurrentPackage").First(&user, userID).Error; err != nil {
+			response.Error(c, 500, "获取用户信息失败")
+			c.Abort()
+			return
+		}
+
+		if user.CurrentPackageID == nil {
+			response.Error(c, 403, "您还没有购买套餐，请先购买套餐")
+			c.Abort()
+			return
+		}
+
+		var permission model.PackagePermission
+		if err := database.DB.Where("package_id = ?", user.CurrentPackageID).First(&permission).Error; err != nil {
+			response.Error(c, 500, "获取套餐权限失败")
+			c.Abort()
+			return
+		}
+
+		maxStorageBytes := int64(permission.MaxStorage) * 1024 * 1024
+		if user.StorageUsed+fileSize > maxStorageBytes {
+			usedMB := float64(user.StorageUsed) / 1024 / 1024
+			maxMB := float64(permission.MaxStorage)
+			response.Error(c, 403, fmt.Sprintf("存储空间不足，已使用 %.2f MB / %.2f MB", usedMB, maxMB))
+			c.Abort()
+			return
+		}
+
+		c.Next()
+	}
+}
+
+func RecordApiUsage() gin.HandlerFunc {
+	return func(c *gin.Context) {
+		startTime := time.Now()
+
+		c.Next()
+
+		userID, exists := c.Get("user_id")
+		if !exists {
+			return
+		}
+
+		appID, _ := c.Get("app_id")
+
+		duration := time.Since(startTime)
+
+		usage := model.ApiUsage{
+			UserID:        userID.(uint),
+			ApplicationID: appID.(uint),
+			Endpoint:      c.Request.URL.Path,
+			Method:        c.Request.Method,
+			IPAddress:     c.ClientIP(),
+			UserAgent:     c.Request.UserAgent(),
+			ResponseTime:  int(duration.Milliseconds()),
+			StatusCode:    c.Writer.Status(),
+			Success:       c.Writer.Status() < 400,
+			CreatedAt:     time.Now(),
+		}
+
+		if len(c.Errors) > 0 {
+			usage.ErrorMessage = c.Errors.String()
+		}
+
+		database.DB.Create(&usage)
+	}
+}
+
+func UpdateStorageUsed(userID uint, fileSize int64, action string) error {
+	var user model.User
+	if err := database.DB.First(&user, userID).Error; err != nil {
+		return err
+	}
+
+	if action == "upload" {
+		user.StorageUsed += fileSize
+	} else if action == "delete" {
+		user.StorageUsed -= fileSize
+		if user.StorageUsed < 0 {
+			user.StorageUsed = 0
+		}
+	}
+
+	return database.DB.Save(&user).Error
+}