fscan/common/state_test.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()
	_ = 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)
	}
}