Fix credential cleanup and explicit tuning flags

This commit is contained in:
ZacharyZcR
2026-06-14 22:23:46 +08:00
parent 8e3cac303d
commit 800cc30794
8 changed files with 205 additions and 86 deletions
+30 -25
View File
@@ -22,20 +22,23 @@ config_struct.go - 配置结构体定义
// Config 扫描器完整配置 - 初始化后只读,可安全共享 // Config 扫描器完整配置 - 初始化后只读,可安全共享
type Config struct { type Config struct {
// 高频访问字段 - 平铺到顶层 // 高频访问字段 - 平铺到顶层
Timeout time.Duration // 通用超时 Timeout time.Duration // 通用超时
ThreadNum int // 主线程数 TimeoutExplicit bool // 用户显式指定了 -time
ThreadNumExplicit bool // 用户显式指定了 -t ThreadNum int // 主线程数
ModuleThreadNum int // 模块线程数 ThreadNumExplicit bool // 用户显式指定了 -t
DisableBrute bool // 禁用暴力破解 ModuleThreadNum int // 模块线程数
DisablePing bool // 禁用Ping检测 ModuleThreadNumExplicit bool // 用户显式指定了 -mt
DisableTcpProbe bool // 禁用TCP补充探测 DisableBrute bool // 禁用暴力破解
DisablePing bool // 禁用Ping检测
DisableTcpProbe bool // 禁用TCP补充探测
// 扫描模式 // 扫描模式
Mode string // 扫描模式 Mode string // 扫描模式
LocalMode bool // 本地模式 LocalMode bool // 本地模式
LocalPlugin string // 本地插件名 LocalPlugin string // 本地插件名
AliveOnly bool // 仅存活检测 AliveOnly bool // 仅存活检测
MaxRetries int // 最大重试次数 MaxRetries int // 最大重试次数
MaxRetriesExplicit bool // 用户显式指定了 -retry
// 高级功能(从AdvancedConfig合并) // 高级功能(从AdvancedConfig合并)
Shellcode string // Shellcode Shellcode string // Shellcode
@@ -81,14 +84,15 @@ type CredentialConfig struct {
// NetworkConfig 网络相关配置 // NetworkConfig 网络相关配置
type NetworkConfig struct { type NetworkConfig struct {
HTTPProxy string HTTPProxy string
Socks5Proxy string Socks5Proxy string
Iface string Iface string
WebTimeout time.Duration WebTimeout time.Duration
MaxRedirects int MaxRedirects int
PacketRateLimit int64 PacketRateLimit int64
MaxPacketCount int64 MaxPacketCount int64
ICMPRate float64 ICMPRate float64
ICMPRateExplicit bool
} }
// OutputConfig 输出相关配置 // OutputConfig 输出相关配置
@@ -107,11 +111,12 @@ type OutputConfig struct {
// POCConfig POC扫描相关配置 // POCConfig POC扫描相关配置
type POCConfig struct { type POCConfig struct {
PocPath string // POC路径 PocPath string // POC路径
PocName string // 指定POC名称 PocName string // 指定POC名称
Full bool // 完整POC扫描 Full bool // 完整POC扫描
Num int // POC并发数 Num int // POC并发数
Disabled bool // 禁用POC扫描 NumExplicit bool // 用户显式指定了 -num
Disabled bool // 禁用POC扫描
} }
// RedisConfig Redis利用相关配置 // RedisConfig Redis利用相关配置
+12 -1
View File
@@ -215,8 +215,19 @@ func Flag(Info *HostInfo) error {
// 检测用户是否显式指定了 -t // 检测用户是否显式指定了 -t
flag.Visit(func(f *flag.Flag) { flag.Visit(func(f *flag.Flag) {
if f.Name == "t" { switch f.Name {
case "t":
fv.ThreadNumExplicit = true 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 PortsFile string
// 扫描控制 // 扫描控制
ScanMode string ScanMode string
ThreadNum int ThreadNum int
ThreadNumExplicit bool // 用户显式指定了 -t ThreadNumExplicit bool // 用户显式指定了 -t
ModuleThreadNum int ModuleThreadNum int
TimeoutSec int64 // 秒,需转换为 time.Duration ModuleThreadNumExplicit bool
GlobalTimeout int64 TimeoutSec int64 // 秒,需转换为 time.Duration
DisablePing bool TimeoutExplicit bool
DisableTcpProbe bool GlobalTimeout int64
LocalPlugin string DisablePing bool
AliveOnly bool DisableTcpProbe bool
DisableBrute bool LocalPlugin string
MaxRetries int AliveOnly bool
DisableBrute bool
MaxRetries int
MaxRetriesExplicit bool
// 认证凭据 // 认证凭据
Username string Username string
@@ -74,6 +77,7 @@ type FlagVars struct {
PocFull bool PocFull bool
DNSLog bool DNSLog bool
PocNum int PocNum int
PocNumExplicit bool
DisablePocScan bool DisablePocScan bool
// Redis利用 // Redis利用
@@ -85,9 +89,10 @@ type FlagVars struct {
DisableRedis bool DisableRedis bool
// 发包频率 // 发包频率
PacketRateLimit int64 PacketRateLimit int64
MaxPacketCount int64 MaxPacketCount int64
ICMPRate float64 ICMPRate float64
ICMPRateExplicit bool
// 输出控制 // 输出控制
Outputfile string Outputfile string
@@ -135,20 +140,23 @@ func GetFlagVars() *FlagVars {
func BuildConfigFromFlags(fv *FlagVars) *Config { func BuildConfigFromFlags(fv *FlagVars) *Config {
return &Config{ return &Config{
// 高频字段 // 高频字段
Timeout: time.Duration(fv.TimeoutSec) * time.Second, Timeout: time.Duration(fv.TimeoutSec) * time.Second,
ThreadNum: fv.ThreadNum, TimeoutExplicit: fv.TimeoutExplicit,
ThreadNumExplicit: fv.ThreadNumExplicit, ThreadNum: fv.ThreadNum,
ModuleThreadNum: fv.ModuleThreadNum, ThreadNumExplicit: fv.ThreadNumExplicit,
DisableBrute: fv.DisableBrute, ModuleThreadNum: fv.ModuleThreadNum,
DisablePing: fv.DisablePing, ModuleThreadNumExplicit: fv.ModuleThreadNumExplicit,
DisableTcpProbe: fv.DisableTcpProbe, DisableBrute: fv.DisableBrute,
DisablePing: fv.DisablePing,
DisableTcpProbe: fv.DisableTcpProbe,
// 扫描模式 // 扫描模式
Mode: fv.ScanMode, Mode: fv.ScanMode,
LocalMode: fv.LocalPlugin != "", LocalMode: fv.LocalPlugin != "",
LocalPlugin: fv.LocalPlugin, LocalPlugin: fv.LocalPlugin,
AliveOnly: fv.AliveOnly, AliveOnly: fv.AliveOnly,
MaxRetries: fv.MaxRetries, MaxRetries: fv.MaxRetries,
MaxRetriesExplicit: fv.MaxRetriesExplicit,
// 高级功能 // 高级功能
Shellcode: fv.Shellcode, Shellcode: fv.Shellcode,
@@ -173,14 +181,15 @@ func BuildConfigFromFlags(fv *FlagVars) *Config {
SSHKeyPath: fv.SSHKeyPath, SSHKeyPath: fv.SSHKeyPath,
}, },
Network: NetworkConfig{ Network: NetworkConfig{
HTTPProxy: fv.HTTPProxy, HTTPProxy: fv.HTTPProxy,
Socks5Proxy: fv.Socks5Proxy, Socks5Proxy: fv.Socks5Proxy,
Iface: fv.Iface, Iface: fv.Iface,
WebTimeout: time.Duration(fv.WebTimeout) * time.Second, WebTimeout: time.Duration(fv.WebTimeout) * time.Second,
MaxRedirects: fv.MaxRedirects, MaxRedirects: fv.MaxRedirects,
PacketRateLimit: fv.PacketRateLimit, PacketRateLimit: fv.PacketRateLimit,
MaxPacketCount: fv.MaxPacketCount, MaxPacketCount: fv.MaxPacketCount,
ICMPRate: fv.ICMPRate, ICMPRate: fv.ICMPRate,
ICMPRateExplicit: fv.ICMPRateExplicit,
}, },
Output: OutputConfig{ Output: OutputConfig{
File: fv.Outputfile, File: fv.Outputfile,
@@ -195,11 +204,12 @@ func BuildConfigFromFlags(fv *FlagVars) *Config {
PerfStats: fv.PerfStats, PerfStats: fv.PerfStats,
}, },
POC: POCConfig{ POC: POCConfig{
PocPath: fv.PocPath, PocPath: fv.PocPath,
PocName: fv.PocName, PocName: fv.PocName,
Full: fv.PocFull, Full: fv.PocFull,
Num: fv.PocNum, Num: fv.PocNum,
Disabled: fv.DisablePocScan, NumExplicit: fv.PocNumExplicit,
Disabled: fv.DisablePocScan,
}, },
Redis: RedisConfig{ Redis: RedisConfig{
Disabled: fv.DisableRedis, Disabled: fv.DisableRedis,
+15 -3
View File
@@ -575,9 +575,9 @@ func TestClampInt(t *testing.T) {
{0, 1, 10, 1}, {0, 1, 10, 1},
{15, 1, 10, 10}, {15, 1, 10, 10},
{-5, -10, -1, -5}, {-5, -10, -1, -5},
{5, 5, 5, 5}, // min == max == v {5, 5, 5, 5}, // min == max == v
{3, 5, 5, 5}, // v < min == max {3, 5, 5, 5}, // v < min == max
{10, 5, 5, 5}, // v > min == max {10, 5, 5, 5}, // v > min == max
} }
for _, tt := range tests { for _, tt := range tests {
@@ -636,4 +636,16 @@ func TestIsExplicit(t *testing.T) {
if !isExplicit(config, "t") { if !isExplicit(config, "t") {
t.Error("ThreadNumExplicit=true 应视为显式") 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 return base
} }
// isExplicit 检查参数是否被用户显式指定 // isExplicit 检查参数是否被用户显式指定。
// 目前只有 ThreadNum 有 explicit 标记,其他参数通过检查是否为默认值来判断 // 显式标记来自 CLI flag.Visit;值比较保留 SDK/测试里直接构造 Config 的旧行为。
func isExplicit(config *common.Config, flagName string) bool { func isExplicit(config *common.Config, flagName string) bool {
switch flagName { switch flagName {
case "t": case "t":
return config.ThreadNumExplicit return config.ThreadNumExplicit
case "time": case "time":
return config.Timeout != 3*time.Second // 默认值 return config.TimeoutExplicit || config.Timeout != 3*time.Second
case "mt": case "mt":
return config.ModuleThreadNum != 20 // 默认值 return config.ModuleThreadNumExplicit || config.ModuleThreadNum != 20
case "retry": case "retry":
return config.MaxRetries != 3 // 默认值 return config.MaxRetriesExplicit || config.MaxRetries != 3
case "icmp-rate": case "icmp-rate":
return config.Network.ICMPRate != 0.1 // 默认值 return config.Network.ICMPRateExplicit || config.Network.ICMPRate != 0.1
case "num": case "num":
return config.POC.Num != 20 // 默认值 return config.POC.NumExplicit || config.POC.Num != 20
} }
return false return false
} }
+44 -5
View File
@@ -214,11 +214,11 @@ func TestTuneConfig_SlowLossy(t *testing.T) {
func TestTuneConfig_ExplicitOverride(t *testing.T) { func TestTuneConfig_ExplicitOverride(t *testing.T) {
config := makeDefaultConfig() config := makeDefaultConfig()
config.Timeout = 5 * time.Second // 用户设了 -time 5 config.Timeout = 5 * time.Second // 用户设了 -time 5
config.ModuleThreadNum = 50 // 用户设了 -mt 50 config.ModuleThreadNum = 50 // 用户设了 -mt 50
config.MaxRetries = 1 // 用户设了 -retry 1 config.MaxRetries = 1 // 用户设了 -retry 1
config.Network.ICMPRate = 0.8 // 用户设了 -icmp-rate 0.8 config.Network.ICMPRate = 0.8 // 用户设了 -icmp-rate 0.8
config.POC.Num = 100 // 用户设了 -num 100 config.POC.Num = 100 // 用户设了 -num 100
session := makeTestSession(config) session := makeTestSession(config)
ep := &EnvironmentProfile{ 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 约束 // 集成测试:fd limit 约束
// ============================================================================= // =============================================================================
+12 -5
View File
@@ -61,6 +61,8 @@ type AuthFunc func(ctx context.Context, cred Credential) *AuthResult
// ErrorClassifier 错误分类函数 // ErrorClassifier 错误分类函数
type ErrorClassifier func(err error) ErrorType type ErrorClassifier func(err error) ErrorType
var authCleanupWait = 2 * time.Second
// ============================================================================= // =============================================================================
// 单凭据测试(解决 goroutine 泄漏) // 单凭据测试(解决 goroutine 泄漏)
// ============================================================================= // =============================================================================
@@ -79,12 +81,17 @@ func TestSingleCredential(ctx context.Context, cred Credential, authFn AuthFunc)
case result := <-resultChan: case result := <-resultChan:
return result return result
case <-ctx.Done(): case <-ctx.Done():
// context 被取消,等待 authFn goroutine 返回并清理连接 // context 被取消后只做有界等待,避免 authFn 卡死时清理 goroutine 也永久泄漏。
// 各插件的 authFn 应在 context 取消时关闭底层连接使 goroutine 快速退出
go func() { go func() {
result := <-resultChan timer := time.NewTimer(authCleanupWait)
if result != nil && result.Conn != nil { defer timer.Stop()
_ = result.Conn.Close()
select {
case result := <-resultChan:
if result != nil && result.Conn != nil {
_ = result.Conn.Close()
}
case <-timer.C:
} }
}() }()
return &AuthResult{ return &AuthResult{
@@ -4,6 +4,7 @@ import (
"context" "context"
"errors" "errors"
"io" "io"
"runtime"
"sync/atomic" "sync/atomic"
"testing" "testing"
"time" "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)
}
}
// ============================================================================= // =============================================================================
// 重试逻辑测试 // 重试逻辑测试
// ============================================================================= // =============================================================================