goprotobuf: Fix merging of duplicate message set entries.

LGTM=crawshaw
R=crawshaw
CC=golang-codereviews
https://codereview.appspot.com/118400044

2 files changed, 84 insertions, 5 deletions

diff --git a/proto/message_set.go b/proto/message_set.go
index 6ddcc30..1a17809 100644
--- a/proto/message_set.go
+++ b/proto/message_set.go
@@ -188,12 +188,25 @@ func UnmarshalMessageSet(buf []byte, m map[int32]Extension) error {
 		return err
 	}
 	for _, item := range ms.Item {
-		// restore wire type and field number varint, plus length varint.
-		b := EncodeVarint(uint64(*item.TypeId)<<3 | WireBytes)
-		b = append(b, EncodeVarint(uint64(len(item.Message)))...)
-		b = append(b, item.Message...)
+		id := *item.TypeId
+		msg := item.Message
+
+		// Restore wire type and field number varint, plus length varint.
+		// Be careful to preserve duplicate items.
+		b := EncodeVarint(uint64(id)<<3 | WireBytes)
+		if ext, ok := m[id]; ok {
+			// Existing data; rip off the tag and length varint
+			// so we join the new data correctly.
+			// We can assume that ext.enc is set because we are unmarshaling.
+			o := ext.enc[len(b):]   // skip wire type and field number
+			_, n := DecodeVarint(o) // calculate length of length varint
+			o = o[n:]               // skip length varint
+			msg = append(o, msg...) // join old data and new data
+		}
+		b = append(b, EncodeVarint(uint64(len(msg)))...)
+		b = append(b, msg...)
 
-		m[*item.TypeId] = Extension{enc: b}
+		m[id] = Extension{enc: b}
 	}
 	return nil
 }
diff --git a/proto/message_set_test.go b/proto/message_set_test.go
new file mode 100644
index 0000000..bb311bc
--- /dev/null
+++ b/proto/message_set_test.go
@@ -0,0 +1,66 @@
+// Go support for Protocol Buffers - Google's data interchange format
+//
+// Copyright 2014 The Go Authors.  All rights reserved.
+// http://code.google.com/p/goprotobuf/
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are
+// met:
+//
+//     * Redistributions of source code must retain the above copyright
+// notice, this list of conditions and the following disclaimer.
+//     * Redistributions in binary form must reproduce the above
+// copyright notice, this list of conditions and the following disclaimer
+// in the documentation and/or other materials provided with the
+// distribution.
+//     * Neither the name of Google Inc. nor the names of its
+// contributors may be used to endorse or promote products derived from
+// this software without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+package proto
+
+import (
+	"bytes"
+	"testing"
+)
+
+func TestUnmarshalMessageSetWithDuplicate(t *testing.T) {
+	// Check that a repeated message set entry will be concatenated.
+	in := &MessageSet{
+		Item: []*_MessageSet_Item{
+			{TypeId: Int32(12345), Message: []byte("hoo")},
+			{TypeId: Int32(12345), Message: []byte("hah")},
+		},
+	}
+	b, err := Marshal(in)
+	if err != nil {
+		t.Fatalf("Marshal: %v", err)
+	}
+	t.Logf("Marshaled bytes: %q", b)
+
+	m := make(map[int32]Extension)
+	if err := UnmarshalMessageSet(b, m); err != nil {
+		t.Fatalf("UnmarshalMessageSet: %v", err)
+	}
+	ext, ok := m[12345]
+	if !ok {
+		t.Fatalf("Didn't retrieve extension 12345; map is %v", m)
+	}
+	// Skip wire type/field number and length varints.
+	got := skipVarint(skipVarint(ext.enc))
+	if want := []byte("hoohah"); !bytes.Equal(got, want) {
+		t.Errorf("Combined extension is %q, want %q", got, want)
+	}
+}