[ratelimiter] Add IP allowlist

common/ratelimit/limiter.go

@@ -13,24 +13,31 @@ type rateLimiter struct {
 	globalRateParams common.GlobalRateParams
 
 	bucketStore BucketStore
+	allowlist   []string
 
 	logger common.Logger
 }
 
-func NewRateLimiter(rateParams common.GlobalRateParams, bucketStore BucketStore, logger common.Logger) common.RateLimiter {
+func NewRateLimiter(rateParams common.GlobalRateParams, bucketStore BucketStore, allowlist []string, logger common.Logger) common.RateLimiter {
 	return &rateLimiter{
 		globalRateParams: rateParams,
 		bucketStore:      bucketStore,
+		allowlist:        allowlist,
 		logger:           logger,
 	}
 }
 
 // Checks whether a request from the given requesterID is allowed
 func (d *rateLimiter) AllowRequest(ctx context.Context, requesterID common.RequesterID, blobSize uint, rate common.RateParam) (bool, error) {
+	for _, id := range d.allowlist {
+		if id == requesterID {
+			// Allow 10x the rate for allowlisted IDs
+			rate = rate * 10
+		}
+	}
 
 	// Retrieve bucket params for the requester ID
 	// This will be from dynamo for Disperser and from local storage for DA node
-
 	bucketParams, err := d.bucketStore.GetItem(ctx, requesterID)
 	if err != nil {
 

common/ratelimit/limiter_cli.go

@@ -15,12 +15,14 @@ const (
 	BucketMultipliersFlagName = "bucket-multipliers"
 	CountFailedFlagName       = "count-failed"
 	BucketStoreSizeFlagName   = "bucket-store-size"
+	AllowlistFlagName         = "allowlist"
 )
 
 type Config struct {
 	common.GlobalRateParams
 	BucketStoreSize  int
 	UniformRateParam common.RateParam
+	Allowlist        []string
 }
 
 func RatelimiterCLIFlags(envPrefix string, flagPrefix string) []cli.Flag {
@@ -52,6 +54,13 @@ func RatelimiterCLIFlags(envPrefix string, flagPrefix string) []cli.Flag {
 			EnvVar:   common.PrefixEnvVar(envPrefix, "BUCKET_STORE_SIZE"),
 			Required: false,
 		},
+		cli.StringSliceFlag{
+			Name:     common.PrefixFlag(flagPrefix, AllowlistFlagName),
+			Usage:    "Allowlist of IPs to bypass rate limiting",
+			EnvVar:   common.PrefixEnvVar(envPrefix, "ALLOWLIST"),
+			Required: false,
+			Value:    &cli.StringSlice{},
+		},
 	}
 }
 
@@ -97,6 +106,7 @@ func ReadCLIConfig(ctx *cli.Context, flagPrefix string) (Config, error) {
 	cfg.Multipliers = multipliers
 	cfg.GlobalRateParams.CountFailed = ctx.Bool(common.PrefixFlag(flagPrefix, CountFailedFlagName))
 	cfg.BucketStoreSize = ctx.Int(common.PrefixFlag(flagPrefix, BucketStoreSizeFlagName))
+	cfg.Allowlist = ctx.StringSlice(common.PrefixFlag(flagPrefix, AllowlistFlagName))
 
 	err := validateConfig(cfg)
 	if err != nil {

common/ratelimit/ratelimit_test.go

@@ -25,14 +25,13 @@ func makeTestRatelimiter() (common.RateLimiter, error) {
 		return nil, err
 	}
 
-	ratelimiter := ratelimit.NewRateLimiter(globalParams, bucketStore, &mock.Logger{})
+	ratelimiter := ratelimit.NewRateLimiter(globalParams, bucketStore, []string{"testRetriever2"}, &mock.Logger{})
 
 	return ratelimiter, nil
 
 }
 
 func TestRatelimit(t *testing.T) {
-
 	ratelimiter, err := makeTestRatelimiter()
 	assert.NoError(t, err)
 
@@ -50,3 +49,23 @@ func TestRatelimit(t *testing.T) {
 	assert.NoError(t, err)
 	assert.Equal(t, false, allow)
 }
+
+func TestRatelimitAllowlist(t *testing.T) {
+	ratelimiter, err := makeTestRatelimiter()
+	assert.NoError(t, err)
+
+	ctx := context.Background()
+
+	retreiverID := "testRetriever2"
+
+	// 10x more requests allowed for allowlisted IDs
+	for i := 0; i < 100; i++ {
+		allow, err := ratelimiter.AllowRequest(ctx, retreiverID, 10, 100)
+		assert.NoError(t, err)
+		assert.Equal(t, true, allow)
+	}
+
+	allow, err := ratelimiter.AllowRequest(ctx, retreiverID, 10, 100)
+	assert.NoError(t, err)
+	assert.Equal(t, false, allow)
+}

disperser/cmd/disperserserver/main.go

@@ -109,7 +109,7 @@ func RunDisperserServer(ctx *cli.Context) error {
 				return err
 			}
 		}
-		ratelimiter = ratelimit.NewRateLimiter(globalParams, bucketStore, logger)
+		ratelimiter = ratelimit.NewRateLimiter(globalParams, bucketStore, config.RatelimiterConfig.Allowlist, logger)
 	}
 
 	// TODO: create a separate metrics for batcher

node/cmd/main.go

@@ -80,7 +80,7 @@ func NodeMain(ctx *cli.Context) error {
 		return err
 	}
 
-	ratelimiter := ratelimit.NewRateLimiter(globalParams, bucketStore, logger)
+	ratelimiter := ratelimit.NewRateLimiter(globalParams, bucketStore, []string{}, logger)
 
 	// Creates the GRPC server.
 	server := grpc.NewServer(config, node, logger, ratelimiter)