package common import ( "sync" "testing" ) /* state_test.go - State 并发安全测试 测试重点: 1. 并发安全性 - 多goroutine同时操作计数器 2. 原子操作一致性 - 增减计数正确 3. Reset功能 - 重置后计数器归零 不测试: - 限速器(需要复杂的时间模拟) - 简单getter/setter */ // TestState_ConcurrentPacketCount 测试并发包计数 func TestState_ConcurrentPacketCount(t *testing.T) { s := NewState() const goroutines = 100 const incrementsPerGoroutine = 1000 var wg sync.WaitGroup wg.Add(goroutines) for i := 0; i < goroutines; i++ { go func() { defer wg.Done() for j := 0; j < incrementsPerGoroutine; j++ { s.IncrementPacketCount() } }() } wg.Wait() expected := int64(goroutines * incrementsPerGoroutine) actual := s.GetPacketCount() if actual != expected { t.Errorf("并发计数不一致: 期望 %d, 实际 %d", expected, actual) } } // TestState_ConcurrentTCPCount 测试并发TCP计数 func TestState_ConcurrentTCPCount(t *testing.T) { s := NewState() const goroutines = 50 const operationsPerGoroutine = 500 var wg sync.WaitGroup wg.Add(goroutines * 2) // 成功和失败各一半 // 成功连接 for i := 0; i < goroutines; i++ { go func() { defer wg.Done() for j := 0; j < operationsPerGoroutine; j++ { s.IncrementTCPSuccessPacketCount() } }() } // 失败连接 for i := 0; i < goroutines; i++ { go func() { defer wg.Done() for j := 0; j < operationsPerGoroutine; j++ { s.IncrementTCPFailedPacketCount() } }() } wg.Wait() expectedTotal := int64(goroutines * operationsPerGoroutine * 2) expectedSuccess := int64(goroutines * operationsPerGoroutine) expectedFailed := int64(goroutines * operationsPerGoroutine) if s.GetPacketCount() != expectedTotal { t.Errorf("总包计数不一致: 期望 %d, 实际 %d", expectedTotal, s.GetPacketCount()) } if s.GetTCPPacketCount() != expectedTotal { t.Errorf("TCP包计数不一致: 期望 %d, 实际 %d", expectedTotal, s.GetTCPPacketCount()) } if s.GetTCPSuccessPacketCount() != expectedSuccess { t.Errorf("TCP成功计数不一致: 期望 %d, 实际 %d", expectedSuccess, s.GetTCPSuccessPacketCount()) } if s.GetTCPFailedPacketCount() != expectedFailed { t.Errorf("TCP失败计数不一致: 期望 %d, 实际 %d", expectedFailed, s.GetTCPFailedPacketCount()) } } // TestState_Reset 测试重置功能 func TestState_Reset(t *testing.T) { s := NewState() // 增加一些计数 for i := 0; i < 100; i++ { s.IncrementTCPSuccessPacketCount() s.IncrementTCPFailedPacketCount() s.IncrementUDPPacketCount() s.IncrementHTTPPacketCount() s.IncrementResourceExhaustedCount() } // 验证有值 if s.GetPacketCount() == 0 { t.Fatal("重置前计数应该非零") } // 重置 s.ResetPacketCounters() // 验证全部归零 if s.GetPacketCount() != 0 { t.Errorf("重置后PacketCount应该为0, 实际 %d", s.GetPacketCount()) } if s.GetTCPPacketCount() != 0 { t.Errorf("重置后TCPPacketCount应该为0, 实际 %d", s.GetTCPPacketCount()) } if s.GetTCPSuccessPacketCount() != 0 { t.Errorf("重置后TCPSuccessPacketCount应该为0, 实际 %d", s.GetTCPSuccessPacketCount()) } if s.GetTCPFailedPacketCount() != 0 { t.Errorf("重置后TCPFailedPacketCount应该为0, 实际 %d", s.GetTCPFailedPacketCount()) } if s.GetUDPPacketCount() != 0 { t.Errorf("重置后UDPPacketCount应该为0, 实际 %d", s.GetUDPPacketCount()) } if s.GetHTTPPacketCount() != 0 { t.Errorf("重置后HTTPPacketCount应该为0, 实际 %d", s.GetHTTPPacketCount()) } if s.GetResourceExhaustedCount() != 0 { t.Errorf("重置后ResourceExhaustedCount应该为0, 实际 %d", s.GetResourceExhaustedCount()) } } // TestState_TaskCounters 测试任务计数器 func TestState_TaskCounters(t *testing.T) { s := NewState() // 初始值应该为0 if s.GetEnd() != 0 || s.GetNum() != 0 { t.Error("初始任务计数器应该为0") } // 设置值 s.SetEnd(100) s.SetNum(50) if s.GetEnd() != 100 { t.Errorf("End应该为100, 实际 %d", s.GetEnd()) } if s.GetNum() != 50 { t.Errorf("Num应该为50, 实际 %d", s.GetNum()) } // 增加值 s.IncrementEnd() s.IncrementNum() if s.GetEnd() != 101 { t.Errorf("IncrementEnd后应该为101, 实际 %d", s.GetEnd()) } if s.GetNum() != 51 { t.Errorf("IncrementNum后应该为51, 实际 %d", s.GetNum()) } } // TestState_ConcurrentTaskCounters 测试并发任务计数 func TestState_ConcurrentTaskCounters(t *testing.T) { s := NewState() const goroutines = 100 const incrementsPerGoroutine = 100 var wg sync.WaitGroup wg.Add(goroutines * 2) // 并发增加End for i := 0; i < goroutines; i++ { go func() { defer wg.Done() for j := 0; j < incrementsPerGoroutine; j++ { s.IncrementEnd() } }() } // 并发增加Num for i := 0; i < goroutines; i++ { go func() { defer wg.Done() for j := 0; j < incrementsPerGoroutine; j++ { s.IncrementNum() } }() } wg.Wait() expected := int64(goroutines * incrementsPerGoroutine) if s.GetEnd() != expected { t.Errorf("End并发计数不一致: 期望 %d, 实际 %d", expected, s.GetEnd()) } if s.GetNum() != expected { t.Errorf("Num并发计数不一致: 期望 %d, 实际 %d", expected, s.GetNum()) } } // TestState_GetOutputMutex 测试获取输出互斥锁指针 func TestState_GetOutputMutex(t *testing.T) { s := NewState() mu := s.GetOutputMutex() if mu == nil { t.Fatal("GetOutputMutex returned nil") } // 验证返回的指针可以正常加锁解锁 mu.Lock() _ = 1 //nolint:staticcheck // SA2001: 故意测试空临界区 mu.Unlock() } // TestState_GetICMPLimiter 测试 ICMP 限速器延迟初始化 func TestState_GetICMPLimiter(t *testing.T) { s := NewState() limiter := s.GetICMPLimiter(0.1) if limiter == nil { t.Fatal("GetICMPLimiter returned nil") } // 再次调用应返回同一个实例(sync.Once 保证) limiter2 := s.GetICMPLimiter(0.5) if limiter != limiter2 { t.Fatal("GetICMPLimiter should return the same instance on repeated calls") } } // TestState_GetICMPLimiterMinRate 测试极低速率下的 ICMP 限速器 func TestState_GetICMPLimiterMinRate(t *testing.T) { s := NewState() // 极低速率(packetsPerSecond < 1)应被钳位到 1 limiter := s.GetICMPLimiter(0.000001) if limiter == nil { t.Fatal("GetICMPLimiter with tiny rate returned nil") } } // TestState_GetPerfStats 测试性能统计数据 func TestState_GetPerfStats(t *testing.T) { s := NewState() // 初始状态:全零 stats := s.GetPerfStats() if stats.TotalPackets != 0 { t.Errorf("初始 TotalPackets 应为 0, 实际 %d", stats.TotalPackets) } if stats.SuccessRate != 0 { t.Errorf("初始 SuccessRate 应为 0, 实际 %f", stats.SuccessRate) } // 增加一些计数后验证统计 s.IncrementTCPSuccessPacketCount() s.IncrementTCPSuccessPacketCount() s.IncrementTCPFailedPacketCount() s.SetNum(3) stats = s.GetPerfStats() if stats.TotalPackets != 3 { t.Errorf("TotalPackets 期望 3, 实际 %d", stats.TotalPackets) } if stats.TCPSuccess != 2 { t.Errorf("TCPSuccess 期望 2, 实际 %d", stats.TCPSuccess) } if stats.TCPFailed != 1 { t.Errorf("TCPFailed 期望 1, 实际 %d", stats.TCPFailed) } if stats.TargetsScanned != 3 { t.Errorf("TargetsScanned 期望 3, 实际 %d", stats.TargetsScanned) } // success rate = 2/3 * 100 ≈ 66.67% if stats.SuccessRate < 66 || stats.SuccessRate > 67 { t.Errorf("SuccessRate 期望约 66.67, 实际 %f", stats.SuccessRate) } } // TestState_GetPerfStatsJSON 测试性能统计 JSON 序列化 func TestState_GetPerfStatsJSON(t *testing.T) { s := NewState() s.IncrementTCPSuccessPacketCount() json := s.GetPerfStatsJSON() if json == "" || json == "{}" { t.Fatalf("GetPerfStatsJSON 返回空: %q", json) } if len(json) < 10 { t.Fatalf("GetPerfStatsJSON 内容过短: %q", json) } // 验证包含关键字段 for _, key := range []string{"total_packets", "tcp_success", "success_rate"} { if !containsStr(json, key) { t.Errorf("GetPerfStatsJSON 缺少字段 %q", key) } } } func containsStr(s, sub string) bool { return len(s) >= len(sub) && (s == sub || len(s) > 0 && stringContains(s, sub)) } func stringContains(s, sub string) bool { for i := 0; i <= len(s)-len(sub); i++ { if s[i:i+len(sub)] == sub { return true } } return false } // TestState_GetPacketLimiter 测试通用发包限速器 func TestState_GetPacketLimiter(t *testing.T) { t.Run("零速率返回nil", func(t *testing.T) { s := NewState() limiter := s.GetPacketLimiter(0) if limiter != nil { t.Fatal("零速率应返回 nil limiter") } }) t.Run("负速率返回nil", func(t *testing.T) { s := NewState() limiter := s.GetPacketLimiter(-1) if limiter != nil { t.Fatal("负速率应返回 nil limiter") } }) t.Run("正速率初始化限速器", func(t *testing.T) { s := NewState() limiter := s.GetPacketLimiter(600) // 600/min = 10/s if limiter == nil { t.Fatal("正速率应返回非 nil limiter") } // 再次调用返回同一实例 limiter2 := s.GetPacketLimiter(1200) if limiter != limiter2 { t.Fatal("GetPacketLimiter 应通过 sync.Once 复用实例") } }) t.Run("低速率被钳位到1pps", func(t *testing.T) { s := NewState() // 1/min < 1/s,应被钳位 limiter := s.GetPacketLimiter(1) if limiter == nil { t.Fatal("低速率钳位后应返回非 nil limiter") } }) } // TestState_CacheService 测试服务识别缓存 func TestState_CacheService(t *testing.T) { s := NewState() // 未缓存时查询返回 false _, ok := s.GetCachedService("192.168.1.1:80") if ok { t.Fatal("未缓存的 key 不应返回 ok=true") } // 缓存并查询 type fakeInfo struct{ Name string } info := &fakeInfo{Name: "http"} s.CacheService("192.168.1.1:80", info) got, ok := s.GetCachedService("192.168.1.1:80") if !ok { t.Fatal("已缓存的 key 应返回 ok=true") } if got != info { t.Fatalf("GetCachedService 返回 %v, 期望 %v", got, info) } // 不同 key 互不干扰 _, ok = s.GetCachedService("192.168.1.1:443") if ok { t.Fatal("不同 key 不应命中缓存") } } // ============================================================================= // CheckAndIncrementPacketRate 测试 // ============================================================================= // TestCheckAndIncrementPacketRate_ZeroLimit 速率为 0 时无限制 func TestCheckAndIncrementPacketRate_ZeroLimit(t *testing.T) { s := NewState() for i := 0; i < 1000; i++ { ok, err := s.CheckAndIncrementPacketRate(0) if !ok || err != nil { t.Fatalf("零速率限制应始终允许: ok=%v err=%v", ok, err) } } } // TestCheckAndIncrementPacketRate_NegativeLimit 负速率等同于无限制 func TestCheckAndIncrementPacketRate_NegativeLimit(t *testing.T) { s := NewState() ok, err := s.CheckAndIncrementPacketRate(-1) if !ok || err != nil { t.Fatalf("负速率应允许: ok=%v err=%v", ok, err) } } // TestCheckAndIncrementPacketRate_AllowsWhenTokensAvailable 有令牌时返回 true func TestCheckAndIncrementPacketRate_AllowsWhenTokensAvailable(t *testing.T) { s := NewState() // 600/min = 10/s,桶容量 20,初始满桶 ok, err := s.CheckAndIncrementPacketRate(600) if !ok || err != nil { t.Fatalf("初始应有令牌: ok=%v err=%v", ok, err) } } // TestCheckAndIncrementPacketRate_RateLimitedAfterExhaustion 耗尽令牌后返回 false 和 PacketLimitError func TestCheckAndIncrementPacketRate_RateLimitedAfterExhaustion(t *testing.T) { s := NewState() // 极低速率:1/min,桶容量为 1(钳位后 packetsPerSecond=1,capacity=2) // 消耗掉所有令牌后应被限速 const limit int64 = 1 // 初始化限速器(第一次调用触发 sync.Once) s.GetPacketLimiter(limit) // 消耗完所有令牌(容量 <= 2) for i := 0; i < 10; i++ { s.CheckAndIncrementPacketRate(limit) //nolint: errcheck } // 此时令牌应已耗尽,下一次调用应被限速 ok, err := s.CheckAndIncrementPacketRate(limit) if ok { // 桶可能还剩令牌(容量 2),多耗几次再判断 for i := 0; i < 20; i++ { ok, err = s.CheckAndIncrementPacketRate(limit) if !ok { break } } } if ok { t.Fatal("令牌耗尽后应返回 ok=false") } if err == nil { t.Fatal("令牌耗尽后应返回 error") } if !isPacketLimitError(err) { t.Errorf("error 类型应为 PacketLimitError, 实际 %T: %v", err, err) } } // isPacketLimitError 检查是否为 PacketLimitError func isPacketLimitError(err error) bool { _, ok := err.(*PacketLimitError) return ok } // TestCheckAndIncrementPacketRate_ErrorUnwrapsToSentinel 验证 error 可 unwrap 到 sentinel func TestCheckAndIncrementPacketRate_ErrorUnwrapsToSentinel(t *testing.T) { s := NewState() const limit int64 = 1 // 耗尽令牌 for i := 0; i < 50; i++ { s.CheckAndIncrementPacketRate(limit) //nolint: errcheck } var lastErr error for i := 0; i < 10; i++ { ok, err := s.CheckAndIncrementPacketRate(limit) if !ok { lastErr = err break } } if lastErr == nil { t.Skip("未能触发限速(可能令牌桶容量较大),跳过 unwrap 测试") } // 验证可 unwrap 到 ErrPacketRateLimited pErr, ok := lastErr.(*PacketLimitError) if !ok { t.Fatalf("期望 *PacketLimitError, 实际 %T", lastErr) } if pErr.Sentinel != ErrPacketRateLimited { t.Errorf("Sentinel = %v, 期望 ErrPacketRateLimited", pErr.Sentinel) } if pErr.Limit != limit { t.Errorf("Limit = %d, 期望 %d", pErr.Limit, limit) } } // TestState_OutputMutex 测试输出互斥锁 func TestState_OutputMutex(t *testing.T) { s := NewState() counter := 0 const goroutines = 100 const incrementsPerGoroutine = 100 var wg sync.WaitGroup wg.Add(goroutines) for i := 0; i < goroutines; i++ { go func() { defer wg.Done() for j := 0; j < incrementsPerGoroutine; j++ { s.LockOutput() counter++ s.UnlockOutput() } }() } wg.Wait() expected := goroutines * incrementsPerGoroutine if counter != expected { t.Errorf("输出互斥锁保护失败: 期望 %d, 实际 %d", expected, counter) } }