Merge branch 'rate-limit' of github.com:CasperGN/nuclei into CasperGN…

…-rate-limit
projectdiscovery · Sep 23, 2020 · 207e78b · 207e78b
2 parents ef49270 + ce051a1
commit 207e78b
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 0 deletions.
diff --git a/v2/go.mod b/v2/go.mod
@@ -17,5 +17,9 @@ require (
 	github.com/projectdiscovery/retryablehttp-go v1.0.1
 	github.com/vbauerster/mpb/v5 v5.3.0
 	golang.org/x/net v0.0.0-20200923182212-328152dc79b1
+	github.com/stretchr/testify v1.5.1
+	github.com/vbauerster/mpb/v5 v5.2.4
+	golang.org/x/net v0.0.0-20200707034311-ab3426394381
+	golang.org/x/time v0.0.0-20200630173020-3af7569d3a1e
 	gopkg.in/yaml.v2 v2.3.0
 )
diff --git a/v2/go.sum b/v2/go.sum
@@ -72,6 +72,8 @@ golang.org/x/sys v0.0.0-20200323222414-85ca7c5b95cd/go.mod h1:h1NjWce9XRLGQEsW7w
 golang.org/x/sys v0.0.0-20200810151505-1b9f1253b3ed h1:WBkVNH1zd9jg/dK4HCM4lNANnmd12EHC9z+LmcCG4ns=
 golang.org/x/sys v0.0.0-20200810151505-1b9f1253b3ed/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
+golang.org/x/time v0.0.0-20200630173020-3af7569d3a1e h1:EHBhcS0mlXEAVwNyO2dLfjToGsyY4j24pTs2ScHnX7s=
+golang.org/x/time v0.0.0-20200630173020-3af7569d3a1e/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
 golang.org/x/tools v0.0.0-20191216052735-49a3e744a425/go.mod h1:TB2adYChydJhpapKDTa4BR/hXlZSLoq2Wpct/0txZ28=
 golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 google.golang.org/appengine v1.1.0/go.mod h1:EbEs0AVv82hx2wNQdGPgUI5lhzA/G0D9YwlJXL52JkM=

diff --git a/v2/internal/runner/options.go b/v2/internal/runner/options.go
@@ -38,6 +38,7 @@ type Options struct {
 	ProxySocksURL      string                 // ProxySocksURL is the URL for the proxy socks server
 	CustomHeaders      requests.CustomHeaders // Custom global headers
 	TemplatesDirectory string                 // TemplatesDirectory is the directory to use for storing templates
+	RateLimit          int                    // Rate-Limit of requests per specified target
 }
 
 type multiStringFlag []string
@@ -78,6 +79,7 @@ func ParseOptions() *Options {
 	flag.BoolVar(&options.JSONRequests, "json-requests", false, "Write requests/responses for matches in JSON output")
 	flag.BoolVar(&options.EnableProgressBar, "pbar", false, "Enable the progress bar")
 	flag.BoolVar(&options.TemplateList, "tl", false, "List available templates")
+	flag.IntVar(&options.RateLimit, "rl", 9999999, "Rate-Limit of requests per specified target") // 9999999 to avoid limiting
 
 	flag.Parse()
 

diff --git a/v2/internal/runner/runner.go b/v2/internal/runner/runner.go
@@ -10,6 +10,9 @@ import (
 	"regexp"
 	"strings"
 	"sync"
+	"time"
+
+	"golang.org/x/time/rate"
 
 	"github.com/logrusorgru/aurora"
 	"github.com/projectdiscovery/gologger"
@@ -238,6 +241,8 @@ func (r *Runner) RunEnumeration() {
 		gologger.Errorf("Could not find any valid input URLs.")
 	} else if totalRequests > 0 || hasWorkflows {
 		ctx := context.Background()
+		// Limiter that will add to the tokenbucket every second and set the max size to -rl flag
+		rateLimit := rate.NewLimiter(rate.Every(1*time.Second), r.options.RateLimit)
 		// tracks global progress and captures stdout/stderr until p.Wait finishes
 		p := r.progress
 		p.InitProgressbar(r.inputCount, templateCount, totalRequests)
@@ -246,6 +251,10 @@ func (r *Runner) RunEnumeration() {
 			wgtemplates.Add(1)
 			go func(template interface{}) {
 				defer wgtemplates.Done()
+				err := rateLimit.Wait(ctx)
+				if err != nil {
+					gologger.Errorf("Issue with rate-limit")
+				}
 				switch tt := template.(type) {
 				case *templates.Template:
 					for _, request := range tt.RequestsDNS {