新增--read-all参数, 用来取消body max read限制

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

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 {

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

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 {

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