refactor: 4项架构优化 — CEL缓存/POC隔离/服务缓存/结果统一

1. CEL 表达式编译缓存
   - 新增 CelProgCache,同一 POC 的所有规则/参数组合共享编译后的 Program
   - clusterpoc 热路径上消除重复的 Compile+Program 调用

2. POC 全局状态消除
   - allPocs/pocLoaded 全局变量改为 pocStore 按 PocPath 缓存
   - 不同 PocPath 的扫描独立加载,Web API 并发场景不再互相覆盖

3. serviceCache 下沉到 per-session State
   - 服务识别缓存从包级全局 map 迁移到 State.serviceCache (sync.Map)
   - BaseScanStrategy 通过 SetState 注入 session state
   - 消除多个并发扫描之间的服务识别缓存串台

4. POC 结果输出路径统一
   - 提取 buildVulnDetails/buildVulnLogMsg/saveVulnResult 三个公共函数
   - CheckMultiPoc 和 recordVulnerabilityResult 共用统一的结果构造逻辑
   - 消除 details 字段名不一致和日志格式差异
This commit is contained in:
ZacharyZcR
2026-06-13 19:24:57 +08:00
parent 6a1636112f
commit 02ad8f5334
10 changed files with 261 additions and 230 deletions
+18
View File
@@ -56,6 +56,10 @@ type State struct {
forwardShellActive int32 // 使用int32以便原子操作
reverseShellActive int32
socks5ProxyActive int32
// 服务识别缓存(per-session,避免跨扫描污染)
// key: "host:port", value: interface{}core.ServiceInfo 指针)
serviceCache sync.Map
}
// NewState 创建新的状态对象
@@ -455,3 +459,17 @@ func (s *State) CheckAndIncrementPacketRate(rateLimit int64) (bool, error) {
return true, nil
}
// =============================================================================
// 服务识别缓存 - per-session,消除跨扫描污染
// =============================================================================
// CacheService 缓存服务信息
func (s *State) CacheService(key string, info interface{}) {
s.serviceCache.Store(key, info)
}
// GetCachedService 获取缓存的服务信息
func (s *State) GetCachedService(key string) (interface{}, bool) {
return s.serviceCache.Load(key)
}
+8 -2
View File
@@ -29,6 +29,7 @@ const (
type BaseScanStrategy struct {
strategyName string
filterType PluginFilterType
state *common.State
}
// NewBaseScanStrategy 创建基础扫描策略
@@ -39,6 +40,11 @@ func NewBaseScanStrategy(name string, filterType PluginFilterType) *BaseScanStra
}
}
// SetState 注入 session state(用于 per-session 服务缓存)
func (b *BaseScanStrategy) SetState(state *common.State) {
b.state = state
}
// GetPlugins 获取插件列表
func (b *BaseScanStrategy) GetPlugins(config *common.Config) ([]string, bool) {
scanMode := config.Mode
@@ -123,7 +129,7 @@ func (b *BaseScanStrategy) isLocalPluginExplicitlySpecified(pluginName string, c
// 匹配策略:端口匹配 → 服务名称匹配(解决非标准端口问题)
func (b *BaseScanStrategy) isPluginApplicableToPortWithHost(pluginName string, targetHost string, targetPort int) bool {
if b.isWebPlugin(pluginName) {
return IsMarkedWebService(targetHost, targetPort)
return IsMarkedWebServiceWithState(b.state, targetHost, targetPort)
}
pluginPorts := b.getPluginPorts(pluginName)
@@ -145,7 +151,7 @@ func (b *BaseScanStrategy) isPluginApplicableToPortWithHost(pluginName string, t
// 端口不匹配时,按指纹识别结果匹配
// 例:8881 端口上识别到 ssh 服务 → ssh 插件应该执行
if targetHost != "" && targetPort > 0 {
if info, ok := GetCachedServiceInfo(targetHost, targetPort); ok && info != nil {
if info, ok := GetCachedServiceInfoWithState(b.state, targetHost, targetPort); ok && info != nil {
if strings.EqualFold(info.Name, pluginName) {
return true
}
+8
View File
@@ -100,6 +100,9 @@ func RunScan(ctx context.Context, info common.HostInfo, session *common.ScanSess
config := session.Config
state := session.State
// 设置全局 State(兼容旧代码路径中未传 state 的调用)
SetGlobalState(state)
// 初始化HTTP客户端(静默,无需日志)
if err := lib.Inithttp(config); err != nil {
session.LogError(i18n.Tr("http_client_init_failed", err))
@@ -190,6 +193,11 @@ func finishScan(session *common.ScanSession) {
func ExecuteScanTasks(ctx context.Context, session *common.ScanSession, targets []common.HostInfo, strategy ScanStrategy, ch chan struct{}, wg *sync.WaitGroup) {
config := session.Config
// 注入 session state 到策略(用于 per-session 服务缓存)
if setter, ok := strategy.(interface{ SetState(*common.State) }); ok {
setter.SetState(session.State)
}
// 获取要执行的插件
pluginsToRun, isCustomMode := strategy.GetPlugins(config)
+3 -3
View File
@@ -4,6 +4,7 @@ import (
"sync"
"testing"
"github.com/shadow1ng/fscan/common"
"github.com/shadow1ng/fscan/plugins"
)
@@ -22,9 +23,8 @@ func registerTestPlugins(t *testing.T) {
}
func clearServiceCache() {
serviceCacheMutex.Lock()
serviceCache = make(map[string]*ServiceInfo)
serviceCacheMutex.Unlock()
state := common.NewState()
SetGlobalState(state)
}
// =============================================================================
+63 -32
View File
@@ -208,12 +208,9 @@ func (w *WebPortDetector) tryHTTP(ctx context.Context, client *http.Client, sess
// 基于服务指纹的Web服务识别
// ===============================
// 服务识别缓存 - 存储所有识别到的服务(不仅限于 Web
// 端口扫描阶段写入,插件匹配阶段读取
var (
serviceCache = make(map[string]*ServiceInfo)
serviceCacheMutex sync.RWMutex
)
// globalState 全局 State 兼容指针(向后兼容不接受 State 的旧调用方
// 新代码应通过 State 方法访问服务缓存
var globalState *common.State
// IsWebServiceByFingerprint 基于服务指纹判断Web服务 - 保持API兼容
// 服务识别规则 - 编译期常量,避免运行时分配
@@ -279,50 +276,84 @@ func isDefinitelyNonWeb(serviceInfo *ServiceInfo) bool {
return false
}
// CacheServiceInfo 缓存识别到的服务信息
func CacheServiceInfo(host string, port int, serviceInfo *ServiceInfo) {
cacheKey := net.JoinHostPort(host, strconv.Itoa(port))
serviceCacheMutex.Lock()
defer serviceCacheMutex.Unlock()
serviceCache[cacheKey] = serviceInfo
// SetGlobalState 设置全局 StateRunScan 入口调用,兼容旧代码路径)
func SetGlobalState(state *common.State) {
globalState = state
}
// MarkAsWebService 标记 Web 服务(兼容旧调用)
func resolveState(state *common.State) *common.State {
if state != nil {
return state
}
return globalState
}
// CacheServiceInfoWithState 缓存服务信息到指定 State
func CacheServiceInfoWithState(state *common.State, host string, port int, serviceInfo *ServiceInfo) {
s := resolveState(state)
if s == nil {
return
}
key := net.JoinHostPort(host, strconv.Itoa(port))
s.CacheService(key, serviceInfo)
}
// CacheServiceInfo 兼容旧调用(使用全局 State)
func CacheServiceInfo(host string, port int, serviceInfo *ServiceInfo) {
CacheServiceInfoWithState(nil, host, port, serviceInfo)
}
// MarkAsWebService 标记 Web 服务
func MarkAsWebService(host string, port int, serviceInfo *ServiceInfo) {
CacheServiceInfo(host, port, serviceInfo)
}
// GetCachedServiceInfo 获取缓存的服务信息
func GetCachedServiceInfo(host string, port int) (*ServiceInfo, bool) {
cacheKey := net.JoinHostPort(host, strconv.Itoa(port))
serviceCacheMutex.RLock()
defer serviceCacheMutex.RUnlock()
serviceInfo, exists := serviceCache[cacheKey]
return serviceInfo, exists
}
// GetWebServiceInfo 获取 Web 服务信息(兼容旧调用)
func GetWebServiceInfo(host string, port int) (*ServiceInfo, bool) {
info, exists := GetCachedServiceInfo(host, port)
if !exists {
// GetCachedServiceInfoWithState 从指定 State 获取缓存的服务信息
func GetCachedServiceInfoWithState(state *common.State, host string, port int) (*ServiceInfo, bool) {
s := resolveState(state)
if s == nil {
return nil, false
}
if !IsWebServiceByFingerprint(info) {
key := net.JoinHostPort(host, strconv.Itoa(port))
val, ok := s.GetCachedService(key)
if !ok {
return nil, false
}
info, ok := val.(*ServiceInfo)
return info, ok
}
// GetCachedServiceInfo 兼容旧调用
func GetCachedServiceInfo(host string, port int) (*ServiceInfo, bool) {
return GetCachedServiceInfoWithState(nil, host, port)
}
// GetWebServiceInfo 获取 Web 服务信息
func GetWebServiceInfo(host string, port int) (*ServiceInfo, bool) {
return GetWebServiceInfoWithState(nil, host, port)
}
// GetWebServiceInfoWithState 从指定 State 获取 Web 服务信息
func GetWebServiceInfoWithState(state *common.State, host string, port int) (*ServiceInfo, bool) {
info, exists := GetCachedServiceInfoWithState(state, host, port)
if !exists || !IsWebServiceByFingerprint(info) {
return nil, false
}
return info, true
}
// IsMarkedWebService 检查是否为 Web 服务
// IsMarkedWebService 检查是否为 Web 服务(使用全局 State
func IsMarkedWebService(host string, port int) bool {
_, exists := GetWebServiceInfo(host, port)
return exists
}
// IsMarkedWebServiceWithState 检查是否为 Web 服务(指定 State)
func IsMarkedWebServiceWithState(state *common.State, host string, port int) bool {
_, exists := GetWebServiceInfoWithState(state, host, port)
return exists
}
// ===============================
// Web扫描策略
// ===============================
+2 -6
View File
@@ -440,9 +440,7 @@ func TestCreateTargetFromURL(t *testing.T) {
// TestWebServiceCache 测试Web服务缓存操作
func TestWebServiceCache(t *testing.T) {
// 清空缓存
serviceCacheMutex.Lock()
serviceCache = make(map[string]*ServiceInfo)
serviceCacheMutex.Unlock()
SetGlobalState(common.NewState())
t.Run("存储和读取", func(t *testing.T) {
serviceInfo := &ServiceInfo{
@@ -517,9 +515,7 @@ func TestWebServiceCache(t *testing.T) {
// TestWebServiceCache_Concurrent 测试并发安全性
func TestWebServiceCache_Concurrent(t *testing.T) {
// 清空缓存
serviceCacheMutex.Lock()
serviceCache = make(map[string]*ServiceInfo)
serviceCacheMutex.Unlock()
SetGlobalState(common.NewState())
t.Run("不同key并发写入", func(t *testing.T) {
var wg sync.WaitGroup
+31 -11
View File
@@ -139,26 +139,46 @@ func MakeVarDecl(key, value string) *exprpb.Decl {
}
}
// Evaluate 评估 CEL 表达式
// CelProgCache 缓存编译后的 CEL Program,避免同一 POC 内重复编译
// 在 executePoc 中创建,同一个 POC 的所有规则/参数组合共享
type CelProgCache map[string]cel.Program
// Evaluate 评估 CEL 表达式(无缓存,用于 Set/Sets 求值等低频路径)
func Evaluate(env *cel.Env, expression string, params map[string]interface{}) (ref.Val, error) {
// 空表达式默认返回 true
return EvaluateCached(env, expression, params, nil)
}
// EvaluateCached 评估 CEL 表达式(带编译缓存,用于规则执行热路径)
func EvaluateCached(env *cel.Env, expression string, params map[string]interface{}, cache CelProgCache) (ref.Val, error) {
if expression == "" {
return types.Bool(true), nil
}
// 编译表达式
ast, issues := env.Compile(expression)
if issues.Err() != nil {
return nil, fmt.Errorf("%s: %w", i18n.GetText("webscan_expression_compile_failed"), issues.Err())
var program cel.Program
if cache != nil {
if cached, ok := cache[expression]; ok {
program = cached
}
}
// 创建程序(使用缓存的程序选项)
program, err := env.Program(ast, GetBaseProgramOptions()...)
if err != nil {
return nil, fmt.Errorf("%s: %w", i18n.GetText("webscan_program_create_failed"), err)
if program == nil {
ast, issues := env.Compile(expression)
if issues.Err() != nil {
return nil, fmt.Errorf("%s: %w", i18n.GetText("webscan_expression_compile_failed"), issues.Err())
}
var err error
program, err = env.Program(ast, GetBaseProgramOptions()...)
if err != nil {
return nil, fmt.Errorf("%s: %w", i18n.GetText("webscan_program_create_failed"), err)
}
if cache != nil {
cache[expression] = program
}
}
// 执行评估
result, _, err := program.Eval(params)
if err != nil {
return nil, fmt.Errorf("%s: %w", i18n.GetText("webscan_expression_eval_failed"), err)
+63 -93
View File
@@ -92,59 +92,7 @@ func CheckMultiPoc(req *http.Request, pocs []*Poc, workers int, pocCtx *POCConte
// 仅当通过普通POC规则(非clusterpoc)检测到漏洞时,才创建结果
// 因为clusterpoc已在内部处理了漏洞输出
if isVulnerable && vulName != "" {
// 构造漏洞详细信息
details := make(map[string]interface{}, 6)
details["vulnerability_type"] = task.Poc.Name
details["vulnerability_name"] = vulName
// 添加作者信息(如果有)
if task.Poc.Detail.Author != "" {
details["author"] = task.Poc.Detail.Author
}
// 添加参考链接(如果有)
if len(task.Poc.Detail.Links) != 0 {
details["references"] = task.Poc.Detail.Links
}
// 添加漏洞描述(如果有)
if task.Poc.Detail.Description != "" {
details["description"] = task.Poc.Detail.Description
}
// 创建并保存扫描结果
result := &output.ScanResult{
Time: time.Now(),
Type: output.TypeVuln,
Target: task.Req.URL.String(),
Status: "vulnerable",
Details: details,
}
_ = pocCtx.Session.SaveResult(result)
// 构造控制台输出的日志信息
logMsg := i18n.Tr("webscan_vuln_detail_header",
task.Req.URL,
task.Poc.Name,
vulName)
// 添加作者信息到日志
if task.Poc.Detail.Author != "" {
logMsg += "\n\t" + i18n.Tr("webscan_vuln_author", task.Poc.Detail.Author)
}
// 添加参考链接到日志
if len(task.Poc.Detail.Links) != 0 {
logMsg += "\n\t" + i18n.Tr("webscan_vuln_references", strings.Join(task.Poc.Detail.Links, "\n"))
}
// 添加描述信息到日志
if task.Poc.Detail.Description != "" {
logMsg += "\n\t" + i18n.Tr("webscan_vuln_description", task.Poc.Detail.Description)
}
// 输出成功日志
pocCtx.Session.LogVuln(logMsg)
saveVulnResult(task.Req.URL.String(), task.Poc, vulName, nil, pocCtx.Session)
}
}
}()
@@ -223,17 +171,20 @@ func executePoc(oReq *http.Request, p *Poc, pocCtx *POCContext) (bool, string, e
}
}
// CEL 编译缓存:同一个 POC 的所有规则/参数组合共享
progCache := make(CelProgCache)
// 处理爆破模式
if len(p.Sets) > 0 {
success, err := clusterpoc(oReq, p, variableMap, req, env, pocCtx)
success, err := clusterpoc(oReq, p, variableMap, req, env, pocCtx, progCache)
return success, "", err
}
return executeRules(oReq, p, variableMap, req, env, pocCtx.Session)
return executeRules(oReq, p, variableMap, req, env, pocCtx.Session, progCache)
}
// executeRules 执行POC规则并返回结果
func executeRules(oReq *http.Request, p *Poc, variableMap map[string]interface{}, req *Request, env *cel.Env, session *common.ScanSession) (bool, string, error) {
func executeRules(oReq *http.Request, p *Poc, variableMap map[string]interface{}, req *Request, env *cel.Env, session *common.ScanSession, progCache CelProgCache) (bool, string, error) {
// 处理单个规则的函数
executeRule := func(rule Rules) (bool, error) {
Headers := cloneMap(rule.Headers)
@@ -301,8 +252,8 @@ func executeRules(oReq *http.Request, p *Poc, variableMap map[string]interface{}
}
}
// 执行表达式
out, err := Evaluate(env, rule.Expression, variableMap)
// 执行表达式(使用编译缓存)
out, err := EvaluateCached(env, rule.Expression, variableMap, progCache)
if err != nil {
return false, err
}
@@ -452,7 +403,7 @@ func newReverse(dnsLog bool) *Reverse {
}
// clusterpoc 执行集群POC检测,支持批量参数组合测试
func clusterpoc(oReq *http.Request, p *Poc, variableMap map[string]interface{}, req *Request, env *cel.Env, pocCtx *POCContext) (success bool, err error) {
func clusterpoc(oReq *http.Request, p *Poc, variableMap map[string]interface{}, req *Request, env *cel.Env, pocCtx *POCContext, progCache CelProgCache) (success bool, err error) {
var strMap StrMap // 存储成功的参数组合
var shiroKeyCount int // shiro key测试计数
@@ -461,7 +412,7 @@ func clusterpoc(oReq *http.Request, p *Poc, variableMap map[string]interface{},
// 检查是否需要进行参数Fuzz测试
if !isFuzz(rule, p.Sets) {
// 不需要Fuzz,直接发送请求
success, err = clustersend(oReq, variableMap, req, env, rule, pocCtx.Session)
success, err = clustersend(oReq, variableMap, req, env, rule, pocCtx.Session, progCache)
if err != nil {
return false, err
}
@@ -527,7 +478,7 @@ func clusterpoc(oReq *http.Request, p *Poc, variableMap map[string]interface{},
ruleHash[ruleMD5] = struct{}{}
// 发送请求并处理结果
success, err = clustersend(oReq, variableMap, req, env, currentRule, pocCtx.Session)
success, err = clustersend(oReq, variableMap, req, env, currentRule, pocCtx.Session, progCache)
if err != nil {
return false, err
}
@@ -651,59 +602,78 @@ func getRuleHash(rule *Rules) string {
return fmt.Sprintf("%x", h.Sum(nil))
}
// recordVulnerabilityResult 记录漏洞检测结果
func recordVulnerabilityResult(targetURL string, pocDef *Poc, params StrMap, skipSave bool, session *common.ScanSession) {
// 构造详细信息
details := make(map[string]interface{})
// buildVulnDetails 构造统一的漏洞详情 map(消除 CheckMultiPoc 和 recordVulnerabilityResult 的重复逻辑)
func buildVulnDetails(pocDef *Poc, vulName string, params StrMap) map[string]interface{} {
details := make(map[string]interface{}, 6)
details["vulnerability_type"] = pocDef.Name
details["vulnerability_name"] = pocDef.Name // 使用POC名称作为漏洞名称
// 添加作者信息(如果有)
details["vulnerability_name"] = vulName
if pocDef.Detail.Author != "" {
details["author"] = pocDef.Detail.Author
}
// 添加参考链接(如果有)
if len(pocDef.Detail.Links) != 0 {
details["references"] = pocDef.Detail.Links
}
// 添加漏洞描述(如果有)
if pocDef.Detail.Description != "" {
details["description"] = pocDef.Detail.Description
}
// 添加参数信息(如果有)
if len(params) > 0 {
paramMap := make(map[string]string)
paramMap := make(map[string]string, len(params))
for _, item := range params {
paramMap[item.Key] = item.Value
}
details["parameters"] = paramMap
}
return details
}
// 保存漏洞结果(除非明确指示跳过)
// buildVulnLogMsg 构造统一的漏洞日志消息
func buildVulnLogMsg(targetURL string, pocDef *Poc, vulName string, params StrMap) string {
var logMsg string
if pocDef.Name == "poc-yaml-backup-file" || pocDef.Name == "poc-yaml-sql-file" {
logMsg = i18n.Tr("webscan_vuln_detected", targetURL, pocDef.Name)
} else if len(params) > 0 {
logMsg = i18n.Tr("webscan_vuln_detected_params", targetURL, pocDef.Name, params)
} else {
logMsg = i18n.Tr("webscan_vuln_detail_header", targetURL, pocDef.Name, vulName)
if pocDef.Detail.Author != "" {
logMsg += "\n\t" + i18n.Tr("webscan_vuln_author", pocDef.Detail.Author)
}
if len(pocDef.Detail.Links) != 0 {
logMsg += "\n\t" + i18n.Tr("webscan_vuln_references", strings.Join(pocDef.Detail.Links, "\n"))
}
if pocDef.Detail.Description != "" {
logMsg += "\n\t" + i18n.Tr("webscan_vuln_description", pocDef.Detail.Description)
}
}
return logMsg
}
// saveVulnResult 统一的漏洞结果保存 + 日志输出
func saveVulnResult(targetURL string, pocDef *Poc, vulName string, params StrMap, session *common.ScanSession) {
details := buildVulnDetails(pocDef, vulName, params)
_ = session.SaveResult(&output.ScanResult{
Time: time.Now(),
Type: output.TypeVuln,
Target: targetURL,
Status: "vulnerable",
Details: details,
})
session.LogVuln(buildVulnLogMsg(targetURL, pocDef, vulName, params))
}
// recordVulnerabilityResult 记录漏洞检测结果(clusterpoc 路径)
func recordVulnerabilityResult(targetURL string, pocDef *Poc, params StrMap, skipSave bool, session *common.ScanSession) {
if !skipSave {
result := &output.ScanResult{
details := buildVulnDetails(pocDef, pocDef.Name, params)
_ = session.SaveResult(&output.ScanResult{
Time: time.Now(),
Type: output.TypeVuln,
Target: targetURL,
Status: "vulnerable",
Details: details,
}
_ = session.SaveResult(result)
})
}
// 生成日志消息
var logMsg string
if pocDef.Name == "poc-yaml-backup-file" || pocDef.Name == "poc-yaml-sql-file" {
logMsg = i18n.Tr("webscan_vuln_detected", targetURL, pocDef.Name)
} else {
logMsg = i18n.Tr("webscan_vuln_detected_params", targetURL, pocDef.Name, params)
}
// 输出成功日志
session.LogVuln(logMsg)
session.LogVuln(buildVulnLogMsg(targetURL, pocDef, pocDef.Name, params))
}
// isFuzz 检查规则是否包含需要Fuzz测试的参数
@@ -773,7 +743,7 @@ func MakeData(base [][]string, nextData []string) [][]string {
}
// clustersend 执行单个规则的HTTP请求和响应检测
func clustersend(oReq *http.Request, variableMap map[string]interface{}, req *Request, env *cel.Env, rule Rules, session *common.ScanSession) (bool, error) {
func clustersend(oReq *http.Request, variableMap map[string]interface{}, req *Request, env *cel.Env, rule Rules, session *common.ScanSession, progCache CelProgCache) (bool, error) {
// 替换请求中的变量
for varName, varValue := range variableMap {
// 跳过map类型的变量
@@ -844,8 +814,8 @@ func clustersend(oReq *http.Request, variableMap map[string]interface{}, req *Re
}
}
// 执行CEL表达式
out, err := Evaluate(env, rule.Expression, variableMap)
// 执行CEL表达式(使用编译缓存)
out, err := EvaluateCached(env, rule.Expression, variableMap, progCache)
if err != nil {
if strings.Contains(err.Error(), "Syntax error") {
common.LogError(i18n.Tr("webscan_cel_syntax_error", rule.Expression, err))
+57 -64
View File
@@ -13,7 +13,6 @@ import (
"strconv"
"strings"
"sync"
"sync/atomic"
"time"
"github.com/shadow1ng/fscan/common"
@@ -42,30 +41,22 @@ var (
//go:embed pocs
var pocsFS embed.FS
var (
pocMu sync.Mutex
pocLoaded atomic.Bool
allPocs []*lib.Poc
cachedPocPath string
)
// pocStore 按 PocPath 缓存已加载的 POC 集合,支持多 session 使用不同 POC 路径
type pocStore struct {
mu sync.Mutex
cache map[string][]*lib.Poc // key: pocPath(空字符串表示内嵌 POC)
}
var globalPocStore = &pocStore{cache: make(map[string][]*lib.Poc)}
// WebScan 执行Web漏洞扫描
func WebScan(ctx context.Context, info *common.HostInfo, cfg *common.Config, session *common.ScanSession) {
// 初始化POC配置(用于CEL回调函数)
lib.InitPOCConfig(cfg.DNSLog)
// 加载POCDCLP: 快速路径无锁,慢路径互斥保护
if !pocLoaded.Load() {
pocMu.Lock()
if !pocLoaded.Load() {
cachedPocPath = cfg.POC.PocPath
initPocs()
if len(allPocs) > 0 {
pocLoaded.Store(true)
}
}
pocMu.Unlock()
}
// 加载POC按 PocPath 缓存,不同路径独立加载
pocs := globalPocStore.getOrLoad(cfg.POC.PocPath)
// 验证输入
if info == nil {
@@ -73,7 +64,7 @@ func WebScan(ctx context.Context, info *common.HostInfo, cfg *common.Config, ses
return
}
if len(allPocs) == 0 {
if len(pocs) == 0 {
session.LogError(i18n.GetText("poc_load_failed"))
return
}
@@ -94,14 +85,11 @@ func WebScan(ctx context.Context, info *common.HostInfo, cfg *common.Config, ses
// 根据扫描策略执行POC
if cfg.POC.PocName == "" && len(info.Info) == 0 {
// 执行所有POC
executePOCs(ctx, config.PocInfo{Target: target}, cfg, session)
executePOCs(ctx, config.PocInfo{Target: target}, cfg, session, pocs)
} else if len(info.Info) > 0 {
// 基于指纹信息执行POC
scanByFingerprints(ctx, target, info.Info, cfg, session)
scanByFingerprints(ctx, target, info.Info, cfg, session, pocs)
} else if cfg.POC.PocName != "" {
// 基于指定POC名称执行
executePOCs(ctx, config.PocInfo{Target: target, PocName: cfg.POC.PocName}, cfg, session)
executePOCs(ctx, config.PocInfo{Target: target, PocName: cfg.POC.PocName}, cfg, session, pocs)
}
}
@@ -178,8 +166,24 @@ func hasMalformedWebURLPort(host string) bool {
return strings.Contains(host, ":")
}
// getOrLoad 获取或加载指定路径的 POC 集合
func (s *pocStore) getOrLoad(pocPath string) []*lib.Poc {
s.mu.Lock()
defer s.mu.Unlock()
if pocs, ok := s.cache[pocPath]; ok {
return pocs
}
pocs := loadPocs(pocPath)
if len(pocs) > 0 {
s.cache[pocPath] = pocs
}
return pocs
}
// scanByFingerprints 根据指纹执行POC
func scanByFingerprints(ctx context.Context, target string, fingerprints []string, cfg *common.Config, session *common.ScanSession) {
func scanByFingerprints(ctx context.Context, target string, fingerprints []string, cfg *common.Config, session *common.ScanSession, pocs []*lib.Poc) {
for _, fingerprint := range fingerprints {
if fingerprint == "" {
continue
@@ -190,12 +194,12 @@ func scanByFingerprints(ctx context.Context, target string, fingerprints []strin
continue
}
executePOCs(ctx, config.PocInfo{Target: target, PocName: pocName}, cfg, session)
executePOCs(ctx, config.PocInfo{Target: target, PocName: pocName}, cfg, session, pocs)
}
}
// executePOCs 执行POC检测
func executePOCs(ctx context.Context, pocInfo config.PocInfo, cfg *common.Config, session *common.ScanSession) {
func executePOCs(ctx context.Context, pocInfo config.PocInfo, cfg *common.Config, session *common.ScanSession, pocs []*lib.Poc) {
// 验证目标
if pocInfo.Target == "" {
session.LogError(ErrEmptyTarget.Error())
@@ -222,7 +226,7 @@ func executePOCs(ctx context.Context, pocInfo config.PocInfo, cfg *common.Config
}
// 筛选POC
matchedPocs := filterPocs(pocInfo.PocName)
matchedPocs := filterPocs(pocInfo.PocName, pocs)
if len(matchedPocs) == 0 {
session.LogDebug(fmt.Sprintf("%v: %s", ErrPocNotFound, pocInfo.PocName))
return
@@ -257,28 +261,22 @@ func createBaseRequest(ctx context.Context, target string, cfg *common.Config) (
return req, nil
}
// initPocs 初始化并加载POC
// 使用cachedPocPath包级变量
func initPocs() {
// 预分配容量避免频繁扩容,典型POC数量在100-500之间
allPocs = make([]*lib.Poc, 0, 256)
if cachedPocPath == "" {
loadEmbeddedPocs()
} else {
loadExternalPocs(cachedPocPath)
// loadPocs 加载指定路径的 POC(空路径表示内嵌 POC)
func loadPocs(pocPath string) []*lib.Poc {
if pocPath == "" {
return loadEmbeddedPocs()
}
return loadExternalPocs(pocPath)
}
// loadEmbeddedPocs 加载内置POC
func loadEmbeddedPocs() {
func loadEmbeddedPocs() []*lib.Poc {
entries, err := pocsFS.ReadDir("pocs")
if err != nil {
common.LogError(i18n.Tr("webscan_builtin_poc_failed", err))
return
return nil
}
// 收集所有POC文件
var pocFiles []string
for _, entry := range entries {
if isPocFile(entry.Name()) {
@@ -286,24 +284,21 @@ func loadEmbeddedPocs() {
}
}
// 并发加载POC文件
loadPocsConcurrently(pocFiles, true, "")
return loadPocsConcurrently(pocFiles, true, "")
}
// loadExternalPocs 从外部路径加载POC
func loadExternalPocs(pocPath string) {
func loadExternalPocs(pocPath string) []*lib.Poc {
if !directoryExists(pocPath) {
common.LogError(i18n.Tr("webscan_poc_dir_not_exist", pocPath))
return
return nil
}
// 收集所有POC文件路径
var pocFiles []string
err := filepath.Walk(pocPath, func(path string, info os.FileInfo, err error) error {
if err != nil || info == nil || info.IsDir() {
return nil
}
if isPocFile(info.Name()) {
pocFiles = append(pocFiles, path)
}
@@ -312,18 +307,17 @@ func loadExternalPocs(pocPath string) {
if err != nil {
common.LogError(i18n.Tr("webscan_poc_dir_walk_failed", err))
return
return nil
}
// 并发加载POC文件
loadPocsConcurrently(pocFiles, false, pocPath)
return loadPocsConcurrently(pocFiles, false, pocPath)
}
// loadPocsConcurrently 并发加载POC文件channel 收集,无锁竞争)
func loadPocsConcurrently(pocFiles []string, isEmbedded bool, pocPath string) {
// loadPocsConcurrently 并发加载POC文件,返回加载结果
func loadPocsConcurrently(pocFiles []string, isEmbedded bool, pocPath string) []*lib.Poc {
pocCount := len(pocFiles)
if pocCount == 0 {
return
return nil
}
var wg sync.WaitGroup
@@ -359,14 +353,14 @@ func loadPocsConcurrently(pocFiles []string, isEmbedded bool, pocPath string) {
close(results)
}()
var successCount int
pocs := make([]*lib.Poc, 0, pocCount)
for poc := range results {
allPocs = append(allPocs, poc)
successCount++
pocs = append(pocs, poc)
}
failCount := pocCount - successCount
common.LogInfo(i18n.Tr("poc_load_complete", pocCount, successCount, failCount))
failCount := pocCount - len(pocs)
common.LogInfo(i18n.Tr("poc_load_complete", pocCount, len(pocs), failCount))
return pocs
}
// directoryExists 检查目录是否存在
@@ -382,16 +376,15 @@ func isPocFile(filename string) bool {
}
// filterPocs 根据POC名称筛选
func filterPocs(pocName string) []*lib.Poc {
func filterPocs(pocName string, pocs []*lib.Poc) []*lib.Poc {
if pocName == "" {
return allPocs
return pocs
}
// 转换为小写以进行不区分大小写的匹配
searchName := strings.ToLower(pocName)
var matchedPocs []*lib.Poc
for _, poc := range allPocs {
for _, poc := range pocs {
if poc != nil && strings.Contains(strings.ToLower(poc.Name), searchName) {
matchedPocs = append(matchedPocs, poc)
}
+8 -19
View File
@@ -338,11 +338,7 @@ func TestFilterPocs(t *testing.T) {
{Name: "Nginx-Path-Traversal"},
}
// 保存原始 allPocs 并在测试后恢复
origAllPocs := allPocs
defer func() { allPocs = origAllPocs }()
allPocs = testPocs
// 直接使用 testPocs 作为输入
tests := []struct {
name string
@@ -408,7 +404,7 @@ func TestFilterPocs(t *testing.T) {
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := filterPocs(tt.pocName)
result := filterPocs(tt.pocName, testPocs)
if len(result) != tt.expectedCount {
t.Errorf("filterPocs(%q) returned %d pocs, want %d", tt.pocName, len(result), tt.expectedCount)
@@ -449,19 +445,14 @@ func TestFilterPocs(t *testing.T) {
}
func TestFilterPocsNilSafety(t *testing.T) {
// 测试全是 nil 的情况
origAllPocs := allPocs
defer func() { allPocs = origAllPocs }()
nilPocs := []*lib.Poc{nil, nil, nil}
allPocs = []*lib.Poc{nil, nil, nil}
result := filterPocs("test")
result := filterPocs("test", nilPocs)
if len(result) != 0 {
t.Errorf("filterPocs with all nil should return empty slice, got %d items", len(result))
}
// 空 pocName 返回所有 POCs(包括 nil
result = filterPocs("")
result = filterPocs("", nilPocs)
if len(result) != 3 {
t.Errorf("filterPocs with empty name should return all pocs (including nil), got %d items, want 3", len(result))
}
@@ -497,12 +488,10 @@ func TestCreateBaseRequestHeaders(t *testing.T) {
func TestExecutePOCsEarlyReturns(t *testing.T) {
cfg := common.NewConfig()
session := common.NewScanSession(cfg, common.NewState(), &common.FlagVars{})
previous := allPocs
allPocs = nil
t.Cleanup(func() { allPocs = previous })
executePOCs(context.Background(), config.PocInfo{}, cfg, session)
executePOCs(context.Background(), config.PocInfo{Target: "http://example.com", PocName: "missing"}, cfg, session)
var emptyPocs []*lib.Poc
executePOCs(context.Background(), config.PocInfo{}, cfg, session, emptyPocs)
executePOCs(context.Background(), config.PocInfo{Target: "http://example.com", PocName: "missing"}, cfg, session, emptyPocs)
}
func TestDirectoryExists(t *testing.T) {