blob: c432817bb9c732f470aab59bc7292eb271c96af2 [file] [log] [blame]
// Copyright 2023 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package dynamicpb
import (
"fmt"
"strings"
"sync"
"sync/atomic"
"google.golang.org/protobuf/internal/errors"
"google.golang.org/protobuf/reflect/protoreflect"
"google.golang.org/protobuf/reflect/protoregistry"
)
type extField struct {
name protoreflect.FullName
number protoreflect.FieldNumber
}
// A Types is a collection of dynamically constructed descriptors.
// Its methods are safe for concurrent use.
//
// Types implements [protoregistry.MessageTypeResolver] and [protoregistry.ExtensionTypeResolver].
// A Types may be used as a [google.golang.org/protobuf/proto.UnmarshalOptions.Resolver].
type Types struct {
// atomicExtFiles is used with sync/atomic and hence must be the first word
// of the struct to guarantee 64-bit alignment.
//
// TODO(stapelberg): once we only support Go 1.19 and newer, switch this
// field to be of type atomic.Uint64 to guarantee alignment on
// stack-allocated values, too.
atomicExtFiles uint64
extMu sync.Mutex
files *protoregistry.Files
extensionsByMessage map[extField]protoreflect.ExtensionDescriptor
}
// NewTypes creates a new Types registry with the provided files.
// The Files registry is retained, and changes to Files will be reflected in Types.
// It is not safe to concurrently change the Files while calling Types methods.
func NewTypes(f *protoregistry.Files) *Types {
return &Types{
files: f,
}
}
// FindEnumByName looks up an enum by its full name;
// e.g., "google.protobuf.Field.Kind".
//
// This returns (nil, [protoregistry.NotFound]) if not found.
func (t *Types) FindEnumByName(name protoreflect.FullName) (protoreflect.EnumType, error) {
d, err := t.files.FindDescriptorByName(name)
if err != nil {
return nil, err
}
ed, ok := d.(protoreflect.EnumDescriptor)
if !ok {
return nil, errors.New("found wrong type: got %v, want enum", descName(d))
}
return NewEnumType(ed), nil
}
// FindExtensionByName looks up an extension field by the field's full name.
// Note that this is the full name of the field as determined by
// where the extension is declared and is unrelated to the full name of the
// message being extended.
//
// This returns (nil, [protoregistry.NotFound]) if not found.
func (t *Types) FindExtensionByName(name protoreflect.FullName) (protoreflect.ExtensionType, error) {
d, err := t.files.FindDescriptorByName(name)
if err != nil {
return nil, err
}
xd, ok := d.(protoreflect.ExtensionDescriptor)
if !ok {
return nil, errors.New("found wrong type: got %v, want extension", descName(d))
}
return NewExtensionType(xd), nil
}
// FindExtensionByNumber looks up an extension field by the field number
// within some parent message, identified by full name.
//
// This returns (nil, [protoregistry.NotFound]) if not found.
func (t *Types) FindExtensionByNumber(message protoreflect.FullName, field protoreflect.FieldNumber) (protoreflect.ExtensionType, error) {
// Construct the extension number map lazily, since not every user will need it.
// Update the map if new files are added to the registry.
if atomic.LoadUint64(&t.atomicExtFiles) != uint64(t.files.NumFiles()) {
t.updateExtensions()
}
xd := t.extensionsByMessage[extField{message, field}]
if xd == nil {
return nil, protoregistry.NotFound
}
return NewExtensionType(xd), nil
}
// FindMessageByName looks up a message by its full name;
// e.g. "google.protobuf.Any".
//
// This returns (nil, [protoregistry.NotFound]) if not found.
func (t *Types) FindMessageByName(name protoreflect.FullName) (protoreflect.MessageType, error) {
d, err := t.files.FindDescriptorByName(name)
if err != nil {
return nil, err
}
md, ok := d.(protoreflect.MessageDescriptor)
if !ok {
return nil, errors.New("found wrong type: got %v, want message", descName(d))
}
return NewMessageType(md), nil
}
// FindMessageByURL looks up a message by a URL identifier.
// See documentation on google.protobuf.Any.type_url for the URL format.
//
// This returns (nil, [protoregistry.NotFound]) if not found.
func (t *Types) FindMessageByURL(url string) (protoreflect.MessageType, error) {
// This function is similar to FindMessageByName but
// truncates anything before and including '/' in the URL.
message := protoreflect.FullName(url)
if i := strings.LastIndexByte(url, '/'); i >= 0 {
message = message[i+len("/"):]
}
return t.FindMessageByName(message)
}
func (t *Types) updateExtensions() {
t.extMu.Lock()
defer t.extMu.Unlock()
if atomic.LoadUint64(&t.atomicExtFiles) == uint64(t.files.NumFiles()) {
return
}
defer atomic.StoreUint64(&t.atomicExtFiles, uint64(t.files.NumFiles()))
t.files.RangeFiles(func(fd protoreflect.FileDescriptor) bool {
t.registerExtensions(fd.Extensions())
t.registerExtensionsInMessages(fd.Messages())
return true
})
}
func (t *Types) registerExtensionsInMessages(mds protoreflect.MessageDescriptors) {
count := mds.Len()
for i := 0; i < count; i++ {
md := mds.Get(i)
t.registerExtensions(md.Extensions())
t.registerExtensionsInMessages(md.Messages())
}
}
func (t *Types) registerExtensions(xds protoreflect.ExtensionDescriptors) {
count := xds.Len()
for i := 0; i < count; i++ {
xd := xds.Get(i)
field := xd.Number()
message := xd.ContainingMessage().FullName()
if t.extensionsByMessage == nil {
t.extensionsByMessage = make(map[extField]protoreflect.ExtensionDescriptor)
}
t.extensionsByMessage[extField{message, field}] = xd
}
}
func descName(d protoreflect.Descriptor) string {
switch d.(type) {
case protoreflect.EnumDescriptor:
return "enum"
case protoreflect.EnumValueDescriptor:
return "enum value"
case protoreflect.MessageDescriptor:
return "message"
case protoreflect.ExtensionDescriptor:
return "extension"
case protoreflect.ServiceDescriptor:
return "service"
default:
return fmt.Sprintf("%T", d)
}
}