diff --git a/.changeset/decodebudget-decoded-bytes.md b/.changeset/decodebudget-decoded-bytes.md new file mode 100644 index 000000000..2d6a7e192 --- /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`, `ValidateProtoJSON`) 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/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..c22b9fb9e --- /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.EstimateProtoJSON(jsonData, desc) + require.NoError(t, err) + 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 new file mode 100644 index 000000000..9b4b9b681 --- /dev/null +++ b/utils/decodebudget/calibration_test.go @@ -0,0 +1,279 @@ +// 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 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 { + 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 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 { + 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{} }}, + {"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 { + 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 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 { + 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{} }}, + {"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 { + t.Run(tc.name, func(t *testing.T) { + 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)) + 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..61e6e506d --- /dev/null +++ b/utils/decodebudget/cost.go @@ -0,0 +1,180 @@ +// 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 +} + +// 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 one at a time. +func appendAllocBytes(n, elem int) int { + total, c := 0, 0 + for l := 0; l < n; l++ { + mem, newCap := appendGrow(l, c, 1, elem) + total += mem + c = newCap + } + 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 + } +} + +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)) + } +} + +// 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 +// 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 depends on earlier occurrences; the caller charges it. + case field.IsExtension(): + cost += extensionCost + case field.ContainingOneof() != nil && !field.ContainingOneof().IsSynthetic(): + cost += oneofWrapperCost(field) + } + return cost +} 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 new file mode 100644 index 000000000..34583537b --- /dev/null +++ b/utils/decodebudget/json.go @@ -0,0 +1,77 @@ +// 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" +) + +// 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/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..557ae8808 --- /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 +) + +// 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) ValidateProtoJSON(data []byte, desc protoreflect.MessageDescriptor) error { + _, err := b.EstimateProtoJSON(data, desc) + return err +} + +// 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) EstimateProtoJSON(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..1c1379d05 --- /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.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)) +} diff --git a/utils/decodebudget/wire.go b/utils/decodebudget/wire.go new file mode 100644 index 000000000..e0decb653 --- /dev/null +++ b/utils/decodebudget/wire.go @@ -0,0 +1,213 @@ +// 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 ( + "errors" + + "google.golang.org/protobuf/encoding/protowire" + "google.golang.org/protobuf/reflect/protoreflect" + "google.golang.org/protobuf/reflect/protoregistry" +) + +// wireWalk walks the protobuf binary encoding. +type wireWalk struct { + *tally +} + +// 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 +} + +type fieldState struct { + field protoreflect.FieldDescriptor + len int + cap int +} + +func (m *messageState) get(field protoreflect.FieldDescriptor) *fieldState { + for i := range m.fields { + if m.fields[i].field == field { + return &m.fields[i] + } + } + m.fields = append(m.fields, fieldState{field: field}) + return &m.fields[len(m.fields)-1] +} + +func (w *wireWalk) message(data []byte, desc protoreflect.MessageDescriptor, depth int) error { + if depth > w.budget.MaxDepth { + return ErrMessageComplexity + } + var state messageState + 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 + } + } + + 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:] + + // 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 + } + 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 && isPackable(field): + err = w.packed(value, field, &state) + default: + err = w.alloc(state.occurrence(field, len(value))) + } + 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 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. Each run reallocates the field's slice at exactly its new +// length, so many small runs cost quadratically. +func (w *wireWalk) packed(value []byte, field protoreflect.FieldDescriptor, state *messageState) 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:] + } + 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 && isPackable(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 isPackable(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 + } +} diff --git a/utils/decodebudget/wire_test.go b/utils/decodebudget/wire_test.go new file mode 100644 index 000000000..94ba3a6d1 --- /dev/null +++ b/utils/decodebudget/wire_test.go @@ -0,0 +1,248 @@ +package decodebudget + +import ( + "bytes" + "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" + "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 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{})) +} + +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 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) +} 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 - } - } -}