blob: d02fbe0c25c1f7066eb1a07a3619a6a641c78dab [file] [log] [blame]
// Copyright 2018 Google LLC
//
// 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 pb reflects over protocol buffer descriptors to generate objects
// that simplify type, enum, and field lookup.
package pb
import (
"google.golang.org/protobuf/proto"
"google.golang.org/protobuf/reflect/protoreflect"
"google.golang.org/protobuf/reflect/protoregistry"
anypb "google.golang.org/protobuf/types/known/anypb"
durpb "google.golang.org/protobuf/types/known/durationpb"
emptypb "google.golang.org/protobuf/types/known/emptypb"
structpb "google.golang.org/protobuf/types/known/structpb"
tspb "google.golang.org/protobuf/types/known/timestamppb"
wrapperspb "google.golang.org/protobuf/types/known/wrapperspb"
)
// Db maps from file / message / enum name to file description.
//
// Each Db is isolated from each other, and while information about protobuf descriptors may be
// fetched from the global protobuf registry, no descriptors are added to this registry, else
// the isolation guarantees of the Db object would be violated.
type Db struct {
revFileDescriptorMap map[string]*FileDescription
// files contains the deduped set of FileDescriptions whose types are contained in the pb.Db.
files []*FileDescription
}
var (
// DefaultDb used at evaluation time or unless overridden at check time.
DefaultDb = &Db{
revFileDescriptorMap: make(map[string]*FileDescription),
files: []*FileDescription{},
}
)
// NewDb creates a new `pb.Db` with an empty type name to file description map.
func NewDb() *Db {
pbdb := &Db{
revFileDescriptorMap: make(map[string]*FileDescription),
files: []*FileDescription{},
}
// The FileDescription objects in the default db contain lazily initialized TypeDescription
// values which may point to the state contained in the DefaultDb irrespective of this shallow
// copy; however, the type graph for a field is idempotently computed, and is guaranteed to
// only be initialized once thanks to atomic values within the TypeDescription objects, so it
// is safe to share these values across instances.
for k, v := range DefaultDb.revFileDescriptorMap {
pbdb.revFileDescriptorMap[k] = v
}
pbdb.files = append(pbdb.files, DefaultDb.files...)
return pbdb
}
// Copy creates a copy of the current database with its own internal descriptor mapping.
func (pbdb *Db) Copy() *Db {
copy := NewDb()
for k, v := range pbdb.revFileDescriptorMap {
copy.revFileDescriptorMap[k] = v
}
for _, f := range pbdb.files {
hasFile := false
for _, f2 := range copy.files {
if f2 == f {
hasFile = true
}
}
if !hasFile {
copy.files = append(copy.files, f)
}
}
return copy
}
// FileDescriptions returns the set of file descriptions associated with this db.
func (pbdb *Db) FileDescriptions() []*FileDescription {
return pbdb.files
}
// RegisterDescriptor produces a `FileDescription` from a `FileDescriptor` and registers the
// message and enum types into the `pb.Db`.
func (pbdb *Db) RegisterDescriptor(fileDesc protoreflect.FileDescriptor) (*FileDescription, error) {
fd, found := pbdb.revFileDescriptorMap[fileDesc.Path()]
if found {
return fd, nil
}
// Make sure to search the global registry to see if a protoreflect.FileDescriptor for
// the file specified has been linked into the binary. If so, use the copy of the descriptor
// from the global cache.
//
// Note: Proto reflection relies on descriptor values being object equal rather than object
// equivalence. This choice means that a FieldDescriptor generated from a FileDescriptorProto
// will be incompatible with the FieldDescriptor in the global registry and any message created
// from that global registry.
globalFD, err := protoregistry.GlobalFiles.FindFileByPath(fileDesc.Path())
if err == nil {
fileDesc = globalFD
}
fd = NewFileDescription(fileDesc, pbdb)
for _, enumValName := range fd.GetEnumNames() {
pbdb.revFileDescriptorMap[enumValName] = fd
}
for _, msgTypeName := range fd.GetTypeNames() {
pbdb.revFileDescriptorMap[msgTypeName] = fd
}
pbdb.revFileDescriptorMap[fileDesc.Path()] = fd
// Return the specific file descriptor registered.
pbdb.files = append(pbdb.files, fd)
return fd, nil
}
// RegisterMessage produces a `FileDescription` from a `message` and registers the message and all
// other definitions within the message file into the `pb.Db`.
func (pbdb *Db) RegisterMessage(message proto.Message) (*FileDescription, error) {
msgDesc := message.ProtoReflect().Descriptor()
msgName := msgDesc.FullName()
typeName := sanitizeProtoName(string(msgName))
if fd, found := pbdb.revFileDescriptorMap[typeName]; found {
return fd, nil
}
return pbdb.RegisterDescriptor(msgDesc.ParentFile())
}
// DescribeEnum takes a qualified enum name and returns an `EnumDescription` if it exists in the
// `pb.Db`.
func (pbdb *Db) DescribeEnum(enumName string) (*EnumValueDescription, bool) {
enumName = sanitizeProtoName(enumName)
if fd, found := pbdb.revFileDescriptorMap[enumName]; found {
return fd.GetEnumDescription(enumName)
}
return nil, false
}
// DescribeType returns a `TypeDescription` for the `typeName` if it exists in the `pb.Db`.
func (pbdb *Db) DescribeType(typeName string) (*TypeDescription, bool) {
typeName = sanitizeProtoName(typeName)
if fd, found := pbdb.revFileDescriptorMap[typeName]; found {
return fd.GetTypeDescription(typeName)
}
return nil, false
}
// CollectFileDescriptorSet builds a file descriptor set associated with the file where the input
// message is declared.
func CollectFileDescriptorSet(message proto.Message) map[string]protoreflect.FileDescriptor {
fdMap := map[string]protoreflect.FileDescriptor{}
parentFile := message.ProtoReflect().Descriptor().ParentFile()
fdMap[parentFile.Path()] = parentFile
// Initialize list of dependencies
deps := make([]protoreflect.FileImport, parentFile.Imports().Len())
for i := 0; i < parentFile.Imports().Len(); i++ {
deps[i] = parentFile.Imports().Get(i)
}
// Expand list for new dependencies
for i := 0; i < len(deps); i++ {
dep := deps[i]
if _, found := fdMap[dep.Path()]; found {
continue
}
fdMap[dep.Path()] = dep.FileDescriptor
for j := 0; j < dep.FileDescriptor.Imports().Len(); j++ {
deps = append(deps, dep.FileDescriptor.Imports().Get(j))
}
}
return fdMap
}
func init() {
// Describe well-known types to ensure they can always be resolved by the check and interpret
// execution phases.
//
// The following subset of message types is enough to ensure that all well-known types can
// resolved in the runtime, since describing the value results in describing the whole file
// where the message is declared.
DefaultDb.RegisterMessage(&anypb.Any{})
DefaultDb.RegisterMessage(&durpb.Duration{})
DefaultDb.RegisterMessage(&emptypb.Empty{})
DefaultDb.RegisterMessage(&tspb.Timestamp{})
DefaultDb.RegisterMessage(&structpb.Value{})
DefaultDb.RegisterMessage(&wrapperspb.BoolValue{})
}