admin/spray
1
0
Fork 0
mirror of https://github.com/chainreactors/spray.git synced 2025-05-06 10:41:21 +00:00
Code Issues Packages Projects Releases Wiki Activity

新增报错阈值

Browse Source
This commit is contained in:
M09Ic 2022-09-23 01:47:24 +08:00
parent 4687a7cf1b
commit 637eb21377
2 changed files with 17 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
internal/baseline.go
View File

@ -9,7 +9,7 @@ import (
"strings" "strings"
) )
func NewBaseline(u *fasthttp.URI, resp *fasthttp.Response) (*baseline, error) { func NewBaseline(u *fasthttp.URI, resp *fasthttp.Response) *baseline {
bl := &baseline{ bl := &baseline{
Url: u, Url: u,
UrlString: u.String(), UrlString: u.String(),
@ -24,7 +24,7 @@ func NewBaseline(u *fasthttp.URI, resp *fasthttp.Response) (*baseline, error) {
bl.HeaderLength = resp.Header.Len() bl.HeaderLength = resp.Header.Len()
bl.RedirectURL = string(resp.Header.Peek("Location")) bl.RedirectURL = string(resp.Header.Peek("Location"))
bl.Raw = append(bl.Header, bl.Body...) bl.Raw = append(bl.Header, bl.Body...)
return bl, nil return bl
} }
func NewInvalidBaseline(u *fasthttp.URI, resp *fasthttp.Response) *baseline { func NewInvalidBaseline(u *fasthttp.URI, resp *fasthttp.Response) *baseline {

20
internal/pool.go
View File

@ -18,11 +18,14 @@ var (
CheckWaf func(*http.Response) bool CheckWaf func(*http.Response) bool
) )
var breakThreshold int = 10
func NewPool(ctx context.Context, config *pkg.Config, outputCh chan *baseline) (*Pool, error) { func NewPool(ctx context.Context, config *pkg.Config, outputCh chan *baseline) (*Pool, error) {
pctx, cancel := context.WithCancel(ctx) pctx, cancel := context.WithCancel(ctx)
pool := &Pool{ pool := &Pool{
Config: config, Config: config,
ctx: pctx, ctx: pctx,
cancel: cancel,
client: pkg.NewClient(config.Thread, 2), client: pkg.NewClient(config.Thread, 2),
worder: words.NewWorder(config.Wordlist), worder: words.NewWorder(config.Wordlist),
outputCh: outputCh, outputCh: outputCh,
@ -57,15 +60,14 @@ func NewPool(ctx context.Context, config *pkg.Config, outputCh chan *baseline) (
//logs.Log.Debugf("%s request error, %s", strurl, err.Error()) //logs.Log.Debugf("%s request error, %s", strurl, err.Error())
pool.errorCount++ pool.errorCount++
bl = &baseline{Err: err} bl = &baseline{Err: err}
} else { } else {
defer fasthttp.ReleaseResponse(resp) defer fasthttp.ReleaseResponse(resp)
defer fasthttp.ReleaseRequest(req) defer fasthttp.ReleaseRequest(req)
//defer resp.Body.Close() // 必须要关闭body ,否则keep-alive无法生效 //defer resp.Body.Close() // 必须要关闭body ,否则keep-alive无法生效
if err = pool.PreCompare(resp); err == nil || unit.source == CheckSource { if err = pool.PreCompare(resp); err == nil || unit.source == CheckSource {
// 通过预对比跳过一些无用数据, 减少性能消耗 // 通过预对比跳过一些无用数据, 减少性能消耗
bl, err = NewBaseline(req.URI(), resp) bl = NewBaseline(req.URI(), resp)
} else if err == ErrWaf {
cancel()
} else { } else {
bl = NewInvalidBaseline(req.URI(), resp) bl = NewInvalidBaseline(req.URI(), resp)
} }
@ -104,12 +106,14 @@ type Pool struct {
pool *ants.PoolWithFunc pool *ants.PoolWithFunc
bar *pkg.Bar bar *pkg.Bar
ctx context.Context ctx context.Context
cancel context.CancelFunc
//baseReq *http.Request //baseReq *http.Request
baseline *baseline baseline *baseline
outputCh chan *baseline outputCh chan *baseline
tempCh chan *baseline tempCh chan *baseline
reqCount int reqCount int
errorCount int errorCount int
failedCount int
checkPeriod int checkPeriod int
errPeriod int errPeriod int
genReq func(s string) (*fasthttp.Request, error) genReq func(s string) (*fasthttp.Request, error)
@ -119,10 +123,15 @@ type Pool struct {
} }
func (p *Pool) check() { func (p *Pool) check() {
p.wg.Add(1) var wg sync.WaitGroup
wg.Add(1)
_ = p.pool.Invoke(newUnit(pkg.RandPath(), CheckSource)) _ = p.pool.Invoke(newUnit(pkg.RandPath(), CheckSource))
//} //}
p.wg.Wait() wg.Wait()
if p.failedCount > breakThreshold {
p.cancel()
}
} }
func (p *Pool) Init() error { func (p *Pool) Init() error {
@ -131,6 +140,7 @@ func (p *Pool) Init() error {
// 检测基本访问能力 // 检测基本访问能力
if p.baseline != nil && p.baseline.Err != nil { if p.baseline != nil && p.baseline.Err != nil {
p.cancel()
return p.baseline.Err return p.baseline.Err
} }