diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 075e12c..780b7a2 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -31,7 +31,7 @@ jobs: - name: Set up Go uses: actions/setup-go@v5 with: - go-version: '1.20.x' + go-version: '1.24.x' cache: true - name: Download dependencies @@ -42,7 +42,7 @@ jobs: - name: Test (race detector) if: runner.os == 'Linux' - run: go test -race -short -count=1 -timeout 300s ./... + run: go test -race -count=1 -timeout 300s ./... - name: Vet run: go vet ./... diff --git a/core/e2e_test.go b/core/e2e_test.go index 6de511b..0c89bef 100644 --- a/core/e2e_test.go +++ b/core/e2e_test.go @@ -74,9 +74,6 @@ func TestE2E_NormalBruteScan(t *testing.T) { // --------------------------------------------------------------------------- func TestE2E_ContextCancellation(t *testing.T) { - if testing.Short() { - t.Skip("skipping context cancellation test in short mode") - } var reqCount int64 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { atomic.AddInt64(&reqCount, 1) diff --git a/core/pool/brutepool.go b/core/pool/brutepool.go index f0d5347..061cdf8 100644 --- a/core/pool/brutepool.go +++ b/core/pool/brutepool.go @@ -19,7 +19,6 @@ import ( "path" "strings" "sync" - "sync/atomic" "time" ) @@ -171,18 +170,18 @@ func (pool *BrutePool) Init() error { } if pool.index.ErrString != "" { logs.Log.Error(pool.index.String()) - return fmt.Errorf(pool.index.ErrString) + return fmt.Errorf("%s", pool.index.ErrString) } if pool.index.Chunked && pool.ClientType == ihttp.FAST { logs.Log.Warn("chunk encoding! buf current client FASTHTTP not support chunk decode") } - logs.Log.Logf(pkg.LogVerbose, "[baseline.index] "+pool.index.Format([]string{"status", "length", "spend", "title", "frame", "redirect"})) + logs.Log.Logf(pkg.LogVerbose, "[baseline.index] %s", pool.index.Format([]string{"status", "length", "spend", "title", "frame", "redirect"})) // 检测基本访问能力 if pool.random.ErrString != "" { logs.Log.Error(pool.index.String()) - return fmt.Errorf(pool.index.ErrString) + return fmt.Errorf("%s", pool.index.ErrString) } - logs.Log.Logf(pkg.LogVerbose, "[baseline.random] "+pool.random.Format([]string{"status", "length", "spend", "title", "frame", "redirect"})) + logs.Log.Logf(pkg.LogVerbose, "[baseline.random] %s", pool.random.Format([]string{"status", "length", "spend", "title", "frame", "redirect"})) // 某些网站http会重定向到https, 如果发现随机目录出现这种情况, 则自定将baseurl升级为https if pool.url.Scheme == "http" { @@ -218,29 +217,29 @@ func (pool *BrutePool) Run(offset, limit int) { pool.launchProducer(pool.doCommonFile) } - // 监控 goroutine: word 流结束 / 超 limit / ctx 取消时, 等 wg 归零后关闭 closeCh - var done atomic.Bool - go func() { - for { - if done.Load() { - pool.wg.Wait() - close(pool.closeCh) - return - } - select { - case <-pool.ctx.Done(): - done.Store(true) - case <-time.After(100 * time.Millisecond): - } - } - }() + // worderCh 指向 Worder 输出; 设为 nil 后 select 不再接受新词, + // 随后 allDone 等待 in-flight 工作全部结束再退出循环. + // 这比 atomic.Bool + 轮询 monitor goroutine 更安全: wg.Add 只发生在 + // worderCh 分支, nil 掉之后不会再有新 Add, 消除了 Add/Wait 竞态. + worderCh := pool.Worder.Output + var allDone <-chan struct{} // nil channel: select 永远不选中 + + var drainOnce sync.Once + startDrain := func() { + drainOnce.Do(func() { + worderCh = nil + ch := make(chan struct{}) + go func() { pool.wg.Wait(); close(ch) }() + allDone = ch + }) + } Loop: for { select { - case w, ok := <-pool.Worder.Output: + case w, ok := <-worderCh: if !ok { - done.Store(true) + startDrain() continue } pool.Statistor.End++ @@ -251,7 +250,7 @@ Loop: } if pool.Statistor.End > limit { - done.Store(true) + startDrain() continue } @@ -267,7 +266,6 @@ Loop: pool.wg.Done() } } else { - // 原样的目录拼接, 输入了几个"/"就是几个, 适配/有语义的中间件 if err := pool.reqPool.Invoke(&Unit{path: pool.safePath(w), source: parsers.WordSource, number: pool.wordOffset}); err != nil { pool.wg.Done() } @@ -288,10 +286,10 @@ Loop: pool.wg.Done() } } - case <-pool.closeCh: + case <-allDone: break Loop case <-pool.ctx.Done(): - done.Store(true) + break Loop } } pool.Close() diff --git a/core/pool/checkpool.go b/core/pool/checkpool.go index 242dbe4..48fd5aa 100644 --- a/core/pool/checkpool.go +++ b/core/pool/checkpool.go @@ -12,7 +12,6 @@ import ( "net/url" "strings" "sync" - "sync/atomic" "time" ) @@ -59,26 +58,25 @@ type CheckPool struct { func (pool *CheckPool) Run(ctx context.Context, offset, limit int) { pool.Worder.Run() - var done atomic.Bool - // 挂起一个监控goroutine, 每100ms判断一次done, 如果已经done, 则关闭closeCh, 然后通过Loop中的select case closeCh去break, 实现退出 - go func() { - for { - if done.Load() { - pool.wg.Wait() - close(pool.closeCh) - return - } - time.Sleep(100 * time.Millisecond) - } - }() + worderCh := pool.Worder.Output + var allDone <-chan struct{} + + var drainOnce sync.Once + startDrain := func() { + drainOnce.Do(func() { + worderCh = nil + ch := make(chan struct{}) + go func() { pool.wg.Wait(); close(ch) }() + allDone = ch + }) + } Loop: for { select { - case u, ok := <-pool.Worder.Output: + case u, ok := <-worderCh: if !ok { - done.Store(true) - time.Sleep(100 * time.Millisecond) + startDrain() continue } @@ -88,6 +86,7 @@ Loop: } if pool.reqCount.Load() > int64(limit) { + startDrain() continue } @@ -102,15 +101,11 @@ Loop: if err := pool.Pool.Invoke(u); err != nil { pool.wg.Done() } - case <-pool.closeCh: + case <-allDone: break Loop case <-ctx.Done(): - // 手动退出，不等待任务完成，直接退出 - done.Store(true) break Loop case <-pool.ctx.Done(): - // 手动退出，不等待任务完成，直接退出 - done.Store(true) break Loop } } diff --git a/core/runner.go b/core/runner.go index 05da1ce..ba543e0 100644 --- a/core/runner.go +++ b/core/runner.go @@ -186,10 +186,14 @@ func (r *Runner) Prepare(ctx context.Context) error { func (r *Runner) RunWithBrute(ctx context.Context) { go func() { + defer close(r.taskCh) for t := range r.Tasks.tasks { - r.taskCh <- t + select { + case r.taskCh <- t: + case <-ctx.Done(): + return + } } - close(r.taskCh) }() if r.Count > 0 { @@ -281,16 +285,6 @@ Loop: for { select { case <-ctx.Done(): - // 如果超过了deadline, 尚未开始的任务都将被记录到stat中 - if len(r.taskCh) > 0 { - for t := range r.taskCh { - stat := r.Option.NewStatistor(t.baseUrl, 0) - r.saveStat(stat.Json()) - } - } - if r.StatFile != nil { - logs.Log.Importantf("already save all stat to %s", r.StatFile.Filename) - } break Loop case t, ok := <-r.taskCh: if !ok { @@ -405,14 +399,16 @@ func (r *Runner) AddRecursive(bl *baseline.Baseline) { } func (r *Runner) AddPool(task *Task) { - // 递归新任务 if _, ok := r.PoolName[task.baseUrl]; ok { logs.Log.Importantf("already added pool, skip %s", task.baseUrl) return } task.depth++ r.poolwg.Add(1) - r.Pools.Invoke(task) + if err := r.Pools.Invoke(task); err != nil { + r.poolwg.Done() + logs.Log.Errorf("submit pool task: %v", err) + } } func (r *Runner) newBar(total int) { diff --git a/pkg/resource_provider.go b/pkg/resource_provider.go index 470a61c..1e77f98 100644 --- a/pkg/resource_provider.go +++ b/pkg/resource_provider.go @@ -31,11 +31,9 @@ func LoadEmbeddedConfig(typ string) []byte { // LoadConfig loads config bytes from the external provider first, then from // the embedded templates kept for standalone spray compatibility. func LoadConfig(typ string) []byte { - provider := func() ResourceProvider { - resourceProvider.RLock() - defer resourceProvider.RUnlock() - return resourceProvider.fn - }() + resourceProvider.RLock() + provider := resourceProvider.fn + resourceProvider.RUnlock() if provider != nil { if data := provider(typ); len(data) > 0 { return data