防爆增加全局黑名单功能

增加防爆测试用例
优化防爆检查逻辑

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.")

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
+}

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"`

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},

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为关闭该功能

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)