diff --git a/internal/consts/consts.go b/internal/consts/consts.go index 0413d3f..c4e419b 100644 --- a/internal/consts/consts.go +++ b/internal/consts/consts.go @@ -174,6 +174,8 @@ const ( // Registration errors CodeRegistrationDisabled = 10077 CodeEmailVerificationDisabled = 10082 + CodeIpRateLimitExceeded = 10090 // IP级别限流 + CodeGlobalRateLimitExceeded = 10091 // 全局限流 // Setup errors CodeSetupCompleted = 10078 @@ -308,6 +310,8 @@ const ( // Registration errors MsgRegistrationDisabled = "注册功能已关闭" MsgEmailVerificationDisabled = "注册邮箱验证未启用" + MsgIpRateLimitExceeded = "该IP注册频率过高,请稍后再试" + MsgGlobalRateLimitExceeded = "系统注册繁忙,请稍后再试" // Setup errors MsgSetupCompleted = "系统已完成初始化" diff --git a/internal/logic/common/register_rate_limit.go b/internal/logic/common/register_rate_limit.go new file mode 100644 index 0000000..f5322a8 --- /dev/null +++ b/internal/logic/common/register_rate_limit.go @@ -0,0 +1,155 @@ +package common + +import ( + "context" + "fmt" + "time" + + "github.com/gogf/gf/v2/frame/g" + "github.com/qianfree/team-api/internal/consts" +) + +// CheckRegisterRateLimit 检查注册速率限制 +// 1. IP级别限流(每小时/每天) +// 2. 全局限流(每分钟) +func CheckRegisterRateLimit(ctx context.Context, ipAddress string) error { + // 1. 检查IP级别限流 + if err := checkIPRateLimit(ctx, ipAddress); err != nil { + return err + } + + // 2. 检查全局限流 + if err := checkGlobalRateLimit(ctx); err != nil { + return err + } + + return nil +} + +// checkIPRateLimit 检查IP级别的注册限流 +func checkIPRateLimit(ctx context.Context, ipAddress string) error { + if ipAddress == "" { + return nil // 无法获取IP时跳过检查 + } + + // 每小时限制 + hourlyLimit := Config().GetInt(ctx, "register_ip_limit_per_hour") + if hourlyLimit > 0 { + hourlyKey := fmt.Sprintf("register:ip:hourly:%s", ipAddress) + count, err := incrementWithExpire(ctx, hourlyKey, 3600) // 1小时过期 + if err != nil { + g.Log().Warningf(ctx, "IP注册限流检查失败: %v", err) + return nil // Redis失败时不阻塞注册 + } + if count > int64(hourlyLimit) { + return NewBusinessError(consts.CodeIpRateLimitExceeded, consts.MsgIpRateLimitExceeded) + } + } + + // 每天限制 + dailyLimit := Config().GetInt(ctx, "register_ip_limit_per_day") + if dailyLimit > 0 { + dailyKey := fmt.Sprintf("register:ip:daily:%s", ipAddress) + count, err := incrementWithExpire(ctx, dailyKey, 86400) // 24小时过期 + if err != nil { + g.Log().Warningf(ctx, "IP注册限流检查失败: %v", err) + return nil // Redis失败时不阻塞注册 + } + if count > int64(dailyLimit) { + return NewBusinessError(consts.CodeIpRateLimitExceeded, consts.MsgIpRateLimitExceeded) + } + } + + return nil +} + +// checkGlobalRateLimit 检查全局注册速率限制 +func checkGlobalRateLimit(ctx context.Context) error { + globalLimit := Config().GetInt(ctx, "register_global_limit_per_minute") + if globalLimit <= 0 { + return nil // 未启用全局限流 + } + + // 使用当前分钟的Unix时间戳作为key,确保每分钟重置 + now := time.Now() + minuteKey := fmt.Sprintf("register:global:%d", now.Unix()/60) // 每分钟一个key + + count, err := incrementWithExpire(ctx, minuteKey, 120) // 2分钟过期(确保当前分钟和下一分钟都能计数) + if err != nil { + g.Log().Warningf(ctx, "全局注册限流检查失败: %v", err) + return nil // Redis失败时不阻塞注册 + } + + if count > int64(globalLimit) { + return NewBusinessError(consts.CodeGlobalRateLimitExceeded, consts.MsgGlobalRateLimitExceeded) + } + + return nil +} + +// incrementWithExpire 原子递增计数器,首次递增时设置过期时间 +// 使用Lua脚本确保原子性:INCR + EXPIRE(仅在第一次) +func incrementWithExpire(ctx context.Context, key string, expireSeconds int) (int64, error) { + // Lua脚本:先递增,如果结果为1则设置过期时间 + luaScript := ` + local count = redis.call("INCR", KEYS[1]) + if count == 1 then + redis.call("EXPIRE", KEYS[1], ARGV[1]) + end + return count + ` + + result, err := g.Redis().Do(ctx, "EVAL", luaScript, 1, key, expireSeconds) + if err != nil { + return 0, err + } + + return result.Int64(), nil +} + +// GetRegisterRateLimitStatus 获取当前IP的注册限流状态(用于前端提示) +func GetRegisterRateLimitStatus(ctx context.Context, ipAddress string) map[string]any { + status := make(map[string]any) + + if ipAddress == "" { + return status + } + + // 查询IP每小时限制 + hourlyLimit := Config().GetInt(ctx, "register_ip_limit_per_hour") + if hourlyLimit > 0 { + hourlyKey := fmt.Sprintf("register:ip:hourly:%s", ipAddress) + hourlyCount, _ := g.Redis().Do(ctx, "GET", hourlyKey) + if !hourlyCount.IsNil() { + remaining := hourlyLimit - hourlyCount.Int() + status["hourly_remaining"] = remaining + status["hourly_limit"] = hourlyLimit + if ttl, err := g.Redis().Do(ctx, "TTL", hourlyKey); err == nil { + status["hourly_reset_seconds"] = ttl.Int() + } + } else { + status["hourly_remaining"] = hourlyLimit + status["hourly_limit"] = hourlyLimit + } + } + + // 查询IP每天限制 + dailyLimit := Config().GetInt(ctx, "register_ip_limit_per_day") + if dailyLimit > 0 { + dailyKey := fmt.Sprintf("register:ip:daily:%s", ipAddress) + dailyCount, _ := g.Redis().Do(ctx, "GET", dailyKey) + if !dailyCount.IsNil() { + remaining := dailyLimit - dailyCount.Int() + status["daily_remaining"] = remaining + status["daily_limit"] = dailyLimit + if ttl, err := g.Redis().Do(ctx, "TTL", dailyKey); err == nil { + status["daily_reset_seconds"] = ttl.Int() + } + } else { + status["daily_remaining"] = dailyLimit + status["daily_limit"] = dailyLimit + } + } + + return status +} diff --git a/internal/logic/common/register_rate_limit_test.go b/internal/logic/common/register_rate_limit_test.go new file mode 100644 index 0000000..3c29ae3 --- /dev/null +++ b/internal/logic/common/register_rate_limit_test.go @@ -0,0 +1,158 @@ +package common + +import ( + "context" + "fmt" + "testing" + "time" + + "github.com/gogf/gf/v2/frame/g" + "github.com/qianfree/team-api/internal/consts" +) + +// TestCheckRegisterRateLimit 测试注册速率限制 +func TestCheckRegisterRateLimit(t *testing.T) { + ctx := context.Background() + + // 测试IP级别限流 + t.Run("IP hourly limit", func(t *testing.T) { + testIP := "192.168.1.100" + + // 清理测试数据 + hourlyKey := "register:ip:hourly:" + testIP + g.Redis().Do(ctx, "DEL", hourlyKey) + + // 设置配置为每小时最多3次 + g.Redis().Do(ctx, "SETEX", "config:register_ip_limit_per_hour", 60, "3") + + // 前3次应该成功 + for i := 0; i < 3; i++ { + err := CheckRegisterRateLimit(ctx, testIP) + if err != nil { + t.Errorf("第 %d 次注册应该成功,但得到错误: %v", i+1, err) + } + } + + // 第4次应该失败 + err := CheckRegisterRateLimit(ctx, testIP) + if err == nil { + t.Error("第4次注册应该被限流,但成功了") + } else if err.Error() != consts.MsgIpRateLimitExceeded { + t.Errorf("第4次注册应该返回IP限流错误,但得到: %v", err) + } + + // 清理测试数据 + g.Redis().Do(ctx, "DEL", hourlyKey) + }) + + t.Run("IP daily limit", func(t *testing.T) { + testIP := "192.168.1.101" + + // 清理测试数据 + dailyKey := "register:ip:daily:" + testIP + g.Redis().Do(ctx, "DEL", dailyKey) + + // 设置配置为每天最多5次 + g.Redis().Do(ctx, "SETEX", "config:register_ip_limit_per_day", 60, "5") + + // 禁用每小时限制 + g.Redis().Do(ctx, "SETEX", "config:register_ip_limit_per_hour", 60, "0") + + // 前5次应该成功 + for i := 0; i < 5; i++ { + err := CheckRegisterRateLimit(ctx, testIP) + if err != nil { + t.Errorf("第 %d 次注册应该成功,但得到错误: %v", i+1, err) + } + } + + // 第6次应该失败 + err := CheckRegisterRateLimit(ctx, testIP) + if err == nil { + t.Error("第6次注册应该被限流,但成功了") + } else if err.Error() != consts.MsgIpRateLimitExceeded { + t.Errorf("第6次注册应该返回IP限流错误,但得到: %v", err) + } + + // 清理测试数据 + g.Redis().Do(ctx, "DEL", dailyKey) + }) +} + +// TestGlobalRateLimit 测试全局限流 +func TestGlobalRateLimit(t *testing.T) { + ctx := context.Background() + + // 清理测试数据 + now := time.Now() + minuteKey := fmt.Sprintf("register:global:%d", now.Unix()/60) + g.Redis().Do(ctx, "DEL", minuteKey) + + // 设置配置为每分钟最多2次 + g.Redis().Do(ctx, "SETEX", "config:register_global_limit_per_minute", 60, "2") + + // 禁用IP限制 + g.Redis().Do(ctx, "SETEX", "config:register_ip_limit_per_hour", 60, "0") + g.Redis().Do(ctx, "SETEX", "config:register_ip_limit_per_day", 60, "0") + + // 前2次应该成功(使用不同IP) + for i := 0; i < 2; i++ { + testIP := fmt.Sprintf("192.168.1.%d", i+100) + err := CheckRegisterRateLimit(ctx, testIP) + if err != nil { + t.Errorf("第 %d 次注册应该成功,但得到错误: %v", i+1, err) + } + } + + // 第3次应该失败(无论使用什么IP) + err := CheckRegisterRateLimit(ctx, "192.168.1.200") + if err == nil { + t.Error("第3次注册应该被全局限流,但成功了") + } else if err.Error() != consts.MsgGlobalRateLimitExceeded { + t.Errorf("第3次注册应该返回全局限流错误,但得到: %v", err) + } + + // 清理测试数据 + g.Redis().Do(ctx, "DEL", minuteKey) +} + +// TestGetRegisterRateLimitStatus 测试获取限流状态 +func TestGetRegisterRateLimitStatus(t *testing.T) { + ctx := context.Background() + testIP := "192.168.1.102" + + // 清理测试数据 + hourlyKey := "register:ip:hourly:" + testIP + dailyKey := "register:ip:daily:" + testIP + g.Redis().Do(ctx, "DEL", hourlyKey) + g.Redis().Do(ctx, "DEL", dailyKey) + + // 设置配置 + g.Redis().Do(ctx, "SETEX", "config:register_ip_limit_per_hour", 60, "5") + g.Redis().Do(ctx, "SETEX", "config:register_ip_limit_per_day", 60, "20") + + // 注册2次 + g.Redis().Do(ctx, "SETEX", hourlyKey, 3600, "2") + g.Redis().Do(ctx, "SETEX", dailyKey, 86400, "2") + + // 获取状态 + status := GetRegisterRateLimitStatus(ctx, testIP) + + // 验证状态 + if hourlyLimit, ok := status["hourly_limit"].(int); !ok || hourlyLimit != 5 { + t.Errorf("期望每小时限制为5,得到: %v", status["hourly_limit"]) + } + if hourlyRemaining, ok := status["hourly_remaining"].(int); !ok || hourlyRemaining != 3 { + t.Errorf("期望每小时剩余3次,得到: %v", status["hourly_remaining"]) + } + if dailyLimit, ok := status["daily_limit"].(int); !ok || dailyLimit != 20 { + t.Errorf("期望每天限制为20,得到: %v", status["daily_limit"]) + } + if dailyRemaining, ok := status["daily_remaining"].(int); !ok || dailyRemaining != 18 { + t.Errorf("期望每天剩余18次,得到: %v", status["daily_remaining"]) + } + + // 清理测试数据 + g.Redis().Do(ctx, "DEL", hourlyKey) + g.Redis().Do(ctx, "DEL", dailyKey) +} diff --git a/internal/logic/common/settings_registry.go b/internal/logic/common/settings_registry.go index 1fd34a8..13effad 100644 --- a/internal/logic/common/settings_registry.go +++ b/internal/logic/common/settings_registry.go @@ -59,6 +59,12 @@ var Registry = []SettingDef{ Label: "开放注册", Description: "是否允许新用户注册", IsPublic: true}, {Key: "register_email_verification", Type: SettingTypeBool, Default: "false", Category: "general", Label: "注册邮箱验证", Description: "注册时是否需要邮箱验证码,关闭时使用滑块验证", IsPublic: true}, + {Key: "register_ip_limit_per_hour", Type: SettingTypeInt, Default: "3", Category: "general", + Label: "IP每小时注册限制", Description: "同一IP每小时最多注册次数(0表示不限制)", Validation: "min:0,max:20"}, + {Key: "register_ip_limit_per_day", Type: SettingTypeInt, Default: "10", Category: "general", + Label: "IP每天注册限制", Description: "同一IP每天最多注册次数(0表示不限制)", Validation: "min:0,max:100"}, + {Key: "register_global_limit_per_minute", Type: SettingTypeInt, Default: "10", Category: "general", + Label: "全局每分钟注册限制", Description: "全系统每分钟最多注册次数(0表示不限制)", Validation: "min:0,max:100"}, {Key: "maintenance_mode", Type: SettingTypeBool, Default: "false", Category: "general", Label: "维护模式", Description: "开启后控制台显示维护提示"}, {Key: "maintenance_message", Type: SettingTypeString, Default: "", Category: "general", diff --git a/internal/logic/tenant/auth.go b/internal/logic/tenant/auth.go index 43ffb13..7b5f957 100644 --- a/internal/logic/tenant/auth.go +++ b/internal/logic/tenant/auth.go @@ -40,6 +40,12 @@ func (s *sTenant) Register(ctx context.Context, req *v1.TenantRegisterReq) (*v1. return nil, common.NewBusinessError(consts.CodeInvalidUsername, err.Error()) } + // Check register rate limit (IP + global) + ipAddress := g.RequestFromCtx(ctx).GetClientIp() + if err := common.CheckRegisterRateLimit(ctx, ipAddress); err != nil { + return nil, err + } + // 根据配置选择验证方式 emailVerificationEnabled := common.Config().GetBool(ctx, "register_email_verification") if emailVerificationEnabled { @@ -80,8 +86,6 @@ func (s *sTenant) Register(ctx context.Context, req *v1.TenantRegisterReq) (*v1. var tenantID, ownerUserID int64 - ipAddress := g.RequestFromCtx(ctx).GetClientIp() - err = dao.TntTenants.Transaction(ctx, func(ctx context.Context, tx gdb.TX) error { // Create tenant tenantResult, err := tx.Model("tnt_tenants").Ctx(ctx).Data(do.TntTenants{