From 3064268cab6666dd0f5811a79ce26e1a02b89dbb Mon Sep 17 00:00:00 2001 From: Paul Wells Date: Sat, 31 Oct 2026 17:39:32 -0700 Subject: [PATCH] decodebudget: share state across merged singular messages The decoder merges each repeated occurrence of a singular message field into the message it already holds, so that message's slices keep growing. The wire walk started every occurrence with fresh state, so a singular field repeated with one packed value each, e.g. ExponentialHistogramDataPoint.positive, was estimated at ~1/1800 of its quadratic decode allocation and passed the budget. Keep a child messageState per singular message field and decode each occurrence into it, charging the struct once. Switching oneof members discards the previous member, so a member that is set again starts over. --- .changeset/decodebudget-merged-messages.md | 5 ++ utils/decodebudget/budget.go | 2 +- utils/decodebudget/calibration_test.go | 35 +++++++++++++ utils/decodebudget/wire.go | 61 +++++++++++++++++++--- utils/decodebudget/wire_test.go | 33 ++++++++++++ 5 files changed, 129 insertions(+), 7 deletions(-) create mode 100644 .changeset/decodebudget-merged-messages.md diff --git a/.changeset/decodebudget-merged-messages.md b/.changeset/decodebudget-merged-messages.md new file mode 100644 index 000000000..5524c48e7 --- /dev/null +++ b/.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 diff --git a/utils/decodebudget/budget.go b/utils/decodebudget/budget.go index bcc7f9596..e98930dfd 100644 --- a/utils/decodebudget/budget.go +++ b/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 } diff --git a/utils/decodebudget/calibration_test.go b/utils/decodebudget/calibration_test.go index 9b4b9b681..b940dc1f9 100644 --- a/utils/decodebudget/calibration_test.go +++ b/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 { diff --git a/utils/decodebudget/wire.go b/utils/decodebudget/wire.go index e0decb653..e92833dc1 100644 --- a/utils/decodebudget/wire.go +++ b/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 } diff --git a/utils/decodebudget/wire_test.go b/utils/decodebudget/wire_test.go index 94ba3a6d1..3b749f1e5 100644 --- a/utils/decodebudget/wire_test.go +++ b/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 +}