mirror of
https://github.com/shadow1ng/fscan.git
synced 2026-09-22 03:10:42 +08:00
541 lines
14 KiB
Go
541 lines
14 KiB
Go
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()
|
||
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)
|
||
}
|
||
}
|