Port DOS protections to our protojson fork

internal/protojson/json_unmarshal.go

@@ -34,6 +34,7 @@ import (
 	"strconv"
 	"strings"
 
+	"google.golang.org/protobuf/encoding/protowire"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/reflect/protoreflect"
 	"google.golang.org/protobuf/reflect/protoregistry"
@@ -71,6 +72,10 @@ type UnmarshalOptions struct {
 		protoregistry.MessageTypeResolver
 		protoregistry.ExtensionTypeResolver
 	}
+
+	// RecursionLimit limits how deeply messages may be nested.
+	// If zero, a default limit is applied.
+	RecursionLimit int
 }
 
 // Unmarshal reads the given []byte and populates the given [proto.Message]
@@ -92,6 +97,10 @@ func (o UnmarshalOptions) unmarshal(b []byte, m proto.Message) error {
 		o.Resolver = protoregistry.GlobalTypes
 	}
 
+	if o.RecursionLimit == 0 {
+		o.RecursionLimit = protowire.DefaultRecursionLimit
+	}
+
 	dec := decoder{json.NewDecoder(b), o}
 	if err := dec.unmarshalMessage(m.ProtoReflect(), false); err != nil {
 		return err
@@ -145,6 +154,11 @@ func (d decoder) syntaxError(pos int, f string, x ...interface{}) error {
 
 // unmarshalMessage unmarshals a message into the given protoreflect.Message.
 func (d decoder) unmarshalMessage(m protoreflect.Message, skipTypeURL bool) error {
+	d.opts.RecursionLimit--
+	if d.opts.RecursionLimit < 0 {
+		return errors.New("exceeded max recursion depth")
+	}
+
 	if jsu, ok := m.Interface().(ProtoJSONMaybeUnmarshaler); ok {
 		if handled, err := jsu.MaybeUnmarshalProtoJSON(d.opts.Metadata, d.Decoder); handled || err != nil {
 			return err

internal/protojson/json_unmarshal_test.go

@@ -42,6 +42,7 @@ import (
 	"google.golang.org/protobuf/types/known/wrapperspb"
 
 	"go.temporal.io/api/internal/protojson"
+	testpb "go.temporal.io/api/internal/protojson/testprotos/test"
 	pb2 "go.temporal.io/api/internal/protojson/testprotos/textpb2"
 	pb3 "go.temporal.io/api/internal/protojson/testprotos/textpb3"
 )
@@ -2421,6 +2422,87 @@ func TestUnmarshal(t *testing.T) {
 				10: 101,
 			},
 		},
+	}, {
+		desc:         "just at recursion limit: nested messages",
+		inputMessage: &testpb.TestAllTypes{},
+		inputText:    `{"optionalNestedMessage":{"corecursive":{"optionalNestedMessage":{"corecursive":{}}}}}`,
+		umo:          protojson.UnmarshalOptions{RecursionLimit: 5},
+	}, {
+		desc:         "exceed recursion limit: nested messages",
+		inputMessage: &testpb.TestAllTypes{},
+		inputText:    `{"optionalNestedMessage":{"corecursive":{"optionalNestedMessage":{"corecursive":{"optionalNestedMessage":{}}}}}}`,
+		umo:          protojson.UnmarshalOptions{RecursionLimit: 5},
+		wantErr:      "exceeded max recursion depth",
+	}, {
+
+		desc:         "just at recursion limit: maps",
+		inputMessage: &testpb.TestAllTypes{},
+		inputText:    `{"mapStringNestedMessage":{"key1":{"corecursive":{"mapStringNestedMessage":{}}}}}`,
+		umo:          protojson.UnmarshalOptions{RecursionLimit: 3},
+	}, {
+		desc:         "exceed recursion limit: maps",
+		inputMessage: &testpb.TestAllTypes{},
+		inputText:    `{"mapStringNestedMessage":{"key1":{"corecursive":{"mapStringNestedMessage":{}}}}}`,
+		umo:          protojson.UnmarshalOptions{RecursionLimit: 2},
+		wantErr:      "exceeded max recursion depth",
+	}, {
+		desc:         "just at recursion limit: arrays",
+		inputMessage: &testpb.TestAllTypes{},
+		inputText:    `{"repeatedNestedMessage":[{"corecursive":{"repeatedInt32":[1,2,3]}}]}`,
+		umo:          protojson.UnmarshalOptions{RecursionLimit: 3},
+	}, {
+		desc:         "exceed recursion limit: arrays",
+		inputMessage: &testpb.TestAllTypes{},
+		inputText:    `{"repeatedNestedMessage":[{"corecursive":{"repeatedNestedMessage":[{}]}}]}`,
+		umo:          protojson.UnmarshalOptions{RecursionLimit: 3},
+		wantErr:      "exceeded max recursion depth",
+	}, {
+		desc:         "just at recursion limit: value",
+		inputMessage: &structpb.Value{},
+		inputText:    `{"a":{"b":{"c":{"d":{}}}}}`,
+		umo:          protojson.UnmarshalOptions{RecursionLimit: 5},
+	}, {
+		desc:         "exceed recursion limit: value",
+		inputMessage: &structpb.Value{},
+		inputText:    `{"a":{"b":{"c":{"d":{"e":[]}}}}}`,
+		umo:          protojson.UnmarshalOptions{RecursionLimit: 5},
+		wantErr:      "exceeded max recursion depth",
+	}, {
+		desc:         "just at recursion limit: list value",
+		inputMessage: &structpb.ListValue{},
+		inputText:    `[[[[[1, 2, 3, 4]]]]]`,
+		// Note: the JSON appears to have recursion of only 5. But it's actually 6 because the
+		// first leaf value (1) is actually a message (google.protobuf.Value), even though the
+		// JSON doesn't use an open brace.
+		umo: protojson.UnmarshalOptions{RecursionLimit: 6},
+	}, {
+		desc:         "exceed recursion limit: list value",
+		inputMessage: &structpb.ListValue{},
+		inputText:    `[[[[[1, 2, 3, 4, ["a", "b"]]]]]]`,
+		umo:          protojson.UnmarshalOptions{RecursionLimit: 6},
+		wantErr:      "exceeded max recursion depth",
+	}, {
+		desc:         "just at recursion limit: struct value",
+		inputMessage: &structpb.Struct{},
+		inputText:    `{"a":{"b":{"c":{"d":{}}}}}`,
+		umo:          protojson.UnmarshalOptions{RecursionLimit: 5},
+	}, {
+		desc:         "exceed recursion limit: struct value",
+		inputMessage: &structpb.Struct{},
+		inputText:    `{"a":{"b":{"c":{"d":{"e":{}]}}}}}`,
+		umo:          protojson.UnmarshalOptions{RecursionLimit: 5},
+		wantErr:      "exceeded max recursion depth",
+	}, {
+		desc:         "just at recursion limit: skip unknown",
+		inputMessage: &testpb.TestAllTypes{},
+		inputText:    `{"foo":{"bar":[{"baz":{}}]}}`,
+		umo:          protojson.UnmarshalOptions{RecursionLimit: 5, DiscardUnknown: true},
+	}, {
+		desc:         "exceed recursion limit: skip unknown",
+		inputMessage: &testpb.TestAllTypes{},
+		inputText:    `{"foo":{"bar":[{"baz":[{}]]}}`,
+		umo:          protojson.UnmarshalOptions{RecursionLimit: 5, DiscardUnknown: true},
+		wantErr:      "exceeded max recursion depth",
 	}}
 
 	for _, tt := range tests {