From 8014f06bf545e41120cb4e46b6f3159c9e73286b Mon Sep 17 00:00:00 2001
From: M09Ic <m09ician@gmail.com>
Date: Thu, 5 Jan 2023 14:56:23 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96check-pool?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/checkpool.go | 25 +++++++++++++------------
 internal/option.go    |  2 +-
 internal/pool.go      |  1 +
 3 files changed, 15 insertions(+), 13 deletions(-)

diff --git a/internal/checkpool.go b/internal/checkpool.go
index 807e062..07debaf 100644
--- a/internal/checkpool.go
+++ b/internal/checkpool.go
@@ -2,6 +2,7 @@ package internal
 
 import (
 	"context"
+	"fmt"
 	"github.com/chainreactors/logs"
 	"github.com/chainreactors/spray/pkg"
 	"github.com/chainreactors/spray/pkg/ihttp"
@@ -9,6 +10,7 @@ import (
 	"github.com/panjf2000/ants/v2"
 	"github.com/valyala/fasthttp"
 	"sync"
+	"time"
 )
 
 func NewCheckPool(ctx context.Context, config *pkg.Config) (*CheckPool, error) {
@@ -23,17 +25,6 @@ func NewCheckPool(ctx context.Context, config *pkg.Config) (*CheckPool, error) {
 		failedCount: 1,
 	}
 
-	switch config.Mod {
-	case pkg.PathSpray:
-		pool.genReq = func(s string) (*ihttp.Request, error) {
-			return ihttp.BuildPathRequest(pool.ClientType, s, "")
-		}
-	case pkg.HostSpray:
-		pool.genReq = func(s string) (*ihttp.Request, error) {
-			return ihttp.BuildHostRequest(pool.ClientType, s, "")
-		}
-	}
-
 	p, _ := ants.NewPoolWithFunc(config.Thread, func(i interface{}) {
 		unit := i.(*Unit)
 		req, err := pool.genReq(unit.path)
@@ -41,6 +32,7 @@ func NewCheckPool(ctx context.Context, config *pkg.Config) (*CheckPool, error) {
 			logs.Log.Error(err.Error())
 		}
 
+		start := time.Now()
 		var bl *pkg.Baseline
 		resp, reqerr := pool.client.Do(pctx, req)
 		if pool.ClientType == ihttp.FAST {
@@ -56,6 +48,7 @@ func NewCheckPool(ctx context.Context, config *pkg.Config) (*CheckPool, error) {
 			bl.Collect()
 		}
 
+		bl.Spended = time.Since(start).Milliseconds()
 		pool.OutputCh <- bl
 		pool.reqCount++
 		pool.wg.Done()
@@ -75,7 +68,6 @@ type CheckPool struct {
 	cancel      context.CancelFunc
 	reqCount    int
 	failedCount int
-	genReq      func(s string) (*ihttp.Request, error)
 	worder      *words.Worder
 	wg          sync.WaitGroup
 }
@@ -84,6 +76,15 @@ func (p *CheckPool) Close() {
 	p.bar.Close()
 }
 
+func (p *CheckPool) genReq(s string) (*ihttp.Request, error) {
+	if p.Mod == pkg.HostSpray {
+		return ihttp.BuildHostRequest(p.ClientType, p.BaseURL, s)
+	} else if p.Mod == pkg.PathSpray {
+		return ihttp.BuildPathRequest(p.ClientType, p.BaseURL, s)
+	}
+	return nil, fmt.Errorf("unknown mod")
+}
+
 func (p *CheckPool) Run(ctx context.Context, offset, limit int) {
 	p.worder.Run()
 Loop:
diff --git a/internal/option.go b/internal/option.go
index 4d314ea..2d90dea 100644
--- a/internal/option.go
+++ b/internal/option.go
@@ -57,7 +57,7 @@ type OutputOptions struct {
 	Match       string   `long:"match" description:"String, custom match function, e.g.: --match current.Status != 200" json:"match,omitempty"`
 	Filter      string   `long:"filter" description:"String, custom filter function, e.g.: --filter current.Body contains 'hello'" json:"filter,omitempty"`
 	Extracts    []string `long:"extract" description:"String, Multi, extract response, e.g.: --extract js --extract ip --extract version:(.*?)" json:"extracts,omitempty"`
-	OutputFile  string   `short:"f" description:"String, output filename" json:"output_file,omitempty"`
+	OutputFile  string   `short:"f" long:"file" description:"String, output filename" json:"output_file,omitempty"`
 	Format      string   `short:"F" long:"format" description:"String, output format, e.g.: --format 1.json"`
 	FuzzyFile   string   `long:"fuzzy-file" description:"String, fuzzy output filename" json:"fuzzy_file,omitempty"`
 	DumpFile    string   `long:"dump-file" description:"String, dump all request, and write to filename"`
diff --git a/internal/pool.go b/internal/pool.go
index 51b6189..07c3afe 100644
--- a/internal/pool.go
+++ b/internal/pool.go
@@ -293,6 +293,7 @@ func (pool *Pool) genReq(s string) (*ihttp.Request, error) {
 	}
 	return nil, fmt.Errorf("unknown mod")
 }
+
 func (pool *Pool) Run(ctx context.Context, offset, limit int) {
 	pool.worder.RunWithRules()
 	go func() {