mirror of
https://github.com/shadow1ng/fscan.git
synced 2026-09-26 13:11:53 +08:00
refactor: context 穿透扫描生命周期,修复长驻插件阻塞和 Web Stop 无效
- RunScan 接受 context.Context,创建可取消上下文并穿透到所有策略和插件 - 长驻插件(forwardshell/socks5proxy/reverseshell)不再进入 scan WaitGroup, 通过 ctx.Done() 管理生命周期,解除 wg.Wait() 死锁 - Web Stop API 从 stopChan 改为 context.CancelFunc,取消信号真正传播到扫描链路 - ExecuteScanTasks 和 executeScanTask 支持 context 取消检查,停止分发新任务 - CLI 模式传 context.Background(),行为完全不变
This commit is contained in:
+52
-8
@@ -18,7 +18,7 @@ import (
|
||||
|
||||
// ScanStrategy 定义扫描策略接口
|
||||
type ScanStrategy interface {
|
||||
Execute(config *common.Config, state *common.State, info common.HostInfo, ch chan struct{}, wg *sync.WaitGroup)
|
||||
Execute(ctx context.Context, config *common.Config, state *common.State, info common.HostInfo, ch chan struct{}, wg *sync.WaitGroup)
|
||||
GetPlugins(config *common.Config) ([]string, bool)
|
||||
IsPluginApplicableByName(pluginName string, targetHost string, targetPort int, isCustomMode bool, config *common.Config) bool
|
||||
}
|
||||
@@ -73,7 +73,10 @@ func selectStrategy(config *common.Config, state *common.State, info common.Host
|
||||
}
|
||||
|
||||
// RunScan 执行整体扫描流程
|
||||
func RunScan(info common.HostInfo, config *common.Config, state *common.State) {
|
||||
func RunScan(ctx context.Context, info common.HostInfo, config *common.Config, state *common.State) {
|
||||
ctx, cancel := context.WithCancel(ctx)
|
||||
defer cancel()
|
||||
|
||||
// 初始化HTTP客户端(静默,无需日志)
|
||||
if err := lib.Inithttp(config); err != nil {
|
||||
common.LogError(i18n.Tr("http_client_init_failed", err))
|
||||
@@ -88,7 +91,7 @@ func RunScan(info common.HostInfo, config *common.Config, state *common.State) {
|
||||
wg := sync.WaitGroup{}
|
||||
|
||||
// 执行策略
|
||||
strategy.Execute(config, state, info, ch, &wg)
|
||||
strategy.Execute(ctx, config, state, info, ch, &wg)
|
||||
|
||||
// 等待所有扫描完成
|
||||
wg.Wait()
|
||||
@@ -111,6 +114,8 @@ func RunScan(info common.HostInfo, config *common.Config, state *common.State) {
|
||||
signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)
|
||||
<-sigChan
|
||||
common.LogInfo(i18n.GetText("received_exit_signal"))
|
||||
cancel()
|
||||
time.Sleep(500 * time.Millisecond)
|
||||
}
|
||||
|
||||
// 完成扫描
|
||||
@@ -134,7 +139,7 @@ func finishScan(config *common.Config, state *common.State) {
|
||||
}
|
||||
|
||||
// ExecuteScanTasks 任务执行通用框架
|
||||
func ExecuteScanTasks(config *common.Config, state *common.State, targets []common.HostInfo, strategy ScanStrategy, ch chan struct{}, wg *sync.WaitGroup) {
|
||||
func ExecuteScanTasks(ctx context.Context, config *common.Config, state *common.State, targets []common.HostInfo, strategy ScanStrategy, ch chan struct{}, wg *sync.WaitGroup) {
|
||||
// 获取要执行的插件
|
||||
pluginsToRun, isCustomMode := strategy.GetPlugins(config)
|
||||
|
||||
@@ -149,6 +154,13 @@ func ExecuteScanTasks(config *common.Config, state *common.State, targets []comm
|
||||
|
||||
// 流式执行任务,避免预构建大量任务对象
|
||||
for _, target := range targets {
|
||||
// 检查取消
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
default:
|
||||
}
|
||||
|
||||
targetPort := target.Port
|
||||
|
||||
for _, pluginName := range pluginsToRun {
|
||||
@@ -159,7 +171,7 @@ func ExecuteScanTasks(config *common.Config, state *common.State, targets []comm
|
||||
|
||||
// 检查插件是否适用于当前目标
|
||||
if strategy.IsPluginApplicableByName(pluginName, target.Host, targetPort, isCustomMode, config) {
|
||||
executeScanTask(config, state, pluginName, target, ch, wg)
|
||||
executeScanTask(ctx, config, state, pluginName, target, ch, wg)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -182,10 +194,42 @@ func countApplicableTasks(targets []common.HostInfo, pluginsToRun []string, isCu
|
||||
return count
|
||||
}
|
||||
|
||||
// longRunningPlugins 长驻插件,不加入 scan WaitGroup,通过 ctx 取消退出
|
||||
var longRunningPlugins = map[string]bool{
|
||||
"forwardshell": true,
|
||||
"socks5proxy": true,
|
||||
"reverseshell": true,
|
||||
}
|
||||
|
||||
// executeScanTask 执行单个扫描任务
|
||||
func executeScanTask(config *common.Config, state *common.State, pluginName string, target common.HostInfo, ch chan struct{}, wg *sync.WaitGroup) {
|
||||
func executeScanTask(ctx context.Context, config *common.Config, state *common.State, pluginName string, target common.HostInfo, ch chan struct{}, wg *sync.WaitGroup) {
|
||||
// 检查取消
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
default:
|
||||
}
|
||||
|
||||
// 长驻插件不进 WaitGroup,通过 ctx 管理生命周期
|
||||
if longRunningPlugins[pluginName] {
|
||||
go func() {
|
||||
plugin := plugins.Get(pluginName)
|
||||
if plugin != nil {
|
||||
plugin.Scan(ctx, &target, config, state)
|
||||
}
|
||||
}()
|
||||
return
|
||||
}
|
||||
|
||||
wg.Add(1)
|
||||
ch <- struct{}{} // 获取并发槽位
|
||||
|
||||
// 获取并发槽位,支持取消
|
||||
select {
|
||||
case ch <- struct{}{}:
|
||||
case <-ctx.Done():
|
||||
wg.Done()
|
||||
return
|
||||
}
|
||||
|
||||
go func() {
|
||||
// 开始监控插件任务
|
||||
@@ -210,7 +254,7 @@ func executeScanTask(config *common.Config, state *common.State, pluginName stri
|
||||
|
||||
plugin := plugins.Get(pluginName)
|
||||
if plugin != nil {
|
||||
result := plugin.Scan(context.Background(), &target, config, state)
|
||||
result := plugin.Scan(ctx, &target, config, state)
|
||||
if result != nil {
|
||||
if result.Success {
|
||||
// 保存成功的扫描结果到文件
|
||||
|
||||
Reference in New Issue
Block a user