Add editions helper functions for resolving features to protoutil (#283)

These helpers, in particular `protoutil.ResolveFeature` and
`protoutil.ResolveCustomFeature`, will be used from updated checks for
`buf breaking`, which will allow the tool to understand the
features-related semantics of the schema.

This way, it can correctly report issues with incompatible changes to
features in editions source files. And it can also allow changing a
file's syntax (like migrating from proto2 or proto3 to editions) as long
as there are no actual semantic changes to the schema.

Author: jhump

internal/editions/editions.go

@@ -23,9 +23,11 @@ import (
 	"sync"
 
 	"google.golang.org/protobuf/encoding/prototext"
+	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/reflect/protoreflect"
 	"google.golang.org/protobuf/reflect/protoregistry"
 	"google.golang.org/protobuf/types/descriptorpb"
+	"google.golang.org/protobuf/types/dynamicpb"
 )
 
 var (
@@ -77,7 +79,7 @@ var _ HasFeatures = (*descriptorpb.MethodOptions)(nil)
 // override. If there is no overridden value, it returns a zero value.
 func ResolveFeature(
 	element protoreflect.Descriptor,
-	field protoreflect.FieldDescriptor,
+	fields ...protoreflect.FieldDescriptor,
 ) (protoreflect.Value, error) {
 	for {
 		var features *descriptorpb.FeatureSet
@@ -86,9 +88,25 @@ func ResolveFeature(
 			features = withFeatures.GetFeatures()
 		}
 
-		msgRef := features.ProtoReflect()
-		if msgRef.Has(field) {
-			return msgRef.Get(field), nil
+		msgRef, err := adaptFeatureSet(features, fields[0])
+		if err != nil {
+			return protoreflect.Value{}, err
+		}
+		// Navigate the fields to find the value
+		var val protoreflect.Value
+		for i, field := range fields {
+			if i > 0 {
+				msgRef = val.Message()
+			}
+			if !msgRef.Has(field) {
+				val = protoreflect.Value{}
+				break
+			}
+			val = msgRef.Get(field)
+		}
+		if val.IsValid() {
+			// All fields were set!
+			return val, nil
 		}
 
 		parent := element.Parent()
@@ -230,3 +248,94 @@ func GetFeatureDefault(edition descriptorpb.Edition, container protoreflect.Mess
 	}
 	return empty.Get(feature), nil
 }
+
+func adaptFeatureSet(msg *descriptorpb.FeatureSet, field protoreflect.FieldDescriptor) (protoreflect.Message, error) {
+	msgRef := msg.ProtoReflect()
+	if field.IsExtension() {
+		// Extensions can always be used directly with the feature set, even if
+		// field.ContainingMessage() != FeatureSetDescriptor.
+		if msgRef.Has(field) || len(msgRef.GetUnknown()) == 0 {
+			return msgRef, nil
+		}
+		// The field is not present, but the message has unrecognized values. So
+		// let's try to parse the unrecognized bytes, just in case they contain
+		// this extension.
+		temp := &descriptorpb.FeatureSet{}
+		unmarshaler := prototext.UnmarshalOptions{
+			AllowPartial: true,
+			Resolver:     resolverForExtension{field},
+		}
+		if err := unmarshaler.Unmarshal(msgRef.GetUnknown(), temp); err != nil {
+			return nil, fmt.Errorf("failed to parse unrecognized fields of FeatureSet: %w", err)
+		}
+		return temp.ProtoReflect(), nil
+	}
+
+	if field.ContainingMessage() == FeatureSetDescriptor {
+		// Known field, not dynamically generated. Can directly use with the feature set.
+		return msgRef, nil
+	}
+
+	// If we get here, we have a dynamic field descriptor. We want to copy its
+	// value into a dynamic message, which requires marshalling/unmarshalling.
+	msgField := FeatureSetDescriptor.Fields().ByNumber(field.Number())
+	// We only need to copy over the unrecognized bytes (if any)
+	// and the same field (if present).
+	data := msgRef.GetUnknown()
+	if msgField != nil && msgRef.Has(msgField) {
+		subset := &descriptorpb.FeatureSet{}
+		subset.ProtoReflect().Set(msgField, msgRef.Get(msgField))
+		fieldBytes, err := proto.MarshalOptions{AllowPartial: true}.Marshal(subset)
+		if err != nil {
+			return nil, fmt.Errorf("failed to marshal FeatureSet field %s to bytes: %w", field.Name(), err)
+		}
+		data = append(data, fieldBytes...)
+	}
+	if len(data) == 0 {
+		// No relevant data to copy over, so we can just return
+		// a zero value message
+		return dynamicpb.NewMessageType(field.ContainingMessage()).Zero(), nil
+	}
+
+	other := dynamicpb.NewMessage(field.ContainingMessage())
+	// We don't need to use a resolver for this step because we know that
+	// field is not an extension. And features are not allowed to themselves
+	// have extensions.
+	if err := (proto.UnmarshalOptions{AllowPartial: true}).Unmarshal(data, other); err != nil {
+		return nil, fmt.Errorf("failed to marshal FeatureSet field %s to bytes: %w", field.Name(), err)
+	}
+	return other, nil
+}
+
+type resolverForExtension struct {
+	ext protoreflect.ExtensionDescriptor
+}
+
+func (r resolverForExtension) FindMessageByName(_ protoreflect.FullName) (protoreflect.MessageType, error) {
+	return nil, protoregistry.NotFound
+}
+
+func (r resolverForExtension) FindMessageByURL(_ string) (protoreflect.MessageType, error) {
+	return nil, protoregistry.NotFound
+}
+
+func (r resolverForExtension) FindExtensionByName(field protoreflect.FullName) (protoreflect.ExtensionType, error) {
+	if field == r.ext.FullName() {
+		return asExtensionType(r.ext), nil
+	}
+	return nil, protoregistry.NotFound
+}
+
+func (r resolverForExtension) FindExtensionByNumber(message protoreflect.FullName, field protoreflect.FieldNumber) (protoreflect.ExtensionType, error) {
+	if message == r.ext.ContainingMessage().FullName() && field == r.ext.Number() {
+		return asExtensionType(r.ext), nil
+	}
+	return nil, protoregistry.NotFound
+}
+
+func asExtensionType(ext protoreflect.ExtensionDescriptor) protoreflect.ExtensionType {
+	if xtd, ok := ext.(protoreflect.ExtensionTypeDescriptor); ok {
+		return xtd.Type()
+	}
+	return dynamicpb.NewExtensionType(ext)
+}

protoutil/editions.go

@@ -0,0 +1,140 @@
+// Copyright 2020-2024 Buf Technologies, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package protoutil
+
+import (
+	"fmt"
+
+	"google.golang.org/protobuf/reflect/protoreflect"
+	"google.golang.org/protobuf/types/descriptorpb"
+	"google.golang.org/protobuf/types/dynamicpb"
+
+	"github.com/bufbuild/protocompile/internal/editions"
+)
+
+// GetFeatureDefault gets the default value for the given feature and the given
+// edition. The given feature must represent a field of the google.protobuf.FeatureSet
+// message and must not be an extension.
+//
+// If the given field is from a dynamically built descriptor (i.e. it's containing
+// message descriptor is different from the linked-in descriptor for
+// [*descriptorpb.FeatureSet]), the returned value may be a dynamic value. In such
+// cases, the value may not be directly usable using [protoreflect.Message.Set] with
+// an instance of [*descriptorpb.FeatureSet] and must instead be used with a
+// [*dynamicpb.Message].
+//
+// To get the default value of a custom feature, use [GetCustomFeatureDefault]
+// instead.
+func GetFeatureDefault(edition descriptorpb.Edition, feature protoreflect.FieldDescriptor) (protoreflect.Value, error) {
+	if feature.ContainingMessage().FullName() != editions.FeatureSetDescriptor.FullName() {
+		return protoreflect.Value{}, fmt.Errorf("feature %s is a field of %s but should be a field of %s",
+			feature.Name(), feature.ContainingMessage().FullName(), editions.FeatureSetDescriptor.FullName())
+	}
+	var msgType protoreflect.MessageType
+	if feature.ContainingMessage() == editions.FeatureSetDescriptor {
+		msgType = editions.FeatureSetType
+	} else {
+		msgType = dynamicpb.NewMessageType(feature.ContainingMessage())
+	}
+	return editions.GetFeatureDefault(edition, msgType, feature)
+}
+
+// GetCustomFeatureDefault gets the default value for the given custom feature
+// and given edition. A custom feature is a field whose containing message is the
+// type of an extension field of google.protobuf.FeatureSet. The given extension
+// describes that extension field and message type. The given feature must be a
+// field of that extension's message type.
+func GetCustomFeatureDefault(edition descriptorpb.Edition, extension protoreflect.ExtensionType, feature protoreflect.FieldDescriptor) (protoreflect.Value, error) {
+	extDesc := extension.TypeDescriptor()
+	if extDesc.ContainingMessage().FullName() != editions.FeatureSetDescriptor.FullName() {
+		return protoreflect.Value{}, fmt.Errorf("extension %s does not extend %s", extDesc.FullName(), editions.FeatureSetDescriptor.FullName())
+	}
+	if extDesc.Message() == nil {
+		return protoreflect.Value{}, fmt.Errorf("extensions of %s should be messages; %s is instead %s",
+			editions.FeatureSetDescriptor.FullName(), extDesc.FullName(), extDesc.Kind().String())
+	}
+	if feature.IsExtension() {
+		return protoreflect.Value{}, fmt.Errorf("feature %s is an extension, but feature extension %s may not itself have extensions",
+			feature.FullName(), extDesc.FullName())
+	}
+	if feature.ContainingMessage().FullName() != extDesc.Message().FullName() {
+		return protoreflect.Value{}, fmt.Errorf("feature %s is a field of %s but should be a field of %s",
+			feature.Name(), feature.ContainingMessage().FullName(), extDesc.Message().FullName())
+	}
+	if feature.ContainingMessage() != extDesc.Message() {
+		return protoreflect.Value{}, fmt.Errorf("feature %s has a different message descriptor from the given extension type for %s",
+			feature.Name(), extDesc.Message().FullName())
+	}
+	return editions.GetFeatureDefault(edition, extension.Zero().Message().Type(), feature)
+}
+
+// ResolveFeature resolves a feature for the given descriptor.
+//
+// If the given element is in a proto2 or proto3 syntax file, this skips
+// resolution and just returns the relevant default (since such files are not
+// allowed to override features). If neither the given element nor any of its
+// ancestors override the given feature, the relevant default is returned.
+//
+// This has the same caveat as GetFeatureDefault if the given feature is from a
+// dynamically built descriptor.
+func ResolveFeature(element protoreflect.Descriptor, feature protoreflect.FieldDescriptor) (protoreflect.Value, error) {
+	edition := editions.GetEdition(element)
+	defaultVal, err := GetFeatureDefault(edition, feature)
+	if err != nil {
+		return protoreflect.Value{}, err
+	}
+	return resolveFeature(edition, defaultVal, element, feature)
+}
+
+// ResolveCustomFeature resolves a custom feature for the given extension and
+// field descriptor.
+//
+// The given extension must be an extension of google.protobuf.FeatureSet that
+// represents a non-repeated message value. The given feature is a field in
+// that extension's message type.
+//
+// If the given element is in a proto2 or proto3 syntax file, this skips
+// resolution and just returns the relevant default (since such files are not
+// allowed to override features). If neither the given element nor any of its
+// ancestors override the given feature, the relevant default is returned.
+func ResolveCustomFeature(element protoreflect.Descriptor, extension protoreflect.ExtensionType, feature protoreflect.FieldDescriptor) (protoreflect.Value, error) {
+	edition := editions.GetEdition(element)
+	defaultVal, err := GetCustomFeatureDefault(edition, extension, feature)
+	if err != nil {
+		return protoreflect.Value{}, err
+	}
+	return resolveFeature(edition, defaultVal, element, extension.TypeDescriptor(), feature)
+}
+
+func resolveFeature(
+	edition descriptorpb.Edition,
+	defaultVal protoreflect.Value,
+	element protoreflect.Descriptor,
+	fields ...protoreflect.FieldDescriptor,
+) (protoreflect.Value, error) {
+	if edition == descriptorpb.Edition_EDITION_PROTO2 || edition == descriptorpb.Edition_EDITION_PROTO3 {
+		// these syntax levels can't specify features, so we can short-circuit the search
+		// through the descriptor hierarchy for feature overrides
+		return defaultVal, nil
+	}
+	val, err := editions.ResolveFeature(element, fields...)
+	if err != nil {
+		return protoreflect.Value{}, err
+	}
+	if val.IsValid() {
+		return val, nil
+	}
+	return defaultVal, nil
+}