新增--rate-limit, 用来限制单个pool请求速率

go.mod

@@ -32,5 +32,6 @@ require (
 	github.com/twmb/murmur3 v1.1.6 // indirect
 	github.com/valyala/bytebufferpool v1.0.0 // indirect
 	golang.org/x/sys v0.2.0 // indirect
+	golang.org/x/time v0.3.0 // indirect
 	gopkg.in/yaml.v2 v2.4.0 // indirect
 )

go.sum

@@ -103,6 +103,8 @@ golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk=
 golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=
+golang.org/x/time v0.3.0 h1:rg5rLMjNzMS1RkNLzCG38eapWhnYLFYXDXj2gOlr8j4=
+golang.org/x/time v0.3.0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=

internal/option.go

@@ -21,6 +21,7 @@ type Option struct {
 	InputOptions    `group:"Input Options"`
 	FunctionOptions `group:"Function Options"`
 	OutputOptions   `group:"Output Options"`
+	PluginOptions   `group:"Plugin Options"`
 	RequestOptions  `group:"Request Options"`
 	ModeOptions     `group:"Modify Options"`
 	MiscOptions     `group:"Miscellaneous Options"`
@@ -28,15 +29,15 @@ type Option struct {
 
 type InputOptions struct {
 	ResumeFrom   string   `long:"resume"`
-	URL          []string `short:"u" long:"url" description:"String, Multi, input baseurl, e.g.: http://google.com"`
+	URL          []string `short:"u" long:"url" description:"Strings, input baseurl, e.g.: http://google.com"`
 	URLFile      string   `short:"l" long:"list" description:"File, input filename"`
 	Raw          string   `long:"raw" description:"File, input raw request filename"`
 	Offset       int      `long:"offset" description:"Int, wordlist offset"`
 	Limit        int      `long:"limit" description:"Int, wordlist limit, start with offset. e.g.: --offset 1000 --limit 100"`
 	Dictionaries []string `short:"d" long:"dict" description:"Files, Multi,dict files, e.g.: -d 1.txt -d 2.txt"`
 	Word         string   `short:"w" long:"word" description:"String, word generate dsl, e.g.: -w test{?ld#4}"`
-	Rules        []string `short:"r" long:"rules" description:"Files, Multi, rule files, e.g.: -r rule1.txt -r rule2.txt"`
+	Rules        []string `short:"r" long:"rules" description:"Files, rule files, e.g.: -r rule1.txt -r rule2.txt"`
-	AppendRule   []string `long:"append-rule" description:"File, when found valid path , use append rule generator new word with current path"`
+	AppendRule   []string `long:"append-rule" description:"Files, when found valid path , use append rule generator new word with current path"`
 	FilterRule   string   `long:"filter-rule" description:"String, filter rule, e.g.: --rule-filter '>8 <4'"`
 }
 
@@ -46,15 +47,15 @@ type FunctionOptions struct {
 	RemoveExtensions  string            `long:"remove-extension" description:"String, remove extensions (separated by commas), e.g.: --remove-extension jsp,jspx"`
 	Uppercase         bool              `short:"U" long:"uppercase" desvcription:"Bool, upper wordlist, e.g.: --uppercase"`
 	Lowercase         bool              `short:"L" long:"lowercase" description:"Bool, lower wordlist, e.g.: --lowercase"`
-	Prefixes          []string          `long:"prefix" description:"Strings, Multi, add prefix, e.g.: --prefix aaa --prefix bbb"`
+	Prefixes          []string          `long:"prefix" description:"Strings, add prefix, e.g.: --prefix aaa --prefix bbb"`
-	Suffixes          []string          `long:"suffix" description:"Strings, Multi, add suffix, e.g.: --suffix aaa --suffix bbb"`
+	Suffixes          []string          `long:"suffix" description:"Strings, add suffix, e.g.: --suffix aaa --suffix bbb"`
-	Replaces          map[string]string `long:"replace" description:"Strings, Multi, replace string, e.g.: --replace aaa:bbb --replace ccc:ddd"`
+	Replaces          map[string]string `long:"replace" description:"Strings, replace string, e.g.: --replace aaa:bbb --replace ccc:ddd"`
 }
 
 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"`
+	Extracts    []string `long:"extract" description:"Strings, extract response, e.g.: --extract js --extract ip --extract version:(.*?)" json:"extracts,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"`
@@ -66,27 +67,31 @@ type OutputOptions struct {
 }
 
 type RequestOptions struct {
-	Headers         []string `long:"header" description:"String, Multi, custom headers, e.g.: --headers 'Auth: example_auth'"`
+	Headers         []string `long:"header" description:"Strings, custom headers, e.g.: --headers 'Auth: example_auth'"`
 	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"`
+	Cookie          []string `long:"cookie" description:"Strings, 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"`
 }
 
-type ModeOptions struct {
+type PluginOptions struct {
 	Advance    bool   `short:"a" long:"advance" description:"Bool, enable crawl and active"`
 	Active     bool   `long:"active" description:"Bool, enable active finger detect"`
-	Crawl           bool   `long:"crawl" description:"Bool, enable crawl"`
 	Bak        bool   `long:"bak" description:"Bool, enable bak found"`
 	FileBak    bool   `long:"file-bak" description:"Bool, enable valid result bak found, equal --append-rule rule/filebak.txt"`
 	Common     bool   `long:"common" description:"Bool, enable common file found"`
+	Crawl      bool   `long:"crawl" description:"Bool, enable crawl"`
+	CrawlDepth int    `long:"crawl-depth" default:"3" description:"Int, crawl depth"`
+	CrawlScope string `long:"crawl-scope" description:"Int, crawl scope (todo)"`
+}
+
+type ModeOptions struct {
+	RateLimit       int    `long:"rate-limit" default:"0" description:"Int, request rate limit (rate/s), e.g.: --rate-limit 100"`
 	Force           bool   `long:"force" description:"Bool, skip error break"`
 	CheckOnly       bool   `long:"check-only" description:"Bool, check only"`
 	Recursive       string `long:"recursive" default:"current.IsDir()" description:"String,custom recursive rule, e.g.: --recursive current.IsDir()"`
 	Depth           int    `long:"depth" default:"0" description:"Int, recursive depth"`
-	CrawlDepth      int    `long:"crawl-depth" default:"3" description:"Int, crawl depth"`
-	CrawlScope      string `long:"crawl-scope" description:"Int, crawl scope (todo)"`
 	CheckPeriod     int    `long:"check-period" default:"200" description:"Int, check period when request"`
 	ErrPeriod       int    `long:"error-period" default:"10" description:"Int, check period when error"`
 	BreakThreshold  int    `long:"error-threshold" default:"20" description:"Int, break when the error exceeds the threshold "`
@@ -121,6 +126,7 @@ func (opt *Option) PrepareRunner() (*Runner, error) {
 		PoolSize:       opt.PoolSize,
 		Mod:            opt.Mod,
 		Timeout:        opt.Timeout,
+		RateLimit:      opt.RateLimit,
 		Deadline:       opt.Deadline,
 		Headers:        make(map[string]string),
 		Offset:         opt.Offset,

internal/pool.go

@@ -13,6 +13,7 @@ import (
 	"github.com/chainreactors/words/rule"
 	"github.com/panjf2000/ants/v2"
 	"github.com/valyala/fasthttp"
+	"golang.org/x/time/rate"
 	"net/url"
 	"path"
 	"strconv"
@@ -50,8 +51,9 @@ 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),
-		wg:          sync.WaitGroup{},
+		waiter:      sync.WaitGroup{},
 		initwg:      sync.WaitGroup{},
+		limiter:     rate.NewLimiter(rate.Limit(config.RateLimit), 1),
 		reqCount:    1,
 		failedCount: 1,
 	}
@@ -127,7 +129,7 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) {
 
 			// 如果要进行递归判断, 要满足 bl有效, mod为path-spray, 当前深度小于最大递归深度
 			if bl.IsValid {
-				pool.wg.Add(2)
+				pool.waiter.Add(2)
 				pool.doCrawl(bl)
 				pool.doRule(bl)
 				if bl.RecuDepth < MaxRecursion {
@@ -137,7 +139,7 @@ func NewPool(ctx context.Context, config *pkg.Config) (*Pool, error) {
 				}
 			}
 			pool.OutputCh <- bl
-			pool.wg.Done()
+			pool.waiter.Done()
 		}
 
 		pool.analyzeDone = true
@@ -170,8 +172,9 @@ type Pool struct {
 	urls            map[string]struct{}
 	analyzeDone     bool
 	worder          *words.Worder
+	limiter         *rate.Limiter
 	locker          sync.Mutex
-	wg              sync.WaitGroup
+	waiter          sync.WaitGroup
 	initwg          sync.WaitGroup // 初始化用, 之后改成锁
 }
 
@@ -237,17 +240,17 @@ func (pool *Pool) genReq(s string) (*ihttp.Request, error) {
 func (pool *Pool) Run(ctx context.Context, offset, limit int) {
 	pool.worder.RunWithRules()
 	if pool.Active {
-		pool.wg.Add(1)
+		pool.waiter.Add(1)
 		go pool.doActive()
 	}
 
 	if pool.Bak {
-		pool.wg.Add(1)
+		pool.waiter.Add(1)
 		go pool.doBak()
 	}
 
 	if pool.Common {
-		pool.wg.Add(1)
+		pool.waiter.Add(1)
 		go pool.doCommonFile()
 	}
 
@@ -256,7 +259,7 @@ func (pool *Pool) Run(ctx context.Context, offset, limit int) {
 	wait := func() {
 		if !worderDone {
 			worderDone = true
-			pool.wg.Wait()
+			pool.waiter.Wait()
 			close(closeCh)
 		}
 	}
@@ -280,7 +283,7 @@ Loop:
 				continue
 			}
 
-			pool.wg.Add(1)
+			pool.waiter.Add(1)
 			pool.urls[u] = struct{}{}
 			pool.reqPool.Invoke(newUnit(pool.safePath(u), WordSource)) // 原样的目录拼接, 输入了几个"/"就是几个, 适配java的目录解析
 		case source := <-pool.checkCh:
@@ -296,7 +299,7 @@ Loop:
 			}
 			if _, ok := pool.urls[unit.path]; ok {
 				logs.Log.Debugf("[%s] duplicate path: %s, skipped", pkg.GetSourceName(unit.source), pool.base+unit.path)
-				pool.wg.Done()
+				pool.waiter.Done()
 			} else {
 				pool.urls[unit.path] = struct{}{}
 				pool.reqPool.Invoke(unit)
@@ -310,12 +313,16 @@ Loop:
 		}
 	}
 
-	pool.wg.Wait()
+	pool.waiter.Wait()
 	pool.Statistor.EndTime = time.Now().Unix()
 	pool.Close()
 }
 
 func (pool *Pool) Invoke(v interface{}) {
+	if pool.RateLimit != 0 {
+		pool.limiter.Wait(pool.ctx)
+	}
+
 	atomic.AddInt32(&pool.Statistor.ReqTotal, 1)
 	unit := v.(*Unit)
 	req, err := pool.genReq(unit.path)
@@ -355,7 +362,7 @@ func (pool *Pool) Invoke(v interface{}) {
 
 	// 手动处理重定向
 	if bl.IsValid && unit.source != CheckSource && bl.RedirectURL != "" {
-		pool.wg.Add(1)
+		pool.waiter.Add(1)
 		pool.doRedirect(bl, unit.depth)
 	}
 
@@ -378,7 +385,7 @@ func (pool *Pool) Invoke(v interface{}) {
 		pool.locker.Lock()
 		pool.index = bl
 		pool.locker.Unlock()
-		pool.wg.Add(1)
+		pool.waiter.Add(1)
 		pool.doCrawl(bl)
 		if bl.Status == 200 || (bl.Status/100) == 3 {
 			pool.OutputCh <- bl
@@ -524,13 +531,13 @@ func (pool *Pool) Upgrade(bl *pkg.Baseline) error {
 }
 
 func (pool *Pool) doRedirect(bl *pkg.Baseline, depth int) {
-	defer pool.wg.Done()
+	defer pool.waiter.Done()
 	if depth >= MaxRedirect {
 		return
 	}
 	reURL := FormatURL(bl.Url.Path, bl.RedirectURL)
 
-	pool.wg.Add(1)
+	pool.waiter.Add(1)
 	go pool.addAddition(&Unit{
 		path:     reURL,
 		source:   RedirectSource,
@@ -541,24 +548,24 @@ func (pool *Pool) doRedirect(bl *pkg.Baseline, depth int) {
 
 func (pool *Pool) doCrawl(bl *pkg.Baseline) {
 	if !pool.Crawl || bl.ReqDepth >= MaxCrawl {
-		pool.wg.Done()
+		pool.waiter.Done()
 		return
 	}
 	bl.CollectURL()
 	if bl.URLs == nil {
-		pool.wg.Done()
+		pool.waiter.Done()
 		return
 	}
 
 	go func() {
-		defer pool.wg.Done()
+		defer pool.waiter.Done()
 		for _, u := range bl.URLs {
 			if u = FormatURL(bl.Url.Path, u); u == "" {
 				continue
 			}
 
 			// 通过map去重,  只有新的url才会进入到该逻辑
-			pool.wg.Add(1)
+			pool.waiter.Add(1)
 			pool.addAddition(&Unit{
 				path:   u,
 				source: CrawlSource,
@@ -571,18 +578,18 @@ func (pool *Pool) doCrawl(bl *pkg.Baseline) {
 
 func (pool *Pool) doRule(bl *pkg.Baseline) {
 	if pool.AppendRule == nil {
-		pool.wg.Done()
+		pool.waiter.Done()
 		return
 	}
 	if bl.Source == int(RuleSource) {
-		pool.wg.Done()
+		pool.waiter.Done()
 		return
 	}
 
 	go func() {
-		defer pool.wg.Done()
+		defer pool.waiter.Done()
 		for u := range rule.RunAsStream(pool.AppendRule.Expressions, path.Base(bl.Path)) {
-			pool.wg.Add(1)
+			pool.waiter.Add(1)
 			pool.addAddition(&Unit{
 				path:   Dir(bl.Url.Path) + u,
 				source: RuleSource,
@@ -592,9 +599,9 @@ func (pool *Pool) doRule(bl *pkg.Baseline) {
 }
 
 func (pool *Pool) doActive() {
-	defer pool.wg.Done()
+	defer pool.waiter.Done()
 	for _, u := range pkg.ActivePath {
-		pool.wg.Add(1)
+		pool.waiter.Add(1)
 		pool.addAddition(&Unit{
 			path:   pool.dir + u[1:],
 			source: ActiveSource,
@@ -603,14 +610,14 @@ func (pool *Pool) doActive() {
 }
 
 func (pool *Pool) doBak() {
-	defer pool.wg.Done()
+	defer pool.waiter.Done()
 	worder, err := words.NewWorderWithDsl("{?0}.{@bak_ext}", [][]string{pkg.BakGenerator(pool.url.Host)}, nil)
 	if err != nil {
 		return
 	}
 	worder.Run()
 	for w := range worder.C {
-		pool.wg.Add(1)
+		pool.waiter.Add(1)
 		pool.addAddition(&Unit{
 			path:   pool.dir + w,
 			source: BakSource,
@@ -623,7 +630,7 @@ func (pool *Pool) doBak() {
 	}
 	worder.Run()
 	for w := range worder.C {
-		pool.wg.Add(1)
+		pool.waiter.Add(1)
 		pool.addAddition(&Unit{
 			path:   pool.dir + w,
 			source: BakSource,
@@ -632,9 +639,9 @@ func (pool *Pool) doBak() {
 }
 
 func (pool *Pool) doCommonFile() {
-	defer pool.wg.Done()
+	defer pool.waiter.Done()
 	for _, u := range mask.SpecialWords["common_file"] {
-		pool.wg.Add(1)
+		pool.waiter.Add(1)
 		pool.addAddition(&Unit{
 			path:   pool.dir + u,
 			source: CommonFileSource,
@@ -665,7 +672,7 @@ func (pool *Pool) addAddition(u *Unit) {
 func (pool *Pool) addFuzzyBaseline(bl *pkg.Baseline) {
 	if _, ok := pool.baselines[bl.Status]; !ok && pkg.IntsContains(FuzzyStatus, bl.Status) {
 		bl.Collect()
-		pool.wg.Add(1)
+		pool.waiter.Add(1)
 		pool.doCrawl(bl)
 		pool.baselines[bl.Status] = bl
 		logs.Log.Infof("[baseline.%dinit] %s", bl.Status, bl.Format([]string{"status", "length", "spend", "title", "frame", "redirect"}))

internal/runner.go

@@ -64,6 +64,7 @@ type Runner struct {
 	Progress       *uiprogress.Progress
 	Offset         int
 	Limit          int
+	RateLimit      int
 	Total          int
 	Deadline       int
 	CheckPeriod    int
@@ -83,6 +84,7 @@ func (r *Runner) PrepareConfig() *pkg.Config {
 	config := &pkg.Config{
 		Thread:         r.Threads,
 		Timeout:        r.Timeout,
+		RateLimit:      r.RateLimit,
 		Headers:        r.Headers,
 		Mod:            pkg.ModMap[r.Mod],
 		OutputCh:       r.OutputCh,

pkg/config.go

@@ -24,6 +24,7 @@ type Config struct {
 	Thread         int
 	Wordlist       []string
 	Timeout        int
+	RateLimit      int
 	CheckPeriod    int
 	ErrPeriod      int
 	BreakThreshold int