codec_extension.go

  1// Copyright 2019 The Go Authors. All rights reserved.
  2// Use of this source code is governed by a BSD-style
  3// license that can be found in the LICENSE file.
  4
  5package impl
  6
  7import (
  8	"sync"
  9	"sync/atomic"
 10
 11	"google.golang.org/protobuf/encoding/protowire"
 12	"google.golang.org/protobuf/internal/errors"
 13	"google.golang.org/protobuf/reflect/protoreflect"
 14)
 15
 16type extensionFieldInfo struct {
 17	wiretag             uint64
 18	tagsize             int
 19	unmarshalNeedsValue bool
 20	funcs               valueCoderFuncs
 21	validation          validationInfo
 22}
 23
 24func getExtensionFieldInfo(xt protoreflect.ExtensionType) *extensionFieldInfo {
 25	if xi, ok := xt.(*ExtensionInfo); ok {
 26		xi.lazyInit()
 27		return xi.info
 28	}
 29	// Ideally we'd cache the resulting *extensionFieldInfo so we don't have to
 30	// recompute this metadata repeatedly. But without support for something like
 31	// weak references, such a cache would pin temporary values (like dynamic
 32	// extension types, constructed for the duration of a user request) to the
 33	// heap forever, causing memory usage of the cache to grow unbounded.
 34	// See discussion in https://github.com/golang/protobuf/issues/1521.
 35	return makeExtensionFieldInfo(xt.TypeDescriptor())
 36}
 37
 38func makeExtensionFieldInfo(xd protoreflect.ExtensionDescriptor) *extensionFieldInfo {
 39	var wiretag uint64
 40	if !xd.IsPacked() {
 41		wiretag = protowire.EncodeTag(xd.Number(), wireTypes[xd.Kind()])
 42	} else {
 43		wiretag = protowire.EncodeTag(xd.Number(), protowire.BytesType)
 44	}
 45	e := &extensionFieldInfo{
 46		wiretag: wiretag,
 47		tagsize: protowire.SizeVarint(wiretag),
 48		funcs:   encoderFuncsForValue(xd),
 49	}
 50	// Does the unmarshal function need a value passed to it?
 51	// This is true for composite types, where we pass in a message, list, or map to fill in,
 52	// and for enums, where we pass in a prototype value to specify the concrete enum type.
 53	switch xd.Kind() {
 54	case protoreflect.MessageKind, protoreflect.GroupKind, protoreflect.EnumKind:
 55		e.unmarshalNeedsValue = true
 56	default:
 57		if xd.Cardinality() == protoreflect.Repeated {
 58			e.unmarshalNeedsValue = true
 59		}
 60	}
 61	return e
 62}
 63
 64type lazyExtensionValue struct {
 65	atomicOnce uint32 // atomically set if value is valid
 66	mu         sync.Mutex
 67	xi         *extensionFieldInfo
 68	value      protoreflect.Value
 69	b          []byte
 70}
 71
 72type ExtensionField struct {
 73	typ protoreflect.ExtensionType
 74
 75	// value is either the value of GetValue,
 76	// or a *lazyExtensionValue that then returns the value of GetValue.
 77	value protoreflect.Value
 78	lazy  *lazyExtensionValue
 79}
 80
 81func (f *ExtensionField) appendLazyBytes(xt protoreflect.ExtensionType, xi *extensionFieldInfo, num protowire.Number, wtyp protowire.Type, b []byte) {
 82	if f.lazy == nil {
 83		f.lazy = &lazyExtensionValue{xi: xi}
 84	}
 85	f.typ = xt
 86	f.lazy.xi = xi
 87	f.lazy.b = protowire.AppendTag(f.lazy.b, num, wtyp)
 88	f.lazy.b = append(f.lazy.b, b...)
 89}
 90
 91func (f *ExtensionField) canLazy(xt protoreflect.ExtensionType) bool {
 92	if f.typ == nil {
 93		return true
 94	}
 95	if f.typ == xt && f.lazy != nil && atomic.LoadUint32(&f.lazy.atomicOnce) == 0 {
 96		return true
 97	}
 98	return false
 99}
100
101// isUnexpandedLazy returns true if the ExensionField is lazy and not
102// yet expanded, which means it's present and already checked for
103// initialized required fields.
104func (f *ExtensionField) isUnexpandedLazy() bool {
105	return f.lazy != nil && atomic.LoadUint32(&f.lazy.atomicOnce) == 0
106}
107
108// lazyBuffer retrieves the buffer for a lazy extension if it's not yet expanded.
109//
110// The returned buffer has to be kept over whatever operation we're planning,
111// as re-retrieving it will fail after the message is lazily decoded.
112func (f *ExtensionField) lazyBuffer() []byte {
113	// This function might be in the critical path, so check the atomic without
114	// taking a look first, then only take the lock if needed.
115	if !f.isUnexpandedLazy() {
116		return nil
117	}
118	f.lazy.mu.Lock()
119	defer f.lazy.mu.Unlock()
120	return f.lazy.b
121}
122
123func (f *ExtensionField) lazyInit() {
124	f.lazy.mu.Lock()
125	defer f.lazy.mu.Unlock()
126	if atomic.LoadUint32(&f.lazy.atomicOnce) == 1 {
127		return
128	}
129	if f.lazy.xi != nil {
130		b := f.lazy.b
131		val := f.typ.New()
132		for len(b) > 0 {
133			var tag uint64
134			if b[0] < 0x80 {
135				tag = uint64(b[0])
136				b = b[1:]
137			} else if len(b) >= 2 && b[1] < 128 {
138				tag = uint64(b[0]&0x7f) + uint64(b[1])<<7
139				b = b[2:]
140			} else {
141				var n int
142				tag, n = protowire.ConsumeVarint(b)
143				if n < 0 {
144					panic(errors.New("bad tag in lazy extension decoding"))
145				}
146				b = b[n:]
147			}
148			num := protowire.Number(tag >> 3)
149			wtyp := protowire.Type(tag & 7)
150			var out unmarshalOutput
151			var err error
152			val, out, err = f.lazy.xi.funcs.unmarshal(b, val, num, wtyp, lazyUnmarshalOptions)
153			if err != nil {
154				panic(errors.New("decode failure in lazy extension decoding: %v", err))
155			}
156			b = b[out.n:]
157		}
158		f.lazy.value = val
159	} else {
160		panic("No support for lazy fns for ExtensionField")
161	}
162	f.lazy.xi = nil
163	f.lazy.b = nil
164	atomic.StoreUint32(&f.lazy.atomicOnce, 1)
165}
166
167// Set sets the type and value of the extension field.
168// This must not be called concurrently.
169func (f *ExtensionField) Set(t protoreflect.ExtensionType, v protoreflect.Value) {
170	f.typ = t
171	f.value = v
172	f.lazy = nil
173}
174
175// Value returns the value of the extension field.
176// This may be called concurrently.
177func (f *ExtensionField) Value() protoreflect.Value {
178	if f.lazy != nil {
179		if atomic.LoadUint32(&f.lazy.atomicOnce) == 0 {
180			f.lazyInit()
181		}
182		return f.lazy.value
183	}
184	return f.value
185}
186
187// Type returns the type of the extension field.
188// This may be called concurrently.
189func (f ExtensionField) Type() protoreflect.ExtensionType {
190	return f.typ
191}
192
193// IsSet returns whether the extension field is set.
194// This may be called concurrently.
195func (f ExtensionField) IsSet() bool {
196	return f.typ != nil
197}
198
199// IsLazy reports whether a field is lazily encoded.
200// It is exported for testing.
201func IsLazy(m protoreflect.Message, fd protoreflect.FieldDescriptor) bool {
202	var mi *MessageInfo
203	var p pointer
204	switch m := m.(type) {
205	case *messageState:
206		mi = m.messageInfo()
207		p = m.pointer()
208	case *messageReflectWrapper:
209		mi = m.messageInfo()
210		p = m.pointer()
211	default:
212		return false
213	}
214	xd, ok := fd.(protoreflect.ExtensionTypeDescriptor)
215	if !ok {
216		return false
217	}
218	xt := xd.Type()
219	ext := mi.extensionMap(p)
220	if ext == nil {
221		return false
222	}
223	f, ok := (*ext)[int32(fd.Number())]
224	if !ok {
225		return false
226	}
227	return f.typ == xt && f.lazy != nil && atomic.LoadUint32(&f.lazy.atomicOnce) == 0
228}