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:
ZacharyZcR
2026-04-27 19:04:26 +08:00
parent 1fe8bb5182
commit d7749c4766
8 changed files with 85 additions and 25 deletions
+2 -1
View File
@@ -1,6 +1,7 @@
package core package core
import ( import (
"context"
"fmt" "fmt"
"sync" "sync"
"time" "time"
@@ -53,7 +54,7 @@ func (s *AliveScanStrategy) Description() string {
} }
// Execute 执行存活探测扫描策略 // Execute 执行存活探测扫描策略
func (s *AliveScanStrategy) Execute(config *common.Config, state *common.State, info common.HostInfo, ch chan struct{}, wg *sync.WaitGroup) { func (s *AliveScanStrategy) Execute(_ context.Context, config *common.Config, state *common.State, info common.HostInfo, ch chan struct{}, wg *sync.WaitGroup) {
// 验证扫描目标(需要同时检查 -h 和 -hf 参数) // 验证扫描目标(需要同时检查 -h 和 -hf 参数)
fv := common.GetFlagVars() fv := common.GetFlagVars()
if info.Host == "" && fv.HostsFile == "" { if info.Host == "" && fv.HostsFile == "" {
+3 -2
View File
@@ -1,6 +1,7 @@
package core package core
import ( import (
"context"
"sync" "sync"
"github.com/shadow1ng/fscan/common" "github.com/shadow1ng/fscan/common"
@@ -41,7 +42,7 @@ func (s *LocalScanStrategy) Description() string {
} }
// Execute 执行本地扫描策略 // Execute 执行本地扫描策略
func (s *LocalScanStrategy) Execute(config *common.Config, state *common.State, info common.HostInfo, ch chan struct{}, wg *sync.WaitGroup) { func (s *LocalScanStrategy) Execute(ctx context.Context, config *common.Config, state *common.State, info common.HostInfo, ch chan struct{}, wg *sync.WaitGroup) {
// 输出扫描开始信息 // 输出扫描开始信息
s.LogScanStart() s.LogScanStart()
@@ -66,7 +67,7 @@ func (s *LocalScanStrategy) Execute(config *common.Config, state *common.State,
targets := s.PrepareTargets(info) targets := s.PrepareTargets(info)
// 执行扫描任务 // 执行扫描任务
ExecuteScanTasks(config, state, targets, s, ch, wg) ExecuteScanTasks(ctx, config, state, targets, s, ch, wg)
} }
// PrepareTargets 准备本地扫描目标 // PrepareTargets 准备本地扫描目标
+52 -8
View File
@@ -18,7 +18,7 @@ import (
// ScanStrategy 定义扫描策略接口 // ScanStrategy 定义扫描策略接口
type ScanStrategy interface { 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) GetPlugins(config *common.Config) ([]string, bool)
IsPluginApplicableByName(pluginName string, targetHost string, targetPort int, isCustomMode bool, config *common.Config) 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 执行整体扫描流程 // 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客户端(静默,无需日志) // 初始化HTTP客户端(静默,无需日志)
if err := lib.Inithttp(config); err != nil { if err := lib.Inithttp(config); err != nil {
common.LogError(i18n.Tr("http_client_init_failed", err)) 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{} wg := sync.WaitGroup{}
// 执行策略 // 执行策略
strategy.Execute(config, state, info, ch, &wg) strategy.Execute(ctx, config, state, info, ch, &wg)
// 等待所有扫描完成 // 等待所有扫描完成
wg.Wait() wg.Wait()
@@ -111,6 +114,8 @@ func RunScan(info common.HostInfo, config *common.Config, state *common.State) {
signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM) signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)
<-sigChan <-sigChan
common.LogInfo(i18n.GetText("received_exit_signal")) 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 任务执行通用框架 // 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) pluginsToRun, isCustomMode := strategy.GetPlugins(config)
@@ -149,6 +154,13 @@ func ExecuteScanTasks(config *common.Config, state *common.State, targets []comm
// 流式执行任务,避免预构建大量任务对象 // 流式执行任务,避免预构建大量任务对象
for _, target := range targets { for _, target := range targets {
// 检查取消
select {
case <-ctx.Done():
return
default:
}
targetPort := target.Port targetPort := target.Port
for _, pluginName := range pluginsToRun { 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) { 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 return count
} }
// longRunningPlugins 长驻插件,不加入 scan WaitGroup,通过 ctx 取消退出
var longRunningPlugins = map[string]bool{
"forwardshell": true,
"socks5proxy": true,
"reverseshell": true,
}
// executeScanTask 执行单个扫描任务 // 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) wg.Add(1)
ch <- struct{}{} // 获取并发槽位
// 获取并发槽位,支持取消
select {
case ch <- struct{}{}:
case <-ctx.Done():
wg.Done()
return
}
go func() { go func() {
// 开始监控插件任务 // 开始监控插件任务
@@ -210,7 +254,7 @@ func executeScanTask(config *common.Config, state *common.State, pluginName stri
plugin := plugins.Get(pluginName) plugin := plugins.Get(pluginName)
if plugin != nil { if plugin != nil {
result := plugin.Scan(context.Background(), &target, config, state) result := plugin.Scan(ctx, &target, config, state)
if result != nil { if result != nil {
if result.Success { if result.Success {
// 保存成功的扫描结果到文件 // 保存成功的扫描结果到文件
+2 -1
View File
@@ -1,6 +1,7 @@
package core package core
import ( import (
"context"
"fmt" "fmt"
"sync" "sync"
"testing" "testing"
@@ -214,7 +215,7 @@ type mockStrategy struct {
applicablePlugins map[string]bool // pluginName -> isApplicable applicablePlugins map[string]bool // pluginName -> isApplicable
} }
func (m *mockStrategy) Execute(config *common.Config, state *common.State, info common.HostInfo, ch chan struct{}, wg *sync.WaitGroup) { func (m *mockStrategy) Execute(_ context.Context, config *common.Config, state *common.State, info common.HostInfo, ch chan struct{}, wg *sync.WaitGroup) {
} }
func (m *mockStrategy) GetPlugins() ([]string, bool) { func (m *mockStrategy) GetPlugins() ([]string, bool) {
+5 -4
View File
@@ -1,6 +1,7 @@
package core package core
import ( import (
"context"
"fmt" "fmt"
"strconv" "strconv"
"strings" "strings"
@@ -112,7 +113,7 @@ func (s *ServiceScanStrategy) Description() string {
} }
// Execute 执行服务扫描策略 // Execute 执行服务扫描策略
func (s *ServiceScanStrategy) Execute(config *common.Config, state *common.State, info common.HostInfo, ch chan struct{}, wg *sync.WaitGroup) { func (s *ServiceScanStrategy) Execute(ctx context.Context, config *common.Config, state *common.State, info common.HostInfo, ch chan struct{}, wg *sync.WaitGroup) {
// 验证扫描目标(需要同时检查 -h 和 -hf 参数) // 验证扫描目标(需要同时检查 -h 和 -hf 参数)
fv := common.GetFlagVars() fv := common.GetFlagVars()
if info.Host == "" && fv.HostsFile == "" { if info.Host == "" && fv.HostsFile == "" {
@@ -133,11 +134,11 @@ func (s *ServiceScanStrategy) Execute(config *common.Config, state *common.State
s.LogPluginInfo(config) s.LogPluginInfo(config)
// 执行主机扫描流程 // 执行主机扫描流程
s.performHostScan(config, state, info, ch, wg) s.performHostScan(ctx, config, state, info, ch, wg)
} }
// performHostScan 执行主机扫描的完整流程 // performHostScan 执行主机扫描的完整流程
func (s *ServiceScanStrategy) performHostScan(config *common.Config, state *common.State, info common.HostInfo, ch chan struct{}, wg *sync.WaitGroup) { func (s *ServiceScanStrategy) performHostScan(ctx context.Context, config *common.Config, state *common.State, info common.HostInfo, ch chan struct{}, wg *sync.WaitGroup) {
// 发现目标主机和端口 // 发现目标主机和端口
targetInfos, err := s.discoverTargets(info.Host, info, config, state) targetInfos, err := s.discoverTargets(info.Host, info, config, state)
if err != nil { if err != nil {
@@ -147,7 +148,7 @@ func (s *ServiceScanStrategy) performHostScan(config *common.Config, state *comm
// 执行漏洞扫描 // 执行漏洞扫描
if len(targetInfos) > 0 { if len(targetInfos) > 0 {
ExecuteScanTasks(config, state, targetInfos, s, ch, wg) ExecuteScanTasks(ctx, config, state, targetInfos, s, ch, wg)
} }
} }
+3 -2
View File
@@ -1,6 +1,7 @@
package core package core
import ( import (
"context"
"crypto/tls" "crypto/tls"
"fmt" "fmt"
"net" "net"
@@ -303,7 +304,7 @@ func (s *WebScanStrategy) Description() string {
} }
// Execute 执行Web扫描策略 // Execute 执行Web扫描策略
func (s *WebScanStrategy) Execute(config *common.Config, state *common.State, info common.HostInfo, ch chan struct{}, wg *sync.WaitGroup) { func (s *WebScanStrategy) Execute(ctx context.Context, config *common.Config, state *common.State, info common.HostInfo, ch chan struct{}, wg *sync.WaitGroup) {
// 输出扫描开始信息 // 输出扫描开始信息
s.LogScanStart() s.LogScanStart()
@@ -320,7 +321,7 @@ func (s *WebScanStrategy) Execute(config *common.Config, state *common.State, in
s.LogPluginInfo(config) s.LogPluginInfo(config)
// 执行扫描任务 // 执行扫描任务
ExecuteScanTasks(config, state, targets, s, ch, wg) ExecuteScanTasks(ctx, config, state, targets, s, ch, wg)
} }
// PrepareTargets 准备URL目标列表 // PrepareTargets 准备URL目标列表
+2 -1
View File
@@ -1,6 +1,7 @@
package main package main
import ( import (
"context"
"os" "os"
"os/signal" "os/signal"
"syscall" "syscall"
@@ -66,5 +67,5 @@ func main() {
defer func() { _ = common.Cleanup() }() defer func() { _ = common.Cleanup() }()
// 执行扫描 // 执行扫描
core.RunScan(*result.Info, result.Config, result.State) core.RunScan(context.Background(), *result.Info, result.Config, result.State)
} }
+16 -6
View File
@@ -3,6 +3,7 @@
package api package api
import ( import (
"context"
"encoding/json" "encoding/json"
"net/http" "net/http"
"sync" "sync"
@@ -73,7 +74,7 @@ type ScanHandler struct {
hub *ws.Hub hub *ws.Hub
state int32 state int32
startTime time.Time startTime time.Time
stopChan chan struct{} cancelFn context.CancelFunc
mu sync.RWMutex mu sync.RWMutex
results *ResultStore results *ResultStore
} }
@@ -122,7 +123,6 @@ func (h *ScanHandler) Start(w http.ResponseWriter, r *http.Request) {
h.mu.Lock() h.mu.Lock()
h.startTime = time.Now() h.startTime = time.Now()
h.stopChan = make(chan struct{})
h.mu.Unlock() h.mu.Unlock()
// 清空旧结果 // 清空旧结果
@@ -145,8 +145,18 @@ func (h *ScanHandler) Start(w http.ResponseWriter, r *http.Request) {
// runScan 执行扫描 // runScan 执行扫描
func (h *ScanHandler) runScan(req ScanRequest) { func (h *ScanHandler) runScan(req ScanRequest) {
ctx, cancel := context.WithCancel(context.Background())
h.mu.Lock()
h.cancelFn = cancel
h.mu.Unlock()
defer func() { defer func() {
common.ClearResultCallback() // 清除回调 cancel()
h.mu.Lock()
h.cancelFn = nil
h.mu.Unlock()
common.ClearResultCallback()
atomic.StoreInt32(&h.state, int32(ScanStateIdle)) atomic.StoreInt32(&h.state, int32(ScanStateIdle))
h.hub.Broadcast(ws.MsgScanCompleted, map[string]interface{}{ h.hub.Broadcast(ws.MsgScanCompleted, map[string]interface{}{
"duration": time.Since(h.startTime).Seconds(), "duration": time.Since(h.startTime).Seconds(),
@@ -209,7 +219,7 @@ func (h *ScanHandler) runScan(req ScanRequest) {
}) })
// 执行扫描 // 执行扫描
core.RunScan(info, config, state) core.RunScan(ctx, info, config, state)
} }
// Stop 停止扫描 // Stop 停止扫描
@@ -229,8 +239,8 @@ func (h *ScanHandler) Stop(w http.ResponseWriter, r *http.Request) {
atomic.StoreInt32(&h.state, int32(ScanStateStopping)) atomic.StoreInt32(&h.state, int32(ScanStateStopping))
h.mu.Lock() h.mu.Lock()
if h.stopChan != nil { if h.cancelFn != nil {
close(h.stopChan) h.cancelFn()
} }
h.mu.Unlock() h.mu.Unlock()