...

Source file src/google.golang.org/protobuf/internal/impl/codec_extension.go

Documentation: google.golang.org/protobuf/internal/impl

     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  
     5  package impl
     6  
     7  import (
     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  
    16  type extensionFieldInfo struct {
    17  	wiretag             uint64
    18  	tagsize             int
    19  	unmarshalNeedsValue bool
    20  	funcs               valueCoderFuncs
    21  	validation          validationInfo
    22  }
    23  
    24  func 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  
    38  func 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  
    64  type 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  	fn         func() protoreflect.Value
    71  }
    72  
    73  type ExtensionField struct {
    74  	typ protoreflect.ExtensionType
    75  
    76  	// value is either the value of GetValue,
    77  	// or a *lazyExtensionValue that then returns the value of GetValue.
    78  	value protoreflect.Value
    79  	lazy  *lazyExtensionValue
    80  }
    81  
    82  func (f *ExtensionField) appendLazyBytes(xt protoreflect.ExtensionType, xi *extensionFieldInfo, num protowire.Number, wtyp protowire.Type, b []byte) {
    83  	if f.lazy == nil {
    84  		f.lazy = &lazyExtensionValue{xi: xi}
    85  	}
    86  	f.typ = xt
    87  	f.lazy.xi = xi
    88  	f.lazy.b = protowire.AppendTag(f.lazy.b, num, wtyp)
    89  	f.lazy.b = append(f.lazy.b, b...)
    90  }
    91  
    92  func (f *ExtensionField) canLazy(xt protoreflect.ExtensionType) bool {
    93  	if f.typ == nil {
    94  		return true
    95  	}
    96  	if f.typ == xt && f.lazy != nil && atomic.LoadUint32(&f.lazy.atomicOnce) == 0 {
    97  		return true
    98  	}
    99  	return false
   100  }
   101  
   102  func (f *ExtensionField) lazyInit() {
   103  	f.lazy.mu.Lock()
   104  	defer f.lazy.mu.Unlock()
   105  	if atomic.LoadUint32(&f.lazy.atomicOnce) == 1 {
   106  		return
   107  	}
   108  	if f.lazy.xi != nil {
   109  		b := f.lazy.b
   110  		val := f.typ.New()
   111  		for len(b) > 0 {
   112  			var tag uint64
   113  			if b[0] < 0x80 {
   114  				tag = uint64(b[0])
   115  				b = b[1:]
   116  			} else if len(b) >= 2 && b[1] < 128 {
   117  				tag = uint64(b[0]&0x7f) + uint64(b[1])<<7
   118  				b = b[2:]
   119  			} else {
   120  				var n int
   121  				tag, n = protowire.ConsumeVarint(b)
   122  				if n < 0 {
   123  					panic(errors.New("bad tag in lazy extension decoding"))
   124  				}
   125  				b = b[n:]
   126  			}
   127  			num := protowire.Number(tag >> 3)
   128  			wtyp := protowire.Type(tag & 7)
   129  			var out unmarshalOutput
   130  			var err error
   131  			val, out, err = f.lazy.xi.funcs.unmarshal(b, val, num, wtyp, lazyUnmarshalOptions)
   132  			if err != nil {
   133  				panic(errors.New("decode failure in lazy extension decoding: %v", err))
   134  			}
   135  			b = b[out.n:]
   136  		}
   137  		f.lazy.value = val
   138  	} else {
   139  		f.lazy.value = f.lazy.fn()
   140  	}
   141  	f.lazy.xi = nil
   142  	f.lazy.fn = nil
   143  	f.lazy.b = nil
   144  	atomic.StoreUint32(&f.lazy.atomicOnce, 1)
   145  }
   146  
   147  // Set sets the type and value of the extension field.
   148  // This must not be called concurrently.
   149  func (f *ExtensionField) Set(t protoreflect.ExtensionType, v protoreflect.Value) {
   150  	f.typ = t
   151  	f.value = v
   152  	f.lazy = nil
   153  }
   154  
   155  // SetLazy sets the type and a value that is to be lazily evaluated upon first use.
   156  // This must not be called concurrently.
   157  func (f *ExtensionField) SetLazy(t protoreflect.ExtensionType, fn func() protoreflect.Value) {
   158  	f.typ = t
   159  	f.lazy = &lazyExtensionValue{fn: fn}
   160  }
   161  
   162  // Value returns the value of the extension field.
   163  // This may be called concurrently.
   164  func (f *ExtensionField) Value() protoreflect.Value {
   165  	if f.lazy != nil {
   166  		if atomic.LoadUint32(&f.lazy.atomicOnce) == 0 {
   167  			f.lazyInit()
   168  		}
   169  		return f.lazy.value
   170  	}
   171  	return f.value
   172  }
   173  
   174  // Type returns the type of the extension field.
   175  // This may be called concurrently.
   176  func (f ExtensionField) Type() protoreflect.ExtensionType {
   177  	return f.typ
   178  }
   179  
   180  // IsSet returns whether the extension field is set.
   181  // This may be called concurrently.
   182  func (f ExtensionField) IsSet() bool {
   183  	return f.typ != nil
   184  }
   185  
   186  // IsLazy reports whether a field is lazily encoded.
   187  // It is exported for testing.
   188  func IsLazy(m protoreflect.Message, fd protoreflect.FieldDescriptor) bool {
   189  	var mi *MessageInfo
   190  	var p pointer
   191  	switch m := m.(type) {
   192  	case *messageState:
   193  		mi = m.messageInfo()
   194  		p = m.pointer()
   195  	case *messageReflectWrapper:
   196  		mi = m.messageInfo()
   197  		p = m.pointer()
   198  	default:
   199  		return false
   200  	}
   201  	xd, ok := fd.(protoreflect.ExtensionTypeDescriptor)
   202  	if !ok {
   203  		return false
   204  	}
   205  	xt := xd.Type()
   206  	ext := mi.extensionMap(p)
   207  	if ext == nil {
   208  		return false
   209  	}
   210  	f, ok := (*ext)[int32(fd.Number())]
   211  	if !ok {
   212  		return false
   213  	}
   214  	return f.typ == xt && f.lazy != nil && atomic.LoadUint32(&f.lazy.atomicOnce) == 0
   215  }
   216  

View as plain text