Add support for extensions

decode.go

@@ -57,6 +57,11 @@ type UnmarshalOptions struct {
 	}
 }
 
+type protoResolver interface {
+	protoregistry.MessageTypeResolver
+	protoregistry.ExtensionTypeResolver
+}
+
 // Unmarshal a Protobuf message from the given YAML data.
 func Unmarshal(data []byte, message proto.Message) error {
 	return (UnmarshalOptions{}).Unmarshal(data, message)
@@ -100,7 +105,7 @@ func (o UnmarshalOptions) unmarshalNode(node *yaml.Node, message proto.Message,
 		case err == nil: // Valid.
 		case errors.As(err, &verr):
 			for _, violation := range verr.Violations {
-				closest := nodeClosestToPath(node, message.ProtoReflect().Descriptor(), violation.GetFieldPath(), violation.GetForKey())
+				closest := unm.nodeClosestToPath(node, message.ProtoReflect().Descriptor(), violation.GetFieldPath(), violation.GetForKey())
 				unm.addError(closest, &violationError{
 					Violation: violation,
 				})
@@ -165,13 +170,7 @@ func (u *unmarshaler) findAnyTypeURL(node *yaml.Node) string {
 
 func (u *unmarshaler) resolveAnyType(typeURL string) (protoreflect.MessageType, error) {
 	// Get the message type.
-	var msgType protoreflect.MessageType
-	var err error
-	if u.options.Resolver != nil {
-		msgType, err = u.options.Resolver.FindMessageByURL(typeURL)
-	} else { // Use the global registry.
-		msgType, err = protoregistry.GlobalTypes.FindMessageByURL(typeURL)
-	}
+	msgType, err := u.getResolver().FindMessageByURL(typeURL)
 	if err != nil {
 		return nil, err
 	}
@@ -446,22 +445,42 @@ func parseIntLiteral(value string) (intLit, error) {
 	return lit, err
 }
 
-// Searches for the field with the given 'key' first by Name, then by JSONName,
+func (u *unmarshaler) getResolver() protoResolver {
+	if u.options.Resolver != nil {
+		return u.options.Resolver
+	}
+	return protoregistry.GlobalTypes
+}
+
+// findField searches for the field with the given 'key' by extension type, JSONName, TextName,
 // and finally by Number.
-func findField(key string, fields protoreflect.FieldDescriptors) protoreflect.FieldDescriptor {
-	if field := fields.ByName(protoreflect.Name(key)); field != nil {
-		return field
+func (u *unmarshaler) findField(key string, msgDesc protoreflect.MessageDescriptor) (protoreflect.FieldDescriptor, error) {
+	fields := msgDesc.Fields()
+	if strings.HasPrefix(key, "[") && strings.HasSuffix(key, "]") {
+		extName := protoreflect.FullName(key[1 : len(key)-1])
+		extType, err := u.getResolver().FindExtensionByName(extName)
+		if err != nil {
+			return nil, err
+		}
+		result := extType.TypeDescriptor()
+		if !msgDesc.ExtensionRanges().Has(result.Number()) || result.ContainingMessage().FullName() != msgDesc.FullName() {
+			return nil, fmt.Errorf("message %v cannot be extended by %v", msgDesc.FullName(), result.FullName())
+		}
+		return result, nil
 	}
 	if field := fields.ByJSONName(key); field != nil {
-		return field
+		return field, nil
+	}
+	if field := fields.ByTextName(key); field != nil {
+		return field, nil
 	}
 	num, err := strconv.ParseInt(key, 10, 32)
 	if err == nil {
 		if field := fields.ByNumber(protoreflect.FieldNumber(num)); field != nil {
-			return field
+			return field, nil
 		}
 	}
-	return nil
+	return nil, protoregistry.NotFound
 }
 
 // Unmarshal a field, handling isList/isMap.
@@ -552,7 +571,6 @@ func (u *unmarshaler) findNodeForCustom(node *yaml.Node, forAny bool) *yaml.Node
 // Unmarshal the given yaml node into the given proto.Message.
 func (u *unmarshaler) unmarshalMessage(node *yaml.Node, message proto.Message, forAny bool) {
 	// Check for a custom unmarshaler
-
 	if custom, ok := u.custom[message.ProtoReflect().Descriptor().FullName()]; ok {
 		valueNode := u.findNodeForCustom(node, forAny)
 		if valueNode == nil {
@@ -570,20 +588,36 @@ func (u *unmarshaler) unmarshalMessage(node *yaml.Node, message proto.Message, f
 		return
 	}
 	// Decode the fields
-	fields := message.ProtoReflect().Descriptor().Fields()
+	msgDesc := message.ProtoReflect().Descriptor()
 	for i := 0; i < len(node.Content); i += 2 {
 		keyNode := node.Content[i]
-		if u.checkKind(keyNode, yaml.ScalarNode) {
-			if forAny && keyNode.Value == atTypeFieldName {
-				continue // Skip the @type field for Any messages
-			}
-			// Get the field Name, JSONName, or Number
-			if field := findField(keyNode.Value, fields); field != nil {
-				valueNode := node.Content[i+1]
-				u.unmarshalField(valueNode, field, message)
-			} else {
-				u.addErrorf(keyNode, "unknown field %#v, expended one of %v", keyNode.Value, getFieldNames(fields))
+		var key string
+		switch keyNode.Kind {
+		case yaml.ScalarNode:
+			key = keyNode.Value
+		case yaml.SequenceNode:
+			if len(keyNode.Content) == 1 && keyNode.Content[0].Kind == yaml.ScalarNode {
+				key = "[" + keyNode.Content[0].Value + "]"
+				break
 			}
+			fallthrough
+		default:
+			u.checkKind(keyNode, yaml.ScalarNode) // Always an error.
+			continue
+		}
+
+		if forAny && key == atTypeFieldName {
+			continue // Skip the @type field for Any messages
+		}
+		field, err := u.findField(key, msgDesc)
+		switch {
+		case errors.Is(err, protoregistry.NotFound):
+			u.addErrorf(keyNode, "unknown field %#v, expected one of %v", key, getFieldNames(msgDesc.Fields()))
+		case err != nil:
+			u.addError(keyNode, err)
+		default:
+			valueNode := node.Content[i+1]
+			u.unmarshalField(valueNode, field, message)
 		}
 	}
 }
@@ -1007,12 +1041,12 @@ func (u *unmarshaler) unmarshalScalarFloat(node *yaml.Node, value *structpb.Valu
 //   - 'foo[0]' -> the first element of the repeated field foo or the map entry with key '0'
 //   - 'foo.bar' -> the field bar in the message field foo
 //   - 'foo["bar"]' -> the map entry with key 'bar' in the map field foo
-func nodeClosestToPath(root *yaml.Node, msgDesc protoreflect.MessageDescriptor, path string, toKey bool) *yaml.Node {
+func (u *unmarshaler) nodeClosestToPath(root *yaml.Node, msgDesc protoreflect.MessageDescriptor, path string, toKey bool) *yaml.Node {
 	parsedPath, err := parseFieldPath(path)
 	if err != nil {
 		return root
 	}
-	return findNodeByPath(root, msgDesc, parsedPath, toKey)
+	return u.findNodeByPath(root, msgDesc, parsedPath, toKey)
 }
 
 func parseFieldPath(path string) ([]string, error) {
@@ -1076,16 +1110,16 @@ func parseNextValue(path string) (string, string) {
 }
 
 // Returns the node as close to the given path as possible.
-func findNodeByPath(root *yaml.Node, msgDesc protoreflect.MessageDescriptor, path []string, toKey bool) *yaml.Node {
+func (u *unmarshaler) findNodeByPath(root *yaml.Node, msgDesc protoreflect.MessageDescriptor, path []string, toKey bool) *yaml.Node {
 	cur := root
 	curMsg := msgDesc
 	var curMap protoreflect.FieldDescriptor
 	for i, key := range path {
 		switch cur.Kind {
 		case yaml.MappingNode:
 			if curMsg != nil {
-				field := findField(key, curMsg.Fields())
-				if field == nil {
+				field, err := u.findField(key, curMsg)
+				if err != nil {
 					return cur
 				}
 				var found bool

decode_test.go

@@ -24,6 +24,7 @@ import (
 	"github.com/bufbuild/protovalidate-go"
 	testv1 "github.com/bufbuild/protoyaml-go/internal/gen/proto/buf/protoyaml/test/v1"
 	"github.com/google/go-cmp/cmp"
+	"github.com/stretchr/testify/require"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/testing/protocmp"
 	"google.golang.org/protobuf/types/known/durationpb"
@@ -37,7 +38,10 @@ func TestGoldenFiles(t *testing.T) {
 			return err
 		}
 		if !info.IsDir() && strings.HasSuffix(path, ".yaml") {
-			testRunYAMLFile(t, path)
+			t.Run(path, func(t *testing.T) {
+				t.Parallel()
+				testRunYAMLFile(t, path)
+			})
 		}
 		return nil
 	}); err != nil {
@@ -97,6 +101,15 @@ func TestParseDuration(t *testing.T) {
 	}
 }
 
+func TestExtension(t *testing.T) {
+	t.Parallel()
+
+	actual := &testv1.Proto2Test{}
+	err := Unmarshal([]byte(`[buf.protoyaml.test.v1.p2t_string_ext]: hi`), actual)
+	require.NoError(t, err)
+	require.Equal(t, "hi", proto.GetExtension(actual, testv1.E_P2TStringExt))
+}
+
 func testRunYAML(path string, msg proto.Message) error {
 	// Read the test file
 	file, err := os.Open(path)
@@ -125,6 +138,8 @@ func testRunYAMLFile(t *testing.T, testFile string) {
 
 	var err error
 	switch {
+	case strings.HasSuffix(testFile, ".proto2test.yaml"):
+		err = testRunYAML(testFile, &testv1.Proto2Test{})
 	case strings.HasSuffix(testFile, ".proto3test.yaml"):
 		err = testRunYAML(testFile, &testv1.Proto3Test{})
 	case strings.HasSuffix(testFile, ".const.yaml"):

go.mod

@@ -3,22 +3,22 @@ module github.com/bufbuild/protoyaml-go
 go 1.20
 
 require (
-	buf.build/gen/go/bufbuild/protovalidate/protocolbuffers/go v1.33.0-20240221180331-f05a6f4403ce.1
-	github.com/bufbuild/protovalidate-go v0.6.0
-	github.com/google/go-cmp v0.5.9
-	github.com/stretchr/testify v1.8.4
+	buf.build/gen/go/bufbuild/protovalidate/protocolbuffers/go v1.33.0-20240401165935-b983156c5e99.1
+	github.com/bufbuild/protovalidate-go v0.6.1
+	github.com/google/go-cmp v0.6.0
+	github.com/stretchr/testify v1.9.0
 	google.golang.org/protobuf v1.33.0
 	gopkg.in/yaml.v3 v3.0.1
 )
 
 require (
 	github.com/antlr4-go/antlr/v4 v4.13.0 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
-	github.com/google/cel-go v0.20.0 // indirect
+	github.com/google/cel-go v0.20.1 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/stoewer/go-strcase v1.3.0 // indirect
-	golang.org/x/exp v0.0.0-20230905200255-921286631fa9 // indirect
+	golang.org/x/exp v0.0.0-20240325151524-a685a6edb6d8 // indirect
 	golang.org/x/text v0.14.0 // indirect
-	google.golang.org/genproto/googleapis/api v0.0.0-20230920204549-e6e6cdab5c13 // indirect
-	google.golang.org/genproto/googleapis/rpc v0.0.0-20230920204549-e6e6cdab5c13 // indirect
+	google.golang.org/genproto/googleapis/api v0.0.0-20240401170217-c3f982113cda // indirect
+	google.golang.org/genproto/googleapis/rpc v0.0.0-20240401170217-c3f982113cda // indirect
 )

go.sum

@@ -1,19 +1,19 @@
-buf.build/gen/go/bufbuild/protovalidate/protocolbuffers/go v1.33.0-20240221180331-f05a6f4403ce.1 h1:0nWhrRcnkgw1kwJ7xibIO8bqfOA7pBzBjGCDBxIHch8=
-buf.build/gen/go/bufbuild/protovalidate/protocolbuffers/go v1.33.0-20240221180331-f05a6f4403ce.1/go.mod h1:Tgn5bgL220vkFOI0KPStlcClPeOJzAv4uT+V8JXGUnw=
+buf.build/gen/go/bufbuild/protovalidate/protocolbuffers/go v1.33.0-20240401165935-b983156c5e99.1 h1:2IGhRovxlsOIQgx2ekZWo4wTPAYpck41+18ICxs37is=
+buf.build/gen/go/bufbuild/protovalidate/protocolbuffers/go v1.33.0-20240401165935-b983156c5e99.1/go.mod h1:Tgn5bgL220vkFOI0KPStlcClPeOJzAv4uT+V8JXGUnw=
 github.com/antlr4-go/antlr/v4 v4.13.0 h1:lxCg3LAv+EUK6t1i0y1V6/SLeUi0eKEKdhQAlS8TVTI=
 github.com/antlr4-go/antlr/v4 v4.13.0/go.mod h1:pfChB/xh/Unjila75QW7+VU4TSnWnnk9UTnmpPaOR2g=
-github.com/bufbuild/protovalidate-go v0.6.0 h1:Jgs1kFuZ2LHvvdj8SpCLA1W/+pXS8QSM3F/E2l3InPY=
-github.com/bufbuild/protovalidate-go v0.6.0/go.mod h1:1LamgoYHZ2NdIQH0XGczGTc6Z8YrTHjcJVmiBaar4t4=
+github.com/bufbuild/protovalidate-go v0.6.1 h1:uzW8r0CDvqApUChNj87VzZVoQSKhcVdw5UWOE605UIw=
+github.com/bufbuild/protovalidate-go v0.6.1/go.mod h1:4BR3rKEJiUiTy+sqsusFn2ladOf0kYmA2Reo6BHSBgQ=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/envoyproxy/protoc-gen-validate v1.0.4 h1:gVPz/FMfvh57HdSJQyvBtF00j8JU4zdyUgIUNhlgg0A=
 github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk=
-github.com/google/cel-go v0.20.0 h1:h4n6DOCppEMpWERzllyNkntl7JrDyxoE543KWS6BLpc=
-github.com/google/cel-go v0.20.0/go.mod h1:kWcIzTsPX0zmQ+H3TirHstLLf9ep5QTsZBN9u4dOYLg=
+github.com/google/cel-go v0.20.1 h1:nDx9r8S3L4pE61eDdt8igGj8rf5kjYR3ILxWIpWNi84=
+github.com/google/cel-go v0.20.1/go.mod h1:kWcIzTsPX0zmQ+H3TirHstLLf9ep5QTsZBN9u4dOYLg=
 github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
-github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38=
-github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
+github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
+github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
 github.com/kr/pretty v0.1.0 h1:L/CwN0zerZDmRFUapSPitk6f+Q3+0za1rQkzVuMiMFI=
 github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
@@ -26,17 +26,17 @@ github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpE
 github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
 github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
-github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=
-github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
-golang.org/x/exp v0.0.0-20230905200255-921286631fa9 h1:GoHiUyI/Tp2nVkLI2mCxVkOjsbSXD66ic0XW0js0R9g=
-golang.org/x/exp v0.0.0-20230905200255-921286631fa9/go.mod h1:S2oDrQGGwySpoQPVqRShND87VCbxmc6bL1Yd2oYrm6k=
+github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
+github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
+golang.org/x/exp v0.0.0-20240325151524-a685a6edb6d8 h1:aAcj0Da7eBAtrTp03QXWvm88pSyOt+UgdZw2BFZ+lEw=
+golang.org/x/exp v0.0.0-20240325151524-a685a6edb6d8/go.mod h1:CQ1k9gNrJ50XIzaKCRR2hssIjF07kZFEiieALBM/ARQ=
 golang.org/x/text v0.14.0 h1:ScX5w1eTa3QqT8oi6+ziP7dTV1S2+ALU0bI+0zXKWiQ=
 golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
-google.golang.org/genproto/googleapis/api v0.0.0-20230920204549-e6e6cdab5c13 h1:U7+wNaVuSTaUqNvK2+osJ9ejEZxbjHHk8F2b6Hpx0AE=
-google.golang.org/genproto/googleapis/api v0.0.0-20230920204549-e6e6cdab5c13/go.mod h1:RdyHbowztCGQySiCvQPgWQWgWhGnouTdCflKoDBt32U=
-google.golang.org/genproto/googleapis/rpc v0.0.0-20230920204549-e6e6cdab5c13 h1:N3bU/SQDCDyD6R528GJ/PwW9KjYcJA3dgyH+MovAkIM=
-google.golang.org/genproto/googleapis/rpc v0.0.0-20230920204549-e6e6cdab5c13/go.mod h1:KSqppvjFjtoCI+KGd4PELB0qLNxdJHRGqRI09mB6pQA=
+google.golang.org/genproto/googleapis/api v0.0.0-20240401170217-c3f982113cda h1:b6F6WIV4xHHD0FA4oIyzU6mHWg2WI2X1RBehwa5QN38=
+google.golang.org/genproto/googleapis/api v0.0.0-20240401170217-c3f982113cda/go.mod h1:AHcE/gZH76Bk/ROZhQphlRoWo5xKDEtz3eVEO1LfA8c=
+google.golang.org/genproto/googleapis/rpc v0.0.0-20240401170217-c3f982113cda h1:LI5DOvAxUPMv/50agcLLoo+AdWc1irS9Rzz4vPuD1V4=
+google.golang.org/genproto/googleapis/rpc v0.0.0-20240401170217-c3f982113cda/go.mod h1:WtryC6hu0hhx87FDGxWCDptyssuo68sk10vYjF+T9fY=
 google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw=
 google.golang.org/protobuf v1.33.0 h1:uNO2rsAINq/JlFpSdYEKIZ0uKD/R9cpdv0T+yoGwGmI=
 google.golang.org/protobuf v1.33.0/go.mod h1:c6P6GXX6sHbq/GpV6MGZEdwhWPcYBgnhAHhKbcUYpos=

internal/cmd/generate-txt-testdata/main.go

@@ -97,6 +97,9 @@ func tryParse(filePath string) (string, error) {
 	}
 
 	switch {
+	case strings.HasSuffix(filePath, ".proto2test.yaml"):
+		testCase := &testv1.Proto2Test{}
+		err = options.Unmarshal(data, testCase)
 	case strings.HasSuffix(filePath, ".proto3test.yaml"):
 		testCase := &testv1.Proto3Test{}
 		err = options.Unmarshal(data, testCase)
@@ -109,5 +112,8 @@ func tryParse(filePath string) (string, error) {
 	default:
 		return "", fmt.Errorf("unknown file type: %s", filePath)
 	}
-	return err.Error(), nil
+	if err != nil {
+		return err.Error(), nil //nolint:nilerr
+	}
+	return "", nil
 }