From 148e786ba0e3026ad7daef74859106bee93883e9 Mon Sep 17 00:00:00 2001
From: Tian Qin <tian@dydx.exchange>
Date: Thu, 23 May 2024 13:56:40 -0400
Subject: [PATCH] validate that vault deposit does not exceed max uint64

---
 .../vault/keeper/msg_server_deposit_to_vault.go  |  9 +++++++--
 .../keeper/msg_server_deposit_to_vault_test.go   | 16 +++++++++++++++-
 protocol/x/vault/types/msg_deposit_to_vault.go   |  3 ++-
 3 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/protocol/x/vault/keeper/msg_server_deposit_to_vault.go b/protocol/x/vault/keeper/msg_server_deposit_to_vault.go
index da776f2a73..1e67c86f65 100644
--- a/protocol/x/vault/keeper/msg_server_deposit_to_vault.go
+++ b/protocol/x/vault/keeper/msg_server_deposit_to_vault.go
@@ -3,6 +3,7 @@ package keeper
 import (
 	"context"
 
+	errorsmod "cosmossdk.io/errors"
 	"github.com/dydxprotocol/v4-chain/protocol/lib"
 	"github.com/dydxprotocol/v4-chain/protocol/lib/log"
 	"github.com/dydxprotocol/v4-chain/protocol/lib/metrics"
@@ -17,13 +18,14 @@ func (k msgServer) DepositToVault(
 	msg *types.MsgDepositToVault,
 ) (*types.MsgDepositToVaultResponse, error) {
 	ctx := lib.UnwrapSDKContext(goCtx, types.ModuleName)
+	quoteQuantums := msg.QuoteQuantums.BigInt()
 
 	// Mint shares for the vault.
 	err := k.MintShares(
 		ctx,
 		*msg.VaultId,
 		msg.SubaccountId.Owner,
-		msg.QuoteQuantums.BigInt(),
+		quoteQuantums,
 	)
 	if err != nil {
 		return nil, err
@@ -32,13 +34,16 @@ func (k msgServer) DepositToVault(
 	// Transfer from sender subaccount to vault.
 	// Note: Transfer should take place after minting shares for
 	// shares calculation to be correct.
+	if !quoteQuantums.IsUint64() {
+		return nil, errorsmod.Wrap(types.ErrInvalidDepositAmount, "quote quantums must be strictly less than 2^64")
+	}
 	err = k.sendingKeeper.ProcessTransfer(
 		ctx,
 		&sendingtypes.Transfer{
 			Sender:    *msg.SubaccountId,
 			Recipient: *msg.VaultId.ToSubaccountId(),
 			AssetId:   assettypes.AssetUsdc.Id,
-			Amount:    msg.QuoteQuantums.BigInt().Uint64(),
+			Amount:    quoteQuantums.Uint64(),
 		},
 	)
 	if err != nil {
diff --git a/protocol/x/vault/keeper/msg_server_deposit_to_vault_test.go b/protocol/x/vault/keeper/msg_server_deposit_to_vault_test.go
index b803769a86..af5c72f70f 100644
--- a/protocol/x/vault/keeper/msg_server_deposit_to_vault_test.go
+++ b/protocol/x/vault/keeper/msg_server_deposit_to_vault_test.go
@@ -2,6 +2,7 @@ package keeper_test
 
 import (
 	"bytes"
+	"math"
 	"math/big"
 	"testing"
 
@@ -198,7 +199,7 @@ func TestMsgDepositToVault(t *testing.T) {
 				big.NewInt(1_000),
 			},
 		},
-		"Two failed deposits due to non-positive amounts": {
+		"Three failed deposits due to invalid deposit amount": {
 			vaultId: constants.Vault_Clob_1,
 			depositorSetups: []DepositorSetup{
 				{
@@ -227,14 +228,27 @@ func TestMsgDepositToVault(t *testing.T) {
 					checkTxResponseContains: "Deposit amount is invalid",
 					expectedOwnerShares:     nil,
 				},
+				{
+					depositor: constants.Bob_Num0,
+					depositAmount: new(big.Int).Add(
+						new(big.Int).SetUint64(math.MaxUint64),
+						big.NewInt(1),
+					),
+					msgSigner:               constants.Bob_Num0.Owner,
+					checkTxFails:            true,
+					checkTxResponseContains: "Deposit amount is invalid",
+					expectedOwnerShares:     nil,
+				},
 			},
 			totalSharesHistory: []*big.Int{
 				big.NewInt(0),
 				big.NewInt(0),
+				big.NewInt(0),
 			},
 			vaultEquityHistory: []*big.Int{
 				big.NewInt(0),
 				big.NewInt(0),
+				big.NewInt(0),
 			},
 		},
 	}
diff --git a/protocol/x/vault/types/msg_deposit_to_vault.go b/protocol/x/vault/types/msg_deposit_to_vault.go
index eb3d091402..5255873f3d 100644
--- a/protocol/x/vault/types/msg_deposit_to_vault.go
+++ b/protocol/x/vault/types/msg_deposit_to_vault.go
@@ -1,6 +1,7 @@
 package types
 
 import (
+	errorsmod "cosmossdk.io/errors"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/dydxprotocol/v4-chain/protocol/dtypes"
 )
@@ -16,7 +17,7 @@ func (msg *MsgDepositToVault) ValidateBasic() error {
 
 	// Validate that quote quantums is positive.
 	if msg.QuoteQuantums.Cmp(dtypes.NewInt(0)) <= 0 {
-		return ErrInvalidDepositAmount
+		return errorsmod.Wrap(ErrInvalidDepositAmount, "quote quantums must be strictly positive")
 	}
 
 	return nil