diff --git a/v2/cmd/integration-test/code.go b/v2/cmd/integration-test/code.go index 1e68549c6..4171efeba 100644 --- a/v2/cmd/integration-test/code.go +++ b/v2/cmd/integration-test/code.go @@ -30,7 +30,7 @@ import ( "github.com/projectdiscovery/nuclei/v2/pkg/reporting" "github.com/projectdiscovery/nuclei/v2/pkg/testutils" "github.com/projectdiscovery/nuclei/v2/pkg/types" - "github.com/projectdiscovery/nuclei/v2/pkg/utils/ratelimit" + "github.com/projectdiscovery/ratelimit" ) var codeTestcases = map[string]testutils.TestCase{ diff --git a/v2/internal/runner/runner.go b/v2/internal/runner/runner.go index 6543e9111..8d6337d59 100644 --- a/v2/internal/runner/runner.go +++ b/v2/internal/runner/runner.go @@ -17,7 +17,7 @@ import ( "github.com/blang/semver" "github.com/logrusorgru/aurora" "github.com/pkg/errors" - "github.com/projectdiscovery/nuclei/v2/pkg/utils/ratelimit" + "github.com/projectdiscovery/ratelimit" "go.uber.org/atomic" "github.com/projectdiscovery/gologger" diff --git a/v2/pkg/protocols/protocols.go b/v2/pkg/protocols/protocols.go index fecc38b12..e9486dc0f 100644 --- a/v2/pkg/protocols/protocols.go +++ b/v2/pkg/protocols/protocols.go @@ -1,7 +1,7 @@ package protocols import ( - "github.com/projectdiscovery/nuclei/v2/pkg/utils/ratelimit" + "github.com/projectdiscovery/ratelimit" "github.com/logrusorgru/aurora" diff --git a/v2/pkg/testutils/testutils.go b/v2/pkg/testutils/testutils.go index 09024b288..8551607cd 100644 --- a/v2/pkg/testutils/testutils.go +++ b/v2/pkg/testutils/testutils.go @@ -4,7 +4,7 @@ import ( "context" "time" - "github.com/projectdiscovery/nuclei/v2/pkg/utils/ratelimit" + "github.com/projectdiscovery/ratelimit" "github.com/logrusorgru/aurora" diff --git a/v2/pkg/utils/ratelimit/ratelimit.go b/v2/pkg/utils/ratelimit/ratelimit.go deleted file mode 100644 index 9b789bc99..000000000 --- a/v2/pkg/utils/ratelimit/ratelimit.go +++ /dev/null @@ -1,68 +0,0 @@ -package ratelimit - -import ( - "context" - "math" - "time" -) - -// Limiter allows a burst of request during the defined duration -type Limiter struct { - maxCount int64 - count int64 - ticker *time.Ticker - tokens chan struct{} - ctx context.Context -} - -func (limiter *Limiter) run() { - for { - if limiter.count <= 0 { - <-limiter.ticker.C - limiter.count = limiter.maxCount - } - - select { - case <-limiter.ctx.Done(): - limiter.ticker.Stop() - return - case limiter.tokens <- struct{}{}: - limiter.count-- - case <-limiter.ticker.C: - limiter.count = limiter.maxCount - } - } -} - -// Take one token from the bucket -func (rateLimiter *Limiter) Take() { - <-rateLimiter.tokens -} - -// New creates a new limiter instance with the tokens amount and the interval -func New(ctx context.Context, max int64, duration time.Duration) *Limiter { - limiter := &Limiter{ - maxCount: max, - count: max, - ticker: time.NewTicker(duration), - tokens: make(chan struct{}), - ctx: ctx, - } - go limiter.run() - - return limiter -} - -// NewUnlimited create a bucket with approximated unlimited tokens -func NewUnlimited(ctx context.Context) *Limiter { - limiter := &Limiter{ - maxCount: math.MaxInt64, - count: math.MaxInt64, - ticker: time.NewTicker(time.Millisecond), - tokens: make(chan struct{}), - ctx: ctx, - } - go limiter.run() - - return limiter -} diff --git a/v2/pkg/utils/ratelimit/ratelimit_test.go b/v2/pkg/utils/ratelimit/ratelimit_test.go deleted file mode 100644 index 0dc14638d..000000000 --- a/v2/pkg/utils/ratelimit/ratelimit_test.go +++ /dev/null @@ -1,44 +0,0 @@ -package ratelimit - -import ( - "context" - "testing" - "time" - - "github.com/stretchr/testify/require" -) - -func TestRateLimit(t *testing.T) { - t.Run("Standard Rate Limit", func(t *testing.T) { - expected := time.Duration(15 * time.Second) - limiter := New(context.Background(), 10, expected) - require.NotNil(t, limiter) - var count int - start := time.Now() - for i := 0; i < 10; i++ { - limiter.Take() - count++ - } - took := time.Since(start) - require.Equal(t, count, 10) - require.True(t, took < expected) - // take another one above max - limiter.Take() - took = time.Since(start) - require.True(t, took >= expected) - }) - - t.Run("Unlimited Rate Limit", func(t *testing.T) { - limiter := NewUnlimited(context.Background()) - require.NotNil(t, limiter) - var count int - start := time.Now() - for i := 0; i < 1000; i++ { - limiter.Take() - count++ - } - took := time.Since(start) - require.Equal(t, count, 1000) - require.True(t, took < time.Duration(1*time.Second)) - }) -}