forked from jhump/protoreflect
/
message_factory.go
201 lines (178 loc) · 6.97 KB
/
message_factory.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
package dynamic
import (
"reflect"
"sync"
"github.com/golang/protobuf/proto"
"github.com/phpstudyer/protoreflect/desc"
)
// MessageFactory can be used to create new empty message objects. A default instance
// (without extension registry or known-type registry specified) will always return
// dynamic messages (e.g. type will be *dynamic.Message) except for "well-known" types.
// The well-known types include primitive wrapper types and a handful of other special
// types defined in standard protobuf definitions, like Any, Duration, and Timestamp.
type MessageFactory struct {
er *ExtensionRegistry
ktr *KnownTypeRegistry
}
// NewMessageFactoryWithExtensionRegistry creates a new message factory where any
// dynamic messages produced will use the given extension registry to recognize and
// parse extension fields.
func NewMessageFactoryWithExtensionRegistry(er *ExtensionRegistry) *MessageFactory {
return NewMessageFactoryWithRegistries(er, nil)
}
// NewMessageFactoryWithKnownTypeRegistry creates a new message factory where the
// known types, per the given registry, will be returned as normal protobuf messages
// (e.g. generated structs, instead of dynamic messages).
func NewMessageFactoryWithKnownTypeRegistry(ktr *KnownTypeRegistry) *MessageFactory {
return NewMessageFactoryWithRegistries(nil, ktr)
}
// NewMessageFactoryWithDefaults creates a new message factory where all "default" types
// (those for which protoc-generated code is statically linked into the Go program) are
// known types. If any dynamic messages are produced, they will recognize and parse all
// "default" extension fields. This is the equivalent of:
// NewMessageFactoryWithRegistries(
// NewExtensionRegistryWithDefaults(),
// NewKnownTypeRegistryWithDefaults())
func NewMessageFactoryWithDefaults() *MessageFactory {
return NewMessageFactoryWithRegistries(NewExtensionRegistryWithDefaults(), NewKnownTypeRegistryWithDefaults())
}
// NewMessageFactoryWithRegistries creates a new message factory with the given extension
// and known type registries.
func NewMessageFactoryWithRegistries(er *ExtensionRegistry, ktr *KnownTypeRegistry) *MessageFactory {
return &MessageFactory{
er: er,
ktr: ktr,
}
}
// NewMessage creates a new empty message that corresponds to the given descriptor.
// If the given descriptor describes a "known type" then that type is instantiated.
// Otherwise, an empty dynamic message is returned.
func (f *MessageFactory) NewMessage(md *desc.MessageDescriptor) proto.Message {
var ktr *KnownTypeRegistry
if f != nil {
ktr = f.ktr
}
if m := ktr.CreateIfKnown(md.GetFullyQualifiedName()); m != nil {
return m
}
return NewMessageWithMessageFactory(md, f)
}
// NewDynamicMessage creates a new empty dynamic message that corresponds to the given
// descriptor. This is like f.NewMessage(md) except the known type registry is not
// consulted so the return value is always a dynamic message.
//
// This is also like dynamic.NewMessage(md) except that the returned message will use
// this factory when creating other messages, like during de-serialization of fields
// that are themselves message types.
func (f *MessageFactory) NewDynamicMessage(md *desc.MessageDescriptor) *Message {
return NewMessageWithMessageFactory(md, f)
}
// GetKnownTypeRegistry returns the known type registry that this factory uses to
// instantiate known (e.g. generated) message types.
func (f *MessageFactory) GetKnownTypeRegistry() *KnownTypeRegistry {
if f == nil {
return nil
}
return f.ktr
}
// GetExtensionRegistry returns the extension registry that this factory uses to
// create dynamic messages. The registry is used by dynamic messages to recognize
// and parse extension fields during de-serialization.
func (f *MessageFactory) GetExtensionRegistry() *ExtensionRegistry {
if f == nil {
return nil
}
return f.er
}
type wkt interface {
XXX_WellKnownType() string
}
var typeOfWkt = reflect.TypeOf((*wkt)(nil)).Elem()
// KnownTypeRegistry is a registry of known message types, as identified by their
// fully-qualified name. A known message type is one for which a protoc-generated
// struct exists, so a dynamic message is not necessary to represent it. A
// MessageFactory uses a KnownTypeRegistry to decide whether to create a generated
// struct or a dynamic message. The zero-value registry (including the behavior of
// a nil pointer) only knows about the "well-known types" in protobuf. These
// include only the wrapper types and a handful of other special types like Any,
// Duration, and Timestamp.
type KnownTypeRegistry struct {
excludeWkt bool
includeDefault bool
mu sync.RWMutex
types map[string]reflect.Type
}
// NewKnownTypeRegistryWithDefaults creates a new registry that knows about all
// "default" types (those for which protoc-generated code is statically linked
// into the Go program).
func NewKnownTypeRegistryWithDefaults() *KnownTypeRegistry {
return &KnownTypeRegistry{includeDefault: true}
}
// NewKnownTypeRegistryWithoutWellKnownTypes creates a new registry that does *not*
// include the "well-known types" in protobuf. So even well-known types would be
// represented by a dynamic message.
func NewKnownTypeRegistryWithoutWellKnownTypes() *KnownTypeRegistry {
return &KnownTypeRegistry{excludeWkt: true}
}
// AddKnownType adds the types of the given messages as known types.
func (r *KnownTypeRegistry) AddKnownType(kts ...proto.Message) {
r.mu.Lock()
defer r.mu.Unlock()
if r.types == nil {
r.types = map[string]reflect.Type{}
}
for _, kt := range kts {
r.types[proto.MessageName(kt)] = reflect.TypeOf(kt)
}
}
// CreateIfKnown will construct an instance of the given message if it is a known type.
// If the given name is unknown, nil is returned.
func (r *KnownTypeRegistry) CreateIfKnown(messageName string) proto.Message {
msgType := r.GetKnownType(messageName)
if msgType == nil {
return nil
}
if msgType.Kind() == reflect.Ptr {
return reflect.New(msgType.Elem()).Interface().(proto.Message)
} else {
return reflect.New(msgType).Elem().Interface().(proto.Message)
}
}
func isWellKnownType(t reflect.Type) bool {
if t.Implements(typeOfWkt) {
return true
}
if msg, ok := reflect.Zero(t).Interface().(proto.Message); ok {
name := proto.MessageName(msg)
_, ok := wellKnownTypeNames[name]
return ok
}
return false
}
// GetKnownType will return the reflect.Type for the given message name if it is
// known. If it is not known, nil is returned.
func (r *KnownTypeRegistry) GetKnownType(messageName string) reflect.Type {
var msgType reflect.Type
if r == nil {
// a nil registry behaves the same as zero value instance: only know of well-known types
t := proto.MessageType(messageName)
if t != nil && isWellKnownType(t) {
msgType = t
}
} else {
if r.includeDefault {
msgType = proto.MessageType(messageName)
} else if !r.excludeWkt {
t := proto.MessageType(messageName)
if t != nil && isWellKnownType(t) {
msgType = t
}
}
if msgType == nil {
r.mu.RLock()
msgType = r.types[messageName]
r.mu.RUnlock()
}
}
return msgType
}