goprotobuf: Ensure GetExtension returns the same result when called on a freshly unmarshaled message.

LGTM=djd
R=djd
CC=golang-codereviews
https://codereview.appspot.com/178920043

2 files changed, 37 insertions, 1 deletions

diff --git a/proto/extensions.go b/proto/extensions.go
index e592053..fca7bde 100644
--- a/proto/extensions.go
+++ b/proto/extensions.go
@@ -227,7 +227,8 @@ func GetExtension(pb extendableProto, extension *ExtensionDesc) (interface{}, er
 		return nil, err
 	}
 
-	e, ok := pb.ExtensionMap()[extension.Field]
+	emap := pb.ExtensionMap()
+	e, ok := emap[extension.Field]
 	if !ok {
 		return nil, ErrMissingExtension
 	}
@@ -252,6 +253,7 @@ func GetExtension(pb extendableProto, extension *ExtensionDesc) (interface{}, er
 	e.value = v
 	e.desc = extension
 	e.enc = nil
+	emap[extension.Field] = e
 	return e.value, nil
 }
 
diff --git a/proto/extensions_test.go b/proto/extensions_test.go
index 03f756b..f8ed150 100644
--- a/proto/extensions_test.go
+++ b/proto/extensions_test.go
@@ -58,3 +58,37 @@ func TestGetExtensionsWithMissingExtensions(t *testing.T) {
 		t.Errorf("ext2 in returned extensions: %T %v", exts[1], exts[1])
 	}
 }
+
+func TestGetExtensionStability(t *testing.T) {
+	check := func(m *pb.MyMessage) bool {
+		ext1, err := proto.GetExtension(m, pb.E_Ext_More)
+		if err != nil {
+			t.Fatalf("GetExtension() failed: %s", err)
+		}
+		ext2, err := proto.GetExtension(m, pb.E_Ext_More)
+		if err != nil {
+			t.Fatalf("GetExtension() failed: %s", err)
+		}
+		return ext1 == ext2
+	}
+	msg := &pb.MyMessage{Count: proto.Int32(4)}
+	ext0 := &pb.Ext{}
+	if err := proto.SetExtension(msg, pb.E_Ext_More, ext0); err != nil {
+		t.Fatalf("Could not set ext1: %s", ext0)
+	}
+	if !check(msg) {
+		t.Errorf("GetExtension() not stable before marshaling")
+	}
+	bb, err := proto.Marshal(msg)
+	if err != nil {
+		t.Fatalf("Marshal() failed: %s", err)
+	}
+	msg1 := &pb.MyMessage{}
+	err = proto.Unmarshal(bb, msg1)
+	if err != nil {
+		t.Fatalf("Unmarshal() failed: %s", err)
+	}
+	if !check(msg1) {
+		t.Errorf("GetExtension() not stable after unmarshaling")
+	}
+}