anylink/server/admin/lockmanager_test.go

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
}