net/proto2/go: add GetAllExtensionDescs returns a []*ExtensionDesc

// GetAllExtensionDescs returns a slice of all the descriptors
// extensions present in pb.
// // If an extension is not registered, a descriptor with only the
// 'field' value set will
// // be returned instead of a full descriptor.
// // The returned slice is not guaranteed to be in any given order.
// func GetAllExtensionDescs(pb Message) (extensions []*ExtensionDesc,
// err error)

proto/extensions.go

@@ -489,6 +489,34 @@ func GetExtensions(pb Message, es []*ExtensionDesc) (extensions []interface{}, e
 	return
 }
 
+// ExtensionDescs returns a new slice containing pb's extension descriptors, in undefined order.
+// For non-registered extensions, ExtensionDescs returns an incomplete descriptor containing
+// just the Field field, which defines the extension's field number.
+func ExtensionDescs(pb Message) ([]*ExtensionDesc, error) {
+	epb, ok := extendable(pb)
+	if !ok {
+		return nil, fmt.Errorf("proto: %T is not an extendable proto.Message", pb)
+	}
+	registeredExtensions := RegisteredExtensions(pb)
+
+	emap, mu := epb.extensionsRead()
+	mu.Lock()
+	defer mu.Unlock()
+	extensions := make([]*ExtensionDesc, 0, len(emap))
+	for extid, e := range emap {
+		desc := e.desc
+		if desc == nil {
+			desc = registeredExtensions[extid]
+			if desc == nil {
+				desc = &ExtensionDesc{Field: extid}
+			}
+		}
+
+		extensions = append(extensions, desc)
+	}
+	return extensions, nil
+}
+
 // SetExtension sets the specified extension of pb to the specified value.
 func SetExtension(pb Message, extension *ExtensionDesc, value interface{}) error {
 	epb, ok := extendable(pb)

proto/extensions_test.go

@@ -35,6 +35,7 @@ import (
 	"bytes"
 	"fmt"
 	"reflect"
+	"sort"
 	"testing"
 
 	"github.com/golang/protobuf/proto"
@@ -45,7 +46,7 @@ func TestGetExtensionsWithMissingExtensions(t *testing.T) {
 	msg := &pb.MyMessage{}
 	ext1 := &pb.Ext{}
 	if err := proto.SetExtension(msg, pb.E_Ext_More, ext1); err != nil {
-		t.Fatalf("Could not set ext1: %s", ext1)
+		t.Fatalf("Could not set ext1: %s", err)
 	}
 	exts, err := proto.GetExtensions(msg, []*proto.ExtensionDesc{
 		pb.E_Ext_More,
@@ -62,6 +63,54 @@ func TestGetExtensionsWithMissingExtensions(t *testing.T) {
 	}
 }
 
+func TestExtensionDescsWithMissingExtensions(t *testing.T) {
+	msg := &pb.MyMessage{Count: proto.Int32(0)}
+	extdesc1 := pb.E_Ext_More
+	ext1 := &pb.Ext{}
+	if err := proto.SetExtension(msg, extdesc1, ext1); err != nil {
+		t.Fatalf("Could not set ext1: %s", err)
+	}
+	extdesc2 := &proto.ExtensionDesc{
+		ExtendedType:  (*pb.MyMessage)(nil),
+		ExtensionType: (*bool)(nil),
+		Field:         123456789,
+		Name:          "a.b",
+		Tag:           "varint,123456789,opt",
+	}
+	ext2 := proto.Bool(false)
+	if err := proto.SetExtension(msg, extdesc2, ext2); err != nil {
+		t.Fatalf("Could not set ext2: %s", err)
+	}
+
+	b, err := proto.Marshal(msg)
+	if err != nil {
+		t.Fatalf("Could not marshal msg: %v", err)
+	}
+	if err := proto.Unmarshal(b, msg); err != nil {
+		t.Fatalf("Could not unmarshal into msg: %v", err)
+	}
+
+	descs, err := proto.ExtensionDescs(msg)
+	if err != nil {
+		t.Fatalf("proto.ExtensionDescs: got error %v", err)
+	}
+	sortExtDescs(descs)
+	wantDescs := []*proto.ExtensionDesc{extdesc1, &proto.ExtensionDesc{Field: extdesc2.Field}}
+	if !reflect.DeepEqual(descs, wantDescs) {
+		t.Errorf("proto.ExtensionDescs(msg) sorted extension ids: got %+v, want %+v", descs, wantDescs)
+	}
+}
+
+type ExtensionDescSlice []*proto.ExtensionDesc
+
+func (s ExtensionDescSlice) Len() int           { return len(s) }
+func (s ExtensionDescSlice) Less(i, j int) bool { return s[i].Field < s[j].Field }
+func (s ExtensionDescSlice) Swap(i, j int)      { s[i], s[j] = s[j], s[i] }
+
+func sortExtDescs(s []*proto.ExtensionDesc) {
+	sort.Sort(ExtensionDescSlice(s))
+}
+
 func TestGetExtensionStability(t *testing.T) {
 	check := func(m *pb.MyMessage) bool {
 		ext1, err := proto.GetExtension(m, pb.E_Ext_More)