External/nuclei
1
0
Fork 0
mirror of https://github.com/projectdiscovery/nuclei.git synced 2025-12-18 17:55:25 +00:00
Code Issues Packages Projects Releases Wiki Activity

Appending RateLimiter to th runner

Browse Source 
Signed-off-by: Casper Guldbech Nielsen <whopsec@protonmail.com>
This commit is contained in:
Casper Guldbech Nielsen 2020-09-12 17:30:54 +02:00
parent 7ceca4daeb
commit 5c67ff9a5d
1 changed files with 9 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
internal/runner/runner.go
View File

@ -10,6 +10,9 @@ import (
"regexp" "regexp"
"strings" "strings"
"sync" "sync"
"time"
"golang.org/x/time/rate"
"github.com/logrusorgru/aurora" "github.com/logrusorgru/aurora"
"github.com/projectdiscovery/gologger" "github.com/projectdiscovery/gologger"
@ -237,6 +240,8 @@ func (r *Runner) RunEnumeration() {
gologger.Errorf("Could not find any valid input URLs.") gologger.Errorf("Could not find any valid input URLs.")
} else if totalRequests > 0 || hasWorkflows { } else if totalRequests > 0 || hasWorkflows {
ctx := context.Background() 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 // tracks global progress and captures stdout/stderr until p.Wait finishes
p := r.progress p := r.progress
p.InitProgressbar(r.inputCount, templateCount, totalRequests) p.InitProgressbar(r.inputCount, templateCount, totalRequests)
@ -245,6 +250,10 @@ func (r *Runner) RunEnumeration() {
wgtemplates.Add(1) wgtemplates.Add(1)
go func(template interface{}) { go func(template interface{}) {
defer wgtemplates.Done() defer wgtemplates.Done()
err := rateLimit.Wait(ctx)
if err != nil {
gologger.Errorf("Issue with rate-limit")
}
switch tt := template.(type) { switch tt := template.(type) {
case *templates.Template: case *templates.Template:
for _, request := range tt.RequestsDNS { for _, request := range tt.RequestsDNS {