From e5be04ca95612bd12308254504489e99a153a80c Mon Sep 17 00:00:00 2001 From: Paul Wells Date: Sat, 31 Oct 2026 16:54:25 -0700 Subject: [PATCH 1/4] decodebudget: estimate decoded allocation Move the protobuf and JSON budgets from utils to utils/decodebudget so callers don't import the utils closure; the package depends only on the protobuf runtime. ProtoBudget/JSONBudget become Proto/JSON. Add Proto.MaxDecodedBytes and Proto.Estimate, which walk the wire format and charge each message's generated Go struct size (cached per type), copied strings and bytes, oneof wrappers, exact append growth for repeated fields, map entries and retained unknown fields. Field counts alone are not a memory bound: an empty submessage is two wire bytes but a few hundred heap bytes. Add Proto.ValidateJSON and Proto.EstimateJSON, which apply the same budget to protojson input in protobuf terms. Calibration tests compare the estimate with measured allocation for realistic OTLP traces, logs and metrics, empty-submessage payloads and livekit types with maps; the estimate stays within [1.00, 1.15] of measured. --- .changeset/decodebudget-decoded-bytes.md | 5 + go.mod | 2 +- utils/decodebudget/calibration_test.go | 205 ++++++++++++++ utils/decodebudget/cost.go | 167 ++++++++++++ utils/decodebudget/json.go | 279 +++++++++++++++++++ utils/decodebudget/proto.go | 329 +++++++++++++++++++++++ utils/decodebudget/proto_test.go | 294 ++++++++++++++++++++ utils/proto_test.go | 247 +---------------- utils/protobudget.go | 214 --------------- 9 files changed, 1282 insertions(+), 460 deletions(-) create mode 100644 .changeset/decodebudget-decoded-bytes.md create mode 100644 utils/decodebudget/calibration_test.go create mode 100644 utils/decodebudget/cost.go create mode 100644 utils/decodebudget/json.go create mode 100644 utils/decodebudget/proto.go create mode 100644 utils/decodebudget/proto_test.go delete mode 100644 utils/protobudget.go diff --git a/.changeset/decodebudget-decoded-bytes.md b/.changeset/decodebudget-decoded-bytes.md new file mode 100644 index 000000000..653ed17f7 --- /dev/null +++ b/.changeset/decodebudget-decoded-bytes.md @@ -0,0 +1,5 @@ +--- +"github.com/livekit/protocol": patch +--- + +decodebudget: move the protobuf/JSON decoding budgets to `utils/decodebudget` and add a calibrated decoded-allocation estimate (`MaxDecodedBytes`, `Estimate`, `ValidateJSON`) diff --git a/go.mod b/go.mod index a685b9cbb..f0560644d 100644 --- a/go.mod +++ b/go.mod @@ -36,6 +36,7 @@ require ( go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.45.0 go.opentelemetry.io/otel/sdk v1.45.0 go.opentelemetry.io/otel/trace v1.47.0 + go.opentelemetry.io/proto/otlp v1.11.0 go.uber.org/atomic v1.11.0 go.uber.org/multierr v1.11.0 go.uber.org/zap v1.28.0 @@ -87,7 +88,6 @@ require ( go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.45.0 // indirect go.opentelemetry.io/otel/log v1.47.0 // indirect go.opentelemetry.io/otel/metric v1.47.0 // indirect - go.opentelemetry.io/proto/otlp v1.11.0 // indirect go.yaml.in/yaml/v3 v3.0.5 // indirect golang.org/x/crypto v0.58.0 // indirect golang.org/x/net v0.61.0 // indirect diff --git a/utils/decodebudget/calibration_test.go b/utils/decodebudget/calibration_test.go new file mode 100644 index 000000000..c62838d16 --- /dev/null +++ b/utils/decodebudget/calibration_test.go @@ -0,0 +1,205 @@ +// The race detector adds allocations of its own, so measured decode sizes only +// match production builds without it. + +//go:build !race + +package decodebudget + +import ( + "fmt" + "runtime" + "strings" + "testing" + + "github.com/stretchr/testify/require" + commonv1 "go.opentelemetry.io/proto/otlp/common/v1" + logsv1 "go.opentelemetry.io/proto/otlp/logs/v1" + metricsv1 "go.opentelemetry.io/proto/otlp/metrics/v1" + resourcev1 "go.opentelemetry.io/proto/otlp/resource/v1" + tracev1 "go.opentelemetry.io/proto/otlp/trace/v1" + "google.golang.org/protobuf/encoding/protowire" + "google.golang.org/protobuf/proto" + + "github.com/livekit/protocol/livekit" + "github.com/livekit/protocol/utils/protojson" +) + +func measureDecode(t *testing.T, data []byte, newMsg func() proto.Message) uint64 { + var best uint64 + for i := 0; i < 3; i++ { + var before, after runtime.MemStats + runtime.GC() + runtime.ReadMemStats(&before) + m := newMsg() + require.NoError(t, proto.Unmarshal(data, m)) + runtime.ReadMemStats(&after) + runtime.KeepAlive(m) + if d := after.TotalAlloc - before.TotalAlloc; best == 0 || d < best { + best = d + } + } + return best +} + +func str(s string) *commonv1.AnyValue { + return &commonv1.AnyValue{Value: &commonv1.AnyValue_StringValue{StringValue: s}} +} + +func calibSpans(n int) *tracev1.TracesData { + ss := &tracev1.ScopeSpans{Scope: &commonv1.InstrumentationScope{Name: "livekit-agents"}} + for i := range n { + attrs := []*commonv1.KeyValue{} + for j := range 20 { + attrs = append(attrs, &commonv1.KeyValue{Key: fmt.Sprintf("lk.attr_%d", j), Value: str(strings.Repeat("v", 10+j))}) + } + attrs = append(attrs, &commonv1.KeyValue{Key: "n", Value: &commonv1.AnyValue{Value: &commonv1.AnyValue_IntValue{IntValue: int64(i)}}}) + ss.Spans = append(ss.Spans, &tracev1.Span{ + TraceId: make([]byte, 16), SpanId: make([]byte, 8), Name: "llm_node", + StartTimeUnixNano: uint64(i), EndTimeUnixNano: uint64(i + 1), Attributes: attrs, + Events: []*tracev1.Span_Event{{Name: "e", Attributes: attrs[:3]}}, + }) + } + return &tracev1.TracesData{ResourceSpans: []*tracev1.ResourceSpans{{ + Resource: &resourcev1.Resource{Attributes: []*commonv1.KeyValue{{Key: "service.name", Value: str("agent")}}}, + ScopeSpans: []*tracev1.ScopeSpans{ss}, + }}} +} + +func calibLogs(n int) *logsv1.LogsData { + sl := &logsv1.ScopeLogs{} + for i := range n { + nested := &commonv1.AnyValue{Value: &commonv1.AnyValue_KvlistValue{KvlistValue: &commonv1.KeyValueList{Values: []*commonv1.KeyValue{ + {Key: "role", Value: str("assistant")}, + {Key: "content", Value: &commonv1.AnyValue{Value: &commonv1.AnyValue_ArrayValue{ArrayValue: &commonv1.ArrayValue{Values: []*commonv1.AnyValue{str(strings.Repeat("hello ", 40))}}}}}, + {Key: "ttft", Value: &commonv1.AnyValue{Value: &commonv1.AnyValue_DoubleValue{DoubleValue: 0.4}}}, + }}}} + sl.LogRecords = append(sl.LogRecords, &logsv1.LogRecord{ + TimeUnixNano: uint64(i), SeverityText: "INFO", Body: str("chat item"), + Attributes: []*commonv1.KeyValue{{Key: "job_id", Value: str("AJ_FuWSAg3gYutS")}, {Key: "chat.item", Value: nested}}, + TraceId: make([]byte, 16), SpanId: make([]byte, 8), + }) + } + return &logsv1.LogsData{ResourceLogs: []*logsv1.ResourceLogs{{ScopeLogs: []*logsv1.ScopeLogs{sl}}}} +} + +func calibHistograms(n int) *metricsv1.MetricsData { + sm := &metricsv1.ScopeMetrics{} + for i := range n { + sm.Metrics = append(sm.Metrics, &metricsv1.Metric{Name: "lat", Data: &metricsv1.Metric_Histogram{Histogram: &metricsv1.Histogram{ + DataPoints: []*metricsv1.HistogramDataPoint{{ + TimeUnixNano: uint64(i), Count: 10, BucketCounts: make([]uint64, 32), ExplicitBounds: make([]float64, 31), + }}, + }}}) + } + return &metricsv1.MetricsData{ResourceMetrics: []*metricsv1.ResourceMetrics{{ScopeMetrics: []*metricsv1.ScopeMetrics{sm}}}} +} + +func calibParticipants() *livekit.ListParticipantsResponse { + participants := &livekit.ListParticipantsResponse{} + for i := range 200 { + p := &livekit.ParticipantInfo{Sid: fmt.Sprintf("PA_%d", i), Identity: "user", Attributes: map[string]string{}} + for j := range 10 { + p.Attributes[fmt.Sprintf("k%d", j)] = "value" + } + p.Tracks = []*livekit.TrackInfo{{Sid: "TR_x", Name: "mic", Layers: []*livekit.VideoLayer{{}, {}}}} + participants.Participants = append(participants.Participants, p) + } + return participants +} + +func emptyChildren(field protowire.Number, n int) []byte { + one := protowire.AppendBytes(protowire.AppendTag(nil, field, protowire.BytesType), nil) + return []byte(strings.Repeat(string(one), n)) +} + +func TestProtoEstimateCalibration(t *testing.T) { + // ScopeSpans{spans: n×{}} inside ResourceSpans inside TracesData. + emptySpans := func(n int) []byte { + scope := emptyChildren(2, n) + rs := protowire.AppendBytes(protowire.AppendTag(nil, 2, protowire.BytesType), scope) + return protowire.AppendBytes(protowire.AppendTag(nil, 1, protowire.BytesType), rs) + } + emptyAttrs := func(n int) []byte { + // LogsData > ResourceLogs > ScopeLogs > LogRecord{attributes: n×{}} + rec := emptyChildren(6, n) + sl := protowire.AppendBytes(protowire.AppendTag(nil, 2, protowire.BytesType), rec) + rl := protowire.AppendBytes(protowire.AppendTag(nil, 2, protowire.BytesType), sl) + return protowire.AppendBytes(protowire.AppendTag(nil, 1, protowire.BytesType), rl) + } + participants := calibParticipants() + cases := []struct { + name string + data []byte + newMsg func() proto.Message + }{ + {"traces 512", must(proto.Marshal(calibSpans(512))), func() proto.Message { return &tracev1.TracesData{} }}, + {"logs 512", must(proto.Marshal(calibLogs(512))), func() proto.Message { return &logsv1.LogsData{} }}, + {"histograms 512", must(proto.Marshal(calibHistograms(512))), func() proto.Message { return &metricsv1.MetricsData{} }}, + {"empty spans 1M", emptySpans(1 << 20), func() proto.Message { return &tracev1.TracesData{} }}, + {"empty attributes 1M", emptyAttrs(1 << 20), func() proto.Message { return &logsv1.LogsData{} }}, + {"participants maps", must(proto.Marshal(participants)), func() proto.Message { return &livekit.ListParticipantsResponse{} }}, + } + budget := Proto{MaxFields: 1 << 30, MaxDepth: 100} + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + est, err := budget.Estimate(tc.data, tc.newMsg().ProtoReflect().Descriptor()) + require.NoError(t, err) + got := measureDecode(t, tc.data, tc.newMsg) + t.Logf("wire=%d measured=%d estimate=%d ratio=%.2f", len(tc.data), got, est, float64(est)/float64(got)) + require.GreaterOrEqual(t, uint64(est), got, "the estimate must not undercount") + require.LessOrEqual(t, float64(est), 1.15*float64(got), "the estimate should stay close enough to admit real payloads") + }) + } +} + +func must[T any](v T, err error) T { + if err != nil { + panic(err) + } + return v +} + +func measureJSONDecode(t *testing.T, data []byte, newMsg func() proto.Message) uint64 { + var best uint64 + for i := 0; i < 3; i++ { + var before, after runtime.MemStats + runtime.GC() + runtime.ReadMemStats(&before) + m := newMsg() + require.NoError(t, protojson.Unmarshal(data, m)) + runtime.ReadMemStats(&after) + runtime.KeepAlive(m) + if d := after.TotalAlloc - before.TotalAlloc; best == 0 || d < best { + best = d + } + } + return best +} + +func TestProtoEstimateJSONCalibration(t *testing.T) { + emptySpans := `{"resourceSpans":[{"scopeSpans":[{"spans":[` + strings.TrimSuffix(strings.Repeat("{},", 1<<18), ",") + `]}]}]}` + emptyValues := `{"resourceLogs":[{"scopeLogs":[{"logRecords":[{"body":{"arrayValue":{"values":[` + strings.TrimSuffix(strings.Repeat("{},", 1<<18), ",") + `]}}}]}]}]}` + cases := []struct { + name string + data []byte + newMsg func() proto.Message + }{ + {"traces 512", must(protojson.Marshal(calibSpans(512))), func() proto.Message { return &tracev1.TracesData{} }}, + {"logs 512", must(protojson.Marshal(calibLogs(512))), func() proto.Message { return &logsv1.LogsData{} }}, + {"histograms 512", must(protojson.Marshal(calibHistograms(512))), func() proto.Message { return &metricsv1.MetricsData{} }}, + {"empty spans", []byte(emptySpans), func() proto.Message { return &tracev1.TracesData{} }}, + {"empty values", []byte(emptyValues), func() proto.Message { return &logsv1.LogsData{} }}, + {"participants maps", must(protojson.Marshal(calibParticipants())), func() proto.Message { return &livekit.ListParticipantsResponse{} }}, + } + budget := Proto{MaxFields: 1 << 30, MaxDepth: 100} + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + est, err := budget.EstimateJSON(tc.data, tc.newMsg().ProtoReflect().Descriptor()) + require.NoError(t, err) + got := measureJSONDecode(t, tc.data, tc.newMsg) + t.Logf("json=%d measured=%d estimate=%d ratio=%.2f", len(tc.data), got, est, float64(est)/float64(got)) + require.GreaterOrEqual(t, uint64(est), got, "the estimate must not undercount") + require.LessOrEqual(t, float64(est), 1.15*float64(got), "the estimate should stay close enough to admit real payloads") + }) + } +} diff --git a/utils/decodebudget/cost.go b/utils/decodebudget/cost.go new file mode 100644 index 000000000..a0e249670 --- /dev/null +++ b/utils/decodebudget/cost.go @@ -0,0 +1,167 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package decodebudget + +import ( + "reflect" + "slices" + "sync" + + "google.golang.org/protobuf/reflect/protoreflect" + "google.golang.org/protobuf/reflect/protoregistry" + "google.golang.org/protobuf/types/dynamicpb" +) + +// The cost model overestimates on purpose: a budget that admits a payload must +// not let its decode allocate more than the estimate. +const ( + // Map tables grow by doubling and are at most 7/8 full. + mapGrowthFactor = 4 + mapEntryOverhead = 16 + extensionCost = 64 + // estimateMarginPercent covers allocations the model does not itemize; the + // calibration tests measured them at under 3%. + estimateMarginPercent = 5 +) + +func withMargin(n int) int { + return n + n*estimateMarginPercent/100 +} + +// Go runtime size classes; an allocation is rounded up to the next one. +var goSizeClasses = []int{ + 8, 16, 24, 32, 48, 64, 80, 96, 112, 128, 144, 160, 176, 192, 208, 224, 240, 256, + 288, 320, 352, 384, 416, 448, 480, 512, 576, 640, 704, 768, 896, 1024, 1152, 1280, + 1408, 1536, 1792, 2048, 2304, 2688, 3072, 3200, 3456, 4096, 4864, 5376, 6144, 6528, + 6784, 6912, 8192, 9472, 9728, 10240, 10880, 12288, 13568, 14336, 16384, 18432, + 19072, 20480, 21760, 24576, 27264, 28672, 32768, +} + +func allocSize(n int) int { + if n <= 0 { + return 0 + } + if i, _ := slices.BinarySearch(goSizeClasses, n); i < len(goSizeClasses) { + return goSizeClasses[i] + } + const page = 8192 + return (n + page - 1) / page * page +} + +// appendAllocBytes is the total size of the backing arrays append allocates +// while growing a nil slice to n elements, following runtime.growslice. +func appendAllocBytes(n, elem int) int { + total, capacity := 0, 0 + for capacity < n { + newLen, newCap := capacity+1, capacity+capacity + switch { + case newLen > newCap: + newCap = newLen + case capacity >= 256: + newCap = capacity + for newCap < newLen { + newCap += (newCap + 3*256) >> 2 + } + } + mem := allocSize(newCap * elem) + total += mem + capacity = mem / elem + } + return total +} + +var protoStructSizes sync.Map // protoreflect.FullName -> int + +// protoStructSize is the allocation for one message of desc's generated Go type. +func protoStructSize(desc protoreflect.MessageDescriptor) int { + if v, ok := protoStructSizes.Load(desc.FullName()); ok { + return v.(int) + } + // A dynamicpb message keeps its fields in a map. + size := 64 + 32*desc.Fields().Len() + if mt, err := protoregistry.GlobalTypes.FindMessageByName(desc.FullName()); err == nil { + if t := reflect.TypeOf(mt.Zero().Interface()); t != reflect.TypeFor[*dynamicpb.Message]() && t.Kind() == reflect.Pointer && t.Elem().Kind() == reflect.Struct { + size = int(t.Elem().Size()) + } + } + size = allocSize(size) + protoStructSizes.Store(desc.FullName(), size) + return size +} + +// goValueSize is the size of one value of field's kind as a slice element, map +// slot or oneof wrapper field. +func goValueSize(field protoreflect.FieldDescriptor) int { + switch field.Kind() { + case protoreflect.BoolKind: + return 1 + case protoreflect.EnumKind, protoreflect.Int32Kind, protoreflect.Sint32Kind, protoreflect.Uint32Kind, + protoreflect.FloatKind, protoreflect.Fixed32Kind, protoreflect.Sfixed32Kind: + return 4 + case protoreflect.StringKind: + return 16 + case protoreflect.BytesKind: + return 24 + default: + // 64-bit scalars and message pointers. + return 8 + } +} + +// packedCost is exact: the decoder sizes a packed run's slice once. +func packedCost(field protoreflect.FieldDescriptor, count int) int { + return allocSize(count * goValueSize(field)) +} + +func mapEntryCost(field protoreflect.FieldDescriptor) int { + return mapGrowthFactor*(goValueSize(field.MapKey())+goValueSize(field.MapValue())) + mapEntryOverhead +} + +func oneofWrapperCost(field protoreflect.FieldDescriptor) int { + switch field.Kind() { + case protoreflect.StringKind, protoreflect.BytesKind, protoreflect.MessageKind, protoreflect.GroupKind: + return allocSize(goValueSize(field)) + default: + return allocSize(goValueSize(field)) + } +} + +// unknownFieldCost covers retaining n raw bytes in a growing unknown-fields slice. +func unknownFieldCost(n int) int { + return 2 * n +} + +// fieldAllocCost is the allocation for one occurrence of field with a string or +// bytes payload of payloadLen. Nested message contents are charged separately. +func fieldAllocCost(field protoreflect.FieldDescriptor, payloadLen int) int { + cost := 0 + if field.Message() != nil && !field.IsMap() { + cost += protoStructSize(field.Message()) + } + if k := field.Kind(); k == protoreflect.StringKind || k == protoreflect.BytesKind { + cost += allocSize(payloadLen) + } + switch { + case field.IsMap(): + cost += mapEntryCost(field) + case field.IsList(): + // Slice growth is charged per message by protoWalk. + case field.IsExtension(): + cost += extensionCost + case field.ContainingOneof() != nil && !field.ContainingOneof().IsSynthetic(): + cost += oneofWrapperCost(field) + } + return cost +} diff --git a/utils/decodebudget/json.go b/utils/decodebudget/json.go new file mode 100644 index 000000000..5bfbcc77e --- /dev/null +++ b/utils/decodebudget/json.go @@ -0,0 +1,279 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package decodebudget + +import ( + "bytes" + "encoding/json" + "errors" + "io" + + "google.golang.org/protobuf/reflect/protoreflect" +) + +// protojson allocates per token on top of the decoded message, and its +// reflection-based list and map setters box each scalar element and map entry. +const ( + jsonTokenCost = 12 + jsonListScalarCost = 24 + jsonMapEntryBoxCost = 32 +) + +// ValidateJSON checks protojson input against the budget in protobuf terms: an +// object key or array element is a field, and nesting is message depth. +func (b Proto) ValidateJSON(data []byte, desc protoreflect.MessageDescriptor) error { + _, err := b.EstimateJSON(data, desc) + return err +} + +// EstimateJSON validates protojson input like ValidateJSON and returns the +// estimated bytes protojson allocates decoding it. Malformed or mistyped input +// is left for the decoder to reject. +func (b Proto) EstimateJSON(data []byte, desc protoreflect.MessageDescriptor) (int, error) { + pw, err := b.newWalk(desc) + if err != nil { + return 0, err + } + dec := json.NewDecoder(bytes.NewReader(data)) + dec.UseNumber() + w := &jsonWalk{protoWalk: pw, dec: dec} + + tok, err := w.next() + if err != nil { + return w.total(), err + } + if err := w.alloc(protoStructSize(desc)); err != nil { + return w.total(), err + } + if tok != json.Delim('{') { + return w.total(), w.skip(tok, 1) + } + if err := w.object(desc, 1); err != nil { + return w.total(), err + } + if _, err := dec.Token(); err != io.EOF { + return w.total(), errors.New("unexpected value after JSON document") + } + return w.total(), nil +} + +type jsonWalk struct { + *protoWalk + dec *json.Decoder +} + +func (w *jsonWalk) next() (json.Token, error) { + tok, err := w.dec.Token() + if err == io.EOF { + return nil, io.ErrUnexpectedEOF + } + if err != nil { + return nil, err + } + return tok, w.alloc(jsonTokenCost) +} + +// object walks the members of a message whose '{' has been consumed. +func (w *jsonWalk) object(desc protoreflect.MessageDescriptor, depth int) error { + if depth > w.budget.MaxDepth { + return ErrMessageComplexity + } + for { + tok, err := w.next() + if err != nil { + return err + } + if tok == json.Delim('}') { + return nil + } + if err := w.count(); err != nil { + return err + } + key, _ := tok.(string) + field := desc.Fields().ByJSONName(key) + if field == nil { + field = desc.Fields().ByTextName(key) + } + tok, err = w.next() + if err != nil { + return err + } + if field == nil { + // Unknown fields and extensions are discarded or rejected by the decoder. + if err := w.skip(tok, depth); err != nil { + return err + } + continue + } + if err := w.field(field, tok, depth); err != nil { + return err + } + } +} + +func (w *jsonWalk) field(field protoreflect.FieldDescriptor, tok json.Token, depth int) error { + switch { + case tok == nil: + return nil + case field.IsMap(): + if tok != json.Delim('{') { + return w.skip(tok, depth) + } + return w.mapEntries(field, depth) + case field.IsList(): + if tok != json.Delim('[') { + return w.skip(tok, depth) + } + for n := 0; ; n++ { + tok, err := w.next() + if err != nil { + return err + } + if tok == json.Delim(']') { + cost := appendAllocBytes(n, goValueSize(field)) + if field.Message() == nil { + cost += n * jsonListScalarCost + } + return w.alloc(cost) + } + if err := w.count(); err != nil { + return err + } + if err := w.value(field, tok, depth); err != nil { + return err + } + } + default: + return w.value(field, tok, depth) + } +} + +func (w *jsonWalk) mapEntries(field protoreflect.FieldDescriptor, depth int) error { + if depth+1 > w.budget.MaxDepth { + return ErrMessageComplexity + } + for { + tok, err := w.next() + if err != nil { + return err + } + if tok == json.Delim('}') { + return nil + } + if err := w.count(); err != nil { + return err + } + key, _ := tok.(string) + cost := mapEntryCost(field) + jsonMapEntryBoxCost + if field.MapKey().Kind() == protoreflect.StringKind { + cost += allocSize(len(key)) + } + if err := w.alloc(cost); err != nil { + return err + } + tok, err = w.next() + if err != nil { + return err + } + if err := w.value(field.MapValue(), tok, depth+1); err != nil { + return err + } + } +} + +// value walks one occurrence of field. Slots, wrappers and map entries were +// charged by the caller only for maps; everything else is charged here. +func (w *jsonWalk) value(field protoreflect.FieldDescriptor, tok json.Token, depth int) error { + if tok == nil { + return nil + } + if !field.IsMap() && field.ContainingMessage().IsMapEntry() { + // Map values carry no slot or wrapper of their own. + if md := field.Message(); md != nil { + return w.message(md, tok, depth) + } + return w.alloc(jsonScalarCost(field, tok)) + } + cost := fieldAllocCost(field, jsonPayloadLen(field, tok)) + if field.Message() != nil { + cost -= protoStructSize(field.Message()) + } + if err := w.alloc(cost); err != nil { + return err + } + if md := field.Message(); md != nil { + return w.message(md, tok, depth) + } + if _, ok := tok.(json.Delim); ok { + return w.skip(tok, depth) + } + return nil +} + +func (w *jsonWalk) message(md protoreflect.MessageDescriptor, tok json.Token, depth int) error { + if err := w.alloc(protoStructSize(md)); err != nil { + return err + } + if tok != json.Delim('{') || md.FullName().Parent() == "google.protobuf" { + // Well-known types have their own JSON forms; walk them generically. + return w.skip(tok, depth+1) + } + return w.object(md, depth+1) +} + +// skip walks a value without a descriptor, counting every token as a field. +func (w *jsonWalk) skip(tok json.Token, depth int) error { + d, ok := tok.(json.Delim) + if !ok || (d != '{' && d != '[') { + return w.count() + } + for open := 1; open > 0; { + if depth+open > w.budget.MaxDepth { + return ErrMessageComplexity + } + tok, err := w.next() + if err != nil { + return err + } + if err := w.count(); err != nil { + return err + } + switch tok { + case json.Delim('{'), json.Delim('['): + open++ + case json.Delim('}'), json.Delim(']'): + open-- + } + } + return nil +} + +func jsonPayloadLen(field protoreflect.FieldDescriptor, tok json.Token) int { + s, ok := tok.(string) + if !ok { + return 0 + } + if field.Kind() == protoreflect.BytesKind { + return len(s) * 3 / 4 + } + return len(s) +} + +func jsonScalarCost(field protoreflect.FieldDescriptor, tok json.Token) int { + if k := field.Kind(); k == protoreflect.StringKind || k == protoreflect.BytesKind { + return allocSize(jsonPayloadLen(field, tok)) + } + return 0 +} diff --git a/utils/decodebudget/proto.go b/utils/decodebudget/proto.go new file mode 100644 index 000000000..43f2fb832 --- /dev/null +++ b/utils/decodebudget/proto.go @@ -0,0 +1,329 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package decodebudget bounds protobuf and protojson input before decoding, +// by field count, nesting depth and estimated decoded allocation. It depends +// only on the protobuf runtime so callers need not import utils. +package decodebudget + +import ( + "bytes" + "encoding/json" + "errors" + "io" + "reflect" + + "google.golang.org/protobuf/encoding/protowire" + "google.golang.org/protobuf/proto" + "google.golang.org/protobuf/reflect/protoreflect" + "google.golang.org/protobuf/reflect/protoregistry" +) + +// ErrMessageComplexity indicates that a message exceeds its structural budget. +var ErrMessageComplexity = errors.New("message exceeds structural complexity limit") + +// Proto limits allocation-driving protobuf structure before decoding. +// Callers must separately limit the serialized byte size. MaxFields and +// MaxDepth must be positive; the zero value does not disable validation. +type Proto struct { + // MaxFields counts every wire-field occurrence and each packed scalar value, + // including occurrences of singular and unknown fields. + MaxFields int + // MaxDepth counts nested messages, including the root message at depth one. + MaxDepth int + // MaxDecodedBytes bounds the estimated bytes the decoder allocates: the Go + // struct of each message, copied strings and bytes, oneof wrappers, and + // slice and map growth. Zero leaves it unchecked. + MaxDecodedBytes int +} + +// Validate checks data against the message descriptor without constructing the +// decoded message. Known extensions use protoregistry.GlobalTypes, matching the +// default protobuf decoder. Legacy protobuf groups are deliberately unsupported. +func (b Proto) Validate(data []byte, desc protoreflect.MessageDescriptor) error { + _, err := b.Estimate(data, desc) + return err +} + +// Estimate validates data like Validate and returns the estimated bytes the +// binary decoder allocates for it, including the root message. +func (b Proto) Estimate(data []byte, desc protoreflect.MessageDescriptor) (int, error) { + w, err := b.newWalk(desc) + if err != nil { + return 0, err + } + if err := w.alloc(protoStructSize(desc)); err != nil { + return w.total(), err + } + err = w.message(data, desc, 1) + return w.total(), err +} + +func (b Proto) newWalk(desc protoreflect.MessageDescriptor) (*protoWalk, error) { + if b.MaxFields <= 0 || b.MaxDepth <= 0 || b.MaxDecodedBytes < 0 { + return nil, errors.New("protobuf budget limits must be positive") + } + if desc == nil { + return nil, errors.New("protobuf descriptor must not be nil") + } + return &protoWalk{budget: b}, nil +} + +type protoWalk struct { + budget Proto + fields int + bytes int +} + +func (w *protoWalk) count() error { + w.fields++ + if w.fields > w.budget.MaxFields { + return ErrMessageComplexity + } + return nil +} + +func (w *protoWalk) total() int { + return withMargin(w.bytes) +} + +func (w *protoWalk) alloc(n int) error { + w.bytes += n + if w.budget.MaxDecodedBytes > 0 && w.total() > w.budget.MaxDecodedBytes { + return ErrMessageComplexity + } + return nil +} + +// repeatedCounts tallies unpacked occurrences of each repeated field in one +// message, so slice growth is charged once for the final length. +type repeatedCounts []struct { + field protoreflect.FieldDescriptor + n int +} + +func (r *repeatedCounts) add(field protoreflect.FieldDescriptor) { + for i := range *r { + if (*r)[i].field == field { + (*r)[i].n++ + return + } + } + *r = append(*r, struct { + field protoreflect.FieldDescriptor + n int + }{field, 1}) +} + +func (w *protoWalk) growth(r repeatedCounts) error { + for _, c := range r { + if err := w.alloc(appendAllocBytes(c.n, goValueSize(c.field))); err != nil { + return err + } + } + return nil +} + +func (w *protoWalk) message(data []byte, desc protoreflect.MessageDescriptor, depth int) error { + if depth > w.budget.MaxDepth { + return ErrMessageComplexity + } + var repeated repeatedCounts + for len(data) > 0 { + if err := w.count(); err != nil { + return err + } + number, kind, n := protowire.ConsumeTag(data) + if n < 0 { + return protowire.ParseError(n) + } + data = data[n:] + if kind == protowire.StartGroupType || kind == protowire.EndGroupType { + return ErrMessageComplexity + } + field := desc.Fields().ByNumber(number) + if field == nil && desc.ExtensionRanges().Has(number) { + extension, err := protoregistry.GlobalTypes.FindExtensionByNumber(desc.FullName(), number) + if err == nil { + field = extension.TypeDescriptor() + } else if !errors.Is(err, protoregistry.NotFound) { + return err + } + } + if kind != protowire.BytesType { + size := protowire.ConsumeFieldValue(number, kind, data) + if size < 0 { + return protowire.ParseError(size) + } + data = data[size:] + cost := unknownFieldCost(n + size) + if field != nil { + cost = fieldAllocCost(field, 0) + if field.IsList() { + repeated.add(field) + } + } + if err := w.alloc(cost); err != nil { + return err + } + continue + } + + value, size := protowire.ConsumeBytes(data) + if size < 0 { + return protowire.ParseError(size) + } + data = data[size:] + switch { + case field == nil: + if err := w.alloc(unknownFieldCost(n + size)); err != nil { + return err + } + case field.Message() != nil: + if err := w.alloc(fieldAllocCost(field, 0)); err != nil { + return err + } + if err := w.message(value, field.Message(), depth+1); err != nil { + return err + } + case isPackableProtoField(field): + if err := w.packed(value, field); err != nil { + return err + } + default: + if err := w.alloc(fieldAllocCost(field, len(value))); err != nil { + return err + } + } + if field != nil && field.IsList() && (!isPackableProtoField(field) || kind != protowire.BytesType) { + repeated.add(field) + } + } + return w.growth(repeated) +} + +// packed walks a packed scalar run; decoders accept one even when the schema +// prefers unpacked. +func (w *protoWalk) packed(value []byte, field protoreflect.FieldDescriptor) error { + count := 0 + for len(value) > 0 { + var consumed int + switch field.Kind() { + case protoreflect.DoubleKind, protoreflect.Fixed64Kind, protoreflect.Sfixed64Kind: + _, consumed = protowire.ConsumeFixed64(value) + case protoreflect.FloatKind, protoreflect.Fixed32Kind, protoreflect.Sfixed32Kind: + _, consumed = protowire.ConsumeFixed32(value) + default: + _, consumed = protowire.ConsumeVarint(value) + } + if consumed < 0 { + return protowire.ParseError(consumed) + } + if err := w.count(); err != nil { + return err + } + count++ + value = value[consumed:] + } + return w.alloc(packedCost(field, count)) +} + +func isPackableProtoField(field protoreflect.FieldDescriptor) bool { + if field == nil || !field.IsList() { + return false + } + switch field.Kind() { + case protoreflect.BoolKind, protoreflect.EnumKind, + protoreflect.Int32Kind, protoreflect.Sint32Kind, protoreflect.Uint32Kind, + protoreflect.Int64Kind, protoreflect.Sint64Kind, protoreflect.Uint64Kind, + protoreflect.FloatKind, protoreflect.Fixed32Kind, protoreflect.Sfixed32Kind, + protoreflect.DoubleKind, protoreflect.Fixed64Kind, protoreflect.Sfixed64Kind: + return true + default: + return false + } +} + +// Unmarshal validates the budget before decoding into a non-nil mutable message. +// It discards unknown fields, uses MaxDepth as the decoder recursion limit, and +// retains the binary decoder's required-field checks. Preflight rejection leaves +// the destination unchanged; other decode errors follow normal protobuf semantics. +func (b Proto) Unmarshal(data []byte, msg proto.Message) error { + if msg == nil || (reflect.ValueOf(msg).Kind() == reflect.Pointer && reflect.ValueOf(msg).IsNil()) { + return errors.New("protobuf message must be non-nil and mutable") + } + reflected := msg.ProtoReflect() + if !reflected.IsValid() { + return errors.New("protobuf message must be non-nil and mutable") + } + if err := b.Validate(data, reflected.Descriptor()); err != nil { + return err + } + return (proto.UnmarshalOptions{DiscardUnknown: true, RecursionLimit: b.MaxDepth}).Unmarshal(data, msg) +} + +// JSON limits JSON structure before a downstream decoder constructs +// collections. Callers must separately limit the serialized byte size. Both +// limits must be positive; the zero value does not disable validation. +type JSON struct { + // MaxTokens includes scalar values, object keys, and opening/closing delimiters. + MaxTokens int + // MaxDepth counts nested arrays/objects, including a root container at depth one. + MaxDepth int +} + +// Validate accepts exactly one complete JSON value, including scalar roots. +func (b JSON) Validate(data []byte) error { + if b.MaxTokens <= 0 || b.MaxDepth <= 0 { + return errors.New("JSON budget limits must be positive") + } + dec := json.NewDecoder(bytes.NewReader(data)) + dec.UseNumber() + depth, tokens := 0, 0 + seen, complete := false, false + for { + token, err := dec.Token() + if err == io.EOF { + if !seen || depth != 0 { + return io.ErrUnexpectedEOF + } + return nil + } + if err != nil { + return err + } + if complete { + return errors.New("unexpected value after JSON document") + } + tokens++ + if tokens > b.MaxTokens { + return ErrMessageComplexity + } + seen = true + if d, ok := token.(json.Delim); ok { + switch d { + case '{', '[': + depth++ + case '}', ']': + depth-- + } + if depth > b.MaxDepth { + return ErrMessageComplexity + } + } + if depth == 0 { + complete = true + } + } +} diff --git a/utils/decodebudget/proto_test.go b/utils/decodebudget/proto_test.go new file mode 100644 index 000000000..60716e67b --- /dev/null +++ b/utils/decodebudget/proto_test.go @@ -0,0 +1,294 @@ +package decodebudget + +import ( + "bytes" + "strings" + "sync" + "testing" + + "github.com/stretchr/testify/require" + "google.golang.org/protobuf/encoding/protowire" + "google.golang.org/protobuf/proto" + "google.golang.org/protobuf/reflect/protodesc" + "google.golang.org/protobuf/reflect/protoreflect" + "google.golang.org/protobuf/reflect/protoregistry" + "google.golang.org/protobuf/types/descriptorpb" + "google.golang.org/protobuf/types/dynamicpb" + + "github.com/livekit/protocol/livekit" + "github.com/livekit/protocol/utils/protojson" +) + +func TestProtoFields(t *testing.T) { + desc := protoTestFile(t).Messages().ByName("Node") + field := desc.Fields().ByName("name") + encoded := protowire.AppendString(protowire.AppendTag(nil, field.Number(), protowire.BytesType), "value") + budget := Proto{MaxFields: 3, MaxDepth: 3} + require.NoError(t, budget.Validate(bytes.Repeat(encoded, 3), desc)) + + msg := dynamicpb.NewMessage(desc) + msg.Set(field, protoreflect.ValueOfString("preserved")) + before := proto.Clone(msg) + require.ErrorIs(t, budget.Unmarshal(bytes.Repeat(encoded, 4), msg), ErrMessageComplexity) + require.True(t, proto.Equal(before, msg)) + + unknown := protowire.AppendBytes(protowire.AppendTag(nil, 9999, protowire.BytesType), []byte("discard")) + require.NoError(t, budget.Unmarshal(unknown, msg)) + require.Empty(t, msg.GetUnknown()) + require.NoError(t, budget.Validate(bytes.Repeat(unknown, 3), desc)) + require.ErrorIs(t, budget.Validate(bytes.Repeat(unknown, 4), desc), ErrMessageComplexity) +} + +func TestProtoPackedFields(t *testing.T) { + desc := protoTestFile(t).Messages().ByName("Repeated") + for _, tc := range []struct { + name string + value []byte + }{ + {name: "integers", value: protowire.AppendVarint(nil, 1)}, + {name: "states", value: protowire.AppendVarint(nil, 1)}, + {name: "fixed32", value: protowire.AppendFixed32(nil, 1)}, + {name: "fixed64", value: protowire.AppendFixed64(nil, 1)}, + } { + t.Run(tc.name, func(t *testing.T) { + field := desc.Fields().ByName(protoreflect.Name(tc.name)) + require.False(t, field.IsPacked()) + packed := func(count int) []byte { + return protowire.AppendBytes(protowire.AppendTag(nil, field.Number(), protowire.BytesType), bytes.Repeat(tc.value, count)) + } + budget := Proto{MaxFields: 3, MaxDepth: 3} + msg := dynamicpb.NewMessage(desc) + // The decoder accepts packed wire values even with packed=false. + require.NoError(t, proto.Unmarshal(packed(3), msg)) + require.Equal(t, 3, msg.Get(field).List().Len()) + before := proto.Clone(msg) + require.ErrorIs(t, budget.Unmarshal(packed(3), msg), ErrMessageComplexity) + require.True(t, proto.Equal(before, msg)) + require.NoError(t, budget.Unmarshal(packed(2), msg)) + require.Equal(t, 2, msg.Get(field).List().Len()) + kind := protowire.VarintType + switch field.Kind() { + case protoreflect.Fixed32Kind: + kind = protowire.Fixed32Type + case protoreflect.Fixed64Kind: + kind = protowire.Fixed64Type + } + unpacked := append(protowire.AppendTag(nil, field.Number(), kind), tc.value...) + require.NoError(t, budget.Unmarshal(bytes.Repeat(unpacked, 3), msg)) + require.Equal(t, 3, msg.Get(field).List().Len()) + }) + } +} + +func TestProtoMaps(t *testing.T) { + input := &livekit.ParticipantInfo{Attributes: map[string]string{"a": "one", "b": "two"}} + data, err := proto.Marshal(input) + require.NoError(t, err) + msg := &livekit.ParticipantInfo{Identity: "preserved"} + before := proto.Clone(msg) + require.ErrorIs(t, (Proto{MaxFields: 5, MaxDepth: 2}).Unmarshal(data, msg), ErrMessageComplexity) + require.True(t, proto.Equal(before, msg)) + require.NoError(t, (Proto{MaxFields: 6, MaxDepth: 2}).Unmarshal(data, msg)) + require.True(t, proto.Equal(input, msg)) + require.ErrorIs(t, (Proto{MaxFields: 6, MaxDepth: 1}).Unmarshal(data, msg), ErrMessageComplexity) +} + +func TestProtoExtensions(t *testing.T) { + file := protoTestFile(t) + desc := file.Messages().ByName("Container") + node := file.Messages().ByName("Node") + extension := file.Extensions().ByName("items") + child := node.Fields().ByName("children") + emptyChild := protowire.AppendBytes(protowire.AppendTag(nil, child.Number(), protowire.BytesType), nil) + encode := func(count int) []byte { + return protowire.AppendBytes(protowire.AppendTag(nil, extension.Number(), protowire.BytesType), bytes.Repeat(emptyChild, count)) + } + msg := dynamicpb.NewMessage(desc) + require.NoError(t, proto.Unmarshal(encode(4), msg)) + xt, err := protoregistry.GlobalTypes.FindExtensionByNumber(desc.FullName(), extension.Number()) + require.NoError(t, err) + require.Equal(t, 4, msg.Get(xt.TypeDescriptor()).List().Get(0).Message().Get(child).List().Len()) + before := proto.Clone(msg) + budget := Proto{MaxFields: 4, MaxDepth: 4} + require.ErrorIs(t, budget.Unmarshal(encode(4), msg), ErrMessageComplexity) + require.True(t, proto.Equal(before, msg)) + require.NoError(t, budget.Unmarshal(encode(1), msg)) + require.Equal(t, 1, msg.Get(xt.TypeDescriptor()).List().Get(0).Message().Get(child).List().Len()) + + packed := file.Extensions().ByName("numbers") + require.False(t, packed.IsPacked()) + data := protowire.AppendBytes(protowire.AppendTag(nil, packed.Number(), protowire.BytesType), bytes.Repeat([]byte{1}, 4)) + require.NoError(t, proto.Unmarshal(data, dynamicpb.NewMessage(desc))) + require.ErrorIs(t, budget.Unmarshal(data, msg), ErrMessageComplexity) + unknown := protowire.AppendBytes(protowire.AppendTag(nil, 150, protowire.BytesType), []byte("discard")) + require.NoError(t, budget.Unmarshal(unknown, msg)) + require.Empty(t, msg.GetUnknown()) +} + +func TestProtoDepth(t *testing.T) { + desc := protoTestFile(t).Messages().ByName("Node") + budget := Proto{MaxFields: 10, MaxDepth: 3} + var data []byte + for depth := 1; depth <= 4; depth++ { + msg := dynamicpb.NewMessage(desc) + if depth <= budget.MaxDepth { + require.NoError(t, budget.Validate(data, desc)) + require.NoError(t, budget.Unmarshal(data, msg)) + } else { + require.ErrorIs(t, budget.Validate(data, desc), ErrMessageComplexity) + require.ErrorIs(t, budget.Unmarshal(data, msg), ErrMessageComplexity) + } + data = protowire.AppendBytes(protowire.AppendTag(nil, 1, protowire.BytesType), data) + } +} + +func TestProtoInvalidInput(t *testing.T) { + desc := protoTestFile(t).Messages().ByName("Node") + budget := Proto{MaxFields: 10, MaxDepth: 3} + for _, data := range [][]byte{{0xff}, {0x12, 5, 'x'}, {0x08, 0x80}} { + require.Error(t, budget.Validate(data, desc)) + } + for _, kind := range []protowire.Type{protowire.StartGroupType, protowire.EndGroupType} { + require.ErrorIs(t, budget.Validate(protowire.AppendTag(nil, 999, kind), desc), ErrMessageComplexity) + } + group := protowire.AppendTag(nil, 999, protowire.StartGroupType) + group = protowire.AppendTag(group, 999, protowire.EndGroupType) + require.NoError(t, proto.Unmarshal(group, dynamicpb.NewMessage(desc))) + require.ErrorIs(t, budget.Validate(group, desc), ErrMessageComplexity) + repeated := protoTestFile(t).Messages().ByName("Repeated") + for _, data := range [][]byte{{0x0a, 1, 0x80}, {0x1a, 3, 1, 2, 3}, {0x22, 7, 1, 2, 3, 4, 5, 6, 7}} { + require.Error(t, budget.Validate(data, repeated)) + } + require.Error(t, budget.Validate(nil, nil)) + require.Error(t, budget.Unmarshal(nil, nil)) + var nilMessage *livekit.SignalRequest + require.Error(t, budget.Unmarshal(nil, nilMessage)) + var nilDynamic *dynamicpb.Message + require.Error(t, budget.Unmarshal(nil, nilDynamic)) + for _, invalid := range []Proto{{}, {MaxFields: 1}, {MaxDepth: 1}, {MaxFields: -1, MaxDepth: 1}} { + require.Error(t, invalid.Validate(nil, desc)) + require.Error(t, invalid.Unmarshal(nil, dynamicpb.NewMessage(desc))) + } + // The bounded binary decoder retains required-field validation. + require.Error(t, budget.Unmarshal(nil, &descriptorpb.UninterpretedOption_NamePart{})) +} + +func TestJSON(t *testing.T) { + budget := JSON{MaxTokens: 20, MaxDepth: 3} + for _, data := range []string{`null`, `true`, `false`, `1e100000`, `"value"`, `{}`, `[]`, `{"name":"camera","layers":[{}]}`} { + t.Run(data, func(t *testing.T) { require.NoError(t, budget.Validate([]byte(data))) }) + } + for _, data := range []string{"", " \n\t", `{} {}`, `true false`, `null 0`, `{"a":`, `[1,]`, `]`} { + t.Run("invalid "+data, func(t *testing.T) { require.Error(t, budget.Validate([]byte(data))) }) + } + require.NoError(t, (JSON{MaxTokens: 3, MaxDepth: 1}).Validate([]byte(`[0]`))) + require.ErrorIs(t, (JSON{MaxTokens: 2, MaxDepth: 1}).Validate([]byte(`[0]`)), ErrMessageComplexity) + require.NoError(t, (JSON{MaxTokens: 4, MaxDepth: 1}).Validate([]byte(`{"a":0}`))) + require.ErrorIs(t, (JSON{MaxTokens: 3, MaxDepth: 1}).Validate([]byte(`{"a":0}`)), ErrMessageComplexity) + require.NoError(t, budget.Validate([]byte(`[[[]]]`))) + require.ErrorIs(t, budget.Validate([]byte(`[[[[]]]]`)), ErrMessageComplexity) + require.ErrorIs(t, budget.Validate([]byte("["+strings.Repeat("{},", 20)+"{}]")), ErrMessageComplexity) + for _, invalid := range []JSON{{}, {MaxTokens: 1}, {MaxDepth: 1}, {MaxTokens: -1, MaxDepth: 1}} { + require.Error(t, invalid.Validate([]byte(`null`))) + } +} + +var protoTestDescriptors = sync.OnceValues(func() (protoreflect.FileDescriptor, error) { + optional := descriptorpb.FieldDescriptorProto_LABEL_OPTIONAL + repeated := descriptorpb.FieldDescriptorProto_LABEL_REPEATED + message := descriptorpb.FieldDescriptorProto_TYPE_MESSAGE + str := descriptorpb.FieldDescriptorProto_TYPE_STRING + int32Kind := descriptorpb.FieldDescriptorProto_TYPE_INT32 + enum := descriptorpb.FieldDescriptorProto_TYPE_ENUM + fixed32 := descriptorpb.FieldDescriptorProto_TYPE_FIXED32 + fixed64 := descriptorpb.FieldDescriptorProto_TYPE_FIXED64 + field := func(name string, number int32, label descriptorpb.FieldDescriptorProto_Label, kind descriptorpb.FieldDescriptorProto_Type, typeName string) *descriptorpb.FieldDescriptorProto { + f := &descriptorpb.FieldDescriptorProto{Name: proto.String(name), Number: proto.Int32(number), Label: &label, Type: &kind} + if typeName != "" { + f.TypeName = proto.String(typeName) + } + return f + } + packed := func(name string, number int32, kind descriptorpb.FieldDescriptorProto_Type, typeName string) *descriptorpb.FieldDescriptorProto { + f := field(name, number, repeated, kind, typeName) + f.Options = &descriptorpb.FieldOptions{Packed: proto.Bool(false)} + return f + } + items := field("items", 100, repeated, message, ".decode_budget_test.Node") + items.Extendee = proto.String(".decode_budget_test.Container") + numbers := packed("numbers", 101, int32Kind, "") + numbers.Extendee = items.Extendee + file, err := protodesc.NewFile(&descriptorpb.FileDescriptorProto{ + Name: proto.String("decode_budget_test.proto"), Package: proto.String("decode_budget_test"), Syntax: proto.String("proto2"), + EnumType: []*descriptorpb.EnumDescriptorProto{{Name: proto.String("State"), Value: []*descriptorpb.EnumValueDescriptorProto{ + {Name: proto.String("ZERO"), Number: proto.Int32(0)}, {Name: proto.String("ONE"), Number: proto.Int32(1)}, + }}}, + MessageType: []*descriptorpb.DescriptorProto{ + {Name: proto.String("Node"), Field: []*descriptorpb.FieldDescriptorProto{ + field("child", 1, optional, message, ".decode_budget_test.Node"), field("name", 2, optional, str, ""), + field("children", 3, repeated, message, ".decode_budget_test.Node"), + }}, + {Name: proto.String("Repeated"), Field: []*descriptorpb.FieldDescriptorProto{ + packed("integers", 1, int32Kind, ""), packed("states", 2, enum, ".decode_budget_test.State"), + packed("fixed32", 3, fixed32, ""), packed("fixed64", 4, fixed64, ""), + }}, + {Name: proto.String("Container"), ExtensionRange: []*descriptorpb.DescriptorProto_ExtensionRange{{Start: proto.Int32(100), End: proto.Int32(200)}}}, + }, + Extension: []*descriptorpb.FieldDescriptorProto{items, numbers}, + }, nil) + if err != nil { + return nil, err + } + for i := 0; i < file.Extensions().Len(); i++ { + if err = protoregistry.GlobalTypes.RegisterExtension(dynamicpb.NewExtensionType(file.Extensions().Get(i))); err != nil { + return nil, err + } + } + return file, nil +}) + +func protoTestFile(t *testing.T) protoreflect.FileDescriptor { + t.Helper() + file, err := protoTestDescriptors() + require.NoError(t, err) + return file +} + +func TestProtoDecodedBytes(t *testing.T) { + input := &livekit.ParticipantInfo{Identity: strings.Repeat("x", 1000), Attributes: map[string]string{"a": "one"}} + data, err := proto.Marshal(input) + require.NoError(t, err) + desc := input.ProtoReflect().Descriptor() + + unlimited := Proto{MaxFields: 100, MaxDepth: 4} + estimate, err := unlimited.Estimate(data, desc) + require.NoError(t, err) + require.Greater(t, estimate, 1000) + + require.NoError(t, (Proto{MaxFields: 100, MaxDepth: 4, MaxDecodedBytes: estimate}).Validate(data, desc)) + require.ErrorIs(t, (Proto{MaxFields: 100, MaxDepth: 4, MaxDecodedBytes: estimate - 1}).Validate(data, desc), ErrMessageComplexity) + require.Error(t, (Proto{MaxFields: 100, MaxDepth: 4, MaxDecodedBytes: -1}).Validate(data, desc)) + + jsonData, err := protojson.Marshal(input) + require.NoError(t, err) + jsonEstimate, err := unlimited.EstimateJSON(jsonData, desc) + require.NoError(t, err) + require.ErrorIs(t, (Proto{MaxFields: 100, MaxDepth: 4, MaxDecodedBytes: jsonEstimate - 1}).ValidateJSON(jsonData, desc), ErrMessageComplexity) +} + +func TestProtoJSONStructure(t *testing.T) { + desc := (&livekit.ParticipantInfo{}).ProtoReflect().Descriptor() + budget := Proto{MaxFields: 4, MaxDepth: 2} + require.NoError(t, budget.ValidateJSON([]byte(`{"identity":"a","tracks":[{"sid":"TR"}]}`), desc)) + require.ErrorIs(t, budget.ValidateJSON([]byte(`{"tracks":[{"sid":"a"},{"sid":"b"},{"sid":"c"}]}`), desc), ErrMessageComplexity) + require.ErrorIs(t, budget.ValidateJSON([]byte(`{"tracks":[{"layers":[{}]}]}`), desc), ErrMessageComplexity) + require.ErrorIs(t, (Proto{MaxFields: 4, MaxDepth: 2}).ValidateJSON([]byte(`{"unknown":[[[[1]]]]}`), desc), ErrMessageComplexity) + require.Error(t, budget.ValidateJSON([]byte(`{} {}`), desc)) + require.Error(t, budget.ValidateJSON([]byte(`{"identity":`), desc)) +} + +func TestAppendAllocBytes(t *testing.T) { + require.Equal(t, 0, appendAllocBytes(0, 8)) + // 8, 16, 32 and 64 byte arrays for caps 1, 2, 4 and 8. + require.Equal(t, 8+16+32+64, appendAllocBytes(5, 8)) +} diff --git a/utils/proto_test.go b/utils/proto_test.go index 8c65cdf40..7459e00e1 100644 --- a/utils/proto_test.go +++ b/utils/proto_test.go @@ -1,20 +1,12 @@ package utils import ( - "bytes" - "strings" - "sync" "testing" - "github.com/livekit/protocol/livekit" "github.com/stretchr/testify/require" - "google.golang.org/protobuf/encoding/protowire" "google.golang.org/protobuf/proto" - "google.golang.org/protobuf/reflect/protodesc" - "google.golang.org/protobuf/reflect/protoreflect" - "google.golang.org/protobuf/reflect/protoregistry" - "google.golang.org/protobuf/types/descriptorpb" - "google.golang.org/protobuf/types/dynamicpb" + + "github.com/livekit/protocol/livekit" ) func TestProtoRedact(t *testing.T) { @@ -51,238 +43,3 @@ func TestProtoRedact(t *testing.T) { } require.True(t, proto.Equal(expectedCloned, cloned)) } - -func TestProtoBudgetFields(t *testing.T) { - desc := protoBudgetTestFile(t).Messages().ByName("Node") - field := desc.Fields().ByName("name") - encoded := protowire.AppendString(protowire.AppendTag(nil, field.Number(), protowire.BytesType), "value") - budget := ProtoBudget{MaxFields: 3, MaxDepth: 3} - require.NoError(t, budget.Validate(bytes.Repeat(encoded, 3), desc)) - - msg := dynamicpb.NewMessage(desc) - msg.Set(field, protoreflect.ValueOfString("preserved")) - before := proto.Clone(msg) - require.ErrorIs(t, budget.Unmarshal(bytes.Repeat(encoded, 4), msg), ErrMessageComplexity) - require.True(t, proto.Equal(before, msg)) - - unknown := protowire.AppendBytes(protowire.AppendTag(nil, 9999, protowire.BytesType), []byte("discard")) - require.NoError(t, budget.Unmarshal(unknown, msg)) - require.Empty(t, msg.GetUnknown()) - require.NoError(t, budget.Validate(bytes.Repeat(unknown, 3), desc)) - require.ErrorIs(t, budget.Validate(bytes.Repeat(unknown, 4), desc), ErrMessageComplexity) -} - -func TestProtoBudgetPackedFields(t *testing.T) { - desc := protoBudgetTestFile(t).Messages().ByName("Repeated") - for _, tc := range []struct { - name string - value []byte - }{ - {name: "integers", value: protowire.AppendVarint(nil, 1)}, - {name: "states", value: protowire.AppendVarint(nil, 1)}, - {name: "fixed32", value: protowire.AppendFixed32(nil, 1)}, - {name: "fixed64", value: protowire.AppendFixed64(nil, 1)}, - } { - t.Run(tc.name, func(t *testing.T) { - field := desc.Fields().ByName(protoreflect.Name(tc.name)) - require.False(t, field.IsPacked()) - packed := func(count int) []byte { - return protowire.AppendBytes(protowire.AppendTag(nil, field.Number(), protowire.BytesType), bytes.Repeat(tc.value, count)) - } - budget := ProtoBudget{MaxFields: 3, MaxDepth: 3} - msg := dynamicpb.NewMessage(desc) - // The decoder accepts packed wire values even with packed=false. - require.NoError(t, proto.Unmarshal(packed(3), msg)) - require.Equal(t, 3, msg.Get(field).List().Len()) - before := proto.Clone(msg) - require.ErrorIs(t, budget.Unmarshal(packed(3), msg), ErrMessageComplexity) - require.True(t, proto.Equal(before, msg)) - require.NoError(t, budget.Unmarshal(packed(2), msg)) - require.Equal(t, 2, msg.Get(field).List().Len()) - kind := protowire.VarintType - switch field.Kind() { - case protoreflect.Fixed32Kind: - kind = protowire.Fixed32Type - case protoreflect.Fixed64Kind: - kind = protowire.Fixed64Type - } - unpacked := append(protowire.AppendTag(nil, field.Number(), kind), tc.value...) - require.NoError(t, budget.Unmarshal(bytes.Repeat(unpacked, 3), msg)) - require.Equal(t, 3, msg.Get(field).List().Len()) - }) - } -} - -func TestProtoBudgetMaps(t *testing.T) { - input := &livekit.ParticipantInfo{Attributes: map[string]string{"a": "one", "b": "two"}} - data, err := proto.Marshal(input) - require.NoError(t, err) - msg := &livekit.ParticipantInfo{Identity: "preserved"} - before := proto.Clone(msg) - require.ErrorIs(t, (ProtoBudget{MaxFields: 5, MaxDepth: 2}).Unmarshal(data, msg), ErrMessageComplexity) - require.True(t, proto.Equal(before, msg)) - require.NoError(t, (ProtoBudget{MaxFields: 6, MaxDepth: 2}).Unmarshal(data, msg)) - require.True(t, proto.Equal(input, msg)) - require.ErrorIs(t, (ProtoBudget{MaxFields: 6, MaxDepth: 1}).Unmarshal(data, msg), ErrMessageComplexity) -} - -func TestProtoBudgetExtensions(t *testing.T) { - file := protoBudgetTestFile(t) - desc := file.Messages().ByName("Container") - node := file.Messages().ByName("Node") - extension := file.Extensions().ByName("items") - child := node.Fields().ByName("children") - emptyChild := protowire.AppendBytes(protowire.AppendTag(nil, child.Number(), protowire.BytesType), nil) - encode := func(count int) []byte { - return protowire.AppendBytes(protowire.AppendTag(nil, extension.Number(), protowire.BytesType), bytes.Repeat(emptyChild, count)) - } - msg := dynamicpb.NewMessage(desc) - require.NoError(t, proto.Unmarshal(encode(4), msg)) - xt, err := protoregistry.GlobalTypes.FindExtensionByNumber(desc.FullName(), extension.Number()) - require.NoError(t, err) - require.Equal(t, 4, msg.Get(xt.TypeDescriptor()).List().Get(0).Message().Get(child).List().Len()) - before := proto.Clone(msg) - budget := ProtoBudget{MaxFields: 4, MaxDepth: 4} - require.ErrorIs(t, budget.Unmarshal(encode(4), msg), ErrMessageComplexity) - require.True(t, proto.Equal(before, msg)) - require.NoError(t, budget.Unmarshal(encode(1), msg)) - require.Equal(t, 1, msg.Get(xt.TypeDescriptor()).List().Get(0).Message().Get(child).List().Len()) - - packed := file.Extensions().ByName("numbers") - require.False(t, packed.IsPacked()) - data := protowire.AppendBytes(protowire.AppendTag(nil, packed.Number(), protowire.BytesType), bytes.Repeat([]byte{1}, 4)) - require.NoError(t, proto.Unmarshal(data, dynamicpb.NewMessage(desc))) - require.ErrorIs(t, budget.Unmarshal(data, msg), ErrMessageComplexity) - unknown := protowire.AppendBytes(protowire.AppendTag(nil, 150, protowire.BytesType), []byte("discard")) - require.NoError(t, budget.Unmarshal(unknown, msg)) - require.Empty(t, msg.GetUnknown()) -} - -func TestProtoBudgetDepth(t *testing.T) { - desc := protoBudgetTestFile(t).Messages().ByName("Node") - budget := ProtoBudget{MaxFields: 10, MaxDepth: 3} - var data []byte - for depth := 1; depth <= 4; depth++ { - msg := dynamicpb.NewMessage(desc) - if depth <= budget.MaxDepth { - require.NoError(t, budget.Validate(data, desc)) - require.NoError(t, budget.Unmarshal(data, msg)) - } else { - require.ErrorIs(t, budget.Validate(data, desc), ErrMessageComplexity) - require.ErrorIs(t, budget.Unmarshal(data, msg), ErrMessageComplexity) - } - data = protowire.AppendBytes(protowire.AppendTag(nil, 1, protowire.BytesType), data) - } -} - -func TestProtoBudgetInvalidInput(t *testing.T) { - desc := protoBudgetTestFile(t).Messages().ByName("Node") - budget := ProtoBudget{MaxFields: 10, MaxDepth: 3} - for _, data := range [][]byte{{0xff}, {0x12, 5, 'x'}, {0x08, 0x80}} { - require.Error(t, budget.Validate(data, desc)) - } - for _, kind := range []protowire.Type{protowire.StartGroupType, protowire.EndGroupType} { - require.ErrorIs(t, budget.Validate(protowire.AppendTag(nil, 999, kind), desc), ErrMessageComplexity) - } - group := protowire.AppendTag(nil, 999, protowire.StartGroupType) - group = protowire.AppendTag(group, 999, protowire.EndGroupType) - require.NoError(t, proto.Unmarshal(group, dynamicpb.NewMessage(desc))) - require.ErrorIs(t, budget.Validate(group, desc), ErrMessageComplexity) - repeated := protoBudgetTestFile(t).Messages().ByName("Repeated") - for _, data := range [][]byte{{0x0a, 1, 0x80}, {0x1a, 3, 1, 2, 3}, {0x22, 7, 1, 2, 3, 4, 5, 6, 7}} { - require.Error(t, budget.Validate(data, repeated)) - } - require.Error(t, budget.Validate(nil, nil)) - require.Error(t, budget.Unmarshal(nil, nil)) - var nilMessage *livekit.SignalRequest - require.Error(t, budget.Unmarshal(nil, nilMessage)) - var nilDynamic *dynamicpb.Message - require.Error(t, budget.Unmarshal(nil, nilDynamic)) - for _, invalid := range []ProtoBudget{{}, {MaxFields: 1}, {MaxDepth: 1}, {MaxFields: -1, MaxDepth: 1}} { - require.Error(t, invalid.Validate(nil, desc)) - require.Error(t, invalid.Unmarshal(nil, dynamicpb.NewMessage(desc))) - } - // The bounded binary decoder retains required-field validation. - require.Error(t, budget.Unmarshal(nil, &descriptorpb.UninterpretedOption_NamePart{})) -} - -func TestJSONBudget(t *testing.T) { - budget := JSONBudget{MaxTokens: 20, MaxDepth: 3} - for _, data := range []string{`null`, `true`, `false`, `1e100000`, `"value"`, `{}`, `[]`, `{"name":"camera","layers":[{}]}`} { - t.Run(data, func(t *testing.T) { require.NoError(t, budget.Validate([]byte(data))) }) - } - for _, data := range []string{"", " \n\t", `{} {}`, `true false`, `null 0`, `{"a":`, `[1,]`, `]`} { - t.Run("invalid "+data, func(t *testing.T) { require.Error(t, budget.Validate([]byte(data))) }) - } - require.NoError(t, (JSONBudget{MaxTokens: 3, MaxDepth: 1}).Validate([]byte(`[0]`))) - require.ErrorIs(t, (JSONBudget{MaxTokens: 2, MaxDepth: 1}).Validate([]byte(`[0]`)), ErrMessageComplexity) - require.NoError(t, (JSONBudget{MaxTokens: 4, MaxDepth: 1}).Validate([]byte(`{"a":0}`))) - require.ErrorIs(t, (JSONBudget{MaxTokens: 3, MaxDepth: 1}).Validate([]byte(`{"a":0}`)), ErrMessageComplexity) - require.NoError(t, budget.Validate([]byte(`[[[]]]`))) - require.ErrorIs(t, budget.Validate([]byte(`[[[[]]]]`)), ErrMessageComplexity) - require.ErrorIs(t, budget.Validate([]byte("["+strings.Repeat("{},", 20)+"{}]")), ErrMessageComplexity) - for _, invalid := range []JSONBudget{{}, {MaxTokens: 1}, {MaxDepth: 1}, {MaxTokens: -1, MaxDepth: 1}} { - require.Error(t, invalid.Validate([]byte(`null`))) - } -} - -var protoBudgetTestDescriptors = sync.OnceValues(func() (protoreflect.FileDescriptor, error) { - optional := descriptorpb.FieldDescriptorProto_LABEL_OPTIONAL - repeated := descriptorpb.FieldDescriptorProto_LABEL_REPEATED - message := descriptorpb.FieldDescriptorProto_TYPE_MESSAGE - str := descriptorpb.FieldDescriptorProto_TYPE_STRING - int32Kind := descriptorpb.FieldDescriptorProto_TYPE_INT32 - enum := descriptorpb.FieldDescriptorProto_TYPE_ENUM - fixed32 := descriptorpb.FieldDescriptorProto_TYPE_FIXED32 - fixed64 := descriptorpb.FieldDescriptorProto_TYPE_FIXED64 - field := func(name string, number int32, label descriptorpb.FieldDescriptorProto_Label, kind descriptorpb.FieldDescriptorProto_Type, typeName string) *descriptorpb.FieldDescriptorProto { - f := &descriptorpb.FieldDescriptorProto{Name: proto.String(name), Number: proto.Int32(number), Label: &label, Type: &kind} - if typeName != "" { - f.TypeName = proto.String(typeName) - } - return f - } - packed := func(name string, number int32, kind descriptorpb.FieldDescriptorProto_Type, typeName string) *descriptorpb.FieldDescriptorProto { - f := field(name, number, repeated, kind, typeName) - f.Options = &descriptorpb.FieldOptions{Packed: proto.Bool(false)} - return f - } - items := field("items", 100, repeated, message, ".decode_budget_test.Node") - items.Extendee = proto.String(".decode_budget_test.Container") - numbers := packed("numbers", 101, int32Kind, "") - numbers.Extendee = items.Extendee - file, err := protodesc.NewFile(&descriptorpb.FileDescriptorProto{ - Name: proto.String("decode_budget_test.proto"), Package: proto.String("decode_budget_test"), Syntax: proto.String("proto2"), - EnumType: []*descriptorpb.EnumDescriptorProto{{Name: proto.String("State"), Value: []*descriptorpb.EnumValueDescriptorProto{ - {Name: proto.String("ZERO"), Number: proto.Int32(0)}, {Name: proto.String("ONE"), Number: proto.Int32(1)}, - }}}, - MessageType: []*descriptorpb.DescriptorProto{ - {Name: proto.String("Node"), Field: []*descriptorpb.FieldDescriptorProto{ - field("child", 1, optional, message, ".decode_budget_test.Node"), field("name", 2, optional, str, ""), - field("children", 3, repeated, message, ".decode_budget_test.Node"), - }}, - {Name: proto.String("Repeated"), Field: []*descriptorpb.FieldDescriptorProto{ - packed("integers", 1, int32Kind, ""), packed("states", 2, enum, ".decode_budget_test.State"), - packed("fixed32", 3, fixed32, ""), packed("fixed64", 4, fixed64, ""), - }}, - {Name: proto.String("Container"), ExtensionRange: []*descriptorpb.DescriptorProto_ExtensionRange{{Start: proto.Int32(100), End: proto.Int32(200)}}}, - }, - Extension: []*descriptorpb.FieldDescriptorProto{items, numbers}, - }, nil) - if err != nil { - return nil, err - } - for i := 0; i < file.Extensions().Len(); i++ { - if err = protoregistry.GlobalTypes.RegisterExtension(dynamicpb.NewExtensionType(file.Extensions().Get(i))); err != nil { - return nil, err - } - } - return file, nil -}) - -func protoBudgetTestFile(t *testing.T) protoreflect.FileDescriptor { - t.Helper() - file, err := protoBudgetTestDescriptors() - require.NoError(t, err) - return file -} diff --git a/utils/protobudget.go b/utils/protobudget.go deleted file mode 100644 index 1f967713d..000000000 --- a/utils/protobudget.go +++ /dev/null @@ -1,214 +0,0 @@ -// Copyright 2026 LiveKit, Inc. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package utils - -import ( - "bytes" - "encoding/json" - "errors" - "io" - "reflect" - - "google.golang.org/protobuf/encoding/protowire" - "google.golang.org/protobuf/proto" - "google.golang.org/protobuf/reflect/protoreflect" - "google.golang.org/protobuf/reflect/protoregistry" -) - -// ErrMessageComplexity indicates that a message exceeds its structural budget. -var ErrMessageComplexity = errors.New("message exceeds structural complexity limit") - -// ProtoBudget limits allocation-driving protobuf structure before decoding. -// Callers must separately limit the serialized byte size. Both limits must be -// positive; the zero value does not disable validation. -type ProtoBudget struct { - // MaxFields counts every wire-field occurrence and each packed scalar value, - // including occurrences of singular and unknown fields. - MaxFields int - // MaxDepth counts nested messages, including the root message at depth one. - MaxDepth int -} - -// Validate checks data against the message descriptor without constructing the -// decoded message. Known extensions use protoregistry.GlobalTypes, matching the -// default protobuf decoder. Legacy protobuf groups are deliberately unsupported. -func (b ProtoBudget) Validate(data []byte, desc protoreflect.MessageDescriptor) error { - if b.MaxFields <= 0 || b.MaxDepth <= 0 { - return errors.New("protobuf budget limits must be positive") - } - if desc == nil { - return errors.New("protobuf descriptor must not be nil") - } - remaining := b.MaxFields - return b.validateFields(data, desc, &remaining, 1) -} - -func (b ProtoBudget) validateFields(data []byte, desc protoreflect.MessageDescriptor, remaining *int, depth int) error { - if depth > b.MaxDepth { - return ErrMessageComplexity - } - for len(data) > 0 { - *remaining-- - if *remaining < 0 { - return ErrMessageComplexity - } - number, kind, n := protowire.ConsumeTag(data) - if n < 0 { - return protowire.ParseError(n) - } - data = data[n:] - if kind == protowire.StartGroupType || kind == protowire.EndGroupType { - return ErrMessageComplexity - } - field := desc.Fields().ByNumber(number) - if field == nil && desc.ExtensionRanges().Has(number) { - extension, err := protoregistry.GlobalTypes.FindExtensionByNumber(desc.FullName(), number) - if err == nil { - field = extension.TypeDescriptor() - } else if !errors.Is(err, protoregistry.NotFound) { - return err - } - } - if kind == protowire.BytesType { - value, size := protowire.ConsumeBytes(data) - if size < 0 { - return protowire.ParseError(size) - } - if field != nil && field.Message() != nil { - if err := b.validateFields(value, field.Message(), remaining, depth+1); err != nil { - return err - } - } else if isPackableProtoField(field) { - // Decoders accept packed input even when the schema prefers unpacked. - for len(value) > 0 { - var consumed int - switch field.Kind() { - case protoreflect.DoubleKind, protoreflect.Fixed64Kind, protoreflect.Sfixed64Kind: - _, consumed = protowire.ConsumeFixed64(value) - case protoreflect.FloatKind, protoreflect.Fixed32Kind, protoreflect.Sfixed32Kind: - _, consumed = protowire.ConsumeFixed32(value) - default: - _, consumed = protowire.ConsumeVarint(value) - } - if consumed < 0 { - return protowire.ParseError(consumed) - } - *remaining-- - if *remaining < 0 { - return ErrMessageComplexity - } - value = value[consumed:] - } - } - data = data[size:] - } else { - size := protowire.ConsumeFieldValue(number, kind, data) - if size < 0 { - return protowire.ParseError(size) - } - data = data[size:] - } - } - return nil -} - -func isPackableProtoField(field protoreflect.FieldDescriptor) bool { - if field == nil || !field.IsList() { - return false - } - switch field.Kind() { - case protoreflect.BoolKind, protoreflect.EnumKind, - protoreflect.Int32Kind, protoreflect.Sint32Kind, protoreflect.Uint32Kind, - protoreflect.Int64Kind, protoreflect.Sint64Kind, protoreflect.Uint64Kind, - protoreflect.FloatKind, protoreflect.Fixed32Kind, protoreflect.Sfixed32Kind, - protoreflect.DoubleKind, protoreflect.Fixed64Kind, protoreflect.Sfixed64Kind: - return true - default: - return false - } -} - -// Unmarshal validates the budget before decoding into a non-nil mutable message. -// It discards unknown fields, uses MaxDepth as the decoder recursion limit, and -// retains the binary decoder's required-field checks. Preflight rejection leaves -// the destination unchanged; other decode errors follow normal protobuf semantics. -func (b ProtoBudget) Unmarshal(data []byte, msg proto.Message) error { - if msg == nil || (reflect.ValueOf(msg).Kind() == reflect.Pointer && reflect.ValueOf(msg).IsNil()) { - return errors.New("protobuf message must be non-nil and mutable") - } - reflected := msg.ProtoReflect() - if !reflected.IsValid() { - return errors.New("protobuf message must be non-nil and mutable") - } - if err := b.Validate(data, reflected.Descriptor()); err != nil { - return err - } - return (proto.UnmarshalOptions{DiscardUnknown: true, RecursionLimit: b.MaxDepth}).Unmarshal(data, msg) -} - -// JSONBudget limits JSON structure before a downstream decoder constructs -// collections. Callers must separately limit the serialized byte size. Both -// limits must be positive; the zero value does not disable validation. -type JSONBudget struct { - // MaxTokens includes scalar values, object keys, and opening/closing delimiters. - MaxTokens int - // MaxDepth counts nested arrays/objects, including a root container at depth one. - MaxDepth int -} - -// Validate accepts exactly one complete JSON value, including scalar roots. -func (b JSONBudget) Validate(data []byte) error { - if b.MaxTokens <= 0 || b.MaxDepth <= 0 { - return errors.New("JSON budget limits must be positive") - } - dec := json.NewDecoder(bytes.NewReader(data)) - dec.UseNumber() - depth, tokens := 0, 0 - seen, complete := false, false - for { - token, err := dec.Token() - if err == io.EOF { - if !seen || depth != 0 { - return io.ErrUnexpectedEOF - } - return nil - } - if err != nil { - return err - } - if complete { - return errors.New("unexpected value after JSON document") - } - tokens++ - if tokens > b.MaxTokens { - return ErrMessageComplexity - } - seen = true - if d, ok := token.(json.Delim); ok { - switch d { - case '{', '[': - depth++ - case '}', ']': - depth-- - } - if depth > b.MaxDepth { - return ErrMessageComplexity - } - } - if depth == 0 { - complete = true - } - } -} From 0fd0ae46e1795345bc7b2a2feeedab0040f28985 Mon Sep 17 00:00:00 2001 From: Paul Wells Date: Sat, 31 Oct 2026 17:18:35 -0700 Subject: [PATCH 2/4] decodebudget: model decoder state that undercounted - Packed runs: each run reallocates the field's slice at exactly its new length, so k small runs cost quadratically. Track each repeated field's length and capacity per message and charge every occurrence against it, packed or appended. - Wire-type mismatches: the decoder keeps a known field with an unexpected wire type as unknown; charge it as unknown. - Unknown fields: charge append growth of the message's unknown-fields slice instead of a flat 2x. - Presence: proto3 optional and proto2 scalars allocate a pointer the first time they are set in a message. - JSON: charge protojson's copy of every string token (keys, enum names, int64 strings, base64), the decoded buffer for bytes, a per-object cost, and presence pointers. Return the full estimate for a non-object root. Calibration adds a probe for each; all stay within [1.00, 1.15]. --- utils/decodebudget/calibration_test.go | 74 +++++++++++ utils/decodebudget/cost.go | 61 ++++++---- utils/decodebudget/json.go | 46 +++---- utils/decodebudget/proto.go | 162 +++++++++++++++---------- utils/decodebudget/proto_test.go | 15 +++ 5 files changed, 245 insertions(+), 113 deletions(-) diff --git a/utils/decodebudget/calibration_test.go b/utils/decodebudget/calibration_test.go index c62838d16..eeac90bfe 100644 --- a/utils/decodebudget/calibration_test.go +++ b/utils/decodebudget/calibration_test.go @@ -94,6 +94,29 @@ func calibHistograms(n int) *metricsv1.MetricsData { return &metricsv1.MetricsData{ResourceMetrics: []*metricsv1.ResourceMetrics{{ScopeMetrics: []*metricsv1.ScopeMetrics{sm}}}} } +func calibEnumBytesLogs(n int) *logsv1.LogsData { + sl := &logsv1.ScopeLogs{} + for i := range n { + sl.LogRecords = append(sl.LogRecords, &logsv1.LogRecord{ + TimeUnixNano: uint64(1763059861000000000 + i), SeverityNumber: logsv1.SeverityNumber_SEVERITY_NUMBER_INFO, + TraceId: []byte("0123456789abcdef"), SpanId: []byte("01234567"), + Body: &commonv1.AnyValue{Value: &commonv1.AnyValue_BytesValue{BytesValue: []byte("some bytes value")}}, + }) + } + return &logsv1.LogsData{ResourceLogs: []*logsv1.ResourceLogs{{ScopeLogs: []*logsv1.ScopeLogs{sl}}}} +} + +func calibOptionalPoints(n int) *metricsv1.MetricsData { + h := &metricsv1.Histogram{} + for i := range n { + v := float64(i) + h.DataPoints = append(h.DataPoints, &metricsv1.HistogramDataPoint{Sum: &v, Min: &v, Max: &v}) + } + return &metricsv1.MetricsData{ResourceMetrics: []*metricsv1.ResourceMetrics{{ScopeMetrics: []*metricsv1.ScopeMetrics{{ + Metrics: []*metricsv1.Metric{{Name: "h", Data: &metricsv1.Metric_Histogram{Histogram: h}}}, + }}}}} +} + func calibParticipants() *livekit.ListParticipantsResponse { participants := &livekit.ListParticipantsResponse{} for i := range 200 { @@ -112,6 +135,49 @@ func emptyChildren(field protowire.Number, n int) []byte { return []byte(strings.Repeat(string(one), n)) } +func field(num protowire.Number, payload []byte) []byte { + return protowire.AppendBytes(protowire.AppendTag(nil, num, protowire.BytesType), payload) +} + +// inLogRecord wraps raw LogRecord fields as LogsData. +func inLogRecord(rec []byte) []byte { + return field(1, field(2, field(2, rec))) +} + +// inHistogram wraps raw HistogramDataPoint messages as MetricsData. +func inHistogram(points []byte) []byte { + return field(1, field(2, field(2, field(9, points)))) +} + +func packedRuns(k int) []byte { + one := field(6, protowire.AppendFixed64(nil, 1)) // bucket_counts + return inHistogram(field(1, []byte(strings.Repeat(string(one), k)))) +} + +func optionalPoints(n int) []byte { + var p []byte + for _, num := range []protowire.Number{5, 11, 12} { // sum, min, max + p = protowire.AppendFixed64(protowire.AppendTag(p, num, protowire.Fixed64Type), 1) + } + return inHistogram([]byte(strings.Repeat(string(field(1, p)), n))) +} + +func mismatchedLength(n int) []byte { + // time_unix_nano is fixed64 on the wire. + return inLogRecord(field(1, make([]byte, n))) +} + +func mismatchedVarints(n int) []byte { + // severity_text is a string. + one := protowire.AppendVarint(protowire.AppendTag(nil, 3, protowire.VarintType), 1) + return inLogRecord([]byte(strings.Repeat(string(one), n))) +} + +func unknownVarints(n int) []byte { + one := protowire.AppendVarint(protowire.AppendTag(nil, 1000, protowire.VarintType), 1) + return inLogRecord([]byte(strings.Repeat(string(one), n))) +} + func TestProtoEstimateCalibration(t *testing.T) { // ScopeSpans{spans: n×{}} inside ResourceSpans inside TracesData. emptySpans := func(n int) []byte { @@ -138,6 +204,11 @@ func TestProtoEstimateCalibration(t *testing.T) { {"empty spans 1M", emptySpans(1 << 20), func() proto.Message { return &tracev1.TracesData{} }}, {"empty attributes 1M", emptyAttrs(1 << 20), func() proto.Message { return &logsv1.LogsData{} }}, {"participants maps", must(proto.Marshal(participants)), func() proto.Message { return &livekit.ListParticipantsResponse{} }}, + {"packed runs 3000", packedRuns(3000), func() proto.Message { return &metricsv1.MetricsData{} }}, + {"optional scalars 64k", optionalPoints(1 << 16), func() proto.Message { return &metricsv1.MetricsData{} }}, + {"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{} }}, } budget := Proto{MaxFields: 1 << 30, MaxDepth: 100} for _, tc := range cases { @@ -190,6 +261,9 @@ func TestProtoEstimateJSONCalibration(t *testing.T) { {"empty spans", []byte(emptySpans), func() proto.Message { return &tracev1.TracesData{} }}, {"empty values", []byte(emptyValues), func() proto.Message { return &logsv1.LogsData{} }}, {"participants maps", must(protojson.Marshal(calibParticipants())), func() proto.Message { return &livekit.ListParticipantsResponse{} }}, + {"enums bytes int64 strings", must(protojson.Marshal(calibEnumBytesLogs(20000))), func() proto.Message { return &logsv1.LogsData{} }}, + {"proto field names", must(protojson.MarshalOptions{UseProtoNames: true}.Marshal(calibLogs(512))), func() proto.Message { return &logsv1.LogsData{} }}, + {"optional scalars", must(protojson.Marshal(calibOptionalPoints(16384))), func() proto.Message { return &metricsv1.MetricsData{} }}, } budget := Proto{MaxFields: 1 << 30, MaxDepth: 100} for _, tc := range cases { diff --git a/utils/decodebudget/cost.go b/utils/decodebudget/cost.go index a0e249670..61e6e506d 100644 --- a/utils/decodebudget/cost.go +++ b/utils/decodebudget/cost.go @@ -60,24 +60,36 @@ func allocSize(n int) int { return (n + page - 1) / page * page } +// appendGrow is the allocation for appending add elements of elem bytes to a +// slice of length l and capacity c, and the resulting capacity, following +// runtime.growslice. +func appendGrow(l, c, add, elem int) (mem, newCap int) { + newLen := l + add + if newLen <= c { + return 0, c + } + newCap = c + c + switch { + case newLen > newCap: + newCap = newLen + case c >= 256: + newCap = c + for newCap < newLen { + newCap += (newCap + 3*256) >> 2 + } + } + mem = allocSize(newCap * elem) + return mem, mem / elem +} + // appendAllocBytes is the total size of the backing arrays append allocates -// while growing a nil slice to n elements, following runtime.growslice. +// while growing a nil slice to n elements one at a time. func appendAllocBytes(n, elem int) int { - total, capacity := 0, 0 - for capacity < n { - newLen, newCap := capacity+1, capacity+capacity - switch { - case newLen > newCap: - newCap = newLen - case capacity >= 256: - newCap = capacity - for newCap < newLen { - newCap += (newCap + 3*256) >> 2 - } - } - mem := allocSize(newCap * elem) + total, c := 0, 0 + for l := 0; l < n; l++ { + mem, newCap := appendGrow(l, c, 1, elem) total += mem - capacity = mem / elem + c = newCap } return total } @@ -120,11 +132,6 @@ func goValueSize(field protoreflect.FieldDescriptor) int { } } -// packedCost is exact: the decoder sizes a packed run's slice once. -func packedCost(field protoreflect.FieldDescriptor, count int) int { - return allocSize(count * goValueSize(field)) -} - func mapEntryCost(field protoreflect.FieldDescriptor) int { return mapGrowthFactor*(goValueSize(field.MapKey())+goValueSize(field.MapValue())) + mapEntryOverhead } @@ -138,9 +145,15 @@ func oneofWrapperCost(field protoreflect.FieldDescriptor) int { } } -// unknownFieldCost covers retaining n raw bytes in a growing unknown-fields slice. -func unknownFieldCost(n int) int { - return 2 * n +// presenceCost is the pointer a generated struct allocates the first time a +// scalar field with explicit presence, such as proto3 optional, is set. +func presenceCost(field protoreflect.FieldDescriptor) int { + if !field.HasPresence() || field.IsList() || field.Message() != nil || field.IsExtension() || + (field.ContainingOneof() != nil && !field.ContainingOneof().IsSynthetic()) || + field.ContainingMessage().IsMapEntry() { + return 0 + } + return allocSize(goValueSize(field)) } // fieldAllocCost is the allocation for one occurrence of field with a string or @@ -157,7 +170,7 @@ func fieldAllocCost(field protoreflect.FieldDescriptor, payloadLen int) int { case field.IsMap(): cost += mapEntryCost(field) case field.IsList(): - // Slice growth is charged per message by protoWalk. + // Slice growth depends on earlier occurrences; the caller charges it. case field.IsExtension(): cost += extensionCost case field.ContainingOneof() != nil && !field.ContainingOneof().IsSynthetic(): diff --git a/utils/decodebudget/json.go b/utils/decodebudget/json.go index 5bfbcc77e..7d0adedfc 100644 --- a/utils/decodebudget/json.go +++ b/utils/decodebudget/json.go @@ -23,12 +23,14 @@ import ( "google.golang.org/protobuf/reflect/protoreflect" ) -// protojson allocates per token on top of the decoded message, and its -// reflection-based list and map setters box each scalar element and map entry. +// protojson allocates per token and per object on top of the decoded message +// and the string copies charged separately, and its reflection-based list and +// map setters box each scalar element and map entry. const ( - jsonTokenCost = 12 + jsonTokenCost = 8 jsonListScalarCost = 24 jsonMapEntryBoxCost = 32 + jsonObjectCost = 8 ) // ValidateJSON checks protojson input against the budget in protobuf terms: an @@ -58,7 +60,8 @@ func (b Proto) EstimateJSON(data []byte, desc protoreflect.MessageDescriptor) (i return w.total(), err } if tok != json.Delim('{') { - return w.total(), w.skip(tok, 1) + err := w.skip(tok, 1) + return w.total(), err } if err := w.object(desc, 1); err != nil { return w.total(), err @@ -82,7 +85,12 @@ func (w *jsonWalk) next() (json.Token, error) { if err != nil { return nil, err } - return tok, w.alloc(jsonTokenCost) + cost := jsonTokenCost + if str, ok := tok.(string); ok { + // protojson's tokenizer copies every string, key or value. + cost += allocSize(len(str)) + } + return tok, w.alloc(cost) } // object walks the members of a message whose '{' has been consumed. @@ -175,12 +183,7 @@ func (w *jsonWalk) mapEntries(field protoreflect.FieldDescriptor, depth int) err if err := w.count(); err != nil { return err } - key, _ := tok.(string) - cost := mapEntryCost(field) + jsonMapEntryBoxCost - if field.MapKey().Kind() == protoreflect.StringKind { - cost += allocSize(len(key)) - } - if err := w.alloc(cost); err != nil { + if err := w.alloc(mapEntryCost(field) + jsonMapEntryBoxCost); err != nil { return err } tok, err = w.next() @@ -206,7 +209,7 @@ func (w *jsonWalk) value(field protoreflect.FieldDescriptor, tok json.Token, dep } return w.alloc(jsonScalarCost(field, tok)) } - cost := fieldAllocCost(field, jsonPayloadLen(field, tok)) + cost := fieldAllocCost(field, 0) + jsonScalarCost(field, tok) + presenceCost(field) if field.Message() != nil { cost -= protoStructSize(field.Message()) } @@ -223,7 +226,7 @@ func (w *jsonWalk) value(field protoreflect.FieldDescriptor, tok json.Token, dep } func (w *jsonWalk) message(md protoreflect.MessageDescriptor, tok json.Token, depth int) error { - if err := w.alloc(protoStructSize(md)); err != nil { + if err := w.alloc(protoStructSize(md) + jsonObjectCost); err != nil { return err } if tok != json.Delim('{') || md.FullName().Parent() == "google.protobuf" { @@ -260,20 +263,11 @@ func (w *jsonWalk) skip(tok json.Token, depth int) error { return nil } -func jsonPayloadLen(field protoreflect.FieldDescriptor, tok json.Token) int { - s, ok := tok.(string) - if !ok { - return 0 - } - if field.Kind() == protoreflect.BytesKind { - return len(s) * 3 / 4 - } - return len(s) -} - +// jsonScalarCost is what decoding a scalar token allocates beyond the copy the +// tokenizer already made: base64 bytes are decoded into a new buffer. func jsonScalarCost(field protoreflect.FieldDescriptor, tok json.Token) int { - if k := field.Kind(); k == protoreflect.StringKind || k == protoreflect.BytesKind { - return allocSize(jsonPayloadLen(field, tok)) + if str, ok := tok.(string); ok && field.Kind() == protoreflect.BytesKind { + return allocSize(len(str) * 3 / 4) } return 0 } diff --git a/utils/decodebudget/proto.go b/utils/decodebudget/proto.go index 43f2fb832..d5f2b38e1 100644 --- a/utils/decodebudget/proto.go +++ b/utils/decodebudget/proto.go @@ -106,40 +106,36 @@ func (w *protoWalk) alloc(n int) error { return nil } -// repeatedCounts tallies unpacked occurrences of each repeated field in one -// message, so slice growth is charged once for the final length. -type repeatedCounts []struct { - field protoreflect.FieldDescriptor - n int +// 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. +type messageState struct { + fields []fieldState + unknownLen int + unknownCap int } -func (r *repeatedCounts) add(field protoreflect.FieldDescriptor) { - for i := range *r { - if (*r)[i].field == field { - (*r)[i].n++ - return - } - } - *r = append(*r, struct { - field protoreflect.FieldDescriptor - n int - }{field, 1}) +type fieldState struct { + field protoreflect.FieldDescriptor + len int + cap int } -func (w *protoWalk) growth(r repeatedCounts) error { - for _, c := range r { - if err := w.alloc(appendAllocBytes(c.n, goValueSize(c.field))); err != nil { - return err +func (m *messageState) get(field protoreflect.FieldDescriptor) *fieldState { + for i := range m.fields { + if m.fields[i].field == field { + return &m.fields[i] } } - return nil + m.fields = append(m.fields, fieldState{field: field}) + return &m.fields[len(m.fields)-1] } func (w *protoWalk) message(data []byte, desc protoreflect.MessageDescriptor, depth int) error { if depth > w.budget.MaxDepth { return ErrMessageComplexity } - var repeated repeatedCounts + var state messageState for len(data) > 0 { if err := w.count(); err != nil { return err @@ -161,61 +157,77 @@ func (w *protoWalk) message(data []byte, desc protoreflect.MessageDescriptor, de return err } } - if kind != protowire.BytesType { - size := protowire.ConsumeFieldValue(number, kind, data) - if size < 0 { - return protowire.ParseError(size) - } - data = data[size:] - cost := unknownFieldCost(n + size) - if field != nil { - cost = fieldAllocCost(field, 0) - if field.IsList() { - repeated.add(field) - } - } - if err := w.alloc(cost); err != nil { - return err - } - continue - } - value, size := protowire.ConsumeBytes(data) + var value []byte + var size int + if kind == protowire.BytesType { + value, size = protowire.ConsumeBytes(data) + } else { + size = protowire.ConsumeFieldValue(number, kind, data) + } if size < 0 { return protowire.ParseError(size) } data = data[size:] - switch { - case field == nil: - if err := w.alloc(unknownFieldCost(n + size)); err != nil { - return err - } - case field.Message() != nil: - if err := w.alloc(fieldAllocCost(field, 0)); err != nil { - return err - } - if err := w.message(value, field.Message(), depth+1); err != nil { + + // The decoder keeps a field with an unexpected wire type as unknown. + if field == nil || !wireTypeMatches(field, kind) { + if err := w.alloc(state.unknown(n + size)); err != nil { return err } - case isPackableProtoField(field): - if err := w.packed(value, field); err != nil { - return err + continue + } + 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) } + case kind == protowire.BytesType && isPackableProtoField(field): + err = w.packed(value, field, &state) default: - if err := w.alloc(fieldAllocCost(field, len(value))); err != nil { - return err - } + err = w.alloc(state.occurrence(field, len(value))) } - if field != nil && field.IsList() && (!isPackableProtoField(field) || kind != protowire.BytesType) { - repeated.add(field) + if err != nil { + return err + } + } + return nil +} + +// 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 { + cost := fieldAllocCost(field, payloadLen) + switch { + case field.IsList(): + fs := m.get(field) + var mem int + mem, fs.cap = appendGrow(fs.len, fs.cap, 1, goValueSize(field)) + fs.len++ + cost += mem + case presenceCost(field) > 0: + if fs := m.get(field); fs.len == 0 { + fs.len = 1 + cost += presenceCost(field) } } - return w.growth(repeated) + return cost +} + +// unknown is the allocation for retaining an n-byte field in the message's +// unknown-fields slice. +func (m *messageState) unknown(n int) int { + mem, newCap := appendGrow(m.unknownLen, m.unknownCap, n, 1) + m.unknownLen += n + m.unknownCap = newCap + return mem } // packed walks a packed scalar run; decoders accept one even when the schema -// prefers unpacked. -func (w *protoWalk) packed(value []byte, field protoreflect.FieldDescriptor) error { +// prefers unpacked. Each run reallocates the field's slice at exactly its new +// length, so many small runs cost quadratically. +func (w *protoWalk) packed(value []byte, field protoreflect.FieldDescriptor, state *messageState) error { count := 0 for len(value) > 0 { var consumed int @@ -236,7 +248,31 @@ func (w *protoWalk) packed(value []byte, field protoreflect.FieldDescriptor) err count++ value = value[consumed:] } - return w.alloc(packedCost(field, count)) + if count == 0 { + return nil + } + fs := state.get(field) + fs.len += count + fs.cap = fs.len + return w.alloc(allocSize(fs.len * goValueSize(field))) +} + +func wireTypeMatches(field protoreflect.FieldDescriptor, kind protowire.Type) bool { + if kind == protowire.BytesType && isPackableProtoField(field) { + return true + } + switch field.Kind() { + case protoreflect.MessageKind, protoreflect.StringKind, protoreflect.BytesKind: + return kind == protowire.BytesType + case protoreflect.GroupKind: + return kind == protowire.StartGroupType + case protoreflect.FloatKind, protoreflect.Fixed32Kind, protoreflect.Sfixed32Kind: + return kind == protowire.Fixed32Type + case protoreflect.DoubleKind, protoreflect.Fixed64Kind, protoreflect.Sfixed64Kind: + return kind == protowire.Fixed64Type + default: + return kind == protowire.VarintType + } } func isPackableProtoField(field protoreflect.FieldDescriptor) bool { diff --git a/utils/decodebudget/proto_test.go b/utils/decodebudget/proto_test.go index 60716e67b..8345d9a31 100644 --- a/utils/decodebudget/proto_test.go +++ b/utils/decodebudget/proto_test.go @@ -5,8 +5,10 @@ import ( "strings" "sync" "testing" + "time" "github.com/stretchr/testify/require" + metricsv1 "go.opentelemetry.io/proto/otlp/metrics/v1" "google.golang.org/protobuf/encoding/protowire" "google.golang.org/protobuf/proto" "google.golang.org/protobuf/reflect/protodesc" @@ -292,3 +294,16 @@ func TestAppendAllocBytes(t *testing.T) { // 8, 16, 32 and 64 byte arrays for caps 1, 2, 4 and 8. require.Equal(t, 8+16+32+64, appendAllocBytes(5, 8)) } + +func TestProtoRejectsQuadraticPackedRuns(t *testing.T) { + // Each packed run reallocates the slice at its new length, so 100k + // one-element runs in 1.2 MB of wire allocate ~40 GB and decode for seconds. + one := protowire.AppendBytes(protowire.AppendTag(nil, 6, protowire.BytesType), protowire.AppendFixed64(nil, 1)) + point := protowire.AppendBytes(protowire.AppendTag(nil, 1, protowire.BytesType), bytes.Repeat(one, 100_000)) + desc := (&metricsv1.Histogram{}).ProtoReflect().Descriptor() + + start := time.Now() + err := Proto{MaxFields: 1 << 24, MaxDepth: 8, MaxDecodedBytes: 512 << 20}.Validate(point, desc) + require.ErrorIs(t, err, ErrMessageComplexity) + require.Less(t, time.Since(start), time.Second) +} From bebbccbed58d6c9edd5485bc4d00af8c31af38dd Mon Sep 17 00:00:00 2001 From: Paul Wells Date: Sat, 31 Oct 2026 17:28:18 -0700 Subject: [PATCH 3/4] decodebudget: split files by encoding budget.go holds the Proto API and the tally shared by both walks, wire.go the binary walk, protojson.go the protojson walk, and json.go the descriptor-free JSON token budget. Tests follow the same split. No behavior change. --- utils/decodebudget/budget.go | 123 ++++++++ utils/decodebudget/budget_test.go | 34 +++ utils/decodebudget/cost_test.go | 13 + utils/decodebudget/json.go | 266 +++-------------- utils/decodebudget/json_test.go | 28 ++ utils/decodebudget/protojson.go | 274 ++++++++++++++++++ utils/decodebudget/protojson_test.go | 20 ++ utils/decodebudget/{proto.go => wire.go} | 168 +---------- .../{proto_test.go => wire_test.go} | 61 ---- 9 files changed, 535 insertions(+), 452 deletions(-) create mode 100644 utils/decodebudget/budget.go create mode 100644 utils/decodebudget/budget_test.go create mode 100644 utils/decodebudget/cost_test.go create mode 100644 utils/decodebudget/json_test.go create mode 100644 utils/decodebudget/protojson.go create mode 100644 utils/decodebudget/protojson_test.go rename utils/decodebudget/{proto.go => wire.go} (52%) rename utils/decodebudget/{proto_test.go => wire_test.go} (77%) diff --git a/utils/decodebudget/budget.go b/utils/decodebudget/budget.go new file mode 100644 index 000000000..bcc7f9596 --- /dev/null +++ b/utils/decodebudget/budget.go @@ -0,0 +1,123 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package decodebudget bounds untrusted input before decoding. Proto checks +// protobuf and protojson input by field count, nesting depth and estimated +// decoded allocation; JSON checks plain JSON by token count and depth. It +// depends only on the protobuf runtime so callers need not import utils. +package decodebudget + +import ( + "errors" + "reflect" + + "google.golang.org/protobuf/proto" + "google.golang.org/protobuf/reflect/protoreflect" +) + +// ErrMessageComplexity indicates that a message exceeds its structural budget. +var ErrMessageComplexity = errors.New("message exceeds structural complexity limit") + +// Proto limits allocation-driving protobuf structure before decoding. +// Callers must separately limit the serialized byte size. MaxFields and +// MaxDepth must be positive; the zero value does not disable validation. +type Proto struct { + // MaxFields counts every wire-field occurrence and each packed scalar value, + // including occurrences of singular and unknown fields. + MaxFields int + // MaxDepth counts nested messages, including the root message at depth one. + MaxDepth int + // MaxDecodedBytes bounds the estimated bytes the decoder allocates: the Go + // struct of each message, copied strings and bytes, oneof wrappers, and + // slice and map growth. Zero leaves it unchecked. + MaxDecodedBytes int +} + +// Validate checks data against the message descriptor without constructing the +// decoded message. Known extensions use protoregistry.GlobalTypes, matching the +// default protobuf decoder. Legacy protobuf groups are deliberately unsupported. +func (b Proto) Validate(data []byte, desc protoreflect.MessageDescriptor) error { + _, err := b.Estimate(data, desc) + return err +} + +// Estimate validates data like Validate and returns the estimated bytes the +// binary decoder allocates for it, including the root message. +func (b Proto) Estimate(data []byte, desc protoreflect.MessageDescriptor) (int, error) { + t, err := b.newTally(desc) + if err != nil { + return 0, err + } + if err := t.alloc(protoStructSize(desc)); err != nil { + return t.total(), err + } + err = (&wireWalk{t}).message(data, desc, 1) + return t.total(), err +} + +func (b Proto) newTally(desc protoreflect.MessageDescriptor) (*tally, error) { + if b.MaxFields <= 0 || b.MaxDepth <= 0 || b.MaxDecodedBytes < 0 { + return nil, errors.New("protobuf budget limits must be positive") + } + if desc == nil { + return nil, errors.New("protobuf descriptor must not be nil") + } + return &tally{budget: b}, nil +} + +// tally is the running field count and allocation estimate shared by the wire +// and protojson walks. +type tally struct { + budget Proto + fields int + bytes int +} + +func (t *tally) count() error { + t.fields++ + if t.fields > t.budget.MaxFields { + return ErrMessageComplexity + } + return nil +} + +func (t *tally) total() int { + return withMargin(t.bytes) +} + +func (t *tally) alloc(n int) error { + t.bytes += n + if t.budget.MaxDecodedBytes > 0 && t.total() > t.budget.MaxDecodedBytes { + return ErrMessageComplexity + } + return nil +} + +// Unmarshal validates the budget before decoding into a non-nil mutable message. +// It discards unknown fields, uses MaxDepth as the decoder recursion limit, and +// retains the binary decoder's required-field checks. Preflight rejection leaves +// the destination unchanged; other decode errors follow normal protobuf semantics. +func (b Proto) Unmarshal(data []byte, msg proto.Message) error { + if msg == nil || (reflect.ValueOf(msg).Kind() == reflect.Pointer && reflect.ValueOf(msg).IsNil()) { + return errors.New("protobuf message must be non-nil and mutable") + } + reflected := msg.ProtoReflect() + if !reflected.IsValid() { + return errors.New("protobuf message must be non-nil and mutable") + } + if err := b.Validate(data, reflected.Descriptor()); err != nil { + return err + } + return (proto.UnmarshalOptions{DiscardUnknown: true, RecursionLimit: b.MaxDepth}).Unmarshal(data, msg) +} diff --git a/utils/decodebudget/budget_test.go b/utils/decodebudget/budget_test.go new file mode 100644 index 000000000..18aac069e --- /dev/null +++ b/utils/decodebudget/budget_test.go @@ -0,0 +1,34 @@ +package decodebudget + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/require" + "google.golang.org/protobuf/proto" + + "github.com/livekit/protocol/livekit" + "github.com/livekit/protocol/utils/protojson" +) + +func TestProtoDecodedBytes(t *testing.T) { + input := &livekit.ParticipantInfo{Identity: strings.Repeat("x", 1000), Attributes: map[string]string{"a": "one"}} + data, err := proto.Marshal(input) + require.NoError(t, err) + desc := input.ProtoReflect().Descriptor() + + unlimited := Proto{MaxFields: 100, MaxDepth: 4} + estimate, err := unlimited.Estimate(data, desc) + require.NoError(t, err) + require.Greater(t, estimate, 1000) + + require.NoError(t, (Proto{MaxFields: 100, MaxDepth: 4, MaxDecodedBytes: estimate}).Validate(data, desc)) + require.ErrorIs(t, (Proto{MaxFields: 100, MaxDepth: 4, MaxDecodedBytes: estimate - 1}).Validate(data, desc), ErrMessageComplexity) + require.Error(t, (Proto{MaxFields: 100, MaxDepth: 4, MaxDecodedBytes: -1}).Validate(data, desc)) + + jsonData, err := protojson.Marshal(input) + require.NoError(t, err) + jsonEstimate, err := unlimited.EstimateJSON(jsonData, desc) + require.NoError(t, err) + require.ErrorIs(t, (Proto{MaxFields: 100, MaxDepth: 4, MaxDecodedBytes: jsonEstimate - 1}).ValidateJSON(jsonData, desc), ErrMessageComplexity) +} diff --git a/utils/decodebudget/cost_test.go b/utils/decodebudget/cost_test.go new file mode 100644 index 000000000..72c887843 --- /dev/null +++ b/utils/decodebudget/cost_test.go @@ -0,0 +1,13 @@ +package decodebudget + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestAppendAllocBytes(t *testing.T) { + require.Equal(t, 0, appendAllocBytes(0, 8)) + // 8, 16, 32 and 64 byte arrays for caps 1, 2, 4 and 8. + require.Equal(t, 8+16+32+64, appendAllocBytes(5, 8)) +} diff --git a/utils/decodebudget/json.go b/utils/decodebudget/json.go index 7d0adedfc..34583537b 100644 --- a/utils/decodebudget/json.go +++ b/utils/decodebudget/json.go @@ -19,255 +19,59 @@ import ( "encoding/json" "errors" "io" - - "google.golang.org/protobuf/reflect/protoreflect" -) - -// protojson allocates per token and per object on top of the decoded message -// and the string copies charged separately, and its reflection-based list and -// map setters box each scalar element and map entry. -const ( - jsonTokenCost = 8 - jsonListScalarCost = 24 - jsonMapEntryBoxCost = 32 - jsonObjectCost = 8 ) -// ValidateJSON checks protojson input against the budget in protobuf terms: an -// object key or array element is a field, and nesting is message depth. -func (b Proto) ValidateJSON(data []byte, desc protoreflect.MessageDescriptor) error { - _, err := b.EstimateJSON(data, desc) - return err +// JSON limits JSON structure before a downstream decoder constructs +// collections. Callers must separately limit the serialized byte size. Both +// limits must be positive; the zero value does not disable validation. +type JSON struct { + // MaxTokens includes scalar values, object keys, and opening/closing delimiters. + MaxTokens int + // MaxDepth counts nested arrays/objects, including a root container at depth one. + MaxDepth int } -// EstimateJSON validates protojson input like ValidateJSON and returns the -// estimated bytes protojson allocates decoding it. Malformed or mistyped input -// is left for the decoder to reject. -func (b Proto) EstimateJSON(data []byte, desc protoreflect.MessageDescriptor) (int, error) { - pw, err := b.newWalk(desc) - if err != nil { - return 0, err +// Validate accepts exactly one complete JSON value, including scalar roots. +func (b JSON) Validate(data []byte) error { + if b.MaxTokens <= 0 || b.MaxDepth <= 0 { + return errors.New("JSON budget limits must be positive") } dec := json.NewDecoder(bytes.NewReader(data)) dec.UseNumber() - w := &jsonWalk{protoWalk: pw, dec: dec} - - tok, err := w.next() - if err != nil { - return w.total(), err - } - if err := w.alloc(protoStructSize(desc)); err != nil { - return w.total(), err - } - if tok != json.Delim('{') { - err := w.skip(tok, 1) - return w.total(), err - } - if err := w.object(desc, 1); err != nil { - return w.total(), err - } - if _, err := dec.Token(); err != io.EOF { - return w.total(), errors.New("unexpected value after JSON document") - } - return w.total(), nil -} - -type jsonWalk struct { - *protoWalk - dec *json.Decoder -} - -func (w *jsonWalk) next() (json.Token, error) { - tok, err := w.dec.Token() - if err == io.EOF { - return nil, io.ErrUnexpectedEOF - } - if err != nil { - return nil, err - } - cost := jsonTokenCost - if str, ok := tok.(string); ok { - // protojson's tokenizer copies every string, key or value. - cost += allocSize(len(str)) - } - return tok, w.alloc(cost) -} - -// object walks the members of a message whose '{' has been consumed. -func (w *jsonWalk) object(desc protoreflect.MessageDescriptor, depth int) error { - if depth > w.budget.MaxDepth { - return ErrMessageComplexity - } + depth, tokens := 0, 0 + seen, complete := false, false for { - tok, err := w.next() - if err != nil { - return err - } - if tok == json.Delim('}') { - return nil - } - if err := w.count(); err != nil { - return err - } - key, _ := tok.(string) - field := desc.Fields().ByJSONName(key) - if field == nil { - field = desc.Fields().ByTextName(key) - } - tok, err = w.next() - if err != nil { - return err - } - if field == nil { - // Unknown fields and extensions are discarded or rejected by the decoder. - if err := w.skip(tok, depth); err != nil { - return err - } - continue - } - if err := w.field(field, tok, depth); err != nil { - return err - } - } -} - -func (w *jsonWalk) field(field protoreflect.FieldDescriptor, tok json.Token, depth int) error { - switch { - case tok == nil: - return nil - case field.IsMap(): - if tok != json.Delim('{') { - return w.skip(tok, depth) - } - return w.mapEntries(field, depth) - case field.IsList(): - if tok != json.Delim('[') { - return w.skip(tok, depth) - } - for n := 0; ; n++ { - tok, err := w.next() - if err != nil { - return err - } - if tok == json.Delim(']') { - cost := appendAllocBytes(n, goValueSize(field)) - if field.Message() == nil { - cost += n * jsonListScalarCost - } - return w.alloc(cost) - } - if err := w.count(); err != nil { - return err + token, err := dec.Token() + if err == io.EOF { + if !seen || depth != 0 { + return io.ErrUnexpectedEOF } - if err := w.value(field, tok, depth); err != nil { - return err - } - } - default: - return w.value(field, tok, depth) - } -} - -func (w *jsonWalk) mapEntries(field protoreflect.FieldDescriptor, depth int) error { - if depth+1 > w.budget.MaxDepth { - return ErrMessageComplexity - } - for { - tok, err := w.next() - if err != nil { - return err - } - if tok == json.Delim('}') { return nil } - if err := w.count(); err != nil { - return err - } - if err := w.alloc(mapEntryCost(field) + jsonMapEntryBoxCost); err != nil { - return err - } - tok, err = w.next() if err != nil { return err } - if err := w.value(field.MapValue(), tok, depth+1); err != nil { - return err - } - } -} - -// value walks one occurrence of field. Slots, wrappers and map entries were -// charged by the caller only for maps; everything else is charged here. -func (w *jsonWalk) value(field protoreflect.FieldDescriptor, tok json.Token, depth int) error { - if tok == nil { - return nil - } - if !field.IsMap() && field.ContainingMessage().IsMapEntry() { - // Map values carry no slot or wrapper of their own. - if md := field.Message(); md != nil { - return w.message(md, tok, depth) + if complete { + return errors.New("unexpected value after JSON document") } - return w.alloc(jsonScalarCost(field, tok)) - } - cost := fieldAllocCost(field, 0) + jsonScalarCost(field, tok) + presenceCost(field) - if field.Message() != nil { - cost -= protoStructSize(field.Message()) - } - if err := w.alloc(cost); err != nil { - return err - } - if md := field.Message(); md != nil { - return w.message(md, tok, depth) - } - if _, ok := tok.(json.Delim); ok { - return w.skip(tok, depth) - } - return nil -} - -func (w *jsonWalk) message(md protoreflect.MessageDescriptor, tok json.Token, depth int) error { - if err := w.alloc(protoStructSize(md) + jsonObjectCost); err != nil { - return err - } - if tok != json.Delim('{') || md.FullName().Parent() == "google.protobuf" { - // Well-known types have their own JSON forms; walk them generically. - return w.skip(tok, depth+1) - } - return w.object(md, depth+1) -} - -// skip walks a value without a descriptor, counting every token as a field. -func (w *jsonWalk) skip(tok json.Token, depth int) error { - d, ok := tok.(json.Delim) - if !ok || (d != '{' && d != '[') { - return w.count() - } - for open := 1; open > 0; { - if depth+open > w.budget.MaxDepth { + tokens++ + if tokens > b.MaxTokens { return ErrMessageComplexity } - tok, err := w.next() - if err != nil { - return err - } - if err := w.count(); err != nil { - return err + seen = true + if d, ok := token.(json.Delim); ok { + switch d { + case '{', '[': + depth++ + case '}', ']': + depth-- + } + if depth > b.MaxDepth { + return ErrMessageComplexity + } } - switch tok { - case json.Delim('{'), json.Delim('['): - open++ - case json.Delim('}'), json.Delim(']'): - open-- + if depth == 0 { + complete = true } } - return nil -} - -// jsonScalarCost is what decoding a scalar token allocates beyond the copy the -// tokenizer already made: base64 bytes are decoded into a new buffer. -func jsonScalarCost(field protoreflect.FieldDescriptor, tok json.Token) int { - if str, ok := tok.(string); ok && field.Kind() == protoreflect.BytesKind { - return allocSize(len(str) * 3 / 4) - } - return 0 } diff --git a/utils/decodebudget/json_test.go b/utils/decodebudget/json_test.go new file mode 100644 index 000000000..6a8f4bd58 --- /dev/null +++ b/utils/decodebudget/json_test.go @@ -0,0 +1,28 @@ +package decodebudget + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestJSON(t *testing.T) { + budget := JSON{MaxTokens: 20, MaxDepth: 3} + for _, data := range []string{`null`, `true`, `false`, `1e100000`, `"value"`, `{}`, `[]`, `{"name":"camera","layers":[{}]}`} { + t.Run(data, func(t *testing.T) { require.NoError(t, budget.Validate([]byte(data))) }) + } + for _, data := range []string{"", " \n\t", `{} {}`, `true false`, `null 0`, `{"a":`, `[1,]`, `]`} { + t.Run("invalid "+data, func(t *testing.T) { require.Error(t, budget.Validate([]byte(data))) }) + } + require.NoError(t, (JSON{MaxTokens: 3, MaxDepth: 1}).Validate([]byte(`[0]`))) + require.ErrorIs(t, (JSON{MaxTokens: 2, MaxDepth: 1}).Validate([]byte(`[0]`)), ErrMessageComplexity) + require.NoError(t, (JSON{MaxTokens: 4, MaxDepth: 1}).Validate([]byte(`{"a":0}`))) + require.ErrorIs(t, (JSON{MaxTokens: 3, MaxDepth: 1}).Validate([]byte(`{"a":0}`)), ErrMessageComplexity) + require.NoError(t, budget.Validate([]byte(`[[[]]]`))) + require.ErrorIs(t, budget.Validate([]byte(`[[[[]]]]`)), ErrMessageComplexity) + require.ErrorIs(t, budget.Validate([]byte("["+strings.Repeat("{},", 20)+"{}]")), ErrMessageComplexity) + for _, invalid := range []JSON{{}, {MaxTokens: 1}, {MaxDepth: 1}, {MaxTokens: -1, MaxDepth: 1}} { + require.Error(t, invalid.Validate([]byte(`null`))) + } +} diff --git a/utils/decodebudget/protojson.go b/utils/decodebudget/protojson.go new file mode 100644 index 000000000..f8c8946b0 --- /dev/null +++ b/utils/decodebudget/protojson.go @@ -0,0 +1,274 @@ +// Copyright 2026 LiveKit, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package decodebudget + +import ( + "bytes" + "encoding/json" + "errors" + "io" + + "google.golang.org/protobuf/reflect/protoreflect" +) + +// protojson allocates per token and per object on top of the decoded message +// and the string copies charged separately, and its reflection-based list and +// map setters box each scalar element and map entry. +const ( + jsonTokenCost = 8 + jsonListScalarCost = 24 + jsonMapEntryBoxCost = 32 + jsonObjectCost = 8 +) + +// ValidateJSON checks protojson input against the budget in protobuf terms: an +// object key or array element is a field, and nesting is message depth. +func (b Proto) ValidateJSON(data []byte, desc protoreflect.MessageDescriptor) error { + _, err := b.EstimateJSON(data, desc) + return err +} + +// EstimateJSON validates protojson input like ValidateJSON and returns the +// estimated bytes protojson allocates decoding it. Malformed or mistyped input +// is left for the decoder to reject. +func (b Proto) EstimateJSON(data []byte, desc protoreflect.MessageDescriptor) (int, error) { + t, err := b.newTally(desc) + if err != nil { + return 0, err + } + dec := json.NewDecoder(bytes.NewReader(data)) + dec.UseNumber() + w := &jsonWalk{tally: t, dec: dec} + + tok, err := w.next() + if err != nil { + return w.total(), err + } + if err := w.alloc(protoStructSize(desc)); err != nil { + return w.total(), err + } + if tok != json.Delim('{') { + err := w.skip(tok, 1) + return w.total(), err + } + if err := w.object(desc, 1); err != nil { + return w.total(), err + } + if _, err := dec.Token(); err != io.EOF { + return w.total(), errors.New("unexpected value after JSON document") + } + return w.total(), nil +} + +// jsonWalk walks protojson input with encoding/json tokens. +type jsonWalk struct { + *tally + dec *json.Decoder +} + +func (w *jsonWalk) next() (json.Token, error) { + tok, err := w.dec.Token() + if err == io.EOF { + return nil, io.ErrUnexpectedEOF + } + if err != nil { + return nil, err + } + cost := jsonTokenCost + if str, ok := tok.(string); ok { + // protojson's tokenizer copies every string, key or value. + cost += allocSize(len(str)) + } + return tok, w.alloc(cost) +} + +// object walks the members of a message whose '{' has been consumed. +func (w *jsonWalk) object(desc protoreflect.MessageDescriptor, depth int) error { + if depth > w.budget.MaxDepth { + return ErrMessageComplexity + } + for { + tok, err := w.next() + if err != nil { + return err + } + if tok == json.Delim('}') { + return nil + } + if err := w.count(); err != nil { + return err + } + key, _ := tok.(string) + field := desc.Fields().ByJSONName(key) + if field == nil { + field = desc.Fields().ByTextName(key) + } + tok, err = w.next() + if err != nil { + return err + } + if field == nil { + // Unknown fields and extensions are discarded or rejected by the decoder. + if err := w.skip(tok, depth); err != nil { + return err + } + continue + } + if err := w.field(field, tok, depth); err != nil { + return err + } + } +} + +func (w *jsonWalk) field(field protoreflect.FieldDescriptor, tok json.Token, depth int) error { + switch { + case tok == nil: + return nil + case field.IsMap(): + if tok != json.Delim('{') { + return w.skip(tok, depth) + } + return w.mapEntries(field, depth) + case field.IsList(): + if tok != json.Delim('[') { + return w.skip(tok, depth) + } + for n := 0; ; n++ { + tok, err := w.next() + if err != nil { + return err + } + if tok == json.Delim(']') { + cost := appendAllocBytes(n, goValueSize(field)) + if field.Message() == nil { + cost += n * jsonListScalarCost + } + return w.alloc(cost) + } + if err := w.count(); err != nil { + return err + } + if err := w.value(field, tok, depth); err != nil { + return err + } + } + default: + return w.value(field, tok, depth) + } +} + +func (w *jsonWalk) mapEntries(field protoreflect.FieldDescriptor, depth int) error { + if depth+1 > w.budget.MaxDepth { + return ErrMessageComplexity + } + for { + tok, err := w.next() + if err != nil { + return err + } + if tok == json.Delim('}') { + return nil + } + if err := w.count(); err != nil { + return err + } + if err := w.alloc(mapEntryCost(field) + jsonMapEntryBoxCost); err != nil { + return err + } + tok, err = w.next() + if err != nil { + return err + } + if err := w.value(field.MapValue(), tok, depth+1); err != nil { + return err + } + } +} + +// value walks one occurrence of field. Slots, wrappers and map entries were +// charged by the caller only for maps; everything else is charged here. +func (w *jsonWalk) value(field protoreflect.FieldDescriptor, tok json.Token, depth int) error { + if tok == nil { + return nil + } + if !field.IsMap() && field.ContainingMessage().IsMapEntry() { + // Map values carry no slot or wrapper of their own. + if md := field.Message(); md != nil { + return w.message(md, tok, depth) + } + return w.alloc(jsonScalarCost(field, tok)) + } + cost := fieldAllocCost(field, 0) + jsonScalarCost(field, tok) + presenceCost(field) + if field.Message() != nil { + cost -= protoStructSize(field.Message()) + } + if err := w.alloc(cost); err != nil { + return err + } + if md := field.Message(); md != nil { + return w.message(md, tok, depth) + } + if _, ok := tok.(json.Delim); ok { + return w.skip(tok, depth) + } + return nil +} + +func (w *jsonWalk) message(md protoreflect.MessageDescriptor, tok json.Token, depth int) error { + if err := w.alloc(protoStructSize(md) + jsonObjectCost); err != nil { + return err + } + if tok != json.Delim('{') || md.FullName().Parent() == "google.protobuf" { + // Well-known types have their own JSON forms; walk them generically. + return w.skip(tok, depth+1) + } + return w.object(md, depth+1) +} + +// skip walks a value without a descriptor, counting every token as a field. +func (w *jsonWalk) skip(tok json.Token, depth int) error { + d, ok := tok.(json.Delim) + if !ok || (d != '{' && d != '[') { + return w.count() + } + for open := 1; open > 0; { + if depth+open > w.budget.MaxDepth { + return ErrMessageComplexity + } + tok, err := w.next() + if err != nil { + return err + } + if err := w.count(); err != nil { + return err + } + switch tok { + case json.Delim('{'), json.Delim('['): + open++ + case json.Delim('}'), json.Delim(']'): + open-- + } + } + return nil +} + +// jsonScalarCost is what decoding a scalar token allocates beyond the copy the +// tokenizer already made: base64 bytes are decoded into a new buffer. +func jsonScalarCost(field protoreflect.FieldDescriptor, tok json.Token) int { + if str, ok := tok.(string); ok && field.Kind() == protoreflect.BytesKind { + return allocSize(len(str) * 3 / 4) + } + return 0 +} diff --git a/utils/decodebudget/protojson_test.go b/utils/decodebudget/protojson_test.go new file mode 100644 index 000000000..22b6908c4 --- /dev/null +++ b/utils/decodebudget/protojson_test.go @@ -0,0 +1,20 @@ +package decodebudget + +import ( + "testing" + + "github.com/stretchr/testify/require" + + "github.com/livekit/protocol/livekit" +) + +func TestProtoJSONStructure(t *testing.T) { + desc := (&livekit.ParticipantInfo{}).ProtoReflect().Descriptor() + budget := Proto{MaxFields: 4, MaxDepth: 2} + require.NoError(t, budget.ValidateJSON([]byte(`{"identity":"a","tracks":[{"sid":"TR"}]}`), desc)) + require.ErrorIs(t, budget.ValidateJSON([]byte(`{"tracks":[{"sid":"a"},{"sid":"b"},{"sid":"c"}]}`), desc), ErrMessageComplexity) + require.ErrorIs(t, budget.ValidateJSON([]byte(`{"tracks":[{"layers":[{}]}]}`), desc), ErrMessageComplexity) + require.ErrorIs(t, (Proto{MaxFields: 4, MaxDepth: 2}).ValidateJSON([]byte(`{"unknown":[[[[1]]]]}`), desc), ErrMessageComplexity) + require.Error(t, budget.ValidateJSON([]byte(`{} {}`), desc)) + require.Error(t, budget.ValidateJSON([]byte(`{"identity":`), desc)) +} diff --git a/utils/decodebudget/proto.go b/utils/decodebudget/wire.go similarity index 52% rename from utils/decodebudget/proto.go rename to utils/decodebudget/wire.go index d5f2b38e1..e0decb653 100644 --- a/utils/decodebudget/proto.go +++ b/utils/decodebudget/wire.go @@ -12,98 +12,19 @@ // See the License for the specific language governing permissions and // limitations under the License. -// Package decodebudget bounds protobuf and protojson input before decoding, -// by field count, nesting depth and estimated decoded allocation. It depends -// only on the protobuf runtime so callers need not import utils. package decodebudget import ( - "bytes" - "encoding/json" "errors" - "io" - "reflect" "google.golang.org/protobuf/encoding/protowire" - "google.golang.org/protobuf/proto" "google.golang.org/protobuf/reflect/protoreflect" "google.golang.org/protobuf/reflect/protoregistry" ) -// ErrMessageComplexity indicates that a message exceeds its structural budget. -var ErrMessageComplexity = errors.New("message exceeds structural complexity limit") - -// Proto limits allocation-driving protobuf structure before decoding. -// Callers must separately limit the serialized byte size. MaxFields and -// MaxDepth must be positive; the zero value does not disable validation. -type Proto struct { - // MaxFields counts every wire-field occurrence and each packed scalar value, - // including occurrences of singular and unknown fields. - MaxFields int - // MaxDepth counts nested messages, including the root message at depth one. - MaxDepth int - // MaxDecodedBytes bounds the estimated bytes the decoder allocates: the Go - // struct of each message, copied strings and bytes, oneof wrappers, and - // slice and map growth. Zero leaves it unchecked. - MaxDecodedBytes int -} - -// Validate checks data against the message descriptor without constructing the -// decoded message. Known extensions use protoregistry.GlobalTypes, matching the -// default protobuf decoder. Legacy protobuf groups are deliberately unsupported. -func (b Proto) Validate(data []byte, desc protoreflect.MessageDescriptor) error { - _, err := b.Estimate(data, desc) - return err -} - -// Estimate validates data like Validate and returns the estimated bytes the -// binary decoder allocates for it, including the root message. -func (b Proto) Estimate(data []byte, desc protoreflect.MessageDescriptor) (int, error) { - w, err := b.newWalk(desc) - if err != nil { - return 0, err - } - if err := w.alloc(protoStructSize(desc)); err != nil { - return w.total(), err - } - err = w.message(data, desc, 1) - return w.total(), err -} - -func (b Proto) newWalk(desc protoreflect.MessageDescriptor) (*protoWalk, error) { - if b.MaxFields <= 0 || b.MaxDepth <= 0 || b.MaxDecodedBytes < 0 { - return nil, errors.New("protobuf budget limits must be positive") - } - if desc == nil { - return nil, errors.New("protobuf descriptor must not be nil") - } - return &protoWalk{budget: b}, nil -} - -type protoWalk struct { - budget Proto - fields int - bytes int -} - -func (w *protoWalk) count() error { - w.fields++ - if w.fields > w.budget.MaxFields { - return ErrMessageComplexity - } - return nil -} - -func (w *protoWalk) total() int { - return withMargin(w.bytes) -} - -func (w *protoWalk) alloc(n int) error { - w.bytes += n - if w.budget.MaxDecodedBytes > 0 && w.total() > w.budget.MaxDecodedBytes { - return ErrMessageComplexity - } - return nil +// wireWalk walks the protobuf binary encoding. +type wireWalk struct { + *tally } // messageState mirrors the decoder state of one message whose cost depends on @@ -131,7 +52,7 @@ func (m *messageState) get(field protoreflect.FieldDescriptor) *fieldState { return &m.fields[len(m.fields)-1] } -func (w *protoWalk) message(data []byte, desc protoreflect.MessageDescriptor, depth int) error { +func (w *wireWalk) message(data []byte, desc protoreflect.MessageDescriptor, depth int) error { if depth > w.budget.MaxDepth { return ErrMessageComplexity } @@ -183,7 +104,7 @@ func (w *protoWalk) message(data []byte, desc protoreflect.MessageDescriptor, de if err = w.alloc(state.occurrence(field, 0)); err == nil { err = w.message(value, field.Message(), depth+1) } - case kind == protowire.BytesType && isPackableProtoField(field): + case kind == protowire.BytesType && isPackable(field): err = w.packed(value, field, &state) default: err = w.alloc(state.occurrence(field, len(value))) @@ -227,7 +148,7 @@ func (m *messageState) unknown(n int) int { // packed walks a packed scalar run; decoders accept one even when the schema // prefers unpacked. Each run reallocates the field's slice at exactly its new // length, so many small runs cost quadratically. -func (w *protoWalk) packed(value []byte, field protoreflect.FieldDescriptor, state *messageState) error { +func (w *wireWalk) packed(value []byte, field protoreflect.FieldDescriptor, state *messageState) error { count := 0 for len(value) > 0 { var consumed int @@ -258,7 +179,7 @@ func (w *protoWalk) packed(value []byte, field protoreflect.FieldDescriptor, sta } func wireTypeMatches(field protoreflect.FieldDescriptor, kind protowire.Type) bool { - if kind == protowire.BytesType && isPackableProtoField(field) { + if kind == protowire.BytesType && isPackable(field) { return true } switch field.Kind() { @@ -275,7 +196,7 @@ func wireTypeMatches(field protoreflect.FieldDescriptor, kind protowire.Type) bo } } -func isPackableProtoField(field protoreflect.FieldDescriptor) bool { +func isPackable(field protoreflect.FieldDescriptor) bool { if field == nil || !field.IsList() { return false } @@ -290,76 +211,3 @@ func isPackableProtoField(field protoreflect.FieldDescriptor) bool { return false } } - -// Unmarshal validates the budget before decoding into a non-nil mutable message. -// It discards unknown fields, uses MaxDepth as the decoder recursion limit, and -// retains the binary decoder's required-field checks. Preflight rejection leaves -// the destination unchanged; other decode errors follow normal protobuf semantics. -func (b Proto) Unmarshal(data []byte, msg proto.Message) error { - if msg == nil || (reflect.ValueOf(msg).Kind() == reflect.Pointer && reflect.ValueOf(msg).IsNil()) { - return errors.New("protobuf message must be non-nil and mutable") - } - reflected := msg.ProtoReflect() - if !reflected.IsValid() { - return errors.New("protobuf message must be non-nil and mutable") - } - if err := b.Validate(data, reflected.Descriptor()); err != nil { - return err - } - return (proto.UnmarshalOptions{DiscardUnknown: true, RecursionLimit: b.MaxDepth}).Unmarshal(data, msg) -} - -// JSON limits JSON structure before a downstream decoder constructs -// collections. Callers must separately limit the serialized byte size. Both -// limits must be positive; the zero value does not disable validation. -type JSON struct { - // MaxTokens includes scalar values, object keys, and opening/closing delimiters. - MaxTokens int - // MaxDepth counts nested arrays/objects, including a root container at depth one. - MaxDepth int -} - -// Validate accepts exactly one complete JSON value, including scalar roots. -func (b JSON) Validate(data []byte) error { - if b.MaxTokens <= 0 || b.MaxDepth <= 0 { - return errors.New("JSON budget limits must be positive") - } - dec := json.NewDecoder(bytes.NewReader(data)) - dec.UseNumber() - depth, tokens := 0, 0 - seen, complete := false, false - for { - token, err := dec.Token() - if err == io.EOF { - if !seen || depth != 0 { - return io.ErrUnexpectedEOF - } - return nil - } - if err != nil { - return err - } - if complete { - return errors.New("unexpected value after JSON document") - } - tokens++ - if tokens > b.MaxTokens { - return ErrMessageComplexity - } - seen = true - if d, ok := token.(json.Delim); ok { - switch d { - case '{', '[': - depth++ - case '}', ']': - depth-- - } - if depth > b.MaxDepth { - return ErrMessageComplexity - } - } - if depth == 0 { - complete = true - } - } -} diff --git a/utils/decodebudget/proto_test.go b/utils/decodebudget/wire_test.go similarity index 77% rename from utils/decodebudget/proto_test.go rename to utils/decodebudget/wire_test.go index 8345d9a31..94ba3a6d1 100644 --- a/utils/decodebudget/proto_test.go +++ b/utils/decodebudget/wire_test.go @@ -2,7 +2,6 @@ package decodebudget import ( "bytes" - "strings" "sync" "testing" "time" @@ -18,7 +17,6 @@ import ( "google.golang.org/protobuf/types/dynamicpb" "github.com/livekit/protocol/livekit" - "github.com/livekit/protocol/utils/protojson" ) func TestProtoFields(t *testing.T) { @@ -175,26 +173,6 @@ func TestProtoInvalidInput(t *testing.T) { require.Error(t, budget.Unmarshal(nil, &descriptorpb.UninterpretedOption_NamePart{})) } -func TestJSON(t *testing.T) { - budget := JSON{MaxTokens: 20, MaxDepth: 3} - for _, data := range []string{`null`, `true`, `false`, `1e100000`, `"value"`, `{}`, `[]`, `{"name":"camera","layers":[{}]}`} { - t.Run(data, func(t *testing.T) { require.NoError(t, budget.Validate([]byte(data))) }) - } - for _, data := range []string{"", " \n\t", `{} {}`, `true false`, `null 0`, `{"a":`, `[1,]`, `]`} { - t.Run("invalid "+data, func(t *testing.T) { require.Error(t, budget.Validate([]byte(data))) }) - } - require.NoError(t, (JSON{MaxTokens: 3, MaxDepth: 1}).Validate([]byte(`[0]`))) - require.ErrorIs(t, (JSON{MaxTokens: 2, MaxDepth: 1}).Validate([]byte(`[0]`)), ErrMessageComplexity) - require.NoError(t, (JSON{MaxTokens: 4, MaxDepth: 1}).Validate([]byte(`{"a":0}`))) - require.ErrorIs(t, (JSON{MaxTokens: 3, MaxDepth: 1}).Validate([]byte(`{"a":0}`)), ErrMessageComplexity) - require.NoError(t, budget.Validate([]byte(`[[[]]]`))) - require.ErrorIs(t, budget.Validate([]byte(`[[[[]]]]`)), ErrMessageComplexity) - require.ErrorIs(t, budget.Validate([]byte("["+strings.Repeat("{},", 20)+"{}]")), ErrMessageComplexity) - for _, invalid := range []JSON{{}, {MaxTokens: 1}, {MaxDepth: 1}, {MaxTokens: -1, MaxDepth: 1}} { - require.Error(t, invalid.Validate([]byte(`null`))) - } -} - var protoTestDescriptors = sync.OnceValues(func() (protoreflect.FileDescriptor, error) { optional := descriptorpb.FieldDescriptorProto_LABEL_OPTIONAL repeated := descriptorpb.FieldDescriptorProto_LABEL_REPEATED @@ -256,45 +234,6 @@ func protoTestFile(t *testing.T) protoreflect.FileDescriptor { return file } -func TestProtoDecodedBytes(t *testing.T) { - input := &livekit.ParticipantInfo{Identity: strings.Repeat("x", 1000), Attributes: map[string]string{"a": "one"}} - data, err := proto.Marshal(input) - require.NoError(t, err) - desc := input.ProtoReflect().Descriptor() - - unlimited := Proto{MaxFields: 100, MaxDepth: 4} - estimate, err := unlimited.Estimate(data, desc) - require.NoError(t, err) - require.Greater(t, estimate, 1000) - - require.NoError(t, (Proto{MaxFields: 100, MaxDepth: 4, MaxDecodedBytes: estimate}).Validate(data, desc)) - require.ErrorIs(t, (Proto{MaxFields: 100, MaxDepth: 4, MaxDecodedBytes: estimate - 1}).Validate(data, desc), ErrMessageComplexity) - require.Error(t, (Proto{MaxFields: 100, MaxDepth: 4, MaxDecodedBytes: -1}).Validate(data, desc)) - - jsonData, err := protojson.Marshal(input) - require.NoError(t, err) - jsonEstimate, err := unlimited.EstimateJSON(jsonData, desc) - require.NoError(t, err) - require.ErrorIs(t, (Proto{MaxFields: 100, MaxDepth: 4, MaxDecodedBytes: jsonEstimate - 1}).ValidateJSON(jsonData, desc), ErrMessageComplexity) -} - -func TestProtoJSONStructure(t *testing.T) { - desc := (&livekit.ParticipantInfo{}).ProtoReflect().Descriptor() - budget := Proto{MaxFields: 4, MaxDepth: 2} - require.NoError(t, budget.ValidateJSON([]byte(`{"identity":"a","tracks":[{"sid":"TR"}]}`), desc)) - require.ErrorIs(t, budget.ValidateJSON([]byte(`{"tracks":[{"sid":"a"},{"sid":"b"},{"sid":"c"}]}`), desc), ErrMessageComplexity) - require.ErrorIs(t, budget.ValidateJSON([]byte(`{"tracks":[{"layers":[{}]}]}`), desc), ErrMessageComplexity) - require.ErrorIs(t, (Proto{MaxFields: 4, MaxDepth: 2}).ValidateJSON([]byte(`{"unknown":[[[[1]]]]}`), desc), ErrMessageComplexity) - require.Error(t, budget.ValidateJSON([]byte(`{} {}`), desc)) - require.Error(t, budget.ValidateJSON([]byte(`{"identity":`), desc)) -} - -func TestAppendAllocBytes(t *testing.T) { - require.Equal(t, 0, appendAllocBytes(0, 8)) - // 8, 16, 32 and 64 byte arrays for caps 1, 2, 4 and 8. - require.Equal(t, 8+16+32+64, appendAllocBytes(5, 8)) -} - func TestProtoRejectsQuadraticPackedRuns(t *testing.T) { // Each packed run reallocates the slice at its new length, so 100k // one-element runs in 1.2 MB of wire allocate ~40 GB and decode for seconds. From adb6ac95363f36ef0ccdde7046ac39e4f8e135a2 Mon Sep 17 00:00:00 2001 From: Paul Wells Date: Sat, 31 Oct 2026 17:29:19 -0700 Subject: [PATCH 4/4] decodebudget: name the protojson methods ValidateProtoJSON/EstimateProtoJSON Keeps them distinct from the plain-JSON budget's Validate. --- .changeset/decodebudget-decoded-bytes.md | 2 +- utils/decodebudget/budget_test.go | 4 ++-- utils/decodebudget/calibration_test.go | 4 ++-- utils/decodebudget/protojson.go | 10 +++++----- utils/decodebudget/protojson_test.go | 12 ++++++------ 5 files changed, 16 insertions(+), 16 deletions(-) diff --git a/.changeset/decodebudget-decoded-bytes.md b/.changeset/decodebudget-decoded-bytes.md index 653ed17f7..2d6a7e192 100644 --- a/.changeset/decodebudget-decoded-bytes.md +++ b/.changeset/decodebudget-decoded-bytes.md @@ -2,4 +2,4 @@ "github.com/livekit/protocol": patch --- -decodebudget: move the protobuf/JSON decoding budgets to `utils/decodebudget` and add a calibrated decoded-allocation estimate (`MaxDecodedBytes`, `Estimate`, `ValidateJSON`) +decodebudget: move the protobuf/JSON decoding budgets to `utils/decodebudget` and add a calibrated decoded-allocation estimate (`MaxDecodedBytes`, `Estimate`, `ValidateProtoJSON`) diff --git a/utils/decodebudget/budget_test.go b/utils/decodebudget/budget_test.go index 18aac069e..c22b9fb9e 100644 --- a/utils/decodebudget/budget_test.go +++ b/utils/decodebudget/budget_test.go @@ -28,7 +28,7 @@ func TestProtoDecodedBytes(t *testing.T) { jsonData, err := protojson.Marshal(input) require.NoError(t, err) - jsonEstimate, err := unlimited.EstimateJSON(jsonData, desc) + jsonEstimate, err := unlimited.EstimateProtoJSON(jsonData, desc) require.NoError(t, err) - require.ErrorIs(t, (Proto{MaxFields: 100, MaxDepth: 4, MaxDecodedBytes: jsonEstimate - 1}).ValidateJSON(jsonData, desc), ErrMessageComplexity) + require.ErrorIs(t, (Proto{MaxFields: 100, MaxDepth: 4, MaxDecodedBytes: jsonEstimate - 1}).ValidateProtoJSON(jsonData, desc), ErrMessageComplexity) } diff --git a/utils/decodebudget/calibration_test.go b/utils/decodebudget/calibration_test.go index eeac90bfe..9b4b9b681 100644 --- a/utils/decodebudget/calibration_test.go +++ b/utils/decodebudget/calibration_test.go @@ -247,7 +247,7 @@ func measureJSONDecode(t *testing.T, data []byte, newMsg func() proto.Message) u return best } -func TestProtoEstimateJSONCalibration(t *testing.T) { +func TestProtoEstimateProtoJSONCalibration(t *testing.T) { emptySpans := `{"resourceSpans":[{"scopeSpans":[{"spans":[` + strings.TrimSuffix(strings.Repeat("{},", 1<<18), ",") + `]}]}]}` emptyValues := `{"resourceLogs":[{"scopeLogs":[{"logRecords":[{"body":{"arrayValue":{"values":[` + strings.TrimSuffix(strings.Repeat("{},", 1<<18), ",") + `]}}}]}]}]}` cases := []struct { @@ -268,7 +268,7 @@ func TestProtoEstimateJSONCalibration(t *testing.T) { budget := Proto{MaxFields: 1 << 30, MaxDepth: 100} for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { - est, err := budget.EstimateJSON(tc.data, tc.newMsg().ProtoReflect().Descriptor()) + est, err := budget.EstimateProtoJSON(tc.data, tc.newMsg().ProtoReflect().Descriptor()) require.NoError(t, err) got := measureJSONDecode(t, tc.data, tc.newMsg) t.Logf("json=%d measured=%d estimate=%d ratio=%.2f", len(tc.data), got, est, float64(est)/float64(got)) diff --git a/utils/decodebudget/protojson.go b/utils/decodebudget/protojson.go index f8c8946b0..557ae8808 100644 --- a/utils/decodebudget/protojson.go +++ b/utils/decodebudget/protojson.go @@ -33,17 +33,17 @@ const ( jsonObjectCost = 8 ) -// ValidateJSON checks protojson input against the budget in protobuf terms: an +// ValidateProtoJSON checks protojson input against the budget in protobuf terms: an // object key or array element is a field, and nesting is message depth. -func (b Proto) ValidateJSON(data []byte, desc protoreflect.MessageDescriptor) error { - _, err := b.EstimateJSON(data, desc) +func (b Proto) ValidateProtoJSON(data []byte, desc protoreflect.MessageDescriptor) error { + _, err := b.EstimateProtoJSON(data, desc) return err } -// EstimateJSON validates protojson input like ValidateJSON and returns the +// EstimateProtoJSON validates protojson input like ValidateProtoJSON and returns the // estimated bytes protojson allocates decoding it. Malformed or mistyped input // is left for the decoder to reject. -func (b Proto) EstimateJSON(data []byte, desc protoreflect.MessageDescriptor) (int, error) { +func (b Proto) EstimateProtoJSON(data []byte, desc protoreflect.MessageDescriptor) (int, error) { t, err := b.newTally(desc) if err != nil { return 0, err diff --git a/utils/decodebudget/protojson_test.go b/utils/decodebudget/protojson_test.go index 22b6908c4..1c1379d05 100644 --- a/utils/decodebudget/protojson_test.go +++ b/utils/decodebudget/protojson_test.go @@ -11,10 +11,10 @@ import ( func TestProtoJSONStructure(t *testing.T) { desc := (&livekit.ParticipantInfo{}).ProtoReflect().Descriptor() budget := Proto{MaxFields: 4, MaxDepth: 2} - require.NoError(t, budget.ValidateJSON([]byte(`{"identity":"a","tracks":[{"sid":"TR"}]}`), desc)) - require.ErrorIs(t, budget.ValidateJSON([]byte(`{"tracks":[{"sid":"a"},{"sid":"b"},{"sid":"c"}]}`), desc), ErrMessageComplexity) - require.ErrorIs(t, budget.ValidateJSON([]byte(`{"tracks":[{"layers":[{}]}]}`), desc), ErrMessageComplexity) - require.ErrorIs(t, (Proto{MaxFields: 4, MaxDepth: 2}).ValidateJSON([]byte(`{"unknown":[[[[1]]]]}`), desc), ErrMessageComplexity) - require.Error(t, budget.ValidateJSON([]byte(`{} {}`), desc)) - require.Error(t, budget.ValidateJSON([]byte(`{"identity":`), desc)) + require.NoError(t, budget.ValidateProtoJSON([]byte(`{"identity":"a","tracks":[{"sid":"TR"}]}`), desc)) + require.ErrorIs(t, budget.ValidateProtoJSON([]byte(`{"tracks":[{"sid":"a"},{"sid":"b"},{"sid":"c"}]}`), desc), ErrMessageComplexity) + require.ErrorIs(t, budget.ValidateProtoJSON([]byte(`{"tracks":[{"layers":[{}]}]}`), desc), ErrMessageComplexity) + require.ErrorIs(t, (Proto{MaxFields: 4, MaxDepth: 2}).ValidateProtoJSON([]byte(`{"unknown":[[[[1]]]]}`), desc), ErrMessageComplexity) + require.Error(t, budget.ValidateProtoJSON([]byte(`{} {}`), desc)) + require.Error(t, budget.ValidateProtoJSON([]byte(`{"identity":`), desc)) }