internal/impl/lazy.go

// Copyright 2024 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package impl

import (
	"fmt"
	"math/bits"
	"os"
	"reflect"
	"sort"
	"sync/atomic"

	"google.golang.org/protobuf/encoding/protowire"
	"google.golang.org/protobuf/internal/errors"
	"google.golang.org/protobuf/internal/protolazy"
	"google.golang.org/protobuf/reflect/protoreflect"
	preg "google.golang.org/protobuf/reflect/protoregistry"
	piface "google.golang.org/protobuf/runtime/protoiface"
)

var enableLazy int32 = func() int32 {
	if os.Getenv("GOPROTODEBUG") == "nolazy" {
		return 0
	}
	return 1
}()

// EnableLazyUnmarshal enables lazy unmarshaling.
func EnableLazyUnmarshal(enable bool) {
	if enable {
		atomic.StoreInt32(&enableLazy, 1)
		return
	}
	atomic.StoreInt32(&enableLazy, 0)
}

// LazyEnabled reports whether lazy unmarshalling is currently enabled.
func LazyEnabled() bool {
	return atomic.LoadInt32(&enableLazy) != 0
}

// UnmarshalField unmarshals a field in a message.
func UnmarshalField(m interface{}, num protowire.Number) {
	switch m := m.(type) {
	case *messageState:
		m.messageInfo().lazyUnmarshal(m.pointer(), num)
	case *messageReflectWrapper:
		m.messageInfo().lazyUnmarshal(m.pointer(), num)
	default:
		panic(fmt.Sprintf("unsupported wrapper type %T", m))
	}
}

func (mi *MessageInfo) lazyUnmarshal(p pointer, num protoreflect.FieldNumber) {
	var f *coderFieldInfo
	if int(num) < len(mi.denseCoderFields) {
		f = mi.denseCoderFields[num]
	} else {
		f = mi.coderFields[num]
	}
	if f == nil {
		panic(fmt.Sprintf("lazyUnmarshal: field info for %v.%v", mi.Desc.FullName(), num))
	}
	lazy := *p.Apply(mi.lazyOffset).LazyInfoPtr()
	start, end, found, _, multipleEntries := lazy.FindFieldInProto(uint32(num))
	if !found && multipleEntries == nil {
		panic(fmt.Sprintf("lazyUnmarshal: can't find field data for %v.%v", mi.Desc.FullName(), num))
	}
	// The actual pointer in the message can not be set until the whole struct is filled in, otherwise we will have races.
	// Create another pointer and set it atomically, if we won the race and the pointer in the original message is still nil.
	fp := pointerOfValue(reflect.New(f.ft))
	if multipleEntries != nil {
		for _, entry := range multipleEntries {
			mi.unmarshalField(lazy.Buffer()[entry.Start:entry.End], fp, f, lazy, lazy.UnmarshalFlags())
		}
	} else {
		mi.unmarshalField(lazy.Buffer()[start:end], fp, f, lazy, lazy.UnmarshalFlags())
	}
	p.Apply(f.offset).AtomicSetPointerIfNil(fp.Elem())
}

func (mi *MessageInfo) unmarshalField(b []byte, p pointer, f *coderFieldInfo, lazyInfo *protolazy.XXX_lazyUnmarshalInfo, flags piface.UnmarshalInputFlags) error {
	opts := lazyUnmarshalOptions
	opts.flags |= flags
	for len(b) > 0 {
		// Parse the tag (field number and wire type).
		var tag uint64
		if b[0] < 0x80 {
			tag = uint64(b[0])
			b = b[1:]
		} else if len(b) >= 2 && b[1] < 128 {
			tag = uint64(b[0]&0x7f) + uint64(b[1])<<7
			b = b[2:]
		} else {
			var n int
			tag, n = protowire.ConsumeVarint(b)
			if n < 0 {
				return errors.New("invalid wire data")
			}
			b = b[n:]
		}
		var num protowire.Number
		if n := tag >> 3; n < uint64(protowire.MinValidNumber) || n > uint64(protowire.MaxValidNumber) {
			return errors.New("invalid wire data")
		} else {
			num = protowire.Number(n)
		}
		wtyp := protowire.Type(tag & 7)
		if num == f.num {
			o, err := f.funcs.unmarshal(b, p, wtyp, f, opts)
			if err == nil {
				b = b[o.n:]
				continue
			}
			if err != errUnknown {
				return err
			}
		}
		n := protowire.ConsumeFieldValue(num, wtyp, b)
		if n < 0 {
			return errors.New("invalid wire data")
		}
		b = b[n:]
	}
	return nil
}

func (mi *MessageInfo) skipField(b []byte, f *coderFieldInfo, wtyp protowire.Type, opts unmarshalOptions) (out unmarshalOutput, _ ValidationStatus) {
	fmi := f.validation.mi
	if fmi == nil {
		fd := mi.Desc.Fields().ByNumber(f.num)
		if fd == nil {
			return out, ValidationUnknown
		}
		messageName := fd.Message().FullName()
		messageType, err := preg.GlobalTypes.FindMessageByName(messageName)
		if err != nil {
			return out, ValidationUnknown
		}
		var ok bool
		fmi, ok = messageType.(*MessageInfo)
		if !ok {
			return out, ValidationUnknown
		}
	}
	fmi.init()
	switch f.validation.typ {
	case validationTypeMessage:
		if wtyp != protowire.BytesType {
			return out, ValidationWrongWireType
		}
		v, n := protowire.ConsumeBytes(b)
		if n < 0 {
			return out, ValidationInvalid
		}
		out, st := fmi.validate(v, 0, opts)
		out.n = n
		return out, st
	case validationTypeGroup:
		if wtyp != protowire.StartGroupType {
			return out, ValidationWrongWireType
		}
		out, st := fmi.validate(b, f.num, opts)
		return out, st
	default:
		return out, ValidationUnknown
	}
}

// unmarshalPointerLazy is similar to unmarshalPointerEager, but it
// specifically handles lazy unmarshalling.  it expects lazyOffset and
// presenceOffset to both be valid.
func (mi *MessageInfo) unmarshalPointerLazy(b []byte, p pointer, groupTag protowire.Number, opts unmarshalOptions) (out unmarshalOutput, err error) {
	initialized := true
	var requiredMask uint64
	var lazy **protolazy.XXX_lazyUnmarshalInfo
	var presence presence
	var lazyIndex []protolazy.IndexEntry
	var lastNum protowire.Number
	outOfOrder := false
	lazyDecode := false
	presence = p.Apply(mi.presenceOffset).PresenceInfo()
	lazy = p.Apply(mi.lazyOffset).LazyInfoPtr()
	if !presence.AnyPresent(mi.presenceSize) {
		if opts.CanBeLazy() {
			// If the message contains existing data, we need to merge into it.
			// Lazy unmarshaling doesn't merge, so only enable it when the
			// message is empty (has no presence bitmap).
			lazyDecode = true
			if *lazy == nil {
				*lazy = &protolazy.XXX_lazyUnmarshalInfo{}
			}
			(*lazy).SetUnmarshalFlags(opts.flags)
			if !opts.AliasBuffer() {
				// Make a copy of the buffer for lazy unmarshaling.
				// Set the AliasBuffer flag so recursive unmarshal
				// operations reuse the copy.
				b = append([]byte{}, b...)
				opts.flags |= piface.UnmarshalAliasBuffer
			}
			(*lazy).SetBuffer(b)
		}
	}
	// Track special handling of lazy fields.
	//
	// In the common case, all fields are lazyValidateOnly (and lazyFields remains nil).
	// In the event that validation for a field fails, this map tracks handling of the field.
	type lazyAction uint8
	const (
		lazyValidateOnly   lazyAction = iota // validate the field only
		lazyUnmarshalNow                     // eagerly unmarshal the field
		lazyUnmarshalLater                   // unmarshal the field after the message is fully processed
	)
	var lazyFields map[*coderFieldInfo]lazyAction
	var exts *map[int32]ExtensionField
	start := len(b)
	pos := 0
	for len(b) > 0 {
		// Parse the tag (field number and wire type).
		var tag uint64
		if b[0] < 0x80 {
			tag = uint64(b[0])
			b = b[1:]
		} else if len(b) >= 2 && b[1] < 128 {
			tag = uint64(b[0]&0x7f) + uint64(b[1])<<7
			b = b[2:]
		} else {
			var n int
			tag, n = protowire.ConsumeVarint(b)
			if n < 0 {
				return out, errDecode
			}
			b = b[n:]
		}
		var num protowire.Number
		if n := tag >> 3; n < uint64(protowire.MinValidNumber) || n > uint64(protowire.MaxValidNumber) {
			return out, errors.New("invalid field number")
		} else {
			num = protowire.Number(n)
		}
		wtyp := protowire.Type(tag & 7)

		if wtyp == protowire.EndGroupType {
			if num != groupTag {
				return out, errors.New("mismatching end group marker")
			}
			groupTag = 0
			break
		}

		var f *coderFieldInfo
		if int(num) < len(mi.denseCoderFields) {
			f = mi.denseCoderFields[num]
		} else {
			f = mi.coderFields[num]
		}
		var n int
		err := errUnknown
		discardUnknown := false
	Field:
		switch {
		case f != nil:
			if f.funcs.unmarshal == nil {
				break
			}
			if f.isLazy && lazyDecode {
				switch {
				case lazyFields == nil || lazyFields[f] == lazyValidateOnly:
					// Attempt to validate this field and leave it for later lazy unmarshaling.
					o, valid := mi.skipField(b, f, wtyp, opts)
					switch valid {
					case ValidationValid:
						// Skip over the valid field and continue.
						err = nil
						presence.SetPresentUnatomic(f.presenceIndex, mi.presenceSize)
						requiredMask |= f.validation.requiredBit
						if !o.initialized {
							initialized = false
						}
						n = o.n
						break Field
					case ValidationInvalid:
						return out, errors.New("invalid proto wire format")
					case ValidationWrongWireType:
						break Field
					case ValidationUnknown:
						if lazyFields == nil {
							lazyFields = make(map[*coderFieldInfo]lazyAction)
						}
						if presence.Present(f.presenceIndex) {
							// We were unable to determine if the field is valid or not,
							// and we've already skipped over at least one instance of this
							// field. Clear the presence bit (so if we stop decoding early,
							// we don't leave a partially-initialized field around) and flag
							// the field for unmarshaling before we return.
							presence.ClearPresent(f.presenceIndex)
							lazyFields[f] = lazyUnmarshalLater
							discardUnknown = true
							break Field
						} else {
							// We were unable to determine if the field is valid or not,
							// but this is the first time we've seen it. Flag it as needing
							// eager unmarshaling and fall through to the eager unmarshal case below.
							lazyFields[f] = lazyUnmarshalNow
						}
					}
				case lazyFields[f] == lazyUnmarshalLater:
					// This field will be unmarshaled in a separate pass below.
					// Skip over it here.
					discardUnknown = true
					break Field
				default:
					// Eagerly unmarshal the field.
				}
			}
			if f.isLazy && !lazyDecode && presence.Present(f.presenceIndex) {
				if p.Apply(f.offset).AtomicGetPointer().IsNil() {
					mi.lazyUnmarshal(p, f.num)
				}
			}
			var o unmarshalOutput
			o, err = f.funcs.unmarshal(b, p.Apply(f.offset), wtyp, f, opts)
			n = o.n
			if err != nil {
				break
			}
			requiredMask |= f.validation.requiredBit
			if f.funcs.isInit != nil && !o.initialized {
				initialized = false
			}
			if f.presenceIndex != noPresence {
				presence.SetPresentUnatomic(f.presenceIndex, mi.presenceSize)
			}
		default:
			// Possible extension.
			if exts == nil && mi.extensionOffset.IsValid() {
				exts = p.Apply(mi.extensionOffset).Extensions()
				if *exts == nil {
					*exts = make(map[int32]ExtensionField)
				}
			}
			if exts == nil {
				break
			}
			var o unmarshalOutput
			o, err = mi.unmarshalExtension(b, num, wtyp, *exts, opts)
			if err != nil {
				break
			}
			n = o.n
			if !o.initialized {
				initialized = false
			}
		}
		if err != nil {
			if err != errUnknown {
				return out, err
			}
			n = protowire.ConsumeFieldValue(num, wtyp, b)
			if n < 0 {
				return out, errDecode
			}
			if !discardUnknown && !opts.DiscardUnknown() && mi.unknownOffset.IsValid() {
				u := mi.mutableUnknownBytes(p)
				*u = protowire.AppendTag(*u, num, wtyp)
				*u = append(*u, b[:n]...)
			}
		}
		b = b[n:]
		end := start - len(b)
		if lazyDecode && f != nil && f.isLazy {
			if num != lastNum {
				lazyIndex = append(lazyIndex, protolazy.IndexEntry{
					FieldNum: uint32(num),
					Start:    uint32(pos),
					End:      uint32(end),
				})
			} else {
				i := len(lazyIndex) - 1
				lazyIndex[i].End = uint32(end)
				lazyIndex[i].MultipleContiguous = true
			}
		}
		if num < lastNum {
			outOfOrder = true
		}
		pos = end
		lastNum = num
	}
	if groupTag != 0 {
		return out, errors.New("missing end group marker")
	}
	if lazyFields != nil {
		// Some fields failed validation, and now need to be unmarshaled.
		for f, action := range lazyFields {
			if action != lazyUnmarshalLater {
				continue
			}
			initialized = false
			if *lazy == nil {
				*lazy = &protolazy.XXX_lazyUnmarshalInfo{}
			}
			if err := mi.unmarshalField((*lazy).Buffer(), p.Apply(f.offset), f, *lazy, opts.flags); err != nil {
				return out, err
			}
			presence.SetPresentUnatomic(f.presenceIndex, mi.presenceSize)
		}
	}
	if lazyDecode {
		if outOfOrder {
			sort.Slice(lazyIndex, func(i, j int) bool {
				return lazyIndex[i].FieldNum < lazyIndex[j].FieldNum ||
					(lazyIndex[i].FieldNum == lazyIndex[j].FieldNum &&
						lazyIndex[i].Start < lazyIndex[j].Start)
			})
		}
		if *lazy == nil {
			*lazy = &protolazy.XXX_lazyUnmarshalInfo{}
		}

		(*lazy).SetIndex(lazyIndex)
	}
	if mi.numRequiredFields > 0 && bits.OnesCount64(requiredMask) != int(mi.numRequiredFields) {
		initialized = false
	}
	if initialized {
		out.initialized = true
	}
	out.n = start - len(b)
	return out, nil
}