From 3595e3d47ba09abebac106ab52c19a53a8d0634a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?F=C3=A9lix=20C=2E=20Morency?=
 <1102868+fmorency@users.noreply.github.com>
Date: Sat, 20 Jul 2024 10:59:36 -0400
Subject: [PATCH] fix(ante): prevent nested withdraw delegator rewards (#201)

---
 ante/disable_withdraw_delegator_rewards.go | 13 +++++++++++
 e2e/helpers/distribution.go                | 15 +++++++++++++
 e2e/poa_test.go                            | 26 ++++++++++++++++++++++
 3 files changed, 54 insertions(+)
 create mode 100644 e2e/helpers/distribution.go

diff --git a/ante/disable_withdraw_delegator_rewards.go b/ante/disable_withdraw_delegator_rewards.go
index a5da585..4cf3ea1 100644
--- a/ante/disable_withdraw_delegator_rewards.go
+++ b/ante/disable_withdraw_delegator_rewards.go
@@ -2,6 +2,7 @@ package poaante
 
 import (
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/x/authz"
 	distrtypes "github.com/cosmos/cosmos-sdk/x/distribution/types"
 
 	"github.com/strangelove-ventures/poa"
@@ -30,6 +31,18 @@ func (mdwr MsgDisableWithdrawDelegatorRewards) AnteHandle(ctx sdk.Context, tx sd
 
 func (mdwr MsgDisableWithdrawDelegatorRewards) hasWithdrawDelegatorRewardsMsg(msgs []sdk.Msg) bool {
 	for _, msg := range msgs {
+		// authz nested message check (recursive)
+		if execMsg, ok := msg.(*authz.MsgExec); ok {
+			msgs, err := execMsg.GetMessages()
+			if err != nil {
+				return true
+			}
+
+			if mdwr.hasWithdrawDelegatorRewardsMsg(msgs) {
+				return true
+			}
+		}
+
 		if _, ok := msg.(*distrtypes.MsgWithdrawDelegatorReward); ok {
 			return true
 		}
diff --git a/e2e/helpers/distribution.go b/e2e/helpers/distribution.go
new file mode 100644
index 0000000..a8d0488
--- /dev/null
+++ b/e2e/helpers/distribution.go
@@ -0,0 +1,15 @@
+package helpers
+
+import (
+	"context"
+	"testing"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/strangelove-ventures/interchaintest/v8/chain/cosmos"
+	"github.com/strangelove-ventures/interchaintest/v8/ibc"
+)
+
+func WithdrawDelegatorRewards(t *testing.T, ctx context.Context, chain *cosmos.CosmosChain, user ibc.Wallet, valoper string) (sdk.TxResponse, error) {
+	cmd := TxCommandBuilder(ctx, chain, []string{"tx", "distribution", "withdraw-rewards", valoper}, user.KeyName())
+	return ExecuteTransaction(ctx, chain, cmd)
+}
diff --git a/e2e/poa_test.go b/e2e/poa_test.go
index 3ee32d7..dda35ab 100644
--- a/e2e/poa_test.go
+++ b/e2e/poa_test.go
@@ -46,6 +46,7 @@ func TestPOABase(t *testing.T) {
 
 	// === Test Cases ===
 	testStakingDisabled(t, ctx, chain, validators, acc0, acc1)
+	testWithdrawDelegatorRewardsDisabled(t, ctx, chain, validators, acc0, acc1)
 	testPowerErrors(t, ctx, chain, validators, incorrectUser, acc0)
 	testRemovePending(t, ctx, chain, acc0)
 }
@@ -74,6 +75,31 @@ func testStakingDisabled(t *testing.T, ctx context.Context, chain *cosmos.Cosmos
 	require.Contains(t, res.RawLog, poa.ErrStakingActionNotAllowed.Error())
 }
 
+func testWithdrawDelegatorRewardsDisabled(t *testing.T, ctx context.Context, chain *cosmos.CosmosChain, validators []string, acc0, acc1 ibc.Wallet) {
+	t.Log("\n===== TEST WITHDRAW DELEGATOR REWARDS DISABLED =====")
+
+	// Normal withdraw delegation rewards execution fails
+	txRes, _ := helpers.WithdrawDelegatorRewards(t, ctx, chain, acc0, validators[0])
+	require.Contains(t, txRes.RawLog, poa.ErrWithdrawDelegatorRewardsNotAllowed.Error())
+
+	granter := acc1
+	grantee := acc0
+
+	// Grant grantee (acc0) the ability to delegate from granter (acc1)
+	res, err := helpers.ExecuteAuthzGrantMsg(t, ctx, chain, granter, grantee, "/cosmos.distribution.v1beta1.MsgWithdrawDelegatorReward")
+	require.NoError(t, err)
+	require.EqualValues(t, res.Code, 0)
+
+	// Generate nested message
+	nested := []string{"tx", "distribution", "withdraw-rewards", validators[0]}
+	nestedCmd := helpers.TxCommandBuilder(ctx, chain, nested, granter.FormattedAddress())
+
+	// Execute nested message via a wrapped Exec
+	res, err = helpers.ExecuteAuthzExecMsg(t, ctx, chain, grantee, nestedCmd)
+	require.NoError(t, err)
+	require.Contains(t, res.RawLog, poa.ErrWithdrawDelegatorRewardsNotAllowed.Error())
+}
+
 func testRemovePending(t *testing.T, ctx context.Context, chain *cosmos.CosmosChain, acc0 ibc.Wallet) {
 	t.Log("\n===== TEST PENDING =====")