From c6095cb82ed81f345b4f47f0396c1ca802f982bc Mon Sep 17 00:00:00 2001
From: Crystal Lemire <crystal@dydx.exchange>
Date: Fri, 20 Oct 2023 09:56:00 -0700
Subject: [PATCH] Address validation on cli query input.

---
 .../clob/client/cli/query_stateful_orders.go  | 11 ++++
 .../client/cli/query_stateful_orders_test.go  | 57 ++++++++++++++-----
 2 files changed, 53 insertions(+), 15 deletions(-)

diff --git a/protocol/x/clob/client/cli/query_stateful_orders.go b/protocol/x/clob/client/cli/query_stateful_orders.go
index 080698dfbd..d051e41af1 100644
--- a/protocol/x/clob/client/cli/query_stateful_orders.go
+++ b/protocol/x/clob/client/cli/query_stateful_orders.go
@@ -2,12 +2,16 @@ package cli
 
 import (
 	"context"
+	"fmt"
 	"github.com/cosmos/cosmos-sdk/client"
 	"github.com/cosmos/cosmos-sdk/client/flags"
+	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/dydxprotocol/v4-chain/protocol/x/clob/types"
 	satypes "github.com/dydxprotocol/v4-chain/protocol/x/subaccounts/types"
 	"github.com/spf13/cast"
 	"github.com/spf13/cobra"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 )
 
 func CmdListStatefulOrders() *cobra.Command {
@@ -59,6 +63,13 @@ func CmdGetStatefulOrderCount() *cobra.Command {
 				return err
 			}
 
+			if _, err := sdk.AccAddressFromBech32(argOwner); err != nil {
+				return status.Error(
+					codes.InvalidArgument,
+					fmt.Sprintf("Invalid owner address: %v", err),
+				)
+			}
+
 			params := &types.QueryStatefulOrderCountRequest{
 				SubaccountId: &satypes.SubaccountId{
 					Owner:  argOwner,
diff --git a/protocol/x/clob/client/cli/query_stateful_orders_test.go b/protocol/x/clob/client/cli/query_stateful_orders_test.go
index 69243df32d..213aa41670 100644
--- a/protocol/x/clob/client/cli/query_stateful_orders_test.go
+++ b/protocol/x/clob/client/cli/query_stateful_orders_test.go
@@ -10,6 +10,8 @@ import (
 	clobcli "github.com/dydxprotocol/v4-chain/protocol/x/clob/client/cli"
 	"github.com/dydxprotocol/v4-chain/protocol/x/clob/types"
 	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 	"testing"
 )
 
@@ -36,23 +38,48 @@ func TestCmdListStatefulOrders(t *testing.T) {
 
 func TestCmdGetStatefulOrderCount(t *testing.T) {
 	net, _ := networkWithClobPairObjects(t, 2)
-
 	ctx := net.Validators[0].ClientCtx
 
-	args := []string{
-		fmt.Sprintf("--%s=json", tmcli.OutputFlag),
-		constants.AliceAccAddress.String(), // owner
-		"0",                                // subaccount number
-	}
-
-	out, err := cli.ExecTestCLICmd(ctx, clobcli.CmdGetStatefulOrderCount(), args)
-	require.NoError(t, err)
+	for name, tc := range map[string]struct {
+		owner       string
+		number      uint32
+		count       uint32
+		expectedErr error
+	}{
+		"valid": {
+			owner:  constants.AliceAccAddress.String(),
+			number: 0,
+			count:  0,
+		},
+		"invalid owner": {
+			owner:  "invalid",
+			number: 0,
+			expectedErr: status.Error(
+				codes.InvalidArgument,
+				"Invalid owner address: decoding bech32 failed: invalid bech32 string length 7",
+			),
+		},
+	} {
+		t.Run(name, func(t *testing.T) {
+			args := []string{
+				fmt.Sprintf("--%s=json", tmcli.OutputFlag),
+				tc.owner,
+				fmt.Sprintf("%d", tc.number),
+			}
+			out, err := cli.ExecTestCLICmd(ctx, clobcli.CmdGetStatefulOrderCount(), args)
+			if tc.expectedErr != nil {
+				require.ErrorIs(t, err, tc.expectedErr)
+			} else {
+				require.NoError(t, err)
 
-	var res types.QueryStatefulOrderCountResponse
-	require.NoError(t, net.Config.Codec.UnmarshalJSON(out.Bytes(), &res))
+				var res types.QueryStatefulOrderCountResponse
+				require.NoError(t, net.Config.Codec.UnmarshalJSON(out.Bytes(), &res))
 
-	// Expect no stateful orders.
-	require.Equal(t, types.QueryStatefulOrderCountResponse{
-		Count: 0,
-	}, res)
+				// Expect no stateful orders.
+				require.Equal(t, types.QueryStatefulOrderCountResponse{
+					Count: tc.count,
+				}, res)
+			}
+		})
+	}
 }