Просмотр исходного кода

If a user passes a nil extension value to SetExtension(), return a useful error message.

The previous behavior was to silently drop all extensions from the encoded message.
David Symonds 11 лет назад
Родитель
Сommit
82b6f0b18a
2 измененных файлов с 25 добавлено и 0 удалено
  1. 9 0
      proto/extensions.go
  2. 16 0
      proto/extensions_test.go

+ 9 - 0
proto/extensions.go

@@ -37,6 +37,7 @@ package proto
 
 import (
 	"errors"
+	"fmt"
 	"reflect"
 	"strconv"
 	"sync"
@@ -321,6 +322,14 @@ func SetExtension(pb extendableProto, extension *ExtensionDesc, value interface{
 	if typ != reflect.TypeOf(value) {
 		return errors.New("proto: bad extension value type")
 	}
+	// nil extension values need to be caught early, because the
+	// encoder can't distinguish an ErrNil due to a nil extension
+	// from an ErrNil due to a missing field. Extensions are
+	// always optional, so the encoder would just swallow the error
+	// and drop all the extensions from the encoded message.
+	if reflect.ValueOf(value).IsNil() {
+		return fmt.Errorf("proto: SetExtension called with nil value of type %T", value)
+	}
 
 	pb.ExtensionMap()[extension.Field] = Extension{desc: extension, value: value}
 	return nil

+ 16 - 0
proto/extensions_test.go

@@ -135,3 +135,19 @@ func TestExtensionsRoundTrip(t *testing.T) {
 		t.Error("expected some sort of type mismatch error, got nil")
 	}
 }
+
+func TestNilExtension(t *testing.T) {
+	msg := &pb.MyMessage{
+		Count: proto.Int32(1),
+	}
+	if err := proto.SetExtension(msg, pb.E_Ext_Text, proto.String("hello")); err != nil {
+		t.Fatal(err)
+	}
+	if err := proto.SetExtension(msg, pb.E_Ext_More, (*pb.Ext)(nil)); err == nil {
+		t.Error("expected SetExtension to fail due to a nil extension")
+	} else if want := "proto: SetExtension called with nil value of type *testdata.Ext"; err.Error() != want {
+		t.Errorf("expected error %v, got %v", want, err)
+	}
+	// Note: if the behavior of Marshal is ever changed to ignore nil extensions, update
+	// this test to verify that E_Ext_Text is properly propagated through marshal->unmarshal.
+}