admin/spray
1
0
Fork 0
mirror of https://github.com/chainreactors/spray.git synced 2025-09-15 11:40:13 +00:00
Code Issues Packages Projects Releases Wiki Activity

优化statistor, 现在是线程安全的

Browse Source
This commit is contained in:
M09Ic 2022-12-11 04:21:42 +08:00
parent 71393bfeb4
commit faf0812858
4 changed files with 19 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
internal/option.go
View File

@ -271,7 +271,7 @@ func (opt *Option) PrepareRunner() (*Runner, error) {
}
taskfrom = "resume " + opt.ResumeFrom
for _, stat := range stats {
tasks = append(tasks, &Task{baseUrl: stat.BaseUrl, offset: stat.Offset + stat.ReqNumber, total: r.Total})
tasks = append(tasks, &Task{baseUrl: stat.BaseUrl, offset: stat.Offset + stat.End, total: r.Total})
}
} else {
var file *os.File

10
internal/pool.go
View File

@ -15,6 +15,7 @@ import (
"strconv"
"strings"
"sync"
"sync/atomic"
"time"
)
@ -37,6 +38,7 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) {
worder: words.NewWorder(config.Wordlist, config.Fns),
baselines: make(map[int]*pkg.Baseline),
tempCh: make(chan *pkg.Baseline, config.Thread),
checkCh: make(chan *Unit),
wg: sync.WaitGroup{},
initwg: sync.WaitGroup{},
reqCount: 1,
@ -45,9 +47,8 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) {
pool.worder.Rules = pool.Rules
pool.worder.RunWithRules()
p, _ := ants.NewPoolWithFunc(config.Thread, func(i interface{}) {
pool.Statistor.Total++
atomic.AddInt32(&pool.Statistor.ReqTotal, 1)
unit := i.(*Unit)
req, err := pool.genReq(unit.path)
if err != nil {
@ -65,7 +66,7 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) {
var bl *pkg.Baseline
if reqerr != nil && reqerr != fasthttp.ErrBodyTooLarge {
pool.failedCount++
pool.Statistor.FailedNumber++
atomic.AddInt32(&pool.Statistor.FailedNumber, 1)
bl = &pkg.Baseline{UrlString: pool.BaseURL + unit.path, IsValid: false, ErrString: reqerr.Error(), Reason: ErrRequestFailed.Error()}
pool.failedBaselines = append(pool.failedBaselines, bl)
} else {
@ -313,7 +314,7 @@ Loop:
break Loop
}
pool.Statistor.End++
if pool.reqCount < offset {
if int(pool.reqCount) < offset {
pool.reqCount++
continue
}
@ -340,7 +341,6 @@ Loop:
}
}
pool.wg.Wait()
pool.Statistor.ReqNumber = pool.reqCount
pool.Statistor.EndTime = time.Now().Unix()
pool.Close()
}

2
internal/runner.go
View File

@ -136,7 +136,7 @@ func (r *Runner) Prepare(ctx context.Context) error {
r.Done()
return
}
pool.Statistor.Total = r.Total
pool.bar = pkg.NewBar(config.BaseURL, t.total-t.offset, r.Progress)
err = pool.Init()
if err != nil {

7
pkg/statistor.go
View File

@ -23,13 +23,14 @@ func NewStatistor(url string) *Statistor {
type Statistor struct {
BaseUrl string `json:"url"`
Counts map[int]int `json:"counts"`
ReqNumber int `json:"req"`
FailedNumber int `json:"failed"`
FailedNumber int32 `json:"failed"`
ReqTotal int32 `json:"req_total"`
CheckNumber int `json:"check"`
FoundNumber int `json:"found"`
FilteredNumber int `json:"filtered"`
FuzzyNumber int `json:"fuzzy"`
WafedNumber int `json:"wafed"`
End int `json:"end"`
Offset int `json:"offset"`
Total int `json:"total"`
@ -42,7 +43,7 @@ type Statistor struct {
func (stat *Statistor) String() string {
var s strings.Builder
s.WriteString(fmt.Sprintf("[stat] %s took %d s, request total: %d, found: %d, check: %d, failed: %d", stat.BaseUrl, stat.EndTime-stat.StartTime, stat.ReqNumber, stat.FoundNumber, stat.CheckNumber, stat.FailedNumber))
s.WriteString(fmt.Sprintf("[stat] %s took %d s, request total: %d, finish: %d/%d, found: %d, check: %d, failed: %d", stat.BaseUrl, stat.EndTime-stat.StartTime, stat.ReqTotal, stat.End, stat.Total, stat.FoundNumber, stat.CheckNumber, stat.FailedNumber))
if stat.FuzzyNumber != 0 {
s.WriteString(", fuzzy: " + strconv.Itoa(stat.FuzzyNumber))