diff --git a/core/format.go b/core/format.go index d753dd5..6d06fff 100644 --- a/core/format.go +++ b/core/format.go @@ -127,7 +127,10 @@ func Format(opts Option) { return } - // Group by host:port instead of just host + // Group by host+path for deduplication. + // Prefer the HTTP Host header when present so that + // different virtual hosts on the same IP:port do not + // get merged together. group := make(map[string]map[string]*baseline.Baseline) for _, line := range bytes.Split(bytes.TrimSpace(content), []byte("\n")) { var result baseline.Baseline @@ -141,26 +144,33 @@ func Format(opts Option) { continue } - // Use host:port as the key - hostPort := result.Url.Host - if result.Url.Port() == "" { - if result.Url.Scheme == "https" { - hostPort += ":443" - } else if result.Url.Scheme == "http" { - hostPort += ":80" + // Determine host key: prefer Host header when available. + hostKey := result.Host + if hostKey == "" { + hostKey = result.Url.Host + // Normalize default ports when URL host does not include one. + if result.Url.Port() == "" { + if result.Url.Scheme == "https" { + hostKey += ":443" + } else if result.Url.Scheme == "http" { + hostKey += ":80" + } } } - - if _, exists := group[hostPort]; !exists { - group[hostPort] = make(map[string]*baseline.Baseline) + if hostKey == "" { + continue } - group[hostPort][result.Path] = &result + + if _, exists := group[hostKey]; !exists { + group[hostKey] = make(map[string]*baseline.Baseline) + } + group[hostKey][result.Path] = &result } // Default to tree mode if not specified outputProbe := opts.OutputProbe if outputProbe == "" { - outputProbe = "tree" + outputProbe = "full" } if outputProbe == "tree" { diff --git a/core/pool/brutepool.go b/core/pool/brutepool.go index 14a88cd..5c1d955 100644 --- a/core/pool/brutepool.go +++ b/core/pool/brutepool.go @@ -173,22 +173,18 @@ func (pool *BrutePool) Init() error { func (pool *BrutePool) Run(offset, limit int) { pool.Worder.Run() if pool.Active { - pool.wg.Add(1) go pool.doActive() } if pool.Bak { - pool.wg.Add(1) go pool.doBak() } if pool.Fuzzuli { - pool.wg.Add(1) go pool.doFuzzuli() } if pool.Common { - pool.wg.Add(1) go pool.doCommonFile() } @@ -430,7 +426,6 @@ func (pool *BrutePool) NoScopeInvoke(v interface{}) { bl.ReqDepth = unit.depth bl.Collect() bl.CollectURL() - pool.wg.Add(1) pool.doScopeCrawl(bl) pool.putToOutput(bl) } @@ -733,9 +728,7 @@ func (pool *BrutePool) doRedirect(bl *baseline.Baseline, depth int) { // return // 不同域名的重定向不处理 //} reURL := pkg.FormatURL(bl.Url.Path, bl.RedirectURL) - pool.wg.Add(1) go func() { - defer pool.wg.Done() pool.addAddition(&Unit{ path: reURL, parent: bl.Number, @@ -758,11 +751,10 @@ func (pool *BrutePool) doCrawl(bl *baseline.Baseline) { return } - pool.wg.Add(2) + pool.wg.Add(1) pool.doScopeCrawl(bl) go func() { - defer pool.wg.Done() for _, u := range bl.URLs { if u = pkg.FormatURL(bl.Url.Path, u); u == "" { continue @@ -782,12 +774,10 @@ func (pool *BrutePool) doCrawl(bl *baseline.Baseline) { func (pool *BrutePool) doScopeCrawl(bl *baseline.Baseline) { if bl.ReqDepth >= pool.MaxCrawlDepth { - pool.wg.Done() return } go func() { - defer pool.wg.Done() for _, u := range bl.URLs { if strings.HasPrefix(u, "http") { if v, _ := url.Parse(u); v == nil || !pkg.MatchWithGlobs(v.Host, pool.Scope) { @@ -812,7 +802,6 @@ func (pool *BrutePool) doScopeCrawl(bl *baseline.Baseline) { } func (pool *BrutePool) doFuzzuli() { - defer pool.wg.Done() if pool.Mod == HostSpray { return } @@ -825,7 +814,6 @@ func (pool *BrutePool) doFuzzuli() { } func (pool *BrutePool) doBak() { - defer pool.wg.Done() if pool.Mod == HostSpray { return } @@ -839,19 +827,16 @@ func (pool *BrutePool) doBak() { } func (pool *BrutePool) doAppend(bl *baseline.Baseline) { - pool.wg.Add(2) pool.doAppendWords(bl) pool.doAppendRule(bl) } func (pool *BrutePool) doAppendRule(bl *baseline.Baseline) { if pool.AppendRule == nil || bl.Source == parsers.AppendRuleSource || bl.ReqDepth >= pool.MaxAppendDepth { - pool.wg.Done() return } go func() { - defer pool.wg.Done() for u := range rule.RunAsStream(pool.AppendRule.Expressions, path.Base(bl.Path)) { pool.addAddition(&Unit{ path: pkg.Dir(bl.Url.Path) + u, @@ -868,13 +853,10 @@ func (pool *BrutePool) doAppendRule(bl *baseline.Baseline) { func (pool *BrutePool) doAppendWords(bl *baseline.Baseline) { if pool.AppendWords == nil || bl.Source == parsers.AppendSource || bl.Source == parsers.RuleSource || bl.ReqDepth >= pool.MaxAppendDepth { // 防止自身递归 - pool.wg.Done() return } go func() { - defer pool.wg.Done() - for u := range NewBruteWords(pool.Config, pool.AppendWords).Output { pool.addAddition(&Unit{ path: pkg.SafePath(bl.Path, u), @@ -889,7 +871,6 @@ func (pool *BrutePool) doAppendWords(bl *baseline.Baseline) { } func (pool *BrutePool) doActive() { - defer pool.wg.Done() if pool.Mod == HostSpray { return } @@ -902,7 +883,6 @@ func (pool *BrutePool) doActive() { } func (pool *BrutePool) doCommonFile() { - defer pool.wg.Done() if pool.Mod == HostSpray { return } diff --git a/core/pool/pool.go b/core/pool/pool.go index 0cf4ec2..523ad9c 100644 --- a/core/pool/pool.go +++ b/core/pool/pool.go @@ -44,16 +44,21 @@ func (pool *BasePool) doRetry(bl *baseline.Baseline) { } func (pool *BasePool) addAddition(u *Unit) { - // 强行屏蔽报错, 防止goroutine泄露 - defer func() { - if err := recover(); err != nil { - } - }() pool.wg.Add(1) select { case pool.additionCh <- u: - case <-pool.ctx.Done(): + default: } + pool.wg.Add(1) + // 强行屏蔽报错, 防止goroutine泄露 + go func() { + select { + case pool.additionCh <- u: + pool.wg.Done() + case <-pool.ctx.Done(): + pool.wg.Done() + } + }() } func (pool *BasePool) putToOutput(bl *baseline.Baseline) { diff --git a/spray.go b/spray.go index afbfe8e..b1a4c6b 100644 --- a/spray.go +++ b/spray.go @@ -5,7 +5,6 @@ import ( "github.com/chainreactors/spray/cmd" "github.com/gookit/config/v2" "github.com/gookit/config/v2/yaml" - //"net/http" //_ "net/http/pprof" ) @@ -18,9 +17,10 @@ func init() { } func main() { - //// 启动 pprof HTTP 服务器 + // 启动 pprof HTTP 服务器，用于调试 goroutine 泄露等问题 //go func() { - // http.ListenAndServe("localhost:6060", nil) + // // 仅监听本地 6060 端口，避免对外暴露 + // _ = http.ListenAndServe("127.0.0.1:6060", nil) //}() cmd.Spray()