decodebudget: share state across merged singular messages

.changeset/decodebudget-merged-messages.md

@@ -0,0 +1,5 @@
+---
+"github.com/livekit/protocol": patch
+---
+
+decodebudget: charge repeated occurrences of a singular message field against the merged message, closing a quadratic undercount

utils/decodebudget/budget.go

@@ -62,7 +62,7 @@ func (b Proto) Estimate(data []byte, desc protoreflect.MessageDescriptor) (int,
 	if err := t.alloc(protoStructSize(desc)); err != nil {
 		return t.total(), err
 	}
-	err = (&wireWalk{t}).message(data, desc, 1)
+	err = (&wireWalk{t}).message(data, desc, 1, &messageState{})
 	return t.total(), err
 }
 

utils/decodebudget/calibration_test.go

@@ -19,6 +19,7 @@ import (
 	tracev1 "go.opentelemetry.io/proto/otlp/trace/v1"
 	"google.golang.org/protobuf/encoding/protowire"
 	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/reflect/protoreflect"
 
 	"github.com/livekit/protocol/livekit"
 	"github.com/livekit/protocol/utils/protojson"
@@ -178,6 +179,37 @@ func unknownVarints(n int) []byte {
 	return inLogRecord([]byte(strings.Repeat(string(one), n)))
 }
 
+func num(m proto.Message, name protoreflect.Name) protowire.Number {
+	return m.ProtoReflect().Descriptor().Fields().ByName(name).Number()
+}
+
+// mergedBuckets repeats a singular Buckets field n times with one packed count
+// each; the decoder merges them, regrowing one slice every time.
+func mergedBuckets(n int) []byte {
+	point := &metricsv1.ExponentialHistogramDataPoint{}
+	buckets := &metricsv1.ExponentialHistogramDataPoint_Buckets{}
+	one := field(num(point, "positive"), field(num(buckets, "bucket_counts"), protowire.AppendVarint(nil, 1)))
+	dp := field(num(&metricsv1.ExponentialHistogram{}, "data_points"), []byte(strings.Repeat(string(one), n)))
+	metric := field(num(&metricsv1.Metric{}, "exponential_histogram"), dp)
+	return field(1, field(2, field(2, metric)))
+}
+
+// mergedResources repeats a singular Resource with one attribute n times.
+func mergedResources(n int) []byte {
+	attr := must(proto.Marshal(&resourcev1.Resource{Attributes: []*commonv1.KeyValue{{Key: "k", Value: str("v")}}}))
+	one := field(num(&logsv1.ResourceLogs{}, "resource"), attr)
+	return field(1, []byte(strings.Repeat(string(one), n)))
+}
+
+// alternatingOneof switches an AnyValue between two message members n times.
+func alternatingOneof(n int) []byte {
+	v := &commonv1.AnyValue{}
+	kv := field(num(v, "kvlist_value"), must(proto.Marshal(&commonv1.KeyValueList{Values: []*commonv1.KeyValue{{Key: "k"}}})))
+	arr := field(num(v, "array_value"), must(proto.Marshal(&commonv1.ArrayValue{Values: []*commonv1.AnyValue{{}}})))
+	body := []byte(strings.Repeat(string(kv)+string(arr), n/2))
+	return inLogRecord(field(num(&logsv1.LogRecord{}, "body"), body))
+}
+
 func TestProtoEstimateCalibration(t *testing.T) {
 	// ScopeSpans{spans: n×{}} inside ResourceSpans inside TracesData.
 	emptySpans := func(n int) []byte {
@@ -209,6 +241,9 @@ func TestProtoEstimateCalibration(t *testing.T) {
 		{"wire type mismatch length", mismatchedLength(1 << 20), func() proto.Message { return &logsv1.LogsData{} }},
 		{"wire type mismatch varints", mismatchedVarints(1 << 18), func() proto.Message { return &logsv1.LogsData{} }},
 		{"unknown varints", unknownVarints(1 << 18), func() proto.Message { return &logsv1.LogsData{} }},
+		{"merged buckets 3000", mergedBuckets(3000), func() proto.Message { return &metricsv1.MetricsData{} }},
+		{"merged resources 64k", mergedResources(1 << 16), func() proto.Message { return &logsv1.LogsData{} }},
+		{"alternating oneof 64k", alternatingOneof(1 << 16), func() proto.Message { return &logsv1.LogsData{} }},
 	}
 	budget := Proto{MaxFields: 1 << 30, MaxDepth: 100}
 	for _, tc := range cases {

utils/decodebudget/wire.go

@@ -29,9 +29,11 @@ type wireWalk struct {
 
 // messageState mirrors the decoder state of one message whose cost depends on
 // earlier occurrences: each repeated field's slice, whether a presence field's
-// pointer exists, and the unknown-fields slice.
+// pointer exists, which member of each oneof is set, and the unknown-fields
+// slice.
 type messageState struct {
 	fields     []fieldState
+	oneofs     []oneofState
 	unknownLen int
 	unknownCap int
 }
@@ -40,6 +42,14 @@ type fieldState struct {
 	field protoreflect.FieldDescriptor
 	len   int
 	cap   int
+	// child is the state of a singular message field. The decoder merges each
+	// later occurrence into the same message, so its slices keep growing.
+	child *messageState
+}
+
+type oneofState struct {
+	oneof protoreflect.OneofDescriptor
+	field protoreflect.FieldDescriptor
 }
 
 func (m *messageState) get(field protoreflect.FieldDescriptor) *fieldState {
@@ -52,11 +62,12 @@ func (m *messageState) get(field protoreflect.FieldDescriptor) *fieldState {
 	return &m.fields[len(m.fields)-1]
 }
 
-func (w *wireWalk) message(data []byte, desc protoreflect.MessageDescriptor, depth int) error {
+// message walks the fields of one message into state, which persists across
+// occurrences that the decoder merges into the same message.
+func (w *wireWalk) message(data []byte, desc protoreflect.MessageDescriptor, depth int, state *messageState) error {
 	if depth > w.budget.MaxDepth {
 		return ErrMessageComplexity
 	}
-	var state messageState
 	for len(data) > 0 {
 		if err := w.count(); err != nil {
 			return err
@@ -101,11 +112,12 @@ func (w *wireWalk) message(data []byte, desc protoreflect.MessageDescriptor, dep
 		var err error
 		switch {
 		case kind == protowire.BytesType && field.Message() != nil:
-			if err = w.alloc(state.occurrence(field, 0)); err == nil {
-				err = w.message(value, field.Message(), depth+1)
+			cost, child := state.messageOccurrence(field)
+			if err = w.alloc(cost); err == nil {
+				err = w.message(value, field.Message(), depth+1, child)
 			}
 		case kind == protowire.BytesType && isPackable(field):
-			err = w.packed(value, field, &state)
+			err = w.packed(value, field, state)
 		default:
 			err = w.alloc(state.occurrence(field, len(value)))
 		}
@@ -116,6 +128,41 @@ func (w *wireWalk) message(data []byte, desc protoreflect.MessageDescriptor, dep
 	return nil
 }
 
+// messageOccurrence is the allocation for one occurrence of a message field,
+// and the state its contents decode into.
+func (m *messageState) messageOccurrence(field protoreflect.FieldDescriptor) (int, *messageState) {
+	if field.IsList() || field.IsMap() {
+		// Every element or map value is a new message.
+		return m.occurrence(field, 0), &messageState{}
+	}
+	fs := m.get(field)
+	m.setOneof(field, fs)
+	if fs.child != nil {
+		return 0, fs.child
+	}
+	fs.child = &messageState{}
+	return fieldAllocCost(field, 0), fs.child
+}
+
+// setOneof records field as its oneof's set member. Switching members discards
+// the previous member's value, so a later return to field starts over.
+func (m *messageState) setOneof(field protoreflect.FieldDescriptor, fs *fieldState) {
+	oneof := field.ContainingOneof()
+	if oneof == nil || oneof.IsSynthetic() {
+		return
+	}
+	for i := range m.oneofs {
+		if m.oneofs[i].oneof == oneof {
+			if m.oneofs[i].field != field {
+				*fs = fieldState{field: field}
+				m.oneofs[i].field = field
+			}
+			return
+		}
+	}
+	m.oneofs = append(m.oneofs, oneofState{oneof: oneof, field: field})
+}
+
 // occurrence is the allocation for one occurrence of field with a string or
 // bytes payload of payloadLen. Nested message contents are charged separately.
 func (m *messageState) occurrence(field protoreflect.FieldDescriptor, payloadLen int) int {
@@ -132,6 +179,8 @@ func (m *messageState) occurrence(field protoreflect.FieldDescriptor, payloadLen
 			fs.len = 1
 			cost += presenceCost(field)
 		}
+	case field.ContainingOneof() != nil:
+		m.setOneof(field, m.get(field))
 	}
 	return cost
 }

utils/decodebudget/wire_test.go

@@ -246,3 +246,36 @@ func TestProtoRejectsQuadraticPackedRuns(t *testing.T) {
 	require.ErrorIs(t, err, ErrMessageComplexity)
 	require.Less(t, time.Since(start), time.Second)
 }
+
+func TestProtoRejectsQuadraticMergedMessages(t *testing.T) {
+	// The decoder merges repeated occurrences of a singular message field, so
+	// one packed count per occurrence regrows a single slice each time.
+	point := &metricsv1.ExponentialHistogramDataPoint{}
+	desc := point.ProtoReflect().Descriptor()
+	positive := desc.Fields().ByName("positive").Number()
+	counts := desc.Fields().ByName("positive").Message().Fields().ByName("bucket_counts").Number()
+	inner := protowire.AppendBytes(protowire.AppendTag(nil, counts, protowire.BytesType), protowire.AppendVarint(nil, 1))
+	one := protowire.AppendBytes(protowire.AppendTag(nil, positive, protowire.BytesType), inner)
+	data := bytes.Repeat(one, 400_000)
+
+	start := time.Now()
+	err := Proto{MaxFields: 1 << 24, MaxDepth: 8, MaxDecodedBytes: 512 << 20}.Validate(data, desc)
+	require.ErrorIs(t, err, ErrMessageComplexity)
+	require.Less(t, time.Since(start), time.Second)
+}
+
+func TestProtoRejectsQuadraticMergedPackedRuns(t *testing.T) {
+	// Each occurrence of the singular field positive (8) merges into one
+	// Buckets message, so its packed bucket_counts (2) slice grows quadratically.
+	var b []byte
+	for range 40_000 {
+		b = protowire.AppendTag(b, 8, protowire.BytesType)
+		b = protowire.AppendVarint(b, 3)
+		b = protowire.AppendTag(b, 2, protowire.BytesType)
+		b = protowire.AppendVarint(b, 1)
+		b = protowire.AppendVarint(b, 1)
+	}
+	desc := (&metricsv1.ExponentialHistogramDataPoint{}).ProtoReflect().Descriptor()
+	err := Proto{MaxFields: 1 << 24, MaxDepth: 8, MaxDecodedBytes: 512 << 20}.Validate(b, desc)
+	require.ErrorIs(t, err, ErrMessageComplexity) // decode allocates ~6.5 GB
+}