diff --git a/pkg/api/interceptor.go b/pkg/api/interceptor.go index f99c051d..8bad0a06 100644 --- 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 index 0d391925..b210ea73 100644 --- 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 {