diff --git a/cmd/cmd.go b/cmd/cmd.go index b2fb38f..0c70573 100644 --- a/cmd/cmd.go +++ b/cmd/cmd.go @@ -106,7 +106,7 @@ func Spray() { signal.Notify(c, os.Interrupt, syscall.SIGTERM) go func() { <-c - fmt.Println("exit signal, save stat and exit") + logs.Log.Important("exit signal, save stat and exit") canceler() }() }() diff --git a/internal/option.go b/internal/option.go index 9d2f039..e827a40 100644 --- a/internal/option.go +++ b/internal/option.go @@ -483,7 +483,7 @@ func (opt *Option) PrepareRunner() (*Runner, error) { if opt.ResumeFrom != "" { r.StatFile, err = files.NewFile(opt.ResumeFrom, false, true, true) } else { - r.StatFile, err = files.NewFile(taskfrom+".stat", false, true, true) + r.StatFile, err = files.NewFile(strings.ReplaceAll(taskfrom, ":", "_")+".stat", false, true, true) } if err != nil { return nil, err diff --git a/internal/pool.go b/internal/pool.go index ebf1b26..4c018f0 100644 --- a/internal/pool.go +++ b/internal/pool.go @@ -40,7 +40,7 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) { pctx, cancel := context.WithCancel(ctx) pool := &Pool{ Config: config, - base: u.Scheme + "://" + u.Hostname(), + base: u.Scheme + "://" + u.Host, isDir: strings.HasSuffix(u.Path, "/"), url: u, ctx: pctx, @@ -185,7 +185,8 @@ func (pool *Pool) Init() error { pool.reqPool.Invoke(newUnit(pool.safePath(pkg.RandPath()), InitRandomSource)) pool.initwg.Wait() if pool.index.ErrString != "" { - return fmt.Errorf(pool.index.String()) + logs.Log.Error(pool.index.String()) + return fmt.Errorf(pool.index.ErrString) } if pool.index.Chunked && pool.ClientType == ihttp.FAST { logs.Log.Warn("chunk encoding! buf current client FASTHTTP not support chunk decode") @@ -193,7 +194,8 @@ func (pool *Pool) Init() error { logs.Log.Info("[baseline.index] " + pool.index.Format([]string{"status", "length", "spend", "title", "frame", "redirect"})) // 检测基本访问能力 if pool.random.ErrString != "" { - return fmt.Errorf(pool.random.String()) + logs.Log.Error(pool.index.String()) + return fmt.Errorf(pool.index.ErrString) } logs.Log.Info("[baseline.random] " + pool.random.Format([]string{"status", "length", "spend", "title", "frame", "redirect"})) @@ -213,30 +215,6 @@ func (pool *Pool) Init() error { return nil } -func (pool *Pool) checkRedirect(redirectURL string) bool { - if pool.random.RedirectURL == "" { - // 如果random的redirectURL为空, 此时该项 - return true - } - - if redirectURL == pool.random.RedirectURL { - // 相同的RedirectURL将被认为是无效数据 - return false - } else { - // path为3xx, 且与baseline中的RedirectURL不同时, 为有效数据 - return true - } -} - -func (pool *Pool) genReq(s string) (*ihttp.Request, error) { - if pool.Mod == pkg.HostSpray { - return ihttp.BuildHostRequest(pool.ClientType, pool.BaseURL, s) - } else if pool.Mod == pkg.PathSpray { - return ihttp.BuildPathRequest(pool.ClientType, pool.base, s) - } - return nil, fmt.Errorf("unknown mod") -} - func (pool *Pool) Run(ctx context.Context, offset, limit int) { pool.worder.RunWithRules() if pool.Active { @@ -313,8 +291,6 @@ Loop: } } - pool.waiter.Wait() - pool.Statistor.EndTime = time.Now().Unix() pool.Close() } @@ -429,6 +405,30 @@ func (pool *Pool) Invoke(v interface{}) { } } +func (pool *Pool) checkRedirect(redirectURL string) bool { + if pool.random.RedirectURL == "" { + // 如果random的redirectURL为空, 此时该项 + return true + } + + if redirectURL == pool.random.RedirectURL { + // 相同的RedirectURL将被认为是无效数据 + return false + } else { + // path为3xx, 且与baseline中的RedirectURL不同时, 为有效数据 + return true + } +} + +func (pool *Pool) genReq(s string) (*ihttp.Request, error) { + if pool.Mod == pkg.HostSpray { + return ihttp.BuildHostRequest(pool.ClientType, pool.BaseURL, s) + } else if pool.Mod == pkg.PathSpray { + return ihttp.BuildPathRequest(pool.ClientType, pool.base, s) + } + return nil, fmt.Errorf("unknown mod") +} + func (pool *Pool) PreCompare(resp *ihttp.Response) error { status := resp.StatusCode() if pkg.IntsContains(WhiteStatus, status) { @@ -705,8 +705,10 @@ func (pool *Pool) Close() { for pool.analyzeDone { time.Sleep(time.Duration(100) * time.Millisecond) } - close(pool.tempCh) + close(pool.additionCh) + close(pool.checkCh) + pool.Statistor.EndTime = time.Now().Unix() pool.bar.Close() } diff --git a/internal/runner.go b/internal/runner.go index 198ebca..13a1fb3 100644 --- a/internal/runner.go +++ b/internal/runner.go @@ -208,10 +208,11 @@ func (r *Runner) Prepare(ctx context.Context) error { pool.bar = pkg.NewBar(config.BaseURL, limit-pool.Statistor.Offset, r.Progress) err = pool.Init() if err != nil { - logs.Log.Error(err.Error()) + pool.Statistor.Error = err.Error() if !r.Force { // 如果没开启force, init失败将会关闭pool - pool.cancel() + pool.Close() + r.PrintStat(pool) r.Done() return } @@ -223,20 +224,7 @@ func (r *Runner) Prepare(ctx context.Context) error { // 如果因为错误积累退出, end将指向第一个错误发生时, 防止resume时跳过大量目标 pool.Statistor.End = pool.failedBaselines[0].Number } - if r.Color { - logs.Log.Important(pool.Statistor.ColorString()) - logs.Log.Important(pool.Statistor.ColorCountString()) - logs.Log.Important(pool.Statistor.ColorSourceString()) - } else { - logs.Log.Important(pool.Statistor.String()) - logs.Log.Important(pool.Statistor.CountString()) - logs.Log.Important(pool.Statistor.SourceString()) - } - - if r.StatFile != nil { - r.StatFile.SafeWrite(pool.Statistor.Json()) - r.StatFile.SafeSync() - } + r.PrintStat(pool) r.Done() }) } @@ -334,6 +322,27 @@ func (r *Runner) Done() { r.poolwg.Done() } +func (r *Runner) PrintStat(pool *Pool) { + if r.Color { + logs.Log.Important(pool.Statistor.ColorString()) + if pool.Statistor.Error == "" { + logs.Log.Important(pool.Statistor.ColorCountString()) + logs.Log.Important(pool.Statistor.ColorSourceString()) + } + } else { + logs.Log.Important(pool.Statistor.String()) + if pool.Statistor.Error == "" { + logs.Log.Important(pool.Statistor.CountString()) + logs.Log.Important(pool.Statistor.SourceString()) + } + } + + if r.StatFile != nil { + r.StatFile.SafeWrite(pool.Statistor.Json()) + r.StatFile.SafeSync() + } +} + func (r *Runner) Outputting() { debugPrint := func(bl *pkg.Baseline) { if r.Color { diff --git a/pkg/statistor.go b/pkg/statistor.go index 01303fa..aab7480 100644 --- a/pkg/statistor.go +++ b/pkg/statistor.go @@ -38,6 +38,7 @@ func NewStatistorFromStat(origin *Statistor) *Statistor { type Statistor struct { BaseUrl string `json:"url"` + Error string `json:"error"` Counts map[int]int `json:"counts"` Sources map[int]int `json:"sources"` FailedNumber int32 `json:"failed"`