From b13903ea9822d4620c79efb72deeb6f016d58622 Mon Sep 17 00:00:00 2001 From: M09Ic Date: Tue, 4 Mar 2025 14:45:31 +0800 Subject: [PATCH 1/3] fix: map panic, https://github.com/chainreactors/spray/issues/93 --- core/ihttp/request.go | 12 ++---------- core/pool/brutepool.go | 9 ++++++++- 2 files changed, 10 insertions(+), 11 deletions(-) diff --git a/core/ihttp/request.go b/core/ihttp/request.go index 93c2949..4e52912 100644 --- a/core/ihttp/request.go +++ b/core/ihttp/request.go @@ -32,16 +32,8 @@ type Request struct { } func (r *Request) SetHeaders(header map[string]string, RandomUA bool) { - if header["User-Agent"] == "" { - if RandomUA { - header["User-Agent"] = pkg.RandomUA() - } else { - header["User-Agent"] = pkg.DefaultUserAgent - } - } - - if header["Accept"] == "" { - header["Accept"] = "*/*" + if RandomUA { + r.SetHeader("User-Agent", pkg.RandomUA()) } if r.StandardRequest != nil { diff --git a/core/pool/brutepool.go b/core/pool/brutepool.go index a7394a7..b051c80 100644 --- a/core/pool/brutepool.go +++ b/core/pool/brutepool.go @@ -107,6 +107,14 @@ type BrutePool struct { } func (pool *BrutePool) Init() error { + if pool.Headers["User-Agent"] == "" { + pool.Headers["User-Agent"] = pkg.DefaultUserAgent + } + + if pool.Headers["Accept"] == "" { + pool.Headers["Accept"] = "*/*" + } + pool.initwg.Add(2) if pool.Index != "/" { logs.Log.Logf(pkg.LogVerbose, "custom index url: %s", pkg.BaseURL(pool.url)+pkg.FormatURL(pkg.BaseURL(pool.url), pool.Index)) @@ -394,7 +402,6 @@ func (pool *BrutePool) NoScopeInvoke(v interface{}) { return } req.SetHeaders(pool.Headers, pool.RandomUserAgent) - req.SetHeader("User-Agent", pkg.RandomUA()) resp, reqerr := pool.client.Do(req) if pool.ClientType == ihttp.FAST { defer fasthttp.ReleaseResponse(resp.FastResponse) From c746c26ff9eba3a4e83b8693fe88557d2a101521 Mon Sep 17 00:00:00 2001 From: M09Ic Date: Tue, 4 Mar 2025 14:56:48 +0800 Subject: [PATCH 2/3] fix: brute init not set header, https://github.com/chainreactors/spray/issues/94 --- core/ihttp/request.go | 10 +++++----- core/option.go | 2 +- core/pool/brutepool.go | 8 ++++---- core/pool/config.go | 3 ++- core/runner.go | 7 ++++++- 5 files changed, 18 insertions(+), 12 deletions(-) diff --git a/core/ihttp/request.go b/core/ihttp/request.go index 4e52912..1302e17 100644 --- a/core/ihttp/request.go +++ b/core/ihttp/request.go @@ -31,18 +31,18 @@ type Request struct { ClientType int } -func (r *Request) SetHeaders(header map[string]string, RandomUA bool) { +func (r *Request) SetHeaders(header http.Header, RandomUA bool) { if RandomUA { r.SetHeader("User-Agent", pkg.RandomUA()) } if r.StandardRequest != nil { - for k, v := range header { - r.StandardRequest.Header.Set(k, v) - } + r.StandardRequest.Header = header } else if r.FastRequest != nil { for k, v := range header { - r.FastRequest.Header.Set(k, v) + for _, i := range v { + r.FastRequest.Header.Set(k, i) + } } } } diff --git a/core/option.go b/core/option.go index 7e790bb..14a38d1 100644 --- a/core/option.go +++ b/core/option.go @@ -99,7 +99,7 @@ type OutputOptions struct { type RequestOptions struct { Method string `short:"X" long:"method" default:"GET" description:"String, request method, e.g.: --method POST" config:"method"` - Headers []string `long:"header" description:"Strings, custom headers, e.g.: --header 'Auth: example_auth'" config:"headers"` + Headers []string `short:"H" long:"header" description:"Strings, custom headers, e.g.: --header 'Auth: example_auth'" config:"headers"` UserAgent string `long:"user-agent" description:"String, custom user-agent, e.g.: --user-agent Custom" config:"useragent"` RandomUserAgent bool `long:"random-agent" description:"Bool, use random with default user-agent" config:"random-useragent"` Cookie []string `long:"cookie" description:"Strings, custom cookie" config:"cookies"` diff --git a/core/pool/brutepool.go b/core/pool/brutepool.go index b051c80..c32701d 100644 --- a/core/pool/brutepool.go +++ b/core/pool/brutepool.go @@ -107,12 +107,12 @@ type BrutePool struct { } func (pool *BrutePool) Init() error { - if pool.Headers["User-Agent"] == "" { - pool.Headers["User-Agent"] = pkg.DefaultUserAgent + if pool.Headers.Get("User-Agent") == "" { + pool.Headers.Set("User-Agent", pkg.DefaultUserAgent) } - if pool.Headers["Accept"] == "" { - pool.Headers["Accept"] = "*/*" + if pool.Headers.Get("Accept") == "" { + pool.Headers.Set("Accept", "*/*") } pool.initwg.Add(2) diff --git a/core/pool/config.go b/core/pool/config.go index 6a0c1f2..5178dfe 100644 --- a/core/pool/config.go +++ b/core/pool/config.go @@ -7,6 +7,7 @@ import ( "github.com/chainreactors/words" "github.com/chainreactors/words/rule" "github.com/expr-lang/expr/vm" + "net/http" "sync" "time" ) @@ -27,7 +28,7 @@ type Config struct { BreakThreshold int32 Method string Mod SprayMod - Headers map[string]string + Headers http.Header ClientType int MatchExpr *vm.Program FilterExpr *vm.Program diff --git a/core/runner.go b/core/runner.go index 168c052..e36821d 100644 --- a/core/runner.go +++ b/core/runner.go @@ -15,6 +15,7 @@ import ( "github.com/panjf2000/ants/v2" "github.com/vbauerster/mpb/v8" "github.com/vbauerster/mpb/v8/decor" + "net/http" "strings" "sync" "time" @@ -67,7 +68,7 @@ func (r *Runner) PrepareConfig() *pool.Config { Thread: r.Threads, Timeout: time.Duration(r.Timeout) * time.Second, RateLimit: r.RateLimit, - Headers: r.Headers, + Headers: make(http.Header), Method: r.Method, Mod: pool.ModMap[r.Mod], OutputCh: r.outputCh, @@ -109,6 +110,10 @@ func (r *Runner) PrepareConfig() *pool.Config { } } + for k, v := range r.Headers { + config.Headers.Set(k, v) + } + return config } From 08ce95b43dc936154863399a67ca2384593188a6 Mon Sep 17 00:00:00 2001 From: M09Ic Date: Tue, 4 Mar 2025 14:58:28 +0800 Subject: [PATCH 3/3] fix: check pool net set headers --- core/pool/brutepool.go | 8 -------- core/pool/checkpool.go | 2 +- core/runner.go | 8 ++++++++ 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/core/pool/brutepool.go b/core/pool/brutepool.go index c32701d..7460a0c 100644 --- a/core/pool/brutepool.go +++ b/core/pool/brutepool.go @@ -107,14 +107,6 @@ type BrutePool struct { } func (pool *BrutePool) Init() error { - if pool.Headers.Get("User-Agent") == "" { - pool.Headers.Set("User-Agent", pkg.DefaultUserAgent) - } - - if pool.Headers.Get("Accept") == "" { - pool.Headers.Set("Accept", "*/*") - } - pool.initwg.Add(2) if pool.Index != "/" { logs.Log.Logf(pkg.LogVerbose, "custom index url: %s", pkg.BaseURL(pool.url)+pkg.FormatURL(pkg.BaseURL(pool.url), pool.Index)) diff --git a/core/pool/checkpool.go b/core/pool/checkpool.go index b6040e2..136d460 100644 --- a/core/pool/checkpool.go +++ b/core/pool/checkpool.go @@ -36,7 +36,7 @@ func NewCheckPool(ctx context.Context, config *Config) (*CheckPool, error) { processCh: make(chan *baseline.Baseline, config.Thread), }, } - pool.Headers = map[string]string{"Connection": "close"} + pool.Headers.Set("Connection", "close") p, _ := ants.NewPoolWithFunc(config.Thread, pool.Invoke) pool.Pool = p diff --git a/core/runner.go b/core/runner.go index e36821d..99b9809 100644 --- a/core/runner.go +++ b/core/runner.go @@ -114,6 +114,14 @@ func (r *Runner) PrepareConfig() *pool.Config { config.Headers.Set(k, v) } + if config.Headers.Get("User-Agent") == "" { + config.Headers.Set("User-Agent", pkg.DefaultUserAgent) + } + + if config.Headers.Get("Accept") == "" { + config.Headers.Set("Accept", "*/*") + } + return config }