From b72d182d2a24df67935e5183a07db6c6d3a06351 Mon Sep 17 00:00:00 2001
From: Alex Peters <alpe@users.noreply.github.com>
Date: Wed, 13 Sep 2023 17:17:30 +0200
Subject: [PATCH] Ensure some contraints and limits on pin/unpin code ids

---
 x/wasm/types/tx.go      | 33 ++++++++++++++++++++++++++++-----
 x/wasm/types/tx_test.go | 37 ++++++++++++++++++++++++++++++++++++-
 2 files changed, 64 insertions(+), 6 deletions(-)

diff --git a/x/wasm/types/tx.go b/x/wasm/types/tx.go
index 888a1a1b3e..2bf96987be 100644
--- a/x/wasm/types/tx.go
+++ b/x/wasm/types/tx.go
@@ -485,12 +485,25 @@ func (msg MsgPinCodes) GetSignBytes() []byte {
 	return sdk.MustSortJSON(ModuleCdc.MustMarshalJSON(&msg))
 }
 
+const maxCodeIDTotal = 50
+
 func (msg MsgPinCodes) ValidateBasic() error {
 	if _, err := sdk.AccAddressFromBech32(msg.Authority); err != nil {
 		return errorsmod.Wrap(err, "authority")
 	}
-	if len(msg.CodeIDs) == 0 {
+	return validateCodeIDs(msg.CodeIDs)
+}
+
+// ensure not empty, not duplicates and not exceeding max number
+func validateCodeIDs(codeIDs []uint64) error {
+	switch n := len(codeIDs); {
+	case n == 0:
 		return errorsmod.Wrap(sdkerrors.ErrInvalidRequest, "empty code ids")
+	case n > maxCodeIDTotal:
+		return errorsmod.Wrapf(sdkerrors.ErrInvalidRequest, "total number of code ids is greater than %d", maxCodeIDTotal)
+	}
+	if hasDuplicates(codeIDs) {
+		return errorsmod.Wrap(sdkerrors.ErrInvalidRequest, "duplicate code ids")
 	}
 	return nil
 }
@@ -519,10 +532,7 @@ func (msg MsgUnpinCodes) ValidateBasic() error {
 	if _, err := sdk.AccAddressFromBech32(msg.Authority); err != nil {
 		return errorsmod.Wrap(err, "authority")
 	}
-	if len(msg.CodeIDs) == 0 {
-		return errorsmod.Wrap(sdkerrors.ErrInvalidRequest, "empty code ids")
-	}
-	return nil
+	return validateCodeIDs(msg.CodeIDs)
 }
 
 func (msg MsgSudoContract) Route() string {
@@ -684,6 +694,7 @@ func (msg MsgRemoveCodeUploadParamsAddresses) ValidateBasic() error {
 func checkDuplicatedAddresses(addresses []string) error {
 	index := map[string]struct{}{}
 	for _, addr := range addresses {
+		addr = strings.ToUpper(addr)
 		if _, err := sdk.AccAddressFromBech32(addr); err != nil {
 			return errorsmod.Wrap(err, "addresses")
 		}
@@ -739,3 +750,15 @@ func (msg MsgStoreAndMigrateContract) ValidateBasic() error {
 	}
 	return nil
 }
+
+// returns true when slice contains any duplicates
+func hasDuplicates[T comparable](s []T) bool {
+	index := make(map[T]struct{}, len(s))
+	for _, v := range s {
+		if _, exists := index[v]; exists {
+			return true
+		}
+		index[v] = struct{}{}
+	}
+	return false
+}
diff --git a/x/wasm/types/tx_test.go b/x/wasm/types/tx_test.go
index cd9d5af38e..e3ea0f610a 100644
--- a/x/wasm/types/tx_test.go
+++ b/x/wasm/types/tx_test.go
@@ -885,7 +885,6 @@ func TestMsgRemoveCodeUploadParamsAddressesValidation(t *testing.T) {
 func TestMsgPinCodesValidation(t *testing.T) {
 	// proper address size
 	goodAddress := sdk.AccAddress(make([]byte, 20)).String()
-
 	specs := map[string]struct {
 		src    MsgPinCodes
 		expErr bool
@@ -915,6 +914,20 @@ func TestMsgPinCodesValidation(t *testing.T) {
 			},
 			expErr: true,
 		},
+		"exceeds max code ids": {
+			src: MsgPinCodes{
+				Authority: goodAddress,
+				CodeIDs:   genCodeIDs(51),
+			},
+			expErr: true,
+		},
+		"duplicate code ids": {
+			src: MsgPinCodes{
+				Authority: goodAddress,
+				CodeIDs:   []uint64{1, 1},
+			},
+			expErr: true,
+		},
 	}
 	for msg, spec := range specs {
 		t.Run(msg, func(t *testing.T) {
@@ -961,6 +974,20 @@ func TestMsgUnpinCodesValidation(t *testing.T) {
 			},
 			expErr: true,
 		},
+		"exceeds max code ids": {
+			src: MsgUnpinCodes{
+				Authority: goodAddress,
+				CodeIDs:   genCodeIDs(51),
+			},
+			expErr: true,
+		},
+		"duplicate code ids": {
+			src: MsgUnpinCodes{
+				Authority: goodAddress,
+				CodeIDs:   []uint64{1, 1},
+			},
+			expErr: true,
+		},
 	}
 	for msg, spec := range specs {
 		t.Run(msg, func(t *testing.T) {
@@ -974,6 +1001,14 @@ func TestMsgUnpinCodesValidation(t *testing.T) {
 	}
 }
 
+func genCodeIDs(max int) []uint64 {
+	r := make([]uint64, max)
+	for i := 0; i < max; i++ {
+		r[i] = uint64(i)
+	}
+	return r
+}
+
 func TestMsgSudoContractValidation(t *testing.T) {
 	badAddress := "abcd"
 	// proper address size