Fix credential cleanup and explicit tuning flags

This commit is contained in:
ZacharyZcR
2026-06-12 15:30:44 +08:00
parent 52f872b8d1
commit 5b7e72e56e
8 changed files with 205 additions and 86 deletions
+30 -25
View File
@@ -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利用相关配置
+12 -1
View File
@@ -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
}
})
+50 -40
View File
@@ -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,
+15 -3
View File
@@ -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 时默认值也应视为显式")
}
}
+7 -7
View File
@@ -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
}
+44 -5
View File
@@ -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 约束
// =============================================================================
+12 -5
View File
@@ -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{
@@ -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)
}
}
// =============================================================================
// 重试逻辑测试
// =============================================================================