现在同时支持http.net与fasthttp两个库, 适用不同的场景

internal/baseline.go

@@ -2,46 +2,48 @@ package internal
 
 import (
 	"encoding/json"
-	"fmt"
 	"github.com/chainreactors/parsers"
 	"github.com/chainreactors/spray/pkg"
-	"github.com/valyala/fasthttp"
+	"github.com/chainreactors/spray/pkg/ihttp"
+	"strconv"
 	"strings"
 )
 
-func NewBaseline(u *fasthttp.URI, resp *fasthttp.Response) *baseline {
+func NewBaseline(u, host string, resp *ihttp.Response) *baseline {
 	bl := &baseline{
-		Url:       u,
-		UrlString: u.String(),
+		//Url:       u,
+		UrlString: u,
+		Host:      host,
 		Status:    resp.StatusCode(),
 		IsValid:   true,
 	}
 
 	bl.Body = resp.Body()
-	bl.BodyLength = resp.Header.ContentLength()
-	bl.Header = resp.Header.Header()
-	bl.HeaderLength = resp.Header.Len()
-	bl.RedirectURL = string(resp.Header.Peek("Location"))
+	bl.BodyLength = resp.ContentLength()
+	bl.Header = resp.Header()
+	bl.HeaderLength = len(bl.Header)
+	bl.RedirectURL = resp.GetHeader("Location")
 	bl.Raw = append(bl.Header, bl.Body...)
 	return bl
 }
 
-func NewInvalidBaseline(u *fasthttp.URI, resp *fasthttp.Response) *baseline {
+func NewInvalidBaseline(u, host string, resp *ihttp.Response) *baseline {
 	bl := &baseline{
-		Url:       u,
-		UrlString: u.String(),
+		//Url:       u,
+		UrlString: u,
+		Host:      host,
 		Status:    resp.StatusCode(),
 		IsValid:   false,
 	}
 
-	bl.RedirectURL = string(resp.Header.Peek("Location"))
+	bl.RedirectURL = string(resp.GetHeader("Location"))
 
 	return bl
 }
 
 type baseline struct {
-	Url          *fasthttp.URI  `json:"-"`
-	UrlString    string         `json:"url_string"`
+	UrlString    string         `json:"url"`
+	Host         string         `json:"host"`
 	Body         []byte         `json:"-"`
 	BodyLength   int            `json:"body_length"`
 	Header       []byte         `json:"-"`
@@ -51,6 +53,7 @@ type baseline struct {
 	Status       int            `json:"status"`
 	IsDynamicUrl bool           `json:"is_dynamic_url"` // 判断是否存在动态的url
 	Spended      int            `json:"spended"`        // 耗时, 毫秒
+	Title        string         `json:"title"`
 	Frameworks   pkg.Frameworks `json:"frameworks"`
 	Extracteds   pkg.Extracteds `json:"extracts"`
 	Err          error          `json:"-"`
@@ -61,6 +64,7 @@ type baseline struct {
 // Collect 深度收集信息
 func (bl *baseline) Collect() {
 	bl.Hashes = parsers.NewHashes(bl.Raw)
+	bl.Title = parsers.MatchTitle(string(bl.Body))
 	// todo extract
 	bl.Extracteds = pkg.Extractors.Extract(string(bl.Raw))
 	// todo 指纹识别
@@ -95,11 +99,21 @@ func (bl *baseline) String() string {
 	var line strings.Builder
 	//line.WriteString("[+] ")
 	line.WriteString(bl.UrlString)
-	line.WriteString(fmt.Sprintf(" - %d - %d ", bl.Status, bl.BodyLength))
+	line.WriteString(" (" + bl.Host + ")")
+	line.WriteString(" - ")
+	line.WriteString(strconv.Itoa(bl.Status))
+	line.WriteString(" - ")
+	line.WriteString(strconv.Itoa(bl.BodyLength))
 	if bl.RedirectURL != "" {
-		line.WriteString("-> ")
+		line.WriteString(" -> ")
 		line.WriteString(bl.RedirectURL)
+		line.WriteString(" ")
 	}
+	line.WriteString(" [" + bl.Title + "]")
+	if bl.Hashes != nil {
+		line.WriteString(" [" + bl.Hashes.BodyMd5 + "]")
+	}
+
 	line.WriteString(bl.Frameworks.ToString())
 	//line.WriteString(bl.Extracteds)
 	//line.WriteString("\n")

internal/pool.go

@@ -4,9 +4,11 @@ import (
 	"context"
 	"github.com/chainreactors/logs"
 	"github.com/chainreactors/spray/pkg"
+	"github.com/chainreactors/spray/pkg/ihttp"
 	"github.com/chainreactors/words"
 	"github.com/panjf2000/ants/v2"
 	"github.com/valyala/fasthttp"
+	"net/http"
 	"sync"
 	"time"
 )
@@ -25,7 +27,7 @@ func NewPool(ctx context.Context, config *pkg.Config, outputCh chan *baseline) (
 		Config:      config,
 		ctx:         pctx,
 		cancel:      cancel,
-		client:      pkg.NewClient(config.Thread, 2),
+		client:      ihttp.NewClient(config.Thread, 2, config.ClientType),
 		worder:      words.NewWorder(config.Wordlist),
 		outputCh:    outputCh,
 		tempCh:      make(chan *baseline, config.Thread),
@@ -37,13 +39,32 @@ func NewPool(ctx context.Context, config *pkg.Config, outputCh chan *baseline) (
 
 	switch config.Mod {
 	case pkg.PathSpray:
-		pool.genReq = func(s string) (*fasthttp.Request, error) {
+		pool.genReq = func(s string) (*ihttp.Request, error) {
 			return pool.buildPathRequest(s)
 		}
+		pool.check = func() {
+			pool.wg.Add(1)
+			_ = pool.pool.Invoke(newUnit(pkg.RandPath(), CheckSource))
+
+			if pool.failedCount > breakThreshold {
+				// 当报错次数超过上限是, 结束任务
+				pool.cancel()
+			}
+		}
 	case pkg.HostSpray:
-		pool.genReq = func(s string) (*fasthttp.Request, error) {
+		pool.genReq = func(s string) (*ihttp.Request, error) {
 			return pool.buildHostRequest(s)
 		}
+
+		pool.check = func() {
+			pool.wg.Add(1)
+			_ = pool.pool.Invoke(newUnit(pkg.RandHost(), CheckSource))
+
+			if pool.failedCount > breakThreshold {
+				// 当报错次数超过上限是, 结束任务
+				pool.cancel()
+			}
+		}
 	}
 
 	p, _ := ants.NewPoolWithFunc(config.Thread, func(i interface{}) {
@@ -56,8 +77,11 @@ func NewPool(ctx context.Context, config *pkg.Config, outputCh chan *baseline) (
 
 		var bl *baseline
 		resp, reqerr := pool.client.Do(pctx, req)
-		defer fasthttp.ReleaseResponse(resp)
-		defer fasthttp.ReleaseRequest(req)
+		if pool.ClientType == ihttp.FAST {
+			defer fasthttp.ReleaseResponse(resp.FastResponse)
+			defer fasthttp.ReleaseRequest(req.FastRequest)
+		}
+
 		if reqerr != nil && reqerr != fasthttp.ErrBodyTooLarge {
 			pool.failedCount++
 			bl = &baseline{UrlString: pool.BaseURL + unit.path, Err: reqerr}
@@ -65,9 +89,9 @@ func NewPool(ctx context.Context, config *pkg.Config, outputCh chan *baseline) (
 			pool.failedCount = 0
 			if err = pool.PreCompare(resp); err == nil || unit.source == CheckSource {
 				// 通过预对比跳过一些无用数据, 减少性能消耗
-				bl = NewBaseline(req.URI(), resp)
+				bl = NewBaseline(req.URI(), req.Host(), resp)
 			} else {
-				bl = NewInvalidBaseline(req.URI(), resp)
+				bl = NewInvalidBaseline(req.URI(), req.Host(), resp)
 			}
 			bl.Err = reqerr
 		}
@@ -108,7 +132,7 @@ func NewPool(ctx context.Context, config *pkg.Config, outputCh chan *baseline) (
 
 type Pool struct {
 	*pkg.Config
-	client *pkg.Client
+	client *ihttp.Client
 	pool   *ants.PoolWithFunc
 	bar    *pkg.Bar
 	ctx    context.Context
@@ -122,22 +146,13 @@ type Pool struct {
 	checkPeriod int
 	errPeriod   int
 	analyzeDone bool
-	genReq      func(s string) (*fasthttp.Request, error)
+	genReq      func(s string) (*ihttp.Request, error)
+	check       func()
 	worder      *words.Worder
 	wg          sync.WaitGroup
 	initwg      sync.WaitGroup // 初始化用, 之后改成锁
 }
 
-func (p *Pool) check() {
-	p.wg.Add(1)
-	_ = p.pool.Invoke(newUnit(pkg.RandPath(), CheckSource))
-
-	if p.failedCount > breakThreshold {
-		// 当报错次数超过上限是, 结束任务
-		p.cancel()
-	}
-}
-
 func (p *Pool) Init() error {
 	p.initwg.Add(1)
 	p.check()
@@ -189,12 +204,12 @@ Loop:
 	p.Close()
 }
 
-func (p *Pool) PreCompare(resp *fasthttp.Response) error {
+func (p *Pool) PreCompare(resp *ihttp.Response) error {
 	if !CheckStatusCode(resp.StatusCode()) {
 		return ErrBadStatus
 	}
 
-	if CheckRedirect != nil && !CheckRedirect(string(resp.Header.Peek("Location"))) {
+	if CheckRedirect != nil && !CheckRedirect(string(resp.GetHeader("Location"))) {
 		return ErrRedirect
 	}
 
@@ -238,15 +253,26 @@ func (p *Pool) Close() {
 	}
 }
 
-func (p *Pool) buildPathRequest(path string) (*fasthttp.Request, error) {
-	req := fasthttp.AcquireRequest()
-	req.SetRequestURI(p.BaseURL + path)
-	return req, nil
+func (p *Pool) buildPathRequest(path string) (*ihttp.Request, error) {
+	if p.Config.ClientType == ihttp.FAST {
+		req := fasthttp.AcquireRequest()
+		req.SetRequestURI(p.BaseURL + path)
+		return &ihttp.Request{FastRequest: req}, nil
+	} else {
+		req, err := http.NewRequest("GET", p.BaseURL+path, nil)
+		return &ihttp.Request{StandardRequest: req}, err
+	}
 }
 
-func (p *Pool) buildHostRequest(host string) (*fasthttp.Request, error) {
-	req := fasthttp.AcquireRequest()
-	req.SetRequestURI(p.BaseURL)
-	req.SetHost(host)
-	return req, nil
+func (p *Pool) buildHostRequest(host string) (*ihttp.Request, error) {
+	if p.Config.ClientType == ihttp.FAST {
+		req := fasthttp.AcquireRequest()
+		req.SetRequestURI(p.BaseURL)
+		req.SetHost(host)
+		return &ihttp.Request{FastRequest: req}, nil
+	} else {
+		req, err := http.NewRequest("GET", p.BaseURL, nil)
+		req.Host = host
+		return &ihttp.Request{StandardRequest: req}, err
+	}
 }

internal/runner.go

@@ -5,7 +5,9 @@ import (
 	"fmt"
 	"github.com/chainreactors/logs"
 	"github.com/chainreactors/spray/pkg"
+	"github.com/chainreactors/spray/pkg/ihttp"
 	"github.com/gosuri/uiprogress"
+	"github.com/panjf2000/ants/v2"
 	"io/ioutil"
 	"net/http"
 	"os"
@@ -27,8 +29,9 @@ type Runner struct {
 	Offset     int         `long:"offset"`
 	Limit      int         `long:"limit"`
 	Threads    int         `short:"t" long:"thread" default:"20"`
-	PoolSize   int         `short:"p" long:"pool"`
-	Pools      map[string]*Pool
+	PoolSize   int         `short:"p" long:"pool" default:"5"`
+	Pools      *ants.PoolWithFunc
+	poolwg     sync.WaitGroup
 	Deadline   int    `long:"deadline" default:"600"` // todo 总的超时时间,适配云函数的deadline
 	Debug      bool   `long:"debug"`
 	Quiet      bool   `short:"q" long:"quiet"`
@@ -108,46 +111,52 @@ func (r *Runner) Prepare() error {
 	}
 
 	r.OutputCh = make(chan *baseline, 100)
-	r.Pools = make(map[string]*Pool)
+	ctx := context.Background()
+
+	r.Pools, err = ants.NewPoolWithFunc(r.PoolSize, func(i interface{}) {
+		u := i.(string)
+		config := &pkg.Config{
+			BaseURL:      u,
+			Wordlist:     r.Wordlist,
+			Thread:       r.Threads,
+			Timeout:      2,
+			Headers:      r.Headers,
+			Mod:          pkg.ModMap[r.Mod],
+			DeadlineTime: r.Deadline,
+		}
+
+		if config.Mod == pkg.PathSpray {
+			config.ClientType = ihttp.FAST
+		} else if config.Mod == pkg.HostSpray {
+			config.ClientType = ihttp.STANDARD
+		}
+
+		pool, err := NewPool(ctx, config, r.OutputCh)
+		if err != nil {
+			logs.Log.Error(err.Error())
+			return
+		}
+		pool.bar = pkg.NewBar(u, len(r.Wordlist), r.Progress)
+		err = pool.Init()
+		if err != nil {
+			logs.Log.Error(err.Error())
+			return
+		}
+		// todo pool 总超时时间
+		pool.Run(ctx)
+		r.poolwg.Done()
+	})
 	go r.Outputting()
 	return nil
 }
 
 func (r *Runner) Run() {
 	// todo pool 结束与并发控制
-	ctx := context.Background()
-	var wg sync.WaitGroup
 	for _, u := range r.URLList {
-		wg.Add(1)
-		u := u
-		go func() {
-			config := &pkg.Config{
-				BaseURL:      u,
-				Wordlist:     r.Wordlist,
-				Thread:       r.Threads,
-				Timeout:      2,
-				Headers:      r.Headers,
-				Mod:          pkg.ModMap[r.Mod],
-				DeadlineTime: r.Deadline,
-			}
-			pool, err := NewPool(ctx, config, r.OutputCh)
-			if err != nil {
-				logs.Log.Error(err.Error())
-				return
-			}
-			pool.bar = pkg.NewBar(u, len(r.Wordlist), r.Progress)
-			err = pool.Init()
-			if err != nil {
-				logs.Log.Error(err.Error())
-				return
-			}
-			r.Pools[u] = pool
-			// todo pool 总超时时间
-			pool.Run(ctx)
-			wg.Done()
-		}()
+		r.poolwg.Add(1)
+		r.Pools.Invoke(u)
 	}
-	wg.Wait()
+	r.poolwg.Wait()
 	for {
 		if len(r.OutputCh) == 0 {
 			close(r.OutputCh)

pkg/client.go

@@ -1,67 +0,0 @@
-package pkg
-
-import (
-	"context"
-	"crypto/tls"
-	"github.com/valyala/fasthttp"
-	"net/http"
-	"time"
-)
-
-var (
-	DefaultMaxBodySize = 1024 * 100 // 100k
-)
-
-func NewClient(thread int, timeout int) *Client {
-	c := &Client{
-		client: &fasthttp.Client{
-			TLSConfig: &tls.Config{
-				Renegotiation:      tls.RenegotiateOnceAsClient,
-				InsecureSkipVerify: true,
-			},
-			//ReadBufferSize:      20480,
-			MaxConnsPerHost:     thread * 2,
-			MaxIdleConnDuration: time.Duration(timeout) * time.Second,
-			MaxConnWaitTimeout:  time.Duration(timeout) * time.Second,
-			ReadTimeout:         time.Duration(timeout) * time.Second,
-			WriteTimeout:        time.Duration(timeout) * time.Second,
-			MaxResponseBodySize: DefaultMaxBodySize,
-		},
-		timeout: time.Duration(timeout) * time.Second,
-	}
-	//c := &Client{
-	//	client: &http.Client{
-	//		Transport:     tr,
-	//		Timeout:       time.Second * time.Duration(timeout),
-	//		CheckRedirect: checkRedirect,
-	//	},
-	//}
-
-	//c.Method = method
-	//c.Headers = Opt.Headers
-	//c.Mod = Opt.Mod
-
-	return c
-}
-
-type Client struct {
-	client  *fasthttp.Client
-	timeout time.Duration
-}
-
-func (c *Client) Do(ctx context.Context, req *fasthttp.Request) (*fasthttp.Response, error) {
-	//if req.Header == nil {
-	//	req.Header = c.Headers
-	//}
-	resp := fasthttp.AcquireResponse()
-	return resp, c.client.Do(req, resp)
-}
-
-var MaxRedirects = 0
-var checkRedirect = func(req *http.Request, via []*http.Request) error {
-	if len(via) > MaxRedirects {
-		return http.ErrUseLastResponse
-	}
-
-	return nil
-}

pkg/config.go

@@ -29,4 +29,5 @@ type Config struct {
 	Headers      http.Header
 	DeadlineTime int
 	EnableFuzzy  bool
+	ClientType   int
 }

pkg/ihttp/client.go

@@ -0,0 +1,94 @@
+package ihttp
+
+import (
+	"context"
+	"crypto/tls"
+	"fmt"
+	"github.com/valyala/fasthttp"
+	"net/http"
+	"time"
+)
+
+var (
+	DefaultMaxBodySize = 1024 * 100 // 100k
+)
+
+const (
+	FAST = iota
+	STANDARD
+)
+
+func NewClient(thread int, timeout int, clientType int) *Client {
+	if clientType == FAST {
+		return &Client{
+			fastClient: &fasthttp.Client{
+				TLSConfig: &tls.Config{
+					Renegotiation:      tls.RenegotiateOnceAsClient,
+					InsecureSkipVerify: true,
+				},
+				//ReadBufferSize:      20480,
+				MaxConnsPerHost:     thread * 2,
+				MaxIdleConnDuration: time.Duration(timeout) * time.Second,
+				MaxConnWaitTimeout:  time.Duration(timeout) * time.Second,
+				ReadTimeout:         time.Duration(timeout) * time.Second,
+				WriteTimeout:        time.Duration(timeout) * time.Second,
+				MaxResponseBodySize: DefaultMaxBodySize,
+			},
+			timeout: time.Duration(timeout) * time.Second,
+		}
+	} else {
+		return &Client{
+			standardClient: &http.Client{
+				Transport: &http.Transport{
+					//Proxy: Proxy,
+					//TLSHandshakeTimeout : delay * time.Second,
+					TLSClientConfig: &tls.Config{
+						Renegotiation:      tls.RenegotiateOnceAsClient,
+						InsecureSkipVerify: true,
+					},
+					MaxConnsPerHost: thread,
+					IdleConnTimeout: time.Duration(timeout) * time.Second,
+				},
+				Timeout:       time.Second * time.Duration(timeout),
+				CheckRedirect: checkRedirect,
+			},
+			timeout: time.Duration(timeout) * time.Second,
+		}
+	}
+}
+
+type Client struct {
+	fastClient     *fasthttp.Client
+	standardClient *http.Client
+	timeout        time.Duration
+}
+
+func (c *Client) FastDo(ctx context.Context, req *fasthttp.Request) (*fasthttp.Response, error) {
+	resp := fasthttp.AcquireResponse()
+	return resp, c.fastClient.Do(req, resp)
+}
+
+func (c *Client) StandardDo(ctx context.Context, req *http.Request) (*http.Response, error) {
+	return c.standardClient.Do(req)
+}
+
+func (c *Client) Do(ctx context.Context, req *Request) (*Response, error) {
+	if c.fastClient != nil {
+		resp, err := c.FastDo(ctx, req.FastRequest)
+		return &Response{FastResponse: resp}, err
+	} else if c.standardClient != nil {
+		resp, err := c.StandardDo(ctx, req.StandardRequest)
+		return &Response{StandardResponse: resp}, err
+	} else {
+		return nil, fmt.Errorf("not found client")
+	}
+}
+
+var MaxRedirects = 0
+var checkRedirect = func(req *http.Request, via []*http.Request) error {
+	if len(via) > MaxRedirects {
+		return http.ErrUseLastResponse
+	}
+
+	return nil
+}

pkg/ihttp/request.go

@@ -0,0 +1,31 @@
+package ihttp
+
+import (
+	"github.com/valyala/fasthttp"
+	"net/http"
+)
+
+type Request struct {
+	StandardRequest *http.Request
+	FastRequest     *fasthttp.Request
+}
+
+func (r *Request) URI() string {
+	if r.FastRequest != nil {
+		return r.FastRequest.URI().String()
+	} else if r.StandardRequest != nil {
+		return r.StandardRequest.URL.String()
+	} else {
+		return ""
+	}
+}
+
+func (r *Request) Host() string {
+	if r.FastRequest != nil {
+		return string(r.FastRequest.Host())
+	} else if r.StandardRequest != nil {
+		return r.StandardRequest.Host
+	} else {
+		return ""
+	}
+}

pkg/ihttp/response.go

@@ -0,0 +1,83 @@
+package ihttp
+
+import (
+	"bytes"
+	"github.com/chainreactors/logs"
+	"github.com/valyala/fasthttp"
+	"io"
+	"net/http"
+)
+
+type Response struct {
+	StandardResponse *http.Response
+	FastResponse     *fasthttp.Response
+}
+
+func (r *Response) StatusCode() int {
+	if r.FastResponse != nil {
+		return r.FastResponse.StatusCode()
+	} else if r.StandardResponse != nil {
+		return r.StandardResponse.StatusCode
+	} else {
+		return 0
+	}
+}
+
+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 {
+			n, err := io.ReadFull(r.StandardResponse.Body, body)
+			_ = r.StandardResponse.Body.Close()
+			if err == nil {
+				return body
+			} else if err == io.ErrUnexpectedEOF {
+				return body[:n]
+			} else {
+				logs.Log.Error("readfull failed" + err.Error())
+				return nil
+			}
+		}
+		return nil
+	} else {
+		return nil
+	}
+}
+
+func (r *Response) ContentLength() int {
+	if r.FastResponse != nil {
+		return r.FastResponse.Header.ContentLength()
+	} else if r.StandardResponse != nil {
+		return int(r.StandardResponse.ContentLength)
+	} else {
+		return 0
+	}
+}
+
+func (r *Response) Header() []byte {
+	if r.FastResponse != nil {
+		return r.FastResponse.Header.Header()
+	} else if r.StandardResponse != nil {
+		var header bytes.Buffer
+		for k, v := range r.StandardResponse.Header {
+			for _, i := range v {
+				header.WriteString(k + ": " + i + "\r\n")
+			}
+		}
+		return header.Bytes()
+	} else {
+		return nil
+	}
+}
+
+func (r *Response) GetHeader(key string) string {
+	if r.FastResponse != nil {
+		return string(r.FastResponse.Header.Peek(key))
+	} else if r.StandardResponse != nil {
+		return r.StandardResponse.Header.Get(key)
+	} else {
+		return ""
+	}
+}

pkg/utils.go

@@ -56,3 +56,23 @@ func RandPath() string {
 	}
 	return *(*string)(unsafe.Pointer(&b))
 }
+
+func RandHost() string {
+	n := 8
+	b := make([]byte, n)
+	// A rand.Int63() generates 63 random bits, enough for letterIdMax letters!
+	for i, cache, remain := n-1, src.Int63(), letterIdMax; i >= 1; {
+		if remain == 0 {
+			cache, remain = src.Int63(), letterIdMax
+		}
+		if idx := int(cache & letterIdMask); idx < len(letters) {
+			b[i] = letters[idx]
+			i--
+		}
+		cache >>= letterIdBits
+		remain--
+	}
+
+	b[5] = byte(0x2e)
+	return *(*string)(unsafe.Pointer(&b))
+}