diff --git a/common/config_struct.go b/common/config_struct.go index 3a6e9ed..164748c 100644 --- a/common/config_struct.go +++ b/common/config_struct.go @@ -22,20 +22,23 @@ config_struct.go - 配置结构体定义 // Config 扫描器完整配置 - 初始化后只读，可安全共享 type Config struct { // 高频访问字段 - 平铺到顶层 - Timeout time.Duration // 通用超时 - ThreadNum int // 主线程数 - ThreadNumExplicit bool // 用户显式指定了 -t - ModuleThreadNum int // 模块线程数 - DisableBrute bool // 禁用暴力破解 - DisablePing bool // 禁用Ping检测 - DisableTcpProbe bool // 禁用TCP补充探测 + Timeout time.Duration // 通用超时 + TimeoutExplicit bool // 用户显式指定了 -time + ThreadNum int // 主线程数 + ThreadNumExplicit bool // 用户显式指定了 -t + ModuleThreadNum int // 模块线程数 + ModuleThreadNumExplicit bool // 用户显式指定了 -mt + DisableBrute bool // 禁用暴力破解 + DisablePing bool // 禁用Ping检测 + DisableTcpProbe bool // 禁用TCP补充探测 // 扫描模式 - Mode string // 扫描模式 - LocalMode bool // 本地模式 - LocalPlugin string // 本地插件名 - AliveOnly bool // 仅存活检测 - MaxRetries int // 最大重试次数 + Mode string // 扫描模式 + LocalMode bool // 本地模式 + LocalPlugin string // 本地插件名 + AliveOnly bool // 仅存活检测 + MaxRetries int // 最大重试次数 + MaxRetriesExplicit bool // 用户显式指定了 -retry // 高级功能（从AdvancedConfig合并） Shellcode string // Shellcode @@ -81,14 +84,15 @@ type CredentialConfig struct { // NetworkConfig 网络相关配置 type NetworkConfig struct { - HTTPProxy string - Socks5Proxy string - Iface string - WebTimeout time.Duration - MaxRedirects int - PacketRateLimit int64 - MaxPacketCount int64 - ICMPRate float64 + HTTPProxy string + Socks5Proxy string + Iface string + WebTimeout time.Duration + MaxRedirects int + PacketRateLimit int64 + MaxPacketCount int64 + ICMPRate float64 + ICMPRateExplicit bool } // OutputConfig 输出相关配置 @@ -107,11 +111,12 @@ type OutputConfig struct { // POCConfig POC扫描相关配置 type POCConfig struct { - PocPath string // POC路径 - PocName string // 指定POC名称 - Full bool // 完整POC扫描 - Num int // POC并发数 - Disabled bool // 禁用POC扫描 + PocPath string // POC路径 + PocName string // 指定POC名称 + Full bool // 完整POC扫描 + Num int // POC并发数 + NumExplicit bool // 用户显式指定了 -num + Disabled bool // 禁用POC扫描 } // RedisConfig Redis利用相关配置 diff --git a/common/flag.go b/common/flag.go index 7a7c82a..4ea91d5 100644 --- a/common/flag.go +++ b/common/flag.go @@ -215,8 +215,19 @@ func Flag(Info *HostInfo) error { // 检测用户是否显式指定了 -t flag.Visit(func(f *flag.Flag) { - if f.Name == "t" { + switch f.Name { + case "t": fv.ThreadNumExplicit = true + case "time": + fv.TimeoutExplicit = true + case "mt": + fv.ModuleThreadNumExplicit = true + case "retry": + fv.MaxRetriesExplicit = true + case "icmp-rate": + fv.ICMPRateExplicit = true + case "num": + fv.PocNumExplicit = true } }) diff --git a/common/flag_config.go b/common/flag_config.go index b147227..085966b 100644 --- a/common/flag_config.go +++ b/common/flag_config.go @@ -30,18 +30,21 @@ type FlagVars struct { PortsFile string // 扫描控制 - ScanMode string - ThreadNum int - ThreadNumExplicit bool // 用户显式指定了 -t - ModuleThreadNum int - TimeoutSec int64 // 秒，需转换为 time.Duration - GlobalTimeout int64 - DisablePing bool - DisableTcpProbe bool - LocalPlugin string - AliveOnly bool - DisableBrute bool - MaxRetries int + ScanMode string + ThreadNum int + ThreadNumExplicit bool // 用户显式指定了 -t + ModuleThreadNum int + ModuleThreadNumExplicit bool + TimeoutSec int64 // 秒，需转换为 time.Duration + TimeoutExplicit bool + GlobalTimeout int64 + DisablePing bool + DisableTcpProbe bool + LocalPlugin string + AliveOnly bool + DisableBrute bool + MaxRetries int + MaxRetriesExplicit bool // 认证凭据 Username string @@ -74,6 +77,7 @@ type FlagVars struct { PocFull bool DNSLog bool PocNum int + PocNumExplicit bool DisablePocScan bool // Redis利用 @@ -85,9 +89,10 @@ type FlagVars struct { DisableRedis bool // 发包频率 - PacketRateLimit int64 - MaxPacketCount int64 - ICMPRate float64 + PacketRateLimit int64 + MaxPacketCount int64 + ICMPRate float64 + ICMPRateExplicit bool // 输出控制 Outputfile string @@ -135,20 +140,23 @@ func GetFlagVars() *FlagVars { func BuildConfigFromFlags(fv *FlagVars) *Config { return &Config{ // 高频字段 - Timeout: time.Duration(fv.TimeoutSec) * time.Second, - ThreadNum: fv.ThreadNum, - ThreadNumExplicit: fv.ThreadNumExplicit, - ModuleThreadNum: fv.ModuleThreadNum, - DisableBrute: fv.DisableBrute, - DisablePing: fv.DisablePing, - DisableTcpProbe: fv.DisableTcpProbe, + Timeout: time.Duration(fv.TimeoutSec) * time.Second, + TimeoutExplicit: fv.TimeoutExplicit, + ThreadNum: fv.ThreadNum, + ThreadNumExplicit: fv.ThreadNumExplicit, + ModuleThreadNum: fv.ModuleThreadNum, + ModuleThreadNumExplicit: fv.ModuleThreadNumExplicit, + DisableBrute: fv.DisableBrute, + DisablePing: fv.DisablePing, + DisableTcpProbe: fv.DisableTcpProbe, // 扫描模式 - Mode: fv.ScanMode, - LocalMode: fv.LocalPlugin != "", - LocalPlugin: fv.LocalPlugin, - AliveOnly: fv.AliveOnly, - MaxRetries: fv.MaxRetries, + Mode: fv.ScanMode, + LocalMode: fv.LocalPlugin != "", + LocalPlugin: fv.LocalPlugin, + AliveOnly: fv.AliveOnly, + MaxRetries: fv.MaxRetries, + MaxRetriesExplicit: fv.MaxRetriesExplicit, // 高级功能 Shellcode: fv.Shellcode, @@ -173,14 +181,15 @@ func BuildConfigFromFlags(fv *FlagVars) *Config { SSHKeyPath: fv.SSHKeyPath, }, Network: NetworkConfig{ - HTTPProxy: fv.HTTPProxy, - Socks5Proxy: fv.Socks5Proxy, - Iface: fv.Iface, - WebTimeout: time.Duration(fv.WebTimeout) * time.Second, - MaxRedirects: fv.MaxRedirects, - PacketRateLimit: fv.PacketRateLimit, - MaxPacketCount: fv.MaxPacketCount, - ICMPRate: fv.ICMPRate, + HTTPProxy: fv.HTTPProxy, + Socks5Proxy: fv.Socks5Proxy, + Iface: fv.Iface, + WebTimeout: time.Duration(fv.WebTimeout) * time.Second, + MaxRedirects: fv.MaxRedirects, + PacketRateLimit: fv.PacketRateLimit, + MaxPacketCount: fv.MaxPacketCount, + ICMPRate: fv.ICMPRate, + ICMPRateExplicit: fv.ICMPRateExplicit, }, Output: OutputConfig{ File: fv.Outputfile, @@ -195,11 +204,12 @@ func BuildConfigFromFlags(fv *FlagVars) *Config { PerfStats: fv.PerfStats, }, POC: POCConfig{ - PocPath: fv.PocPath, - PocName: fv.PocName, - Full: fv.PocFull, - Num: fv.PocNum, - Disabled: fv.DisablePocScan, + PocPath: fv.PocPath, + PocName: fv.PocName, + Full: fv.PocFull, + Num: fv.PocNum, + NumExplicit: fv.PocNumExplicit, + Disabled: fv.DisablePocScan, }, Redis: RedisConfig{ Disabled: fv.DisableRedis, diff --git a/core/edge_cases_test.go b/core/edge_cases_test.go index 854235c..12983fa 100644 --- a/core/edge_cases_test.go +++ b/core/edge_cases_test.go @@ -575,9 +575,9 @@ func TestClampInt(t *testing.T) { {0, 1, 10, 1}, {15, 1, 10, 10}, {-5, -10, -1, -5}, - {5, 5, 5, 5}, // min == max == v - {3, 5, 5, 5}, // v < min == max - {10, 5, 5, 5}, // v > min == max + {5, 5, 5, 5}, // min == max == v + {3, 5, 5, 5}, // v < min == max + {10, 5, 5, 5}, // v > min == max } for _, tt := range tests { @@ -636,4 +636,16 @@ func TestIsExplicit(t *testing.T) { if !isExplicit(config, "t") { t.Error("ThreadNumExplicit=true 应视为显式") } + + config = makeDefaultConfig() + config.TimeoutExplicit = true + config.ModuleThreadNumExplicit = true + config.MaxRetriesExplicit = true + config.Network.ICMPRateExplicit = true + config.POC.NumExplicit = true + if !isExplicit(config, "time") || !isExplicit(config, "mt") || + !isExplicit(config, "retry") || !isExplicit(config, "icmp-rate") || + !isExplicit(config, "num") { + t.Error("显式标记为 true 时默认值也应视为显式") + } } diff --git a/core/env_profiler.go b/core/env_profiler.go index dde2267..40adfe3 100644 --- a/core/env_profiler.go +++ b/core/env_profiler.go @@ -180,22 +180,22 @@ func computeICMPRate(net *NetworkProfile, sys *SystemProfile) float64 { return base } -// isExplicit 检查参数是否被用户显式指定 -// 目前只有 ThreadNum 有 explicit 标记，其他参数通过检查是否为默认值来判断 +// isExplicit 检查参数是否被用户显式指定。 +// 显式标记来自 CLI flag.Visit；值比较保留 SDK/测试里直接构造 Config 的旧行为。 func isExplicit(config *common.Config, flagName string) bool { switch flagName { case "t": return config.ThreadNumExplicit case "time": - return config.Timeout != 3*time.Second // 默认值 + return config.TimeoutExplicit || config.Timeout != 3*time.Second case "mt": - return config.ModuleThreadNum != 20 // 默认值 + return config.ModuleThreadNumExplicit || config.ModuleThreadNum != 20 case "retry": - return config.MaxRetries != 3 // 默认值 + return config.MaxRetriesExplicit || config.MaxRetries != 3 case "icmp-rate": - return config.Network.ICMPRate != 0.1 // 默认值 + return config.Network.ICMPRateExplicit || config.Network.ICMPRate != 0.1 case "num": - return config.POC.Num != 20 // 默认值 + return config.POC.NumExplicit || config.POC.Num != 20 } return false } diff --git a/core/env_profiler_test.go b/core/env_profiler_test.go index eebde01..f0e197e 100644 --- a/core/env_profiler_test.go +++ b/core/env_profiler_test.go @@ -214,11 +214,11 @@ func TestTuneConfig_SlowLossy(t *testing.T) { func TestTuneConfig_ExplicitOverride(t *testing.T) { config := makeDefaultConfig() - config.Timeout = 5 * time.Second // 用户设了 -time 5 - config.ModuleThreadNum = 50 // 用户设了 -mt 50 - config.MaxRetries = 1 // 用户设了 -retry 1 - config.Network.ICMPRate = 0.8 // 用户设了 -icmp-rate 0.8 - config.POC.Num = 100 // 用户设了 -num 100 + config.Timeout = 5 * time.Second // 用户设了 -time 5 + config.ModuleThreadNum = 50 // 用户设了 -mt 50 + config.MaxRetries = 1 // 用户设了 -retry 1 + config.Network.ICMPRate = 0.8 // 用户设了 -icmp-rate 0.8 + config.POC.Num = 100 // 用户设了 -num 100 session := makeTestSession(config) ep := &EnvironmentProfile{ @@ -252,6 +252,45 @@ func TestTuneConfig_ExplicitOverride(t *testing.T) { } } +func TestTuneConfig_ExplicitDefaultValues(t *testing.T) { + config := makeDefaultConfig() + config.TimeoutExplicit = true + config.ModuleThreadNumExplicit = true + config.MaxRetriesExplicit = true + config.Network.ICMPRateExplicit = true + config.POC.NumExplicit = true + session := makeTestSession(config) + + ep := &EnvironmentProfile{ + Net: NetworkProfile{ + Env: EnvLAN, + RTTMedian: 1 * time.Millisecond, + RTTStddev: 500 * time.Microsecond, + LossRate: 0.0, + Samples: 30, + }, + System: SystemProfile{FDLimit: 65536, NumCPU: 8}, + } + + ep.TuneConfig(config, session) + + if config.Timeout != 3*time.Second { + t.Errorf("显式默认 Timeout 被覆盖: %v", config.Timeout) + } + if config.ModuleThreadNum != 20 { + t.Errorf("显式默认 ModuleThreadNum 被覆盖: %d", config.ModuleThreadNum) + } + if config.MaxRetries != 3 { + t.Errorf("显式默认 MaxRetries 被覆盖: %d", config.MaxRetries) + } + if config.Network.ICMPRate != 0.1 { + t.Errorf("显式默认 ICMPRate 被覆盖: %.2f", config.Network.ICMPRate) + } + if config.POC.Num != 20 { + t.Errorf("显式默认 PocNum 被覆盖: %d", config.POC.Num) + } +} + // ============================================================================= // 集成测试：fd limit 约束 // ============================================================================= diff --git a/plugins/services/credential_tester.go b/plugins/services/credential_tester.go index 8066aac..37fed92 100644 --- a/plugins/services/credential_tester.go +++ b/plugins/services/credential_tester.go @@ -61,6 +61,8 @@ type AuthFunc func(ctx context.Context, cred Credential) *AuthResult // ErrorClassifier 错误分类函数 type ErrorClassifier func(err error) ErrorType +var authCleanupWait = 2 * time.Second + // ============================================================================= // 单凭据测试（解决 goroutine 泄漏） // ============================================================================= @@ -79,12 +81,17 @@ func TestSingleCredential(ctx context.Context, cred Credential, authFn AuthFunc) case result := <-resultChan: return result case <-ctx.Done(): - // context 被取消，等待 authFn goroutine 返回并清理连接 - // 各插件的 authFn 应在 context 取消时关闭底层连接使 goroutine 快速退出 + // context 被取消后只做有界等待，避免 authFn 卡死时清理 goroutine 也永久泄漏。 go func() { - result := <-resultChan - if result != nil && result.Conn != nil { - _ = result.Conn.Close() + timer := time.NewTimer(authCleanupWait) + defer timer.Stop() + + select { + case result := <-resultChan: + if result != nil && result.Conn != nil { + _ = result.Conn.Close() + } + case <-timer.C: } }() return &AuthResult{ diff --git a/plugins/services/credential_tester_test.go b/plugins/services/credential_tester_test.go index 1955b11..1e3511e 100644 --- a/plugins/services/credential_tester_test.go +++ b/plugins/services/credential_tester_test.go @@ -4,6 +4,7 @@ import ( "context" "errors" "io" + "runtime" "sync/atomic" "testing" "time" @@ -401,6 +402,40 @@ func TestTestSingleCredential_ContextCancel(t *testing.T) { } } +func TestTestSingleCredential_ContextCancelCleanupIsBounded(t *testing.T) { + oldWait := authCleanupWait + authCleanupWait = 20 * time.Millisecond + defer func() { authCleanupWait = oldWait }() + + authStarted := make(chan struct{}) + releaseAuth := make(chan struct{}) + authFn := func(ctx context.Context, cred Credential) *AuthResult { + close(authStarted) + <-releaseAuth + return nil + } + + ctx, cancel := context.WithCancel(context.Background()) + go func() { + <-authStarted + cancel() + }() + + before := runtime.NumGoroutine() + result := TestSingleCredential(ctx, Credential{Username: "admin", Password: "admin"}, authFn) + if result.Success { + t.Error("context取消后不应该返回成功") + } + + time.Sleep(100 * time.Millisecond) + after := runtime.NumGoroutine() + close(releaseAuth) + + if after > before+1 { + t.Fatalf("清理 goroutine 疑似泄漏: before=%d after=%d", before, after) + } +} + // ============================================================================= // 重试逻辑测试 // =============================================================================