Add V3 rate limits

xmtp · Jun 16, 2024 · d896f0c · d896f0c
1 parent 24971f4
commit d896f0c
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 12 deletions.
diff --git a/pkg/api/interceptor.go b/pkg/api/interceptor.go
@@ -13,7 +13,9 @@ import (
 	"google.golang.org/grpc/status"
 
 	"github.com/xmtp/xmtp-node-go/pkg/logging"
+	identityv1 "github.com/xmtp/xmtp-node-go/pkg/proto/identity/api/v1"
 	messagev1 "github.com/xmtp/xmtp-node-go/pkg/proto/message_api/v1"
+	mlsv1 "github.com/xmtp/xmtp-node-go/pkg/proto/mls/api/v1"
 	"github.com/xmtp/xmtp-node-go/pkg/ratelimiter"
 	"github.com/xmtp/xmtp-node-go/pkg/types"
 )
@@ -160,6 +162,16 @@ func (wa *WalletAuthorizer) applyLimits(ctx context.Context, fullMethod string,
 	cost := 1
 	limitType := ratelimiter.DEFAULT
 	switch req := req.(type) {
+	case *identityv1.PublishIdentityUpdateRequest:
+		limitType = ratelimiter.IDENTITY_PUBLISH
+	case *mlsv1.FetchKeyPackagesRequest, *mlsv1.SubscribeGroupMessagesRequest, *mlsv1.SubscribeWelcomeMessagesRequest, *identityv1.GetIdentityUpdatesRequest, *identityv1.GetInboxIdsRequest:
+		limitType = ratelimiter.V3_DEFAULT
+	case *mlsv1.SendWelcomeMessagesRequest:
+		cost = len(req.Messages)
+		limitType = ratelimiter.V3_PUBLISH
+	case *mlsv1.SendGroupMessagesRequest:
+		cost = len(req.Messages)
+		limitType = ratelimiter.V3_PUBLISH
 	case *messagev1.PublishRequest:
 		cost = len(req.Envelopes)
 		limitType = ratelimiter.PUBLISH

diff --git a/pkg/ratelimiter/rate_limiter.go b/pkg/ratelimiter/rate_limiter.go
@@ -14,16 +14,21 @@ import (
 type LimitType string
 
 const (
-	DEFAULT_PRIORITY_MULTIPLIER = uint16(5)
-	PUBLISH_PRIORITY_MULTIPLIER = uint16(25)
-	DEFAULT_RATE_PER_MINUTE     = uint16(2000)
-	DEFAULT_MAX_TOKENS          = uint16(10000)
-	PUBLISH_RATE_PER_MINUTE     = uint16(200)
-	PUBLISH_MAX_TOKENS          = uint16(1000)
-	MAX_UINT_16                 = 65535
-
-	DEFAULT LimitType = "DEF"
-	PUBLISH LimitType = "PUB"
+	DEFAULT_PRIORITY_MULTIPLIER      = uint16(5)
+	PUBLISH_PRIORITY_MULTIPLIER      = uint16(25)
+	DEFAULT_RATE_PER_MINUTE          = uint16(2000)
+	DEFAULT_MAX_TOKENS               = uint16(10000)
+	PUBLISH_RATE_PER_MINUTE          = uint16(200)
+	PUBLISH_MAX_TOKENS               = uint16(1000)
+	IDENTITY_PUBLISH_RATE_PER_MINUTE = uint16(10)
+	IDENTITY_PUBLISH_MAX_TOKENS      = uint16(50)
+	MAX_UINT_16                      = 65535
+
+	DEFAULT          LimitType = "DEF"
+	V3_DEFAULT       LimitType = "V3DEF"
+	PUBLISH          LimitType = "PUB"
+	V3_PUBLISH       LimitType = "V3PUB"
+	IDENTITY_PUBLISH LimitType = "IDPUB"
 )
 
 type RateLimiter interface {
@@ -91,8 +96,11 @@ func NewTokenBucketRateLimiter(ctx context.Context, log *zap.Logger) *TokenBucke
 	tb.PriorityMultiplier = DEFAULT_PRIORITY_MULTIPLIER
 	tb.PublishPriorityMultiplier = PUBLISH_PRIORITY_MULTIPLIER
 	tb.Limits = map[LimitType]*Limit{
-		DEFAULT: {DEFAULT_MAX_TOKENS, DEFAULT_RATE_PER_MINUTE},
-		PUBLISH: {PUBLISH_MAX_TOKENS, PUBLISH_RATE_PER_MINUTE},
+		DEFAULT:          {DEFAULT_MAX_TOKENS, DEFAULT_RATE_PER_MINUTE},
+		V3_DEFAULT:       {DEFAULT_MAX_TOKENS, DEFAULT_RATE_PER_MINUTE},
+		PUBLISH:          {PUBLISH_MAX_TOKENS, PUBLISH_RATE_PER_MINUTE},
+		V3_PUBLISH:       {PUBLISH_MAX_TOKENS, PUBLISH_RATE_PER_MINUTE},
+		IDENTITY_PUBLISH: {IDENTITY_PUBLISH_MAX_TOKENS, IDENTITY_PUBLISH_RATE_PER_MINUTE},
 	}
 	return tb
 }
@@ -109,6 +117,7 @@ func (rl *TokenBucketRateLimiter) fillAndReturnEntry(limitType LimitType, bucket
 	limit := rl.getLimit(limitType)
 	multiplier := uint16(1)
 	if isPriority {
+		// There are no priority limits for V3
 		if limitType == PUBLISH {
 			multiplier = rl.PublishPriorityMultiplier
 		} else {