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