From a9f02d9cb7447b3da1fb00b4458588db55e9b5fb Mon Sep 17 00:00:00 2001
From: yihuang <huang@crypto.com>
Date: Thu, 29 Sep 2022 09:20:03 +0800
Subject: [PATCH] fix: add gRPC nil/zero check in query (#13352)

---
 CHANGELOG.md               |  1 +
 baseapp/grpcrouter_test.go | 12 ++++++------
 codec/proto_codec.go       |  5 +++++
 codec/proto_codec_test.go  | 27 +++++++++++++++++++++++++++
 4 files changed, 39 insertions(+), 6 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4afd767db5ae..d8b5ca65b4d4 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -172,6 +172,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 * (x/gov) [#13051](https://github.com/cosmos/cosmos-sdk/pull/13051) In SubmitPropsal, when a legacy msg fails it's handler call, wrap the error as ErrInvalidProposalContent (instead of ErrNoProposalHandlerExists).
 * (x/gov) [#13045](https://github.com/cosmos/cosmos-sdk/pull/13045) Fix gov migrations for v3(0.46).
 * (Store) [#13334](https://github.com/cosmos/cosmos-sdk/pull/13334) Call streaming listeners for deliver tx event, it was removed accidentally.
+* (grpc) [#13352](https://github.com/cosmos/cosmos-sdk/pull/13352) fix grpc query panic that could crash the node.
 * (snapshot) [#13400](https://github.com/cosmos/cosmos-sdk/pull/13400) Fix snapshot checksum issue in golang 1.19.
 
 ### Deprecated
diff --git a/baseapp/grpcrouter_test.go b/baseapp/grpcrouter_test.go
index 89b419e18c75..7a672b10b54f 100644
--- a/baseapp/grpcrouter_test.go
+++ b/baseapp/grpcrouter_test.go
@@ -35,9 +35,9 @@ func TestGRPCQueryRouter(t *testing.T) {
 	require.NotNil(t, res)
 	require.Equal(t, "hello", res.Message)
 
-	require.Panics(t, func() {
-		_, _ = client.Echo(context.Background(), nil)
-	})
+	res, err = client.Echo(context.Background(), nil)
+	require.Nil(t, err)
+	require.Empty(t, res.Message)
 
 	res2, err := client.SayHello(context.Background(), &testdata.SayHelloRequest{Name: "Foo"})
 	require.Nil(t, err)
@@ -153,9 +153,9 @@ func testQueryDataRacesSameHandler(t *testing.T, makeClientConn func(*baseapp.GR
 			require.NotNil(t, res)
 			require.Equal(t, "hello", res.Message)
 
-			require.Panics(t, func() {
-				_, _ = client.Echo(context.Background(), nil)
-			})
+			res, err = client.Echo(context.Background(), nil)
+			require.Nil(t, err)
+			require.Empty(t, res.Message)
 
 			res2, err := client.SayHello(context.Background(), &testdata.SayHelloRequest{Name: "Foo"})
 			require.Nil(t, err)
diff --git a/codec/proto_codec.go b/codec/proto_codec.go
index 75815eaa8251..b1d4eb45183e 100644
--- a/codec/proto_codec.go
+++ b/codec/proto_codec.go
@@ -43,6 +43,11 @@ func NewProtoCodec(interfaceRegistry types.InterfaceRegistry) *ProtoCodec {
 // NOTE: this function must be used with a concrete type which
 // implements proto.Message. For interface please use the codec.MarshalInterface
 func (pc *ProtoCodec) Marshal(o ProtoMarshaler) ([]byte, error) {
+	// Size() check can catch the typed nil value.
+	if o == nil || o.Size() == 0 {
+		// return empty bytes instead of nil, because nil has special meaning in places like store.Set
+		return []byte{}, nil
+	}
 	return o.Marshal()
 }
 
diff --git a/codec/proto_codec_test.go b/codec/proto_codec_test.go
index 69e9066733e6..b661db1fe70c 100644
--- a/codec/proto_codec_test.go
+++ b/codec/proto_codec_test.go
@@ -3,15 +3,20 @@ package codec_test
 import (
 	"errors"
 	"fmt"
+	"math"
 	"reflect"
 	"testing"
 
 	"github.com/cosmos/gogoproto/proto"
 	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/encoding"
+	"google.golang.org/grpc/status"
 
 	"github.com/cosmos/cosmos-sdk/codec"
 	"github.com/cosmos/cosmos-sdk/codec/types"
 	"github.com/cosmos/cosmos-sdk/testutil/testdata"
+	banktypes "github.com/cosmos/cosmos-sdk/x/bank/types"
 )
 
 func createTestInterfaceRegistry() types.InterfaceRegistry {
@@ -109,6 +114,28 @@ func TestProtoCodecMarshal(t *testing.T) {
 
 	err = cartoonCdc.UnmarshalInterface(bz, &cartoon)
 	require.NoError(t, err)
+
+	// test typed nil input shouldn't panic
+	var v *banktypes.QueryBalanceResponse
+	bz, err = grpcServerEncode(cartoonCdc.GRPCCodec(), v)
+	require.NoError(t, err)
+	require.Empty(t, bz)
+}
+
+// Emulate grpc server implementation
+// https://github.com/grpc/grpc-go/blob/b1d7f56b81b7902d871111b82dec6ba45f854ede/rpc_util.go#L590
+func grpcServerEncode(c encoding.Codec, msg interface{}) ([]byte, error) {
+	if msg == nil { // NOTE: typed nils will not be caught by this check
+		return nil, nil
+	}
+	b, err := c.Marshal(msg)
+	if err != nil {
+		return nil, status.Errorf(codes.Internal, "grpc: error while marshaling: %v", err.Error())
+	}
+	if uint(len(b)) > math.MaxUint32 {
+		return nil, status.Errorf(codes.ResourceExhausted, "grpc: message too large (%d bytes)", len(b))
+	}
+	return b, nil
 }
 
 func TestProtoCodecUnmarshalLengthPrefixedChecks(t *testing.T) {