Support parsing JSON messages up to a maximum depth

parser/json/options.go

@@ -1,7 +1,10 @@
 package json
 
+import "math"
+
 const (
-	defaultMaxDepth = 3
+	// By default the full JSON message is parsed.
+	defaultMaxDepth = math.MaxInt64
 )
 
 // Options provide a set of parsing options.

parser/json/parser.go

@@ -102,8 +102,9 @@ type parser struct {
 	maxDepth int
 	cache    cache
 
-	str string
-	pos int
+	str   string
+	pos   int
+	depth int
 }
 
 // NewParser creates a JSON parser.
@@ -120,6 +121,7 @@ func (p *parser) Reset() {
 	p.cache.reset()
 	p.str = ""
 	p.pos = 0
+	p.depth = 0
 }
 
 func (p *parser) Parse(str string) (*value.Value, error) {
@@ -206,6 +208,16 @@ func (p *parser) parseObject() (*value.Value, error) {
 		return emptyObjectValue, nil
 	}
 
+	if p.depth >= p.maxDepth {
+		// Skip parsing values due to crossing maximum depth threshold.
+		if err := p.skipObjectValue(); err != nil {
+			return nil, err
+		}
+		// NB: If the parser successfully skips the object value without
+		// encountering any errors, return an empty object value.
+		return emptyObjectValue, nil
+	}
+
 	var kvs *value.KVArray
 	for {
 		p.skipWS()
@@ -229,10 +241,12 @@ func (p *parser) parseObject() (*value.Value, error) {
 
 		// Parse out the value.
 		p.skipWS()
+		p.depth++
 		v, err := p.parseValue()
 		if err != nil {
 			return nil, newParseError("object value", p.pos, err)
 		}
+		p.depth--
 
 		// Consume the separator.
 		p.skipWS()
@@ -261,6 +275,46 @@ func (p *parser) parseObject() (*value.Value, error) {
 	}
 }
 
+// Precondition: The parser has seen a '{', and expects to see '}'.
+func (p *parser) skipObjectValue() error {
+	var (
+		inString      bool
+		numLeftParens = 1
+		d             [20]byte // Temporary buffer to absorb escaped bytes
+	)
+	for !p.eos() {
+		switch p.current() {
+		case '\\':
+			if inString {
+				off, _, err := processEscape(p.str[p.pos:], d[:])
+				if err != nil {
+					return newParseError("escape string", p.pos, err)
+				}
+				p.pos += off
+				continue
+			} else {
+				return newParseError("object", p.pos, errors.New("unexpected escape char"))
+			}
+		case '"':
+			inString = !inString
+		case '{':
+			if !inString {
+				numLeftParens++
+			}
+		case '}':
+			if !inString {
+				numLeftParens--
+			}
+			if numLeftParens == 0 {
+				p.pos++
+				return nil
+			}
+		}
+		p.pos++
+	}
+	return newParseError("object", p.pos, errors.New("missing }"))
+}
+
 func (p *parser) parseArray() (*value.Value, error) {
 	p.skipWS()
 	if p.eos() {

parser/json/parser_benchmark_test.go

@@ -178,20 +178,23 @@ func benchmarkParse(b *testing.B, s string) {
 		benchmarkStdJSONParseEmptyStruct(b, s)
 	})
 	b.Run("fastjson", func(b *testing.B) {
-		benchmarkFastJSONParse(b, s)
+		benchmarkFastJSONParse(b, s, testParserPool)
+	})
+	b.Run("fastjson-max-depth", func(b *testing.B) {
+		benchmarkFastJSONParse(b, s, testParserWithMaxDepthPool)
 	})
 	b.Run("fastjson-get", func(b *testing.B) {
 		benchmarkFastJSONParseGet(b, s)
 	})
 }
 
-func benchmarkFastJSONParse(b *testing.B, s string) {
+func benchmarkFastJSONParse(b *testing.B, s string, parsePool *ParserPool) {
 	b.ReportAllocs()
 	b.SetBytes(int64(len(s)))
 	b.ResetTimer()
 
 	b.RunParallel(func(pb *testing.PB) {
-		p := testParserPool.Get()
+		p := parsePool.Get()
 		for pb.Next() {
 			v, err := p.Parse(s)
 			if err != nil {
@@ -201,7 +204,7 @@ func benchmarkFastJSONParse(b *testing.B, s string) {
 				panic(fmt.Errorf("unexpected value type; got %s; want %s", v.Type(), value.ObjectType))
 			}
 		}
-		testParserPool.Put(p)
+		parsePool.Put(p)
 	})
 }
 
@@ -335,10 +338,21 @@ func s2b(s string) []byte {
 	return b
 }
 
-var testParserPool *ParserPool
+const (
+	testDefaultMaxDepth = 1
+)
+
+var (
+	testParserPool             *ParserPool
+	testParserWithMaxDepthPool *ParserPool
+)
 
 func init() {
 	opts := NewParserPoolOptions().SetSize(32)
 	testParserPool = NewParserPool(opts)
 	testParserPool.Init(func() Parser { return NewParser(nil) })
+
+	parserOpts := NewOptions().SetMaxDepth(testDefaultMaxDepth)
+	testParserWithMaxDepthPool = NewParserPool(opts)
+	testParserWithMaxDepthPool.Init(func() Parser { return NewParser(parserOpts) })
 }

parser/json/parser_test.go

@@ -184,6 +184,93 @@ func TestParserParseIncompleteObject(t *testing.T) {
 	}
 }
 
+func TestParserSkipObjectValue(t *testing.T) {
+	inputs := []struct {
+		str string
+		pos int
+	}{
+		{str: `"foo": "bar"}`, pos: 13},
+		{str: `"foo": "bar", "baz": 123}`, pos: 25},
+		{str: `"foo": "ba}r", "baz": 123}`, pos: 26},
+		{str: `"foo": "ba\"}r", "baz": 123}`, pos: 28},
+		{str: `"foo": "ba\\"}, "baz": 123}`, pos: 14},
+		{str: `"foo": ["bar"], "baz": 123}`, pos: 27},
+		{str: `"foo": {"bar": [{"baz": {"cat\"}": 123}}], "da": "ca}r"}}`, pos: 57},
+	}
+
+	for _, input := range inputs {
+		p := NewParser(NewOptions()).(*parser)
+		p.str = input.str
+		require.NoError(t, p.skipObjectValue())
+		require.Equal(t, input.pos, p.pos)
+		require.Equal(t, 0, p.depth)
+	}
+}
+
+func TestParserSkipObjectValueError(t *testing.T) {
+	inputs := []string{
+		`"foo": "bar"`,
+		`"foo\}": "bar"`,
+		`"foo": "bar}"`,
+		`"foo": \"bar"}`,
+	}
+
+	for _, input := range inputs {
+		p := NewParser(NewOptions()).(*parser)
+		p.str = input
+		require.Error(t, p.skipObjectValue())
+	}
+}
+
+func TestParseMaximumDepth(t *testing.T) {
+	input := `
+	{
+		"foo": 123,
+		"bar": [
+			{
+				"baz": {
+					"cat": 456,
+					"car": 789
+				},
+				"dar": ["bbb"]
+			},
+			666
+		],
+		"rad": ["usa"],
+		"pat": {
+			"qat": {
+				"xw": {
+					"woei": "oiwers",
+					"234": "sdflk"
+				},
+				"bw": 123
+			},
+			"tab": {
+				"enter": "return"
+			},
+			"bzr": 123
+		}
+	}
+`
+
+	expected := []string{
+		`{}`,
+		`{"foo":123,"bar":[{},666],"rad":["usa"],"pat":{}}`,
+		`{"foo":123,"bar":[{"baz":{},"dar":["bbb"]},666],"rad":["usa"],"pat":{"qat":{},"tab":{},"bzr":123}}`,
+		`{"foo":123,"bar":[{"baz":{"cat":456,"car":789},"dar":["bbb"]},666],"rad":["usa"],"pat":{"qat":{"xw":{},"bw":123},"tab":{"enter":"return"},"bzr":123}}`,
+		`{"foo":123,"bar":[{"baz":{"cat":456,"car":789},"dar":["bbb"]},666],"rad":["usa"],"pat":{"qat":{"xw":{"woei":"oiwers","234":"sdflk"},"bw":123},"tab":{"enter":"return"},"bzr":123}}`,
+	}
+
+	opts := NewOptions()
+	for i := 0; i < 5; i++ {
+		opts = opts.SetMaxDepth(i)
+		p := NewParser(opts)
+		v, err := p.Parse(input)
+		require.NoError(t, err)
+		require.Equal(t, expected[i], testMarshalled(t, v))
+	}
+}
+
 func TestParserParseEmptyArray(t *testing.T) {
 	p := NewParser(NewOptions())
 	v, err := p.Parse("[]")