merge v1.2.6

.github/workflows/ci.yml

@@ -0,0 +1,34 @@
+name: CI
+
+on:
+  pull_request:
+    branches: [ master ]
+  push:
+    branches: [ master ]
+
+jobs:
+  test:
+    strategy:
+      matrix:
+        os: [ubuntu-latest, windows-latest]
+    runs-on: ${{ matrix.os }}
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+        with:
+          submodules: recursive
+
+      - name: Set up Go
+        uses: actions/setup-go@v5
+        with:
+          go-version: "1.20"
+
+      - name: Build
+        run: go build ./...
+
+      - name: Test (race)
+        run: go test ./core/pool/ ./core/baseline/ ./core/ihttp/ ./pkg/ -v -race -count=1 -timeout=120s
+
+      - name: Vet
+        run: go vet ./core/... ./pkg/... ./cmd/...
+        continue-on-error: true

core/baseline/baseline.go

@@ -3,6 +3,7 @@
 import (
 	"bytes"
 	"github.com/chainreactors/fingers/common"
+	"github.com/chainreactors/logs"
 	"github.com/chainreactors/parsers"
 	"github.com/chainreactors/spray/core/ihttp"
 	"github.com/chainreactors/spray/pkg"
@@ -53,15 +54,14 @@
 	bl.Raw = append(bl.Header, bl.Body...)
 	bl.Response, err = pkg.ParseRawResponse(bl.Raw)
 	if err != nil {
-		bl.IsValid = false
-		bl.Reason = pkg.ErrResponseError.Error()
-		bl.ErrString = err.Error()
-		return bl
+		// raw 重解析失败不影响 baseline 有效性，live response 已提供所有需要的数据
+		logs.Log.Debugf("ParseRawResponse failed for %s: %s", u, err.Error())
 	}
-	if r := bl.Response.Header.Get("Location"); r != "" {
+	// 始终从 live response 读取 Location
+	if r := resp.GetHeader("Location"); r != "" {
 		bl.RedirectURL = r
 	} else {
-		bl.RedirectURL = bl.Response.Header.Get("location")
+		bl.RedirectURL = resp.GetHeader("location")
 	}
 
 	bl.Dir = bl.IsDir()
@@ -235,7 +235,7 @@
 			return -1
 		}
 	}
 	return -1
 }
 
 func (bl *Baseline) ProbeOutput(format []string) string {

core/pool/brutepool.go

@@ -52,10 +52,11 @@ func NewBrutePool(ctx context.Context, config *Config) (*BrutePool, error) {
 				Timeout:     config.Timeout,
 				ProxyClient: config.ProxyClient,
 			}),
-			additionCh: make(chan *Unit, config.Thread*10),
-			closeCh:    make(chan struct{}),
-			processCh:  make(chan *baseline.Baseline, config.Thread*2),
-			wg:         &sync.WaitGroup{},
+			additionCh:  make(chan *Unit, config.Thread*10),
+			closeCh:     make(chan struct{}),
+			processCh:   make(chan *baseline.Baseline, config.Thread*2),
+			wg:          &sync.WaitGroup{},
+			handlerDone: make(chan struct{}),
 		},
 		base:        u.Scheme + "://" + u.Host,
 		isDir:       strings.HasSuffix(u.Path, "/"),
@@ -105,7 +106,6 @@ type BrutePool struct {
 	urls        sync.Map
 	scopeurls   map[string]struct{}
 	uniques     map[uint16]struct{}
-	analyzeDone bool
 	limiter     *rate.Limiter
 	locker      sync.Mutex
 	scopeLocker sync.Mutex
@@ -197,6 +197,11 @@ func (pool *BrutePool) Run(offset, limit int) {
 				close(pool.closeCh)
 				return
 			}
+			select {
+			case <-pool.ctx.Done():
+				return
+			default:
+			}
 			time.Sleep(100 * time.Millisecond)
 		}
 	}()
@@ -376,7 +381,7 @@ func (pool *BrutePool) Invoke(v interface{}) {
 
 	case parsers.WordSource:
 		// 异步进行性能消耗较大的深度对比
-		pool.processCh <- bl
+		pool.sendProcess(bl)
 		if int(pool.Statistor.ReqTotal)%pool.CheckPeriod == 0 {
 			// 间歇插入check waf的探针
 			pool.doCheck()
@@ -388,9 +393,9 @@ func (pool *BrutePool) Invoke(v interface{}) {
 		pool.Bar.Done()
 	case parsers.RedirectSource:
 		bl.FrontURL = unit.frontUrl
-		pool.processCh <- bl
+		pool.sendProcess(bl)
 	default:
-		pool.processCh <- bl
+		pool.sendProcess(bl)
 	}
 }
 
@@ -432,6 +437,7 @@ func (pool *BrutePool) NoScopeInvoke(v interface{}) {
 }
 
 func (pool *BrutePool) Handler() {
+	defer close(pool.handlerDone)
 	for bl := range pool.processCh {
 		if bl.IsValid {
 			pool.addFuzzyBaseline(bl)
@@ -516,8 +522,6 @@ func (pool *BrutePool) Handler() {
 		}
 		pool.wg.Done()
 	}
-
-	pool.analyzeDone = true
 }
 
 func (pool *BrutePool) checkRedirect(redirectURL string) bool {
@@ -674,15 +678,12 @@ func (pool *BrutePool) fallback() {
 }
 
 func (pool *BrutePool) Close() {
-	for pool.analyzeDone {
-		// 等待缓存的待处理任务完成
-		time.Sleep(time.Duration(100) * time.Millisecond)
-	}
-	close(pool.additionCh) // 关闭addition管道
-	//close(pool.checkCh)    // 关闭check管道
-	pool.Statistor.EndTime = time.Now().Unix()
+	pool.Cancel()
 	pool.reqPool.Release()
 	pool.scopePool.Release()
+	close(pool.processCh)
+	<-pool.handlerDone
+	pool.Statistor.EndTime = time.Now().Unix()
 }
 
 func (pool *BrutePool) safePath(u string) string {
@@ -713,9 +714,15 @@ func (pool *BrutePool) doCheck() {
 	}
 
 	if pool.Mod == HostSpray {
-		pool.checkCh <- struct{}{}
+		select {
+		case pool.checkCh <- struct{}{}:
+		case <-pool.ctx.Done():
+		}
 	} else if pool.Mod == PathSpray {
-		pool.checkCh <- struct{}{}
+		select {
+		case pool.checkCh <- struct{}{}:
+		case <-pool.ctx.Done():
+		}
 	}
 }
 
@@ -755,6 +762,7 @@ func (pool *BrutePool) doCrawl(bl *baseline.Baseline) {
 	pool.doScopeCrawl(bl)
 
 	go func() {
+		defer pool.wg.Done()
 		for _, u := range bl.URLs {
 			if u = pkg.FormatURL(bl.Url.Path, u); u == "" {
 				continue

core/pool/checkpool.go

@@ -34,6 +34,7 @@ func NewCheckPool(ctx context.Context, config *Config) (*CheckPool, error) {
 			additionCh: make(chan *Unit, config.Thread*10),
 			closeCh:    make(chan struct{}),
 			processCh:  make(chan *baseline.Baseline, config.Thread*2),
+			handlerDone: make(chan struct{}),
 		},
 	}
 	pool.Request.Headers.Set("Connection", "close")
@@ -105,6 +106,9 @@ Loop:
 	pool.Close()
 }
 func (pool *CheckPool) Close() {
+	pool.Cancel()
+	close(pool.processCh)
+	<-pool.handlerDone
 	pool.Bar.Close()
 	pool.Pool.Release()
 }
@@ -139,7 +143,7 @@ func (pool *CheckPool) Invoke(v interface{}) {
 				ReqDepth:  unit.depth,
 			},
 		}
-		pool.processCh <- bl
+		pool.sendProcess(bl)
 		return
 	}
 	start := time.Now()
@@ -172,10 +176,11 @@ func (pool *CheckPool) Invoke(v interface{}) {
 	if bl.RedirectURL != "" {
 		pool.doRedirect(bl, bl.ReqDepth)
 	}
-	pool.processCh <- bl
+	pool.sendProcess(bl)
 }
 
 func (pool *CheckPool) Handler() {
+	defer close(pool.handlerDone)
 	for bl := range pool.processCh {
 		if bl.IsValid {
 			params := map[string]interface{}{

core/pool/pool.go

@@ -26,6 +26,7 @@ type BasePool struct {
 	additionCh  chan *Unit
 	closeCh     chan struct{}
 	wg          *sync.WaitGroup
+	handlerDone chan struct{}
 	isFallback  atomic.Bool
 }
 
@@ -44,18 +45,21 @@ func (pool *BasePool) doRetry(bl *baseline.Baseline) {
 }
 
 func (pool *BasePool) addAddition(u *Unit) {
+	if pool.ctx.Err() != nil {
+		return
+	}
 	pool.wg.Add(1)
 	select {
 	case pool.additionCh <- u:
-	default:
-		// 强行屏蔽报错, 防止goroutine泄露
-		go func() {
-			select {
-			case pool.additionCh <- u:
-			case <-pool.ctx.Done():
-				pool.wg.Done()
-			}
-		}()
+	case <-pool.ctx.Done():
+		pool.wg.Done()
+	}
+}
+
+func (pool *BasePool) sendProcess(bl *baseline.Baseline) {
+	select {
+	case pool.processCh <- bl:
+	case <-pool.ctx.Done():
 	}
 }
 
@@ -64,11 +68,19 @@ func (pool *BasePool) putToOutput(bl *baseline.Baseline) {
 		bl.Collect()
 	}
 	pool.Outwg.Add(1)
-	pool.OutputCh <- bl
+	select {
+	case pool.OutputCh <- bl:
+	case <-pool.ctx.Done():
+		pool.Outwg.Done()
+	}
 }
 
 func (pool *BasePool) putToFuzzy(bl *baseline.Baseline) {
 	pool.Outwg.Add(1)
 	bl.IsFuzzy = true
-	pool.FuzzyCh <- bl
+	select {
+	case pool.FuzzyCh <- bl:
+	case <-pool.ctx.Done():
+		pool.Outwg.Done()
+	}
 }