Skip to content

Commit

Permalink
Merge pull request #4986 from projectdiscovery/feat-3072-init-adaptiv…
Browse files Browse the repository at this point in the history
…e-speed

Initial Refactor for speed control
  • Loading branch information
Mzack9999 authored Apr 9, 2024
2 parents 375d1dd + 5fc08ce commit 721ddda
Show file tree
Hide file tree
Showing 34 changed files with 459 additions and 161 deletions.
4 changes: 4 additions & 0 deletions .github/workflows/build-test.yml
Original file line number Diff line number Diff line change
Expand Up @@ -66,3 +66,7 @@ jobs:
- name: Example SDK Advanced
run: go run .
working-directory: examples/advanced/

- name: Example SDK with speed control
run: go run .
working-directory: examples/with_speed_control/
9 changes: 6 additions & 3 deletions cmd/nuclei/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ import (
"github.com/projectdiscovery/utils/auth/pdcp"
"github.com/projectdiscovery/utils/env"
_ "github.com/projectdiscovery/utils/pprof"
stringsutil "github.com/projectdiscovery/utils/strings"

"github.com/projectdiscovery/goflags"
"github.com/projectdiscovery/gologger"
Expand Down Expand Up @@ -329,13 +330,15 @@ on extensive configurability, massive extensibility and ease of use.`)

flagSet.CreateGroup("rate-limit", "Rate-Limit",
flagSet.IntVarP(&options.RateLimit, "rate-limit", "rl", 150, "maximum number of requests to send per second"),
flagSet.IntVarP(&options.RateLimitMinute, "rate-limit-minute", "rlm", 0, "maximum number of requests to send per minute"),
flagSet.DurationVarP(&options.RateLimitDuration, "rate-limit-duration", "rld", time.Second, "maximum number of requests to send per second"),
flagSet.IntVarP(&options.RateLimitMinute, "rate-limit-minute", "rlm", 0, "maximum number of requests to send per minute (DEPRECATED)"),
flagSet.IntVarP(&options.BulkSize, "bulk-size", "bs", 25, "maximum number of hosts to be analyzed in parallel per template"),
flagSet.IntVarP(&options.TemplateThreads, "concurrency", "c", 25, "maximum number of templates to be executed in parallel"),
flagSet.IntVarP(&options.HeadlessBulkSize, "headless-bulk-size", "hbs", 10, "maximum number of headless hosts to be analyzed in parallel per template"),
flagSet.IntVarP(&options.HeadlessTemplateThreads, "headless-concurrency", "headc", 10, "maximum number of headless templates to be executed in parallel"),
flagSet.IntVarP(&options.JsConcurrency, "js-concurrency", "jsc", 120, "maximum number of javascript runtimes to be executed in parallel"),
flagSet.IntVarP(&options.PayloadConcurrency, "payload-concurrency", "pc", 25, "max payload concurrency for each template"),
flagSet.IntVarP(&options.ProbeConcurrency, "probe-concurrency", "prc", 50, "http probe concurrency with httpx"),
)
flagSet.CreateGroup("optimization", "Optimizations",
flagSet.IntVar(&options.Timeout, "timeout", 10, "time to wait in seconds before timeout"),
Expand Down Expand Up @@ -597,10 +600,10 @@ Note: Make sure you have backup of your custom nuclei-templates before proceedin
gologger.Fatal().Msgf("could not read response: %s", err)
}
resp = strings.TrimSpace(resp)
if strings.EqualFold(resp, "y") || strings.EqualFold(resp, "yes") {
if stringsutil.EqualFoldAny(resp, "y", "yes") {
break
}
if strings.EqualFold(resp, "n") || strings.EqualFold(resp, "no") || resp == "" {
if stringsutil.EqualFoldAny(resp, "n", "no", "") {
fmt.Println("Exiting...")
os.Exit(0)
}
Expand Down
7 changes: 5 additions & 2 deletions examples/advanced/advanced.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ package main

import (
nuclei "github.com/projectdiscovery/nuclei/v3/lib"
"github.com/remeh/sizedwaitgroup"
syncutil "github.com/projectdiscovery/utils/sync"
)

func main() {
Expand All @@ -12,7 +12,10 @@ func main() {
panic(err)
}
// setup sizedWaitgroup to handle concurrency
sg := sizedwaitgroup.New(10)
sg, err := syncutil.New(syncutil.WithSize(10))
if err != nil {
panic(err)
}

// scan 1 = run dns templates on scanme.sh
sg.Add()
Expand Down
104 changes: 104 additions & 0 deletions examples/with_speed_control/main.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,104 @@
package main

import (
"log"
"sync"
"time"

nuclei "github.com/projectdiscovery/nuclei/v3/lib"
"github.com/projectdiscovery/nuclei/v3/pkg/templates/types"
)

func main() {
ne, err := initializeNucleiEngine()
if err != nil {
panic(err)
}
defer ne.Close()

ne.LoadTargets([]string{"http://honey.scanme.sh"}, false)

var wg sync.WaitGroup
wg.Add(3)

go testRateLimit(&wg, ne)
go testThreadsAndBulkSize(&wg, ne)
go testPayloadConcurrency(&wg, ne)

err = ne.ExecuteWithCallback(nil)
if err != nil {
panic(err)
}

wg.Wait()
}

func initializeNucleiEngine() (*nuclei.NucleiEngine, error) {
return nuclei.NewNucleiEngine(
nuclei.WithTemplateFilters(nuclei.TemplateFilters{Tags: []string{"oast"}}),
nuclei.EnableStatsWithOpts(nuclei.StatsOptions{MetricServerPort: 6064}),
nuclei.WithGlobalRateLimit(1, time.Second),
nuclei.WithConcurrency(nuclei.Concurrency{
TemplateConcurrency: 1,
HostConcurrency: 1,
HeadlessHostConcurrency: 1,
HeadlessTemplateConcurrency: 1,
JavascriptTemplateConcurrency: 1,
TemplatePayloadConcurrency: 1,
ProbeConcurrency: 1,
}),
)
}

func testRateLimit(wg *sync.WaitGroup, ne *nuclei.NucleiEngine) {
defer wg.Done()
verifyRateLimit(ne, 1, 5000)
}

func testThreadsAndBulkSize(wg *sync.WaitGroup, ne *nuclei.NucleiEngine) {
defer wg.Done()
initialTemplateThreads, initialBulkSize := 1, 1
verifyThreadsAndBulkSize(ne, initialTemplateThreads, initialBulkSize, 25, 25)
}

func testPayloadConcurrency(wg *sync.WaitGroup, ne *nuclei.NucleiEngine) {
defer wg.Done()
verifyPayloadConcurrency(ne, 1, 500)
}

func verifyRateLimit(ne *nuclei.NucleiEngine, initialRate, finalRate int) {
if ne.GetExecuterOptions().RateLimiter.GetLimit() != uint(initialRate) {
panic("wrong initial rate limit")
}
time.Sleep(5 * time.Second)
ne.Options().RateLimit = finalRate
time.Sleep(20 * time.Second)
if ne.GetExecuterOptions().RateLimiter.GetLimit() != uint(finalRate) {
panic("wrong final rate limit")
}
}

func verifyThreadsAndBulkSize(ne *nuclei.NucleiEngine, initialThreads, initialBulk, finalThreads, finalBulk int) {
if ne.Options().TemplateThreads != initialThreads || ne.Options().BulkSize != initialBulk {
panic("wrong initial standard concurrency")
}
time.Sleep(5 * time.Second)
ne.Options().TemplateThreads = finalThreads
ne.Options().BulkSize = finalBulk
time.Sleep(20 * time.Second)
if ne.Engine().GetWorkPool().InputPool(types.HTTPProtocol).Size != finalBulk || ne.Engine().WorkPool().Default.Size != finalThreads {
log.Fatal("wrong final concurrency", ne.Engine().WorkPool().Default.Size, finalThreads, ne.Engine().GetWorkPool().InputPool(types.HTTPProtocol).Size, finalBulk)
}
}

func verifyPayloadConcurrency(ne *nuclei.NucleiEngine, initialPayloadConcurrency, finalPayloadConcurrency int) {
if ne.Options().PayloadConcurrency != initialPayloadConcurrency {
panic("wrong initial payload concurrency")
}
time.Sleep(5 * time.Second)
ne.Options().PayloadConcurrency = finalPayloadConcurrency
time.Sleep(20 * time.Second)
if ne.GetExecuterOptions().GetThreadsForNPayloadRequests(100, 0) != finalPayloadConcurrency {
panic("wrong final payload concurrency")
}
}
12 changes: 7 additions & 5 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@ require (
github.com/weppos/publicsuffix-go v0.30.2-0.20230730094716-a20f9abcc222
github.com/xanzy/go-gitlab v0.84.0
go.uber.org/multierr v1.11.0
golang.org/x/net v0.21.0
golang.org/x/net v0.24.0
golang.org/x/oauth2 v0.11.0
golang.org/x/text v0.14.0
gopkg.in/yaml.v2 v2.4.0
Expand Down Expand Up @@ -94,13 +94,13 @@ require (
github.com/projectdiscovery/tlsx v1.1.6
github.com/projectdiscovery/uncover v1.0.7
github.com/projectdiscovery/useragent v0.0.40
github.com/projectdiscovery/utils v0.0.88
github.com/projectdiscovery/utils v0.0.88-0.20240404181359-663cfe2196d0
github.com/projectdiscovery/wappalyzergo v0.0.116
github.com/redis/go-redis/v9 v9.1.0
github.com/seh-msft/burpxml v1.0.1
github.com/stretchr/testify v1.9.0
github.com/zmap/zgrab2 v0.1.8-0.20230806160807-97ba87c0e706
golang.org/x/term v0.17.0
golang.org/x/term v0.19.0
gopkg.in/yaml.v3 v3.0.1
moul.io/http2curl v1.0.0
)
Expand Down Expand Up @@ -142,6 +142,8 @@ require (
github.com/docker/cli v24.0.5+incompatible // indirect
github.com/docker/docker v24.0.9+incompatible // indirect
github.com/docker/go-connections v0.4.0 // indirect
github.com/eapache/channels v1.1.0 // indirect
github.com/eapache/queue v1.1.0 // indirect
github.com/fatih/color v1.15.0 // indirect
github.com/free5gc/util v1.0.5-0.20230511064842-2e120956883b // indirect
github.com/gabriel-vasile/mimetype v1.4.2 // indirect
Expand Down Expand Up @@ -300,10 +302,10 @@ require (
go.etcd.io/bbolt v1.3.8 // indirect
go.uber.org/zap v1.25.0 // indirect
goftp.io/server/v2 v2.0.1 // indirect
golang.org/x/crypto v0.19.0 // indirect
golang.org/x/crypto v0.22.0 // indirect
golang.org/x/exp v0.0.0-20240119083558-1b970713d09a
golang.org/x/mod v0.14.0 // indirect
golang.org/x/sys v0.17.0 // indirect
golang.org/x/sys v0.19.0 // indirect
golang.org/x/time v0.5.0 // indirect
golang.org/x/tools v0.17.0
google.golang.org/appengine v1.6.7 // indirect
Expand Down
22 changes: 13 additions & 9 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -296,8 +296,11 @@ github.com/dsnet/compress v0.0.2-0.20210315054119-f66993602bf5/go.mod h1:qssHWj6
github.com/dsnet/golib v0.0.0-20171103203638-1ea166775780/go.mod h1:Lj+Z9rebOhdfkVLjJ8T6VcRQv3SXugXy999NBtR9aFY=
github.com/dustin/go-humanize v0.0.0-20171111073723-bb3d318650d4/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk=
github.com/dustin/go-humanize v1.0.0/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk=
github.com/eapache/channels v1.1.0 h1:F1taHcn7/F0i8DYqKXJnyhJcVpp2kgFcNePxXtnyu4k=
github.com/eapache/channels v1.1.0/go.mod h1:jMm2qB5Ubtg9zLd+inMZd2/NUvXgzmWXsDaLyQIGfH0=
github.com/eapache/go-resiliency v1.1.0/go.mod h1:kFI+JgMyC7bLPUVY133qvEBtVayf5mFgVsvEsIPBvNs=
github.com/eapache/go-xerial-snappy v0.0.0-20180814174437-776d5712da21/go.mod h1:+020luEh2TKB4/GOp8oxxtq0Daoen/Cii55CzbTV6DU=
github.com/eapache/queue v1.1.0 h1:YOEu7KNc61ntiQlcEeUIoDTJ2o8mQznoNvUhiigpIqc=
github.com/eapache/queue v1.1.0/go.mod h1:6eCeP0CKFpHLu8blIFXhExK/dRa7WDZfr6jVFPTqq+I=
github.com/edsrzf/mmap-go v1.0.0/go.mod h1:YO35OhQPt3KJa3ryjFM5Bs14WD66h8eGKpfaBNrHW5M=
github.com/elazarl/goproxy v0.0.0-20230808193330-2592e75ae04a h1:mATvB/9r/3gvcejNsXKSkQ6lcIaNec2nyfOdlTBR2lU=
Expand Down Expand Up @@ -883,8 +886,8 @@ github.com/projectdiscovery/uncover v1.0.7 h1:ut+2lTuvmftmveqF5RTjMWAgyLj8ltPQC7
github.com/projectdiscovery/uncover v1.0.7/go.mod h1:HFXgm1sRPuoN0D4oATljPIdmbo/EEh1wVuxQqo/dwFE=
github.com/projectdiscovery/useragent v0.0.40 h1:1LUhReSGPkhqsM5n40OOC9dIoNqMGs1dyGFJcOmg2Fo=
github.com/projectdiscovery/useragent v0.0.40/go.mod h1:EvK1x3s948Gtqb/XOahXcauyejCL/rSgy5d1IAvsKT4=
github.com/projectdiscovery/utils v0.0.88 h1:oYfCXM+8VHNLyH/H6cOibkuDUwHUAOBAMRNPFX6NPrs=
github.com/projectdiscovery/utils v0.0.88/go.mod h1:lAWzFdGXtJRPKdhUu1Z46d8B8JbASTk1Z69WY6H/3kA=
github.com/projectdiscovery/utils v0.0.88-0.20240404181359-663cfe2196d0 h1:2ZR0yiN0cUm/qYEMq79MfcbgM374lJSdftheYhMFxNo=
github.com/projectdiscovery/utils v0.0.88-0.20240404181359-663cfe2196d0/go.mod h1:lAWzFdGXtJRPKdhUu1Z46d8B8JbASTk1Z69WY6H/3kA=
github.com/projectdiscovery/wappalyzergo v0.0.116 h1:xy+mBpwbYo/0PSzmJOQ/RXHomEh0D3nDBcbCxsW69m8=
github.com/projectdiscovery/wappalyzergo v0.0.116/go.mod h1:hc/o+fgM8KtdpFesjfBTmHTwsR+yBd+4kYZW/DGy/x8=
github.com/projectdiscovery/yamldoc-go v1.0.4 h1:eZoESapnMw6WAHiVgRwNqvbJEfNHEH148uthhFbG5jE=
Expand Down Expand Up @@ -1183,8 +1186,8 @@ golang.org/x/crypto v0.6.0/go.mod h1:OFC/31mSvZgRz0V1QTNCzfAI1aIRzbiufJtkMIlEp58
golang.org/x/crypto v0.7.0/go.mod h1:pYwdfH91IfpZVANVyUOhSIPZaFoJGxTFbZhFTx+dXZU=
golang.org/x/crypto v0.11.0/go.mod h1:xgJhtzW8F9jGdVFWZESrid1U1bjeNy4zgy5cRr/CIio=
golang.org/x/crypto v0.17.0/go.mod h1:gCAAfMLgwOJRpTjQ2zCCt2OcSfYMTeZVSRtQlPC7Nq4=
golang.org/x/crypto v0.19.0 h1:ENy+Az/9Y1vSrlrvBSyna3PITt4tiZLf7sgCjZBX7Wo=
golang.org/x/crypto v0.19.0/go.mod h1:Iy9bg/ha4yyC70EfRS8jz+B6ybOBKMaSxLj6P6oBDfU=
golang.org/x/crypto v0.22.0 h1:g1v0xeRhjcugydODzvb3mEM9SQ0HGp9s/nh3COQ/C30=
golang.org/x/crypto v0.22.0/go.mod h1:vr6Su+7cTlO45qkww3VDJlzDn0ctJvRgYbC2NvXHt+M=
golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA=
golang.org/x/exp v0.0.0-20190306152737-a1d7652674e8/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA=
golang.org/x/exp v0.0.0-20190510132918-efd6b22b2522/go.mod h1:ZjyILWgesfNpC6sMxTJOJm9Kp84zZh5NQWvqDGG3Qr8=
Expand Down Expand Up @@ -1277,8 +1280,8 @@ golang.org/x/net v0.8.0/go.mod h1:QVkue5JL9kW//ek3r6jTKnTFis1tRmNAW2P1shuFdJc=
golang.org/x/net v0.9.0/go.mod h1:d48xBJpPfHeWQsugry2m+kC02ZBRGRgulfHnEXEuWns=
golang.org/x/net v0.10.0/go.mod h1:0qNGK6F8kojg2nk9dLZ2mShWaEBan6FAoqfSigmmuDg=
golang.org/x/net v0.12.0/go.mod h1:zEVYFnQC7m/vmpQFELhcD1EWkZlX69l4oqgmer6hfKA=
golang.org/x/net v0.21.0 h1:AQyQV4dYCvJ7vGmJyKki9+PBdyvhkSd8EIx/qb0AYv4=
golang.org/x/net v0.21.0/go.mod h1:bIjVDfnllIU7BJ2DNgfnXvpSvtn8VRwhlsaeUTyUS44=
golang.org/x/net v0.24.0 h1:1PcaxkF854Fu3+lvBIx5SYn9wRlBzzcnHZSiaFFAb0w=
golang.org/x/net v0.24.0/go.mod h1:2Q7sJY5mzlzWjKtYUEXSlBWCdyaioyXzRB2RtU8KVE8=
golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U=
golang.org/x/oauth2 v0.0.0-20190226205417-e64efc72b421/go.mod h1:gOpvHmFTYa4IltrdGE7lF6nIHvwfUNPOp7c8zoXwtLw=
golang.org/x/oauth2 v0.0.0-20190604053449-0f29369cfe45/go.mod h1:gOpvHmFTYa4IltrdGE7lF6nIHvwfUNPOp7c8zoXwtLw=
Expand Down Expand Up @@ -1379,8 +1382,9 @@ golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.10.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.11.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
golang.org/x/sys v0.15.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
golang.org/x/sys v0.17.0 h1:25cE3gD+tdBA7lp7QfhuV+rJiE9YXTcS3VG1SqssI/Y=
golang.org/x/sys v0.17.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
golang.org/x/sys v0.19.0 h1:q5f1RH2jigJ1MoAWp2KTp3gm5zAGFUTarQZ5U386+4o=
golang.org/x/sys v0.19.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
golang.org/x/term v0.0.0-20201117132131-f5c789dd3221/go.mod h1:Nr5EML6q2oocZ2LXRh80K7BxOlk5/8JxuGnuhpl+muw=
golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
golang.org/x/term v0.0.0-20210220032956-6a3ed077a48d/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
Expand All @@ -1392,8 +1396,8 @@ golang.org/x/term v0.7.0/go.mod h1:P32HKFT3hSsZrRxla30E9HqToFYAQPCMs/zFMBUFqPY=
golang.org/x/term v0.8.0/go.mod h1:xPskH00ivmX89bAKVGSKKtLOWNx2+17Eiy94tnKShWo=
golang.org/x/term v0.10.0/go.mod h1:lpqdcUyK/oCiQxvxVrppt5ggO2KCZ5QblwqPnfZ6d5o=
golang.org/x/term v0.15.0/go.mod h1:BDl952bC7+uMoWR75FIrCDx79TPU9oHkTZ9yRbYOrX0=
golang.org/x/term v0.17.0 h1:mkTF7LCd6WGJNL3K1Ad7kwxNfYAW6a8a8QqtMblp/4U=
golang.org/x/term v0.17.0/go.mod h1:lLRBjIVuehSbZlaOtGMbcMncT+aqLLLmKrsjNrUguwk=
golang.org/x/term v0.19.0 h1:+ThwsDv+tYfnJFhF4L8jITxu1tdTWRTZpdsWgEgjL6Q=
golang.org/x/term v0.19.0/go.mod h1:2CuTdWZ7KHSQwUzKva0cbMg6q2DMI3Mmxp+gKJbskEk=
golang.org/x/text v0.0.0-20170915032832-14c0d48ead0c/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
golang.org/x/text v0.3.1-0.20180807135948-17ff2d5776d2/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
Expand Down
26 changes: 17 additions & 9 deletions internal/runner/inputs.go
Original file line number Diff line number Diff line change
Expand Up @@ -12,15 +12,14 @@ import (
"github.com/projectdiscovery/nuclei/v3/pkg/protocols/common/contextargs"
"github.com/projectdiscovery/nuclei/v3/pkg/utils"
stringsutil "github.com/projectdiscovery/utils/strings"
"github.com/remeh/sizedwaitgroup"
syncutil "github.com/projectdiscovery/utils/sync"
)

const probeBulkSize = 50
var GlobalProbeBulkSize = 50

// initializeTemplatesHTTPInput initializes the http form of input
// for any loaded http templates if input is in non-standard format.
func (r *Runner) initializeTemplatesHTTPInput() (*hybrid.HybridMap, error) {

hm, err := hybrid.New(hybrid.DefaultDiskOptions)
if err != nil {
return nil, errors.Wrap(err, "could not create temporary input file")
Expand All @@ -31,8 +30,8 @@ func (r *Runner) initializeTemplatesHTTPInput() (*hybrid.HybridMap, error) {
}
gologger.Info().Msgf("Running httpx on input host")

var bulkSize = probeBulkSize
if r.options.BulkSize > probeBulkSize {
var bulkSize = GlobalProbeBulkSize
if r.options.BulkSize > GlobalProbeBulkSize {
bulkSize = r.options.BulkSize
}

Expand All @@ -44,27 +43,36 @@ func (r *Runner) initializeTemplatesHTTPInput() (*hybrid.HybridMap, error) {
return nil, errors.Wrap(err, "could not create httpx client")
}

shouldFollowGlobalProbeBulkSize := bulkSize == GlobalProbeBulkSize

// Probe the non-standard URLs and store them in cache
swg := sizedwaitgroup.New(bulkSize)
count := int32(0)
swg, err := syncutil.New(syncutil.WithSize(bulkSize))
if err != nil {
return nil, errors.Wrap(err, "could not create adaptive group")
}
var count atomic.Int32
r.inputProvider.Iterate(func(value *contextargs.MetaInput) bool {
if stringsutil.HasPrefixAny(value.Input, "http://", "https://") {
return true
}

if shouldFollowGlobalProbeBulkSize && swg.Size != GlobalProbeBulkSize {
swg.Resize(GlobalProbeBulkSize)
}

swg.Add()
go func(input *contextargs.MetaInput) {
defer swg.Done()

if result := utils.ProbeURL(input.Input, httpxClient); result != "" {
atomic.AddInt32(&count, 1)
count.Add(1)
_ = hm.Set(input.Input, []byte(result))
}
}(value)
return true
})
swg.Wait()

gologger.Info().Msgf("Found %d URL from httpx", atomic.LoadInt32(&count))
gologger.Info().Msgf("Found %d URL from httpx", count.Load())
return hm, nil
}
14 changes: 10 additions & 4 deletions internal/runner/runner.go
Original file line number Diff line number Diff line change
Expand Up @@ -314,11 +314,17 @@ func New(options *types.Options) (*Runner, error) {
}

if options.RateLimitMinute > 0 {
runner.rateLimiter = ratelimit.New(context.Background(), uint(options.RateLimitMinute), time.Minute)
} else if options.RateLimit > 0 {
runner.rateLimiter = ratelimit.New(context.Background(), uint(options.RateLimit), time.Second)
} else {
gologger.Print().Msgf("[%v] %v", aurora.BrightYellow("WRN"), "rate limit per minute is deprecated - use rate-limit-duration")
options.RateLimit = options.RateLimitMinute
options.RateLimitDuration = time.Minute
}
if options.RateLimit > 0 && options.RateLimitDuration == 0 {
options.RateLimitDuration = time.Second
}
if options.RateLimit == 0 && options.RateLimitDuration == 0 {
runner.rateLimiter = ratelimit.NewUnlimited(context.Background())
} else {
runner.rateLimiter = ratelimit.New(context.Background(), uint(options.RateLimit), options.RateLimitDuration)
}

if tmpDir, err := os.MkdirTemp("", "nuclei-tmp-*"); err == nil {
Expand Down
Loading

0 comments on commit 721ddda

Please sign in to comment.