优化统计模块, 现在check失败的任务也会在stat中保留了

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()
 		}()
 	}()

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

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()
 }
 

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 {
+			r.PrintStat(pool)
-				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.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 {

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"`