From 0a25bd7194ee534ee256d79af1f0d8926cc002d7 Mon Sep 17 00:00:00 2001
From: MSalopek <matija.salopek994@gmail.com>
Date: Thu, 1 Feb 2024 10:59:02 +0100
Subject: [PATCH] fix: handle edgecases (limit 0) and add test

---
 x/authz/keeper/keeper.go      |  9 +++--
 x/authz/keeper/keeper_test.go | 69 +++++++++++++++++++++++++++++++++++
 2 files changed, 74 insertions(+), 4 deletions(-)

diff --git a/x/authz/keeper/keeper.go b/x/authz/keeper/keeper.go
index ddac85e2fde5..0056eb9ffd00 100644
--- a/x/authz/keeper/keeper.go
+++ b/x/authz/keeper/keeper.go
@@ -384,6 +384,11 @@ func (k Keeper) DequeueAndDeleteExpiredGrants(ctx sdk.Context, limit int) error
 
 	count := 0
 	for ; iterator.Valid(); iterator.Next() {
+		// limit the amount of iterations to avoid taking too much time
+		if count >= limit {
+			return nil
+		}
+
 		var queueItem authz.GrantQueueItem
 		if err := k.cdc.Unmarshal(iterator.Value(), &queueItem); err != nil {
 			return err
@@ -400,11 +405,7 @@ func (k Keeper) DequeueAndDeleteExpiredGrants(ctx sdk.Context, limit int) error
 			store.Delete(grantStoreKey(grantee, granter, typeURL))
 		}
 
-		// limit the amount of iterations to avoid taking too much time
 		count++
-		if count == limit {
-			return nil
-		}
 	}
 
 	return nil
diff --git a/x/authz/keeper/keeper_test.go b/x/authz/keeper/keeper_test.go
index 737640aa9ad6..2995cfd46cf4 100644
--- a/x/authz/keeper/keeper_test.go
+++ b/x/authz/keeper/keeper_test.go
@@ -370,6 +370,7 @@ func (s *TestSuite) TestDequeueAllGrantsQueue() {
 	require.NoError(err)
 
 	newCtx := s.ctx.WithBlockTime(exp.AddDate(1, 0, 0))
+	// setting a high limit so all grants are dequeued
 	err = s.authzKeeper.DequeueAndDeleteExpiredGrants(newCtx, 200)
 	require.NoError(err)
 
@@ -391,6 +392,74 @@ func (s *TestSuite) TestDequeueAllGrantsQueue() {
 	require.Len(authzs, 1)
 }
 
+func (s *TestSuite) TestDequeueGrantsQueueEdgecases() {
+	require := s.Require()
+	addrs := s.addrs
+	granter := addrs[0]
+	grantee := addrs[1]
+	grantee1 := addrs[2]
+	exp := s.ctx.BlockTime().AddDate(0, 0, 1)
+	a := banktypes.SendAuthorization{SpendLimit: coins100}
+
+	// create few authorizations
+	err := s.authzKeeper.SaveGrant(s.ctx, grantee, granter, &a, &exp)
+	require.NoError(err)
+
+	err = s.authzKeeper.SaveGrant(s.ctx, grantee1, granter, &a, &exp)
+	require.NoError(err)
+
+	exp2 := exp.AddDate(0, 1, 0)
+	err = s.authzKeeper.SaveGrant(s.ctx, granter, grantee1, &a, &exp2)
+	require.NoError(err)
+
+	exp2 = exp.AddDate(2, 0, 0)
+	err = s.authzKeeper.SaveGrant(s.ctx, granter, grantee, &a, &exp2)
+	require.NoError(err)
+
+	newCtx := s.ctx.WithBlockTime(exp.AddDate(1, 0, 0))
+	err = s.authzKeeper.DequeueAndDeleteExpiredGrants(newCtx, 0)
+	require.NoError(err)
+
+	s.T().Log("verify no pruning happens when limit is 0")
+	authzs, err := s.authzKeeper.GetAuthorizations(newCtx, grantee, granter)
+	require.NoError(err)
+	require.Len(authzs, 1)
+
+	authzs, err = s.authzKeeper.GetAuthorizations(newCtx, granter, grantee1)
+	require.NoError(err)
+	require.Len(authzs, 1)
+
+	authzs, err = s.authzKeeper.GetAuthorizations(newCtx, grantee1, granter)
+	require.NoError(err)
+	require.Len(authzs, 1)
+
+	authzs, err = s.authzKeeper.GetAuthorizations(newCtx, granter, grantee)
+	require.NoError(err)
+	require.Len(authzs, 1)
+
+	// expecting to prune 1 record when limit is 1
+	newCtx = s.ctx.WithBlockTime(exp.AddDate(1, 0, 0))
+	err = s.authzKeeper.DequeueAndDeleteExpiredGrants(newCtx, 1)
+	require.NoError(err)
+
+	s.T().Log("verify 1 record is prunded when limit is 1")
+	authzs, err = s.authzKeeper.GetAuthorizations(newCtx, grantee, granter)
+	require.NoError(err)
+	require.Len(authzs, 0)
+
+	authzs, err = s.authzKeeper.GetAuthorizations(newCtx, granter, grantee1)
+	require.NoError(err)
+	require.Len(authzs, 1)
+
+	authzs, err = s.authzKeeper.GetAuthorizations(newCtx, grantee1, granter)
+	require.NoError(err)
+	require.Len(authzs, 1)
+
+	authzs, err = s.authzKeeper.GetAuthorizations(newCtx, granter, grantee)
+	require.NoError(err)
+	require.Len(authzs, 1)
+}
+
 func (s *TestSuite) TestGetAuthorization() {
 	addr1 := s.addrs[3]
 	addr2 := s.addrs[4]