fix: addAddition deadlock

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" {

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
 	}

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) {

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()