From d2a35dcec29f89cd478962373b8466c5a4e7572a Mon Sep 17 00:00:00 2001 From: wsczx Date: Mon, 25 Aug 2025 18:05:07 +0800 Subject: [PATCH] =?UTF-8?q?=E9=98=B2=E7=88=86=E5=A2=9E=E5=8A=A0=E5=85=A8?= =?UTF-8?q?=E5=B1=80=E9=BB=91=E5=90=8D=E5=8D=95=E5=8A=9F=E8=83=BD=20?= =?UTF-8?q?=E5=A2=9E=E5=8A=A0=E9=98=B2=E7=88=86=E6=B5=8B=E8=AF=95=E7=94=A8?= =?UTF-8?q?=E4=BE=8B=20=E4=BC=98=E5=8C=96=E9=98=B2=E7=88=86=E6=A3=80?= =?UTF-8?q?=E6=9F=A5=E9=80=BB=E8=BE=91?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- server/admin/lockmanager.go | 74 ++++-- server/admin/lockmanager_test.go | 433 +++++++++++++++++++++++++++++++ server/base/cfg.go | 3 +- server/base/config.go | 1 + server/conf/server-sample.toml | 2 + server/handler/link_auth.go | 11 +- 6 files changed, 494 insertions(+), 30 deletions(-) create mode 100644 server/admin/lockmanager_test.go diff --git a/server/admin/lockmanager.go b/server/admin/lockmanager.go index af67621..b5d9f0a 100644 --- a/server/admin/lockmanager.go +++ b/server/admin/lockmanager.go @@ -24,7 +24,15 @@ type LockState struct { LockTime time.Time `json:"lock_time"` // 锁定截止时间 LastAttempt time.Time `json:"lastAttempt"` // 最后一次尝试的时间 } -type IPWhitelists struct { + +type IPListType int + +const ( + IPListWhite IPListType = iota + IPListBlack +) + +type IPList struct { IP net.IP CIDR *net.IPNet } @@ -35,7 +43,7 @@ type LockManager struct { ipLocks map[string]*LockState // 全局IP锁定状态 userLocks map[string]*LockState // 全局用户锁定状态 ipUserLocks map[string]map[string]*LockState // 单用户IP锁定状态 - ipWhitelists []IPWhitelists // 全局IP白名单，包含IP地址和CIDR范围 + ipLists map[IPListType][]IPList // 统一的IP列表管理 cleanupTicker *time.Ticker } @@ -46,10 +54,10 @@ func GetLockManager() *LockManager { once.Do(func() { lockmanager = &LockManager{ // LoginStatus: sync.Map{}, - ipLocks: make(map[string]*LockState), - userLocks: make(map[string]*LockState), - ipUserLocks: make(map[string]map[string]*LockState), - ipWhitelists: make([]IPWhitelists, 0), + ipLocks: make(map[string]*LockState), + userLocks: make(map[string]*LockState), + ipUserLocks: make(map[string]map[string]*LockState), + ipLists: make(map[IPListType][]IPList), } }) return lockmanager @@ -64,7 +72,8 @@ func InitLockManager() { base.Cfg.GlobalLockStateExpirationTime = defaultGlobalLockStateExpirationTime } lm.StartCleanupTicker() - lm.InitIPWhitelist() + lm.InitIPList(IPListWhite, base.Cfg.IPWhiteList) + lm.InitIPList(IPListBlack, base.Cfg.IPBlackList) } } @@ -181,40 +190,50 @@ func (lm *LockManager) GetLocksInfo() []LockInfo { return locksInfo } -// 初始化IP白名单 -func (lm *LockManager) InitIPWhitelist() { - ipWhitelist := strings.Split(base.Cfg.IPWhitelist, ",") - for _, ipWhitelist := range ipWhitelist { - ipWhitelist = strings.TrimSpace(ipWhitelist) - if ipWhitelist == "" { +// 初始化IP列表 +func (lm *LockManager) InitIPList(listType IPListType, config string) { + if lm.ipLists == nil { + lm.ipLists = make(map[IPListType][]IPList) + } + + ipList := strings.Split(config, ",") + for _, ipItem := range ipList { + ipItem = strings.TrimSpace(ipItem) + if ipItem == "" { continue } - _, ipNet, err := net.ParseCIDR(ipWhitelist) + _, ipNet, err := net.ParseCIDR(ipItem) if err == nil { - lm.ipWhitelists = append(lm.ipWhitelists, IPWhitelists{CIDR: ipNet}) + lm.ipLists[listType] = append(lm.ipLists[listType], IPList{CIDR: ipNet}) continue } - ip := net.ParseIP(ipWhitelist) + ip := net.ParseIP(ipItem) if ip != nil { - lm.ipWhitelists = append(lm.ipWhitelists, IPWhitelists{IP: ip}) + lm.ipLists[listType] = append(lm.ipLists[listType], IPList{IP: ip}) continue } } } -// 检查 IP 是否在白名单中 -func (lm *LockManager) IsWhitelisted(ip string) bool { +// 检查 IP 列表 +func (lm *LockManager) IsInIPList(ip string, listType IPListType) bool { clientIP := net.ParseIP(ip) if clientIP == nil { return false } - for _, ipWhitelist := range lm.ipWhitelists { - if ipWhitelist.CIDR != nil && ipWhitelist.CIDR.Contains(clientIP) { + + ipList, exists := lm.ipLists[listType] + if !exists { + return false + } + + for _, ipItem := range ipList { + if ipItem.CIDR != nil && ipItem.CIDR.Contains(clientIP) { return true } - if ipWhitelist.IP != nil && ipWhitelist.IP.Equal(clientIP) { + if ipItem.IP != nil && ipItem.IP.Equal(clientIP) { return true } } @@ -372,6 +391,9 @@ func (lm *LockManager) UpdateUserIPLock(username, ip string, now time.Time, succ // 更新锁定状态 func (lm *LockManager) UpdateLockState(state *LockState, now time.Time, success bool, maxBanCount, lockTime int) { + if state.Locked { + return // 已经锁定，不处理 + } if success { lm.Unlock(state) // 成功登录后解锁 } else { @@ -424,10 +446,16 @@ func (lm *LockManager) CheckLocked(username, ipaddr string) bool { now := time.Now() // 检查IP是否在白名单中 - if lm.IsWhitelisted(ip) { + if lm.IsInIPList(ip, IPListWhite) { return true } + // 检查IP是否在黑名单中 + if lm.IsInIPList(ip, IPListBlack) { + base.Warn("IP", ip, "is blacklisted. Access denied.") + return false + } + // 检查全局 IP 锁定 if base.Cfg.MaxGlobalIPBanCount > 0 && lm.CheckGlobalIPLock(ip, now) { base.Warn("IP", ip, "is globally locked. Try again later.") diff --git a/server/admin/lockmanager_test.go b/server/admin/lockmanager_test.go new file mode 100644 index 0000000..4ff2de6 --- /dev/null +++ b/server/admin/lockmanager_test.go @@ -0,0 +1,433 @@ +package admin + +import ( + "fmt" + "net" + "sync" + "testing" + "time" + + "github.com/bjdgyc/anylink/base" + "github.com/stretchr/testify/assert" +) + +// Helper function to reset the singleton for test isolation +func resetLockManager() { + once = sync.Once{} + lockmanager = nil +} + +// 测试 GetLockManager 函数 +func TestGetLockManager(t *testing.T) { + resetLockManager() + base.Test() + setupTestConfig() + + t.Run("Singleton_Pattern", func(t *testing.T) { + lm1 := GetLockManager() + lm2 := GetLockManager() + assert.Same(t, lm1, lm2, "GetLockManager应该返回同一个实例") + }) +} + +// 测试并发竞争条件 +func TestLockManager_RaceConditions(t *testing.T) { + resetLockManager() + base.Test() + setupTestConfig() + + lm := GetLockManager() + + t.Run("Concurrent_CheckAndUpdate", func(t *testing.T) { + username := "raceuser" + ipaddr := "192.168.1.10:12345" + + var wg sync.WaitGroup + results := make([]bool, 20) + + // 并发执行检查和更新操作 + for i := 0; i < 20; i++ { + wg.Add(1) + go func(index int) { + defer wg.Done() + // 交替执行检查和更新操作 + if index%2 == 0 { + results[index] = lm.CheckLocked(username, ipaddr) + } else { + lm.UpdateLoginStatus(username, ipaddr, false) + } + }(i) + } + + wg.Wait() + + // 验证最终状态的一致性 + finalResult := lm.CheckLocked(username, ipaddr) + assert.False(t, finalResult, "高并发后应该被锁定") + }) + + t.Run("Concurrent_MultipleUsers", func(t *testing.T) { + ipaddr := "192.168.1.11:12345" + var wg sync.WaitGroup + + // 多个用户同时从同一IP进行攻击 + for i := 0; i < 50; i++ { + wg.Add(1) + go func(userIndex int) { + defer wg.Done() + username := fmt.Sprintf("user%d", userIndex) + lm.UpdateLoginStatus(username, ipaddr, false) + }(i) + } + + wg.Wait() + + // 验证全局IP锁定是否正确触发 + result := lm.CheckLocked("newuser", ipaddr) + assert.False(t, result, "多用户并发攻击后IP应该被全局锁定") + }) + + t.Run("Concurrent_CleanupAndUpdate", func(t *testing.T) { + username := "cleanuprace" + ipaddr := "192.168.1.12:12345" + + // 先创建一些状态 + lm.UpdateLoginStatus(username, ipaddr, false) + + var wg sync.WaitGroup + + // 并发执行清理和更新操作 + wg.Add(2) + go func() { + defer wg.Done() + lm.CleanupExpiredLocks() + }() + + go func() { + defer wg.Done() + lm.UpdateLoginStatus(username, ipaddr, false) + }() + + wg.Wait() + + // 验证状态一致性 + result := lm.CheckLocked(username, ipaddr) + // 结果可能是true或false，但不应该panic + _ = result + }) +} + +// 测试 CheckLocked 函数 +func TestCheckLocked(t *testing.T) { + resetLockManager() + base.Test() + setupTestConfig() + + lm := GetLockManager() + + t.Run("InitialState_AllowLogin", func(t *testing.T) { + result := lm.CheckLocked("testuser", "192.168.1.1:12345") + assert.True(t, result, "初始状态应该允许登录") + }) + + t.Run("LockedState_DenyLogin", func(t *testing.T) { + username := "testuser" + ipaddr := "192.168.1.1:12345" + + // 模拟5次失败登录 + for i := 0; i < 5; i++ { + lm.UpdateLoginStatus(username, ipaddr, false) + } + + result := lm.CheckLocked(username, ipaddr) + assert.False(t, result, "5次失败后应该被锁定") + }) +} + +// 测试 UpdateLoginStatus 函数 +func TestUpdateLoginStatus(t *testing.T) { + resetLockManager() + base.Test() + setupTestConfig() + + lm := GetLockManager() + + t.Run("FailureCount_Increment", func(t *testing.T) { + username := "testuser" + ipaddr := "192.168.1.1:12345" + + // 测试失败计数递增 + for i := 1; i <= 3; i++ { + lm.UpdateLoginStatus(username, ipaddr, false) + + lm.mu.Lock() + userIPMap := lm.ipUserLocks[username] + ip, _, _ := net.SplitHostPort(ipaddr) + state := userIPMap[ip] + assert.Equal(t, i, state.FailureCount, fmt.Sprintf("第%d次失败后计数应该为%d", i, i)) + lm.mu.Unlock() + } + }) + + t.Run("SuccessLogin_ResetCount", func(t *testing.T) { + username := "successuser" + ipaddr := "192.168.1.2:12345" + + // 先失败几次 + for i := 0; i < 3; i++ { + lm.UpdateLoginStatus(username, ipaddr, false) + } + + // 成功登录 + lm.UpdateLoginStatus(username, ipaddr, true) + + // 验证计数被重置 + lm.mu.Lock() + userIPMap := lm.ipUserLocks[username] + ip, _, _ := net.SplitHostPort(ipaddr) + state := userIPMap[ip] + assert.Equal(t, 0, state.FailureCount, "成功登录应该重置失败计数") + lm.mu.Unlock() + }) +} + +// 测试 UpdateLockState 函数 +func TestUpdateLockState(t *testing.T) { + resetLockManager() + base.Test() + setupTestConfig() + + lm := GetLockManager() + + t.Run("LockedState_NoUpdate", func(t *testing.T) { + // 测试已锁定状态不会被重复更新 + username := "lockeduser" + ipaddr := "192.168.1.2:12345" + + // 先锁定用户 + for i := 0; i < 5; i++ { + lm.UpdateLoginStatus(username, ipaddr, false) + } + + // 获取当前状态 + lm.mu.Lock() + userIPMap := lm.ipUserLocks[username] + ip, _, _ := net.SplitHostPort(ipaddr) + state := userIPMap[ip] + originalCount := state.FailureCount + originalLastAttempt := state.LastAttempt + lm.mu.Unlock() + + // 尝试再次更新 + lm.UpdateLoginStatus(username, ipaddr, false) + + // 验证状态没有改变 + lm.mu.Lock() + newState := lm.ipUserLocks[username][ip] + assert.Equal(t, originalCount, newState.FailureCount, "已锁定状态的失败计数不应该改变") + assert.Equal(t, originalLastAttempt, newState.LastAttempt, "已锁定状态的时间戳不应该改变") + lm.mu.Unlock() + }) +} + +// 测试 CheckGlobalIPLock 函数 +func TestCheckGlobalIPLock(t *testing.T) { + resetLockManager() + base.Test() + setupTestConfig() + + lm := GetLockManager() + + t.Run("GlobalIP_Protection", func(t *testing.T) { + ipaddr := "192.168.1.3:12345" + ip, _, _ := net.SplitHostPort(ipaddr) + + // 使用不同用户名从同一IP进行攻击 + for i := 0; i < 40; i++ { + username := fmt.Sprintf("user%d", i) + lm.UpdateLoginStatus(username, ipaddr, false) + } + + // 验证全局IP锁定检查 + result := lm.CheckGlobalIPLock(ip, time.Now()) + assert.True(t, result, "IP应该被全局锁定") + }) +} + +// 测试 CheckGlobalUserLock 函数 +func TestCheckGlobalUserLock(t *testing.T) { + resetLockManager() + base.Test() + setupTestConfig() + + lm := GetLockManager() + + t.Run("GlobalUser_Protection", func(t *testing.T) { + username := "globaluser" + + // 同一用户从不同IP进行攻击 + for i := 0; i < 20; i++ { + ipaddr := fmt.Sprintf("192.168.1.%d:12345", 100+i) + lm.UpdateLoginStatus(username, ipaddr, false) + } + + // 验证全局用户锁定检查 + result := lm.CheckGlobalUserLock(username, time.Now()) + assert.True(t, result, "用户应该被全局锁定") + }) +} + +// 测试 CheckUserIPLock 函数 +func TestCheckUserIPLock(t *testing.T) { + resetLockManager() + base.Test() + setupTestConfig() + + lm := GetLockManager() + + t.Run("UserIP_Protection", func(t *testing.T) { + username := "useripuser" + ipaddr := "192.168.1.4:12345" + ip, _, _ := net.SplitHostPort(ipaddr) + + // 单用户IP锁定测试 + for i := 0; i < 5; i++ { + lm.UpdateLoginStatus(username, ipaddr, false) + } + + // 验证单用户IP锁定检查 + result := lm.CheckUserIPLock(username, ip, time.Now()) + assert.True(t, result, "单用户IP应该被锁定") + }) +} + +// 测试 InitIPList 和 IsInIPList 函数 +func TestInitIPList_IsInIPList(t *testing.T) { + resetLockManager() + base.Test() + setupTestConfig() + + lm := GetLockManager() + + t.Run("Whitelist_Functionality", func(t *testing.T) { + // 手动初始化IP列表 + lm.InitIPList(IPListWhite, base.Cfg.IPWhiteList) + + // 测试白名单检查 + result := lm.IsInIPList("192.168.90.1", IPListWhite) + assert.True(t, result, "192.168.90.1应该在白名单中") + + // 测试CIDR范围 + result2 := lm.IsInIPList("172.16.0.100", IPListWhite) + assert.True(t, result2, "172.16.0.100应该在CIDR范围内") + }) + + t.Run("Blacklist_Functionality", func(t *testing.T) { + // 手动初始化黑名单 + lm.InitIPList(IPListBlack, base.Cfg.IPBlackList) + + // 测试黑名单检查 + result := lm.IsInIPList("10.0.0.1", IPListBlack) + assert.True(t, result, "10.0.0.1应该在黑名单中") + }) +} + +// 测试 GetLocksInfo 函数 +func TestGetLocksInfo(t *testing.T) { + resetLockManager() + base.Test() + setupTestConfig() + + lm := GetLockManager() + + t.Run("EmptyState", func(t *testing.T) { + locksInfo := lm.GetLocksInfo() + assert.Empty(t, locksInfo, "初始状态应该没有锁定信息") + }) + + t.Run("WithLocks", func(t *testing.T) { + // 创建锁定状态 + username := "testuser" + ipaddr := "192.168.1.5:12345" + + for i := 0; i < 5; i++ { + lm.UpdateLoginStatus(username, ipaddr, false) + } + + locksInfo := lm.GetLocksInfo() + assert.NotEmpty(t, locksInfo, "应该有锁定信息") + }) +} + +// 测试 CleanupExpiredLocks 函数 +func TestCleanupExpiredLocks(t *testing.T) { + resetLockManager() + base.Test() + setupTestConfig() + + lm := GetLockManager() + + t.Run("ExpiredLocks_Cleanup", func(t *testing.T) { + username := "cleanupuser" + ipaddr := "192.168.1.6:12345" + + // 创建锁定状态 + lm.UpdateLoginStatus(username, ipaddr, false) + + // 模拟过期状态 + lm.mu.Lock() + userIPMap := lm.ipUserLocks[username] + ip, _, _ := net.SplitHostPort(ipaddr) + state := userIPMap[ip] + state.LastAttempt = time.Now().Add(-7200 * time.Second) // 2小时前 + lm.mu.Unlock() + + // 执行清理 + lm.CleanupExpiredLocks() + + // 验证过期状态被清理 + lm.mu.Lock() + _, exists := lm.ipUserLocks[username] + lm.mu.Unlock() + assert.False(t, exists, "过期的锁定状态应该被清理") + }) +} + +// 测试 CheckLockState 函数 +func TestCheckLockState(t *testing.T) { + resetLockManager() + base.Test() + setupTestConfig() + + lm := GetLockManager() + + t.Run("TimeWindow_Reset", func(t *testing.T) { + state := &LockState{ + FailureCount: 3, + LastAttempt: time.Now().Add(-700 * time.Second), // 700秒前 + } + + // 检查状态（应该重置计数） + result := lm.CheckLockState(state, time.Now(), 600) // 600秒重置时间 + + assert.False(t, result, "超过重置时间应该返回false") + assert.Equal(t, 0, state.FailureCount, "失败计数应该被重置") + }) +} + +// 辅助函数 +func setupTestConfig() { + base.Cfg.AntiBruteForce = true + base.Cfg.IPWhiteList = "192.168.90.1,172.16.0.0/24" + base.Cfg.IPBlackList = "10.0.0.1" + base.Cfg.MaxBanCount = 5 + base.Cfg.BanResetTime = 600 + base.Cfg.LockTime = 300 + base.Cfg.MaxGlobalUserBanCount = 20 + base.Cfg.GlobalUserBanResetTime = 600 + base.Cfg.GlobalUserLockTime = 300 + base.Cfg.MaxGlobalIPBanCount = 40 + base.Cfg.GlobalIPBanResetTime = 1200 + base.Cfg.GlobalIPLockTime = 300 + base.Cfg.GlobalLockStateExpirationTime = 3600 +} diff --git a/server/base/cfg.go b/server/base/cfg.go index 390b018..dfee3bd 100644 --- a/server/base/cfg.go +++ b/server/base/cfg.go @@ -97,7 +97,8 @@ type ServerConfig struct { EncryptionPassword bool `json:"encryption_password"` AntiBruteForce bool `json:"anti_brute_force"` - IPWhitelist string `json:"ip_whitelist"` + IPWhiteList string `json:"ip_whitelist"` + IPBlackList string `json:"ip_blacklist"` MaxBanCount int `json:"max_ban_score"` BanResetTime int `json:"ban_reset_time"` diff --git a/server/base/config.go b/server/base/config.go index 4903ce2..1028f77 100644 --- a/server/base/config.go +++ b/server/base/config.go @@ -82,6 +82,7 @@ var configs = []config{ {Typ: cfgBool, Name: "anti_brute_force", Usage: "是否开启防爆功能", ValBool: true}, {Typ: cfgStr, Name: "ip_whitelist", Usage: "全局IP白名单,多个用逗号分隔，支持单IP和CIDR范围", ValStr: "192.168.90.1,172.16.0.0/24"}, + {Typ: cfgStr, Name: "ip_blacklist", Usage: "全局IP黑名单,多个用逗号分隔，支持单IP和CIDR范围", ValStr: ""}, {Typ: cfgInt, Name: "max_ban_score", Usage: "单位时间内最大尝试次数，0为关闭该功能", ValInt: 5}, {Typ: cfgInt, Name: "ban_reset_time", Usage: "设置单位时间(秒)，超过则重置计数", ValInt: 10}, diff --git a/server/conf/server-sample.toml b/server/conf/server-sample.toml index dda032f..c0f83f8 100644 --- a/server/conf/server-sample.toml +++ b/server/conf/server-sample.toml @@ -139,6 +139,8 @@ encryption_password = false anti_brute_force = true #全局IP白名单,多个用逗号分隔，支持单IP和CIDR范围 ip_whitelist = "192.168.90.1,172.16.0.0/24" +#全局IP黑名单,多个用逗号分隔，支持单IP和CIDR范围 +ip_blacklist = "" #锁定时间最好不要超过单位时间 #单位时间内最大尝试次数，0为关闭该功能 diff --git a/server/handler/link_auth.go b/server/handler/link_auth.go index c288c44..caa8b93 100644 --- a/server/handler/link_auth.go +++ b/server/handler/link_auth.go @@ -68,6 +68,11 @@ func LinkAuth(w http.ResponseWriter, r *http.Request) { ClientRequest: cr, UserActLog: ua, } + // 锁定状态判断 + if !lockManager.CheckLocked(cr.Auth.Username, r.RemoteAddr) { + w.WriteHeader(http.StatusTooManyRequests) + return + } // setCommonHeader(w) if cr.Type == "logout" { // 退出删除session信息 @@ -136,12 +141,6 @@ func LinkAuth(w http.ResponseWriter, r *http.Request) { return } - // 锁定状态判断 - if !lockManager.CheckLocked(cr.Auth.Username, r.RemoteAddr) { - w.WriteHeader(http.StatusTooManyRequests) - return - } - // TODO 用户密码校验 ext := map[string]any{"mac_addr": cr.MacAddressList.MacAddress} err = dbdata.CheckUser(cr.Auth.Username, cr.Auth.Password, cr.GroupSelect, ext)