From bf1b73fcd00f42641c56ff8f05b0d754ce3afd33 Mon Sep 17 00:00:00 2001
From: Mohammed Affan <affan@dydx.exchange>
Date: Wed, 11 Sep 2024 17:06:19 -0400
Subject: [PATCH] [OTE-730] add test to check register affiliate auth (#2235)

Co-authored-by: Teddy Ding <teddy@dydx.exchange>
---
 .../affiliates/e2e/register_affiliate_test.go | 73 +++++++++++++++++++
 protocol/x/affiliates/keeper/msg_server.go    |  5 +-
 2 files changed, 77 insertions(+), 1 deletion(-)
 create mode 100644 protocol/x/affiliates/e2e/register_affiliate_test.go

diff --git a/protocol/x/affiliates/e2e/register_affiliate_test.go b/protocol/x/affiliates/e2e/register_affiliate_test.go
new file mode 100644
index 00000000000..5408d56b02d
--- /dev/null
+++ b/protocol/x/affiliates/e2e/register_affiliate_test.go
@@ -0,0 +1,73 @@
+package affiliate_test
+
+import (
+	"testing"
+
+	testapp "github.com/dydxprotocol/v4-chain/protocol/testutil/app"
+	constants "github.com/dydxprotocol/v4-chain/protocol/testutil/constants"
+	"github.com/dydxprotocol/v4-chain/protocol/x/affiliates/types"
+	"github.com/stretchr/testify/require"
+)
+
+func TestRegisterAffiliateInvalidSigner(t *testing.T) {
+	tApp := testapp.NewTestAppBuilder(t).Build()
+	ctx := tApp.InitChain()
+
+	testCases := []struct {
+		name          string
+		referee       string
+		affiliate     string
+		signer        string
+		expectSuccess bool
+	}{
+		{
+			name:          "Valid signer (referee)",
+			referee:       constants.BobAccAddress.String(),
+			affiliate:     constants.AliceAccAddress.String(),
+			signer:        constants.BobAccAddress.String(),
+			expectSuccess: true,
+		},
+		{
+			name:          "Invalid signer (affiliate)",
+			referee:       constants.BobAccAddress.String(),
+			affiliate:     constants.AliceAccAddress.String(),
+			signer:        constants.AliceAccAddress.String(),
+			expectSuccess: false,
+		},
+		{
+			name:          "Invalid signer (non-related address)",
+			referee:       constants.BobAccAddress.String(),
+			affiliate:     constants.AliceAccAddress.String(),
+			signer:        constants.CarlAccAddress.String(),
+			expectSuccess: false,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			msgRegisterAffiliate := types.MsgRegisterAffiliate{
+				Referee:   tc.referee,
+				Affiliate: tc.affiliate,
+			}
+
+			checkTxMsgRegisterAffiliate := testapp.MustMakeCheckTx(
+				ctx,
+				tApp.App,
+				testapp.MustMakeCheckTxOptions{
+					AccAddressForSigning: tc.signer,
+					Gas:                  constants.TestGasLimit,
+					FeeAmt:               constants.TestFeeCoins_5Cents,
+				},
+				&msgRegisterAffiliate,
+			)
+			checkTxResp := tApp.CheckTx(checkTxMsgRegisterAffiliate)
+
+			if tc.expectSuccess {
+				require.True(t, checkTxResp.IsOK(), "Expected CheckTx to succeed with valid signer")
+			} else {
+				require.True(t, checkTxResp.IsErr(), "Expected CheckTx to fail with invalid signer")
+				require.Contains(t, checkTxResp.Log, "pubKey does not match signer address")
+			}
+		})
+	}
+}
diff --git a/protocol/x/affiliates/keeper/msg_server.go b/protocol/x/affiliates/keeper/msg_server.go
index 00780ed7fbb..b21094d20ba 100644
--- a/protocol/x/affiliates/keeper/msg_server.go
+++ b/protocol/x/affiliates/keeper/msg_server.go
@@ -14,7 +14,10 @@ type msgServer struct {
 }
 
 // RegisterAffiliate implements types.MsgServer.
-// This is only valid if a referee signs the message.
+// This is only valid if a referee signs the message
+// since the referee field is annotated with cosmos.msg.v1.signer
+// in protos. This ensures that only referee is returned
+// as a signer when GetSigners is called for authentication.
 // For example, if Alice is the referee and Bob is the affiliate,
 // then only Alice can register Bob as an affiliate. Any
 // other signer that sends this message will be rejected.