jsonpb: Implementing marshaling of proto2 extensions.

Signed-off-by: David Symonds <dsymonds@golang.org>
diff --git a/jsonpb/jsonpb.go b/jsonpb/jsonpb.go
index c9afada..57b58fd 100644
--- a/jsonpb/jsonpb.go
+++ b/jsonpb/jsonpb.go
@@ -84,6 +84,13 @@
 	return buf.String(), nil
 }
 
+type int32Slice []int32
+
+// For sorting extensions ids to ensure stable output.
+func (s int32Slice) Len() int           { return len(s) }
+func (s int32Slice) Less(i, j int) bool { return s[i] < s[j] }
+func (s int32Slice) Swap(i, j int)      { s[i], s[j] = s[j], s[i] }
+
 // marshalObject writes a struct to the Writer.
 func (m *Marshaler) marshalObject(out *errWriter, v proto.Message, indent string) error {
 	out.write("{")
@@ -92,14 +99,13 @@
 	}
 
 	s := reflect.ValueOf(v).Elem()
-	writeBeforeField := ""
+	firstField := true
 	for i := 0; i < s.NumField(); i++ {
 		value := s.Field(i)
 		valueField := s.Type().Field(i)
 		if strings.HasPrefix(valueField.Name, "XXX_") {
 			continue
 		}
-		fieldName := jsonFieldName(valueField)
 
 		// TODO: proto3 objects should have default values omitted.
 
@@ -117,33 +123,50 @@
 			sv := value.Elem().Elem() // interface -> *T -> T
 			value = sv.Field(0)
 			valueField = sv.Type().Field(0)
-
-			var p proto.Properties
-			p.Parse(sv.Type().Field(0).Tag.Get("protobuf"))
-			fieldName = p.OrigName
 		}
-
-		out.write(writeBeforeField)
-		if m.Indent != "" {
-			out.write(indent)
-			out.write(m.Indent)
+		prop := jsonProperties(valueField)
+		if !firstField {
+			m.writeSep(out)
 		}
-		out.write(`"`)
-		out.write(fieldName)
-		out.write(`":`)
-		if m.Indent != "" {
-			out.write(" ")
-		}
-
-		if err := m.marshalValue(out, value, valueField, indent); err != nil {
+		if err := m.marshalField(out, prop, value, indent); err != nil {
 			return err
 		}
+		firstField = false
+	}
 
-		if m.Indent != "" {
-			writeBeforeField = ",\n"
-		} else {
-			writeBeforeField = ","
+	// Handle proto2 extensions.
+	if ep, ok := v.(extendableProto); ok {
+		extensions := proto.RegisteredExtensions(v)
+		extensionMap := ep.ExtensionMap()
+		// Sort extensions for stable output.
+		ids := make([]int32, 0, len(extensionMap))
+		for id := range extensionMap {
+			ids = append(ids, id)
 		}
+		sort.Sort(int32Slice(ids))
+		for _, id := range ids {
+			desc := extensions[id]
+			if desc == nil {
+				// unknown extension
+				continue
+			}
+			ext, extErr := proto.GetExtension(ep, desc)
+			if extErr != nil {
+				return extErr
+			}
+			value := reflect.ValueOf(ext)
+			var prop proto.Properties
+			prop.Parse(desc.Tag)
+			prop.OrigName = fmt.Sprintf("[%s]", desc.Name)
+			if !firstField {
+				m.writeSep(out)
+			}
+			if err := m.marshalField(out, &prop, value, indent); err != nil {
+				return err
+			}
+			firstField = false
+		}
+
 	}
 
 	if m.Indent != "" {
@@ -154,9 +177,34 @@
 	return out.err
 }
 
+func (m *Marshaler) writeSep(out *errWriter) {
+	if m.Indent != "" {
+		out.write(",\n")
+	} else {
+		out.write(",")
+	}
+}
+
+// marshalField writes field description and value to the Writer.
+func (m *Marshaler) marshalField(out *errWriter, prop *proto.Properties, v reflect.Value, indent string) error {
+	if m.Indent != "" {
+		out.write(indent)
+		out.write(m.Indent)
+	}
+	out.write(`"`)
+	out.write(prop.OrigName)
+	out.write(`":`)
+	if m.Indent != "" {
+		out.write(" ")
+	}
+	if err := m.marshalValue(out, prop, v, indent); err != nil {
+		return err
+	}
+	return nil
+}
+
 // marshalValue writes the value to the Writer.
-func (m *Marshaler) marshalValue(out *errWriter, v reflect.Value,
-	structField reflect.StructField, indent string) error {
+func (m *Marshaler) marshalValue(out *errWriter, prop *proto.Properties, v reflect.Value, indent string) error {
 
 	var err error
 	v = reflect.Indirect(v)
@@ -174,7 +222,7 @@
 				out.write(m.Indent)
 				out.write(m.Indent)
 			}
-			m.marshalValue(out, sliceVal, structField, indent+m.Indent)
+			m.marshalValue(out, prop, sliceVal, indent+m.Indent)
 			comma = ","
 		}
 		if m.Indent != "" {
@@ -187,8 +235,7 @@
 	}
 
 	// Handle enumerations.
-	protoInfo := structField.Tag.Get("protobuf")
-	if !m.EnumsAsInts && strings.Contains(protoInfo, ",enum=") {
+	if !m.EnumsAsInts && prop.Enum != "" {
 		// Unknown enum values will are stringified by the proto library as their
 		// value. Such values should _not_ be quoted or they will be interpreted
 		// as an enum string instead of their value.
@@ -253,7 +300,7 @@
 				out.write(` `)
 			}
 
-			if err := m.marshalValue(out, v.MapIndex(k), structField, indent+m.Indent); err != nil {
+			if err := m.marshalValue(out, prop, v.MapIndex(k), indent+m.Indent); err != nil {
 				return err
 			}
 		}
@@ -323,7 +370,7 @@
 			if strings.HasPrefix(ft.Name, "XXX_") {
 				continue
 			}
-			fieldName := jsonFieldName(ft)
+			fieldName := jsonProperties(ft).OrigName
 
 			valueForField, ok := jsonFields[fieldName]
 			if !ok {
@@ -438,11 +485,18 @@
 	return json.Unmarshal(inputValue, target.Addr().Interface())
 }
 
-// jsonFieldName returns the field name to use.
-func jsonFieldName(f reflect.StructField) string {
+// jsonProperties returns parsed proto.Properties for the field.
+func jsonProperties(f reflect.StructField) *proto.Properties {
 	var prop proto.Properties
 	prop.Init(f.Type, f.Name, f.Tag.Get("protobuf"), &f)
-	return prop.OrigName
+	return &prop
+}
+
+// extendableProto is an interface implemented by any protocol buffer that may be extended.
+type extendableProto interface {
+	proto.Message
+	ExtensionRangeArray() []proto.ExtensionRange
+	ExtensionMap() map[int32]proto.Extension
 }
 
 // Writer wrapper inspired by https://blog.golang.org/errors-are-values
diff --git a/jsonpb/jsonpb_test.go b/jsonpb/jsonpb_test.go
index 9e5b064..0a5b4ec 100644
--- a/jsonpb/jsonpb_test.go
+++ b/jsonpb/jsonpb_test.go
@@ -254,8 +254,25 @@
     }
   }
 }`
+	realNumber     = &pb.Real{Value: proto.Float64(3.14159265359)}
+	realNumberName = "Pi"
+	complexNumber  = &pb.Complex{Imaginary: proto.Float64(0.5772156649)}
+	realNumberJSON = `{` +
+		`"value":3.14159265359,` +
+		`"[jsonpb.Complex.real_extension]":{"imaginary":0.5772156649},` +
+		`"[jsonpb.name]":"Pi"` +
+		`}`
 )
 
+func init() {
+	if err := proto.SetExtension(realNumber, pb.E_Name, &realNumberName); err != nil {
+		panic(err)
+	}
+	if err := proto.SetExtension(realNumber, pb.E_Complex_RealExtension, complexNumber); err != nil {
+		panic(err)
+	}
+}
+
 var marshalingTests = []struct {
 	desc      string
 	marshaler Marshaler
@@ -294,6 +311,7 @@
 		`{"m_bool_simple":{"true":{"o_int32":1}}}`},
 	{"oneof, not set", marshaler, &pb.MsgWithOneof{}, `{}`},
 	{"oneof, set", marshaler, &pb.MsgWithOneof{Union: &pb.MsgWithOneof_Title{"Grand Poobah"}}, `{"title":"Grand Poobah"}`},
+	{"proto2 extension", marshaler, realNumber, realNumberJSON},
 }
 
 func TestMarshaling(t *testing.T) {
diff --git a/jsonpb/jsonpb_test_proto/more_test_objects.pb.go b/jsonpb/jsonpb_test_proto/more_test_objects.pb.go
index e9a0976..a61d539 100644
--- a/jsonpb/jsonpb_test_proto/more_test_objects.pb.go
+++ b/jsonpb/jsonpb_test_proto/more_test_objects.pb.go
@@ -17,6 +17,8 @@
 	Widget
 	Maps
 	MsgWithOneof
+	Real
+	Complex
 */
 package jsonpb
 
diff --git a/jsonpb/jsonpb_test_proto/test_objects.pb.go b/jsonpb/jsonpb_test_proto/test_objects.pb.go
index fd45129..8acd9a7 100644
--- a/jsonpb/jsonpb_test_proto/test_objects.pb.go
+++ b/jsonpb/jsonpb_test_proto/test_objects.pb.go
@@ -418,6 +418,86 @@
 	}
 }
 
+type Real struct {
+	Value            *float64                  `protobuf:"fixed64,1,opt,name=value" json:"value,omitempty"`
+	XXX_extensions   map[int32]proto.Extension `json:"-"`
+	XXX_unrecognized []byte                    `json:"-"`
+}
+
+func (m *Real) Reset()         { *m = Real{} }
+func (m *Real) String() string { return proto.CompactTextString(m) }
+func (*Real) ProtoMessage()    {}
+
+var extRange_Real = []proto.ExtensionRange{
+	{100, 536870911},
+}
+
+func (*Real) ExtensionRangeArray() []proto.ExtensionRange {
+	return extRange_Real
+}
+func (m *Real) ExtensionMap() map[int32]proto.Extension {
+	if m.XXX_extensions == nil {
+		m.XXX_extensions = make(map[int32]proto.Extension)
+	}
+	return m.XXX_extensions
+}
+
+func (m *Real) GetValue() float64 {
+	if m != nil && m.Value != nil {
+		return *m.Value
+	}
+	return 0
+}
+
+type Complex struct {
+	Imaginary        *float64                  `protobuf:"fixed64,1,opt,name=imaginary" json:"imaginary,omitempty"`
+	XXX_extensions   map[int32]proto.Extension `json:"-"`
+	XXX_unrecognized []byte                    `json:"-"`
+}
+
+func (m *Complex) Reset()         { *m = Complex{} }
+func (m *Complex) String() string { return proto.CompactTextString(m) }
+func (*Complex) ProtoMessage()    {}
+
+var extRange_Complex = []proto.ExtensionRange{
+	{100, 536870911},
+}
+
+func (*Complex) ExtensionRangeArray() []proto.ExtensionRange {
+	return extRange_Complex
+}
+func (m *Complex) ExtensionMap() map[int32]proto.Extension {
+	if m.XXX_extensions == nil {
+		m.XXX_extensions = make(map[int32]proto.Extension)
+	}
+	return m.XXX_extensions
+}
+
+func (m *Complex) GetImaginary() float64 {
+	if m != nil && m.Imaginary != nil {
+		return *m.Imaginary
+	}
+	return 0
+}
+
+var E_Complex_RealExtension = &proto.ExtensionDesc{
+	ExtendedType:  (*Real)(nil),
+	ExtensionType: (*Complex)(nil),
+	Field:         123,
+	Name:          "jsonpb.Complex.real_extension",
+	Tag:           "bytes,123,opt,name=real_extension",
+}
+
+var E_Name = &proto.ExtensionDesc{
+	ExtendedType:  (*Real)(nil),
+	ExtensionType: (*string)(nil),
+	Field:         124,
+	Name:          "jsonpb.name",
+	Tag:           "bytes,124,opt,name=name",
+}
+
 func init() {
 	proto.RegisterEnum("jsonpb.Widget_Color", Widget_Color_name, Widget_Color_value)
+	proto.RegisterExtension(E_Complex_RealExtension)
+	proto.RegisterExtension(E_Name)
 }
diff --git a/jsonpb/jsonpb_test_proto/test_objects.proto b/jsonpb/jsonpb_test_proto/test_objects.proto
index 85700bf..77f7fba 100644
--- a/jsonpb/jsonpb_test_proto/test_objects.proto
+++ b/jsonpb/jsonpb_test_proto/test_objects.proto
@@ -91,3 +91,20 @@
     int64 salary = 2;
   }
 }
+
+message Real {
+  optional double value = 1;
+  extensions 100 to max;
+}
+
+extend Real {
+  optional string name = 124;
+}
+
+message Complex {
+  extend Real {
+    optional Complex real_extension = 123;
+  }
+  optional double imaginary = 1;
+  extensions 100 to max;
+}