From b019324383ebf00c2a8f3f5e30c51cce1ac50f1e Mon Sep 17 00:00:00 2001 From: M09Ic Date: Mon, 9 Jan 2023 21:47:06 +0800 Subject: [PATCH] =?UTF-8?q?=E6=96=B0=E5=A2=9E--read-all=E5=8F=82=E6=95=B0,?= =?UTF-8?q?=20=E7=94=A8=E6=9D=A5=E5=8F=96=E6=B6=88body=20max=20read?= =?UTF-8?q?=E9=99=90=E5=88=B6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- cmd/cmd.go | 22 ++++++++++++++++++++++ internal/option.go | 20 ++------------------ internal/pool.go | 2 -- pkg/ihttp/client.go | 3 ++- pkg/ihttp/response.go | 17 +++++++++++++++-- 5 files changed, 41 insertions(+), 23 deletions(-) diff --git a/cmd/cmd.go b/cmd/cmd.go index ce95161..ff65a95 100644 --- a/cmd/cmd.go +++ b/cmd/cmd.go @@ -3,13 +3,16 @@ package cmd import ( "context" "fmt" + "github.com/chainreactors/gogo/v2/pkg/fingers" "github.com/chainreactors/gogo/v2/pkg/utils" "github.com/chainreactors/logs" "github.com/chainreactors/spray/internal" "github.com/chainreactors/spray/pkg" + "github.com/chainreactors/spray/pkg/ihttp" "github.com/jessevdk/go-flags" "os" "os/signal" + "regexp" "syscall" "time" ) @@ -56,6 +59,25 @@ func Spray() { utils.Fatal(err.Error()) } + if option.Extracts != nil { + for _, e := range option.Extracts { + if reg, ok := fingers.PresetExtracts[e]; ok { + pkg.Extractors[e] = reg + } else { + pkg.Extractors[e] = regexp.MustCompile(e) + } + } + } + // 一些全局变量初始化 + if option.Debug { + logs.Log.Level = logs.Debug + } + + pkg.Distance = uint8(option.SimhashDistance) + ihttp.DefaultMaxBodySize = option.MaxBodyLength * 1024 + if option.ReadAll { + ihttp.DefaultMaxBodySize = 0 + } var runner *internal.Runner if option.ResumeFrom != "" { runner, err = option.PrepareRunner() diff --git a/internal/option.go b/internal/option.go index 64b5839..9be39f5 100644 --- a/internal/option.go +++ b/internal/option.go @@ -4,17 +4,14 @@ import ( "fmt" "github.com/antonmedv/expr" "github.com/chainreactors/files" - "github.com/chainreactors/gogo/v2/pkg/fingers" "github.com/chainreactors/logs" "github.com/chainreactors/spray/pkg" - "github.com/chainreactors/spray/pkg/ihttp" "github.com/chainreactors/words/mask" "github.com/chainreactors/words/rule" "github.com/gosuri/uiprogress" "io/ioutil" "net/url" "os" - "regexp" "strconv" "strings" ) @@ -72,6 +69,7 @@ type RequestOptions struct { UserAgent string `long:"user-agent" description:"String, custom user-agent, e.g.: --user-agent Custom"` RandomUserAgent bool `long:"random-agent" description:"Bool, use random with default user-agent"` Cookie []string `long:"cookie" description:"String, Multi, custom cookie"` + ReadAll bool `long:"read-all" description:"Bool, read all response body"` MaxBodyLength int `long:"max-length" default:"100" description:"Int, max response body length (kb), default 100k, e.g. -max-length 1000"` } @@ -140,25 +138,13 @@ func (opt *Option) PrepareRunner() (*Runner, error) { Common: opt.Common, } - if opt.Extracts != nil { - for _, e := range opt.Extracts { - if reg, ok := fingers.PresetExtracts[e]; ok { - pkg.Extractors[e] = reg - } else { - pkg.Extractors[e] = regexp.MustCompile(e) - } - } - } - // 一些全局变量初始化 + // log and bar if !opt.NoColor { logs.Log.Color = true logs.DefaultColorMap[logs.Info] = logs.PurpleBold logs.DefaultColorMap[logs.Important] = logs.Green r.Color = true } - if opt.Debug { - logs.Log.Level = logs.Debug - } if opt.Quiet { logs.Log.Quiet = true logs.Log.Color = false @@ -168,8 +154,6 @@ func (opt *Option) PrepareRunner() (*Runner, error) { r.Progress.Start() logs.Log.Writer = r.Progress.Bypass() } - pkg.Distance = uint8(opt.SimhashDistance) - ihttp.DefaultMaxBodySize = opt.MaxBodyLength * 1024 // configuration if opt.Force { diff --git a/internal/pool.go b/internal/pool.go index 51ecd18..eb78674 100644 --- a/internal/pool.go +++ b/internal/pool.go @@ -42,7 +42,6 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) { tempCh: make(chan *pkg.Baseline, config.Thread), checkCh: make(chan int), additionCh: make(chan *Unit, 100), - closeCh: make(chan struct{}), wg: sync.WaitGroup{}, initwg: sync.WaitGroup{}, reqCount: 1, @@ -134,7 +133,6 @@ type Pool struct { tempCh chan *pkg.Baseline // 待处理的baseline checkCh chan int // 独立的check管道， 防止与redirect/crawl冲突 additionCh chan *Unit - closeCh chan struct{} reqCount int failedCount int isFailed bool diff --git a/pkg/ihttp/client.go b/pkg/ihttp/client.go index b45b2b9..6ee6950 100644 --- a/pkg/ihttp/client.go +++ b/pkg/ihttp/client.go @@ -31,7 +31,7 @@ func NewClient(thread int, timeout int, clientType int) *Client { MaxConnWaitTimeout: time.Duration(timeout) * time.Second, ReadTimeout: time.Duration(timeout) * time.Second, WriteTimeout: time.Duration(timeout) * time.Second, - ReadBufferSize: 16384, + ReadBufferSize: 16384, // 16k MaxResponseBodySize: DefaultMaxBodySize, NoDefaultUserAgentHeader: true, DisablePathNormalizing: true, @@ -52,6 +52,7 @@ func NewClient(thread int, timeout int, clientType int) *Client { }, MaxConnsPerHost: thread * 3 / 2, IdleConnTimeout: time.Duration(timeout) * time.Second, + ReadBufferSize: 16384, // 16k }, Timeout: time.Second * time.Duration(timeout), CheckRedirect: func(req *http.Request, via []*http.Request) error { diff --git a/pkg/ihttp/response.go b/pkg/ihttp/response.go index c835657..399c935 100644 --- a/pkg/ihttp/response.go +++ b/pkg/ihttp/response.go @@ -29,8 +29,20 @@ func (r *Response) Body() []byte { if r.FastResponse != nil { return r.FastResponse.Body() } else if r.StandardResponse != nil { - body := make([]byte, 20480) - if r.StandardResponse.ContentLength > 0 { + if DefaultMaxBodySize == 0 { + body, err := io.ReadAll(r.StandardResponse.Body) + if err != nil { + return nil + } + return body + } else { + var body []byte + if r.StandardResponse.ContentLength > 0 && r.StandardResponse.ContentLength < int64(DefaultMaxBodySize) { + body = make([]byte, r.StandardResponse.ContentLength) + } else { + body = make([]byte, DefaultMaxBodySize) + } + n, err := io.ReadFull(r.StandardResponse.Body, body) _ = r.StandardResponse.Body.Close() if err == nil { @@ -40,6 +52,7 @@ func (r *Response) Body() []byte { } else { logs.Log.Error("readfull failed" + err.Error()) return nil + } } return nil