ratelimit library

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{

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"

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"
 

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"
 

v2/pkg/utils/ratelimit/ratelimit.go

@@ -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
-}

v2/pkg/utils/ratelimit/ratelimit_test.go

@@ -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))
-	})
-}