diff --git a/compiler/cpp/src/thrift/generate/t_go_generator.cc b/compiler/cpp/src/thrift/generate/t_go_generator.cc index d52a6e261b..87524215cb 100644 --- a/compiler/cpp/src/thrift/generate/t_go_generator.cc +++ b/compiler/cpp/src/thrift/generate/t_go_generator.cc @@ -1768,6 +1768,11 @@ void t_go_generator::generate_countsetfields_helper(ostream& out, out << indent() << "func (p *" << tstruct_name << ") CountSetFields" << tstruct_name << "() int {" << '\n'; indent_up(); + out << indent() << "if p == nil {" << '\n'; + indent_up(); + out << indent() << "return 0" << '\n'; + indent_down(); + out << indent() << "}" << '\n'; out << indent() << "count := 0" << '\n'; for (f_iter = fields.begin(); f_iter != fields.end(); ++f_iter) { if ((*f_iter)->get_req() == t_field::T_REQUIRED) diff --git a/lib/go/test/UnionDefaultValueTest.thrift b/lib/go/test/UnionDefaultValueTest.thrift index 4c9348003e..9ecb631b6c 100644 --- a/lib/go/test/UnionDefaultValueTest.thrift +++ b/lib/go/test/UnionDefaultValueTest.thrift @@ -32,3 +32,8 @@ union Descendant { struct TestStruct { 1: optional Descendant descendant = { "option1": {}} } + +struct StructWithUnsetUnion { + 1: required bool f_1 = 1, + 2: Descendant f_2, +} diff --git a/lib/go/test/tests/union_default_value_test.go b/lib/go/test/tests/union_default_value_test.go index a02569a6e9..f90a0e0ed0 100644 --- a/lib/go/test/tests/union_default_value_test.go +++ b/lib/go/test/tests/union_default_value_test.go @@ -20,9 +20,12 @@ package tests import ( + "context" + "strings" "testing" "github.com/apache/thrift/lib/go/test/gopath/src/uniondefaultvaluetest" + "github.com/apache/thrift/lib/go/thrift" ) func TestUnionDefaultValue(t *testing.T) { @@ -32,3 +35,33 @@ func TestUnionDefaultValue(t *testing.T) { t.Error("Default Union value not set!") } } + +func TestNilUnion(t *testing.T) { + var d *uniondefaultvaluetest.Descendant + if count := d.CountSetFieldsDescendant(); count != 0 { + t.Errorf("Expected 0 set fields for nil union, got %d", count) + } + + proto := thrift.NewTBinaryProtocolConf(thrift.NewTMemoryBuffer(), nil) + err := d.Write(context.Background(), proto) + if err == nil { + t.Error("Expected error when writing nil union, got nil") + } +} + +func TestStructWithUnsetUnion(t *testing.T) { + s := uniondefaultvaluetest.NewStructWithUnsetUnion() + buf := thrift.NewTMemoryBuffer() + proto := thrift.NewTBinaryProtocolConf(buf, nil) + + // Writing a struct whose union field is nil used to dereference the nil + // receiver inside CountSetFields and panic. It now reports which union + // could not be written instead. + err := s.Write(context.Background(), proto) + if err == nil { + t.Fatal("Expected an error writing a struct with an unset union, got nil") + } + if !strings.Contains(err.Error(), "exactly one field must be set") { + t.Errorf("Expected the union arity error, got %v", err) + } +}