From 0e74440a3435bf86933e9c839ccfdd2e293093f5 Mon Sep 17 00:00:00 2001 From: Steven Danna Date: Wed, 17 Feb 2021 22:32:10 +0000 Subject: [PATCH] rpc: install custom proto codec This change enables future upgrades of google.golang.org/grpc. Recent changes to google.golang.org/grpc have changed the default proto codec such that it is incompatible with gogoproto: https://github.com/grpc/grpc-go/pull/3958 While `proto.Marshal` does do a check for the Marshal/Unmarshal it appears to do this through reflection methods that will panic on fields using gogoproto.nullable. To avoid this, we install a proto codec that does an interface check before dispatching to proto.Marshal/proto.Unmarshal. Note that I haven't replicated the protobuffer caching that previously existed in google.golang.org/grpc. The referenced change claims removing the caching may be a small performance win; however, they also removed the typecheck that we need to do, so it is unclear. Fixes #60531 Release note: None --- pkg/rpc/BUILD.bazel | 7 ++++- pkg/rpc/codec.go | 48 +++++++++++++++++++++++++++++++++ pkg/rpc/codec_test.go | 44 ++++++++++++++++++++++++++++++ pkg/rpc/context.go | 4 +-- pkg/testutils/lint/lint_test.go | 7 +++++ 5 files changed, 106 insertions(+), 4 deletions(-) create mode 100644 pkg/rpc/codec.go create mode 100644 pkg/rpc/codec_test.go diff --git a/pkg/rpc/BUILD.bazel b/pkg/rpc/BUILD.bazel index 257b6d1323ce..740b931c04f3 100644 --- a/pkg/rpc/BUILD.bazel +++ b/pkg/rpc/BUILD.bazel @@ -9,6 +9,7 @@ go_library( "auth_tenant.go", "breaker.go", "clock_offset.go", + "codec.go", "connection_class.go", "context.go", "context_testutils.go", @@ -49,6 +50,8 @@ go_library( "@com_github_cockroachdb_errors//:errors", "@com_github_cockroachdb_logtags//:logtags", "@com_github_facebookgo_clock//:clock", + "@com_github_gogo_protobuf//proto", + "@com_github_golang_protobuf//proto:go_default_library", "@com_github_golang_snappy//:snappy", "@com_github_montanaflynn_stats//:stats", "@com_github_vividcortex_ewma//:ewma", @@ -57,7 +60,6 @@ go_library( "@org_golang_google_grpc//codes", "@org_golang_google_grpc//credentials", "@org_golang_google_grpc//encoding", - "@org_golang_google_grpc//encoding/proto", "@org_golang_google_grpc//keepalive", "@org_golang_google_grpc//metadata", "@org_golang_google_grpc//peer", @@ -73,6 +75,7 @@ go_test( srcs = [ "auth_test.go", "clock_offset_test.go", + "codec_test.go", "context_test.go", "heartbeat_test.go", "main_test.go", @@ -105,9 +108,11 @@ go_test( "@com_github_cockroachdb_errors//:errors", "@com_github_cockroachdb_logtags//:logtags", "@com_github_stretchr_testify//require", + "@io_etcd_go_etcd_raft_v3//raftpb", "@org_golang_google_grpc//:go_default_library", "@org_golang_google_grpc//codes", "@org_golang_google_grpc//credentials", + "@org_golang_google_grpc//health/grpc_health_v1", "@org_golang_google_grpc//keepalive", "@org_golang_google_grpc//metadata", "@org_golang_google_grpc//peer", diff --git a/pkg/rpc/codec.go b/pkg/rpc/codec.go new file mode 100644 index 000000000000..8211f602ebf3 --- /dev/null +++ b/pkg/rpc/codec.go @@ -0,0 +1,48 @@ +// Copyright 2020 The Cockroach Authors. +// +// Use of this software is governed by the Business Source License +// included in the file licenses/BSL.txt. +// +// As of the Change Date specified in that file, in accordance with +// the Business Source License, use of this software will be governed +// by the Apache License, Version 2.0, included in the file +// licenses/APL.txt. + +package rpc + +import ( + "github.com/gogo/protobuf/proto" + // Used instead of gogo/protobuf/proto for the fallback case + // to match the behavior of the upstream codec in + // google.golang.org/grpc/encoding/proto that we are + // replacing: + // + // https://github.com/grpc/grpc-go/blob/7b167fd6eca1ab8f05ec14085d63197cacd41438/encoding/proto/proto.go + // + gproto "github.com/golang/protobuf/proto" + "google.golang.org/grpc/encoding" +) + +const name = "proto" + +type codec struct{} + +var _ encoding.Codec = codec{} + +func (codec) Marshal(v interface{}) ([]byte, error) { + if pm, ok := v.(proto.Marshaler); ok { + return pm.Marshal() + } + return gproto.Marshal(v.(gproto.Message)) +} + +func (codec) Unmarshal(data []byte, v interface{}) error { + if pm, ok := v.(proto.Unmarshaler); ok { + return pm.Unmarshal(data) + } + return gproto.Unmarshal(data, v.(gproto.Message)) +} + +func (codec) Name() string { + return name +} diff --git a/pkg/rpc/codec_test.go b/pkg/rpc/codec_test.go new file mode 100644 index 000000000000..102f9a1e7b58 --- /dev/null +++ b/pkg/rpc/codec_test.go @@ -0,0 +1,44 @@ +// Copyright 2020 The Cockroach Authors. +// +// Use of this software is governed by the Business Source License +// included in the file licenses/BSL.txt. +// +// As of the Change Date specified in that file, in accordance with +// the Business Source License, use of this software will be governed +// by the Apache License, Version 2.0, included in the file +// licenses/APL.txt. + +package rpc + +import ( + "fmt" + "testing" + + "github.com/cockroachdb/cockroach/pkg/roachpb" + "github.com/cockroachdb/cockroach/pkg/util/leaktest" + "github.com/stretchr/testify/require" + "go.etcd.io/etcd/raft/v3/raftpb" + "google.golang.org/grpc/health/grpc_health_v1" +) + +func TestCodec(t *testing.T) { + defer leaktest.AfterTest(t)() + + testCodec := codec{} + for _, test := range []struct { + name string + msgBuilder func() interface{} + }{ + {"rpc.PingRequest", func() interface{} { return &PingRequest{} }}, + {"raftpb.Message", func() interface{} { return &raftpb.Message{} }}, + {"grpc_health_v1.HealthCheckRequest", func() interface{} { return &grpc_health_v1.HealthCheckRequest{} }}, + {"roachpb.GetRequest", func() interface{} { return &roachpb.GetRequest{} }}, + } { + t.Run(fmt.Sprintf("can call marshal and unmarhsal on %s", test.name), func(t *testing.T) { + marshaled, err := testCodec.Marshal(test.msgBuilder()) + require.NoError(t, err, "marshal failed") + err = testCodec.Unmarshal(marshaled, test.msgBuilder()) + require.NoError(t, err, "unmarshal failed") + }) + } +} diff --git a/pkg/rpc/context.go b/pkg/rpc/context.go index 0e5f0f81d7c9..865f9de9fcdc 100644 --- a/pkg/rpc/context.go +++ b/pkg/rpc/context.go @@ -46,7 +46,6 @@ import ( "google.golang.org/grpc/codes" "google.golang.org/grpc/credentials" "google.golang.org/grpc/encoding" - encodingproto "google.golang.org/grpc/encoding/proto" "google.golang.org/grpc/metadata" grpcstatus "google.golang.org/grpc/status" ) @@ -794,8 +793,7 @@ func (c growStackCodec) Unmarshal(data []byte, v interface{}) error { // Install the growStackCodec over the default proto codec in order to grow the // stack for BatchRequest RPCs prior to unmarshaling. func init() { - protoCodec := encoding.GetCodec(encodingproto.Name) - encoding.RegisterCodec(growStackCodec{Codec: protoCodec}) + encoding.RegisterCodec(growStackCodec{Codec: codec{}}) } // onlyOnceDialer implements the grpc.WithDialer interface but only diff --git a/pkg/testutils/lint/lint_test.go b/pkg/testutils/lint/lint_test.go index 8e77384acebc..463123cd5553 100644 --- a/pkg/testutils/lint/lint_test.go +++ b/pkg/testutils/lint/lint_test.go @@ -998,6 +998,8 @@ func TestLint(t *testing.T) { ":!sql/*.pb.go", ":!util/protoutil/marshal.go", ":!util/protoutil/marshaler.go", + ":!rpc/codec.go", + ":!rpc/codec_test.go", ":!settings/settings_test.go", ":!sql/types/types_jsonpb.go", ":!sql/schemachanger/scbuild/builder_test.go", @@ -1042,6 +1044,8 @@ func TestLint(t *testing.T) { ":!util/protoutil/marshaler.go", ":!util/encoding/encoding.go", ":!util/hlc/timestamp.go", + ":!rpc/codec.go", + ":!rpc/codec_test.go", ":!sql/types/types_jsonpb.go", ) if err != nil { @@ -1481,6 +1485,7 @@ func TestLint(t *testing.T) { stream.GrepNot(`cockroach/pkg/util/log: github\.com/pkg/errors$`), stream.GrepNot(`cockroach/pkg/(base|release|security|util/(log|randutil|stop)): log$`), stream.GrepNot(`cockroach/pkg/(server/serverpb|ts/tspb): github\.com/golang/protobuf/proto$`), + stream.GrepNot(`cockroachdb/cockroach/pkg/rpc: github\.com/golang/protobuf/proto$`), stream.GrepNot(`cockroachdb/cockroach/pkg/sql/lex/allkeywords: log$`), stream.GrepNot(`cockroachdb/cockroach/pkg/util/timeutil/gen: log$`), stream.GrepNot(`cockroachdb/cockroach/pkg/roachpb/gen: log$`), @@ -1646,6 +1651,8 @@ func TestLint(t *testing.T) { stream.GrepNot(`pkg/.*.go:.* func .*\.Cause is unused`), // Using deprecated WireLength call. stream.GrepNot(`pkg/rpc/stats_handler.go:.*v.WireLength is deprecated: This field is never set.*`), + // rpc/codec.go imports the same proto package that grpc-go imports (as of crdb@dd87d1145 and grpc-go@7b167fd6) + stream.GrepNot(`pkg/rpc/codec.go:.*package github.com/golang/protobuf/proto is deprecated: Use the "google.golang.org/protobuf/proto" package instead.`), ), func(s string) { t.Errorf("\n%s", s) }); err != nil {