From d296277d5cb0cf0491048742cede7266d4f464fc Mon Sep 17 00:00:00 2001
From: Ben Brooks <ben@bbrks.me>
Date: Mon, 11 Nov 2019 16:13:57 +0000
Subject: [PATCH] Adds MaxDepth config option

Defaults to 10,000 to match the existing maxDepth constant everywhetre,
except when using `ConfigCompatibleWithStandardLibrary` - which retains
the limitless depth (and causes a stack overflow).

Added tests for the new config, and also up to jsoniter's stack overflow limit.
---
 api_tests/config_test.go | 49 +++++++++++++++++++++++++++++++++++++++-
 config.go                | 10 ++++++++
 iter.go                  |  5 +---
 3 files changed, 59 insertions(+), 5 deletions(-)

diff --git a/api_tests/config_test.go b/api_tests/config_test.go
index 2179e7c..e5eb351 100644
--- a/api_tests/config_test.go
+++ b/api_tests/config_test.go
@@ -2,9 +2,11 @@ package test
 
 import (
 	"encoding/json"
+	"fmt"
+	"strings"
 	"testing"
 
-	"github.com/json-iterator/go"
+	jsoniter "github.com/json-iterator/go"
 	"github.com/stretchr/testify/require"
 )
 
@@ -24,6 +26,51 @@ func Test_customize_float_marshal(t *testing.T) {
 	should.Equal("1.234568", str)
 }
 
+func Test_max_depth(t *testing.T) {
+	deepJSON := func(depth int) []byte {
+		return []byte(strings.Repeat(`[`, depth) + strings.Repeat(`]`, depth))
+	}
+
+	tests := []struct {
+		jsonDepth   int
+		cfgMaxDepth int
+		expectedErr string
+	}{
+		// Test the default depth
+		{jsonDepth: 10000, cfgMaxDepth: 0},
+		{jsonDepth: 10001, cfgMaxDepth: 0, expectedErr: "max depth"},
+		// Test max depth logic
+		{jsonDepth: 5, cfgMaxDepth: 6},
+		{jsonDepth: 5, cfgMaxDepth: 5},
+		{jsonDepth: 5, cfgMaxDepth: 4, expectedErr: "max depth"},
+		// Now try some larger values to figure out the limit
+		{jsonDepth: 128000, cfgMaxDepth: -1},
+		{jsonDepth: 512000, cfgMaxDepth: -1},
+		{jsonDepth: 768000, cfgMaxDepth: -1},
+		{jsonDepth: 860367, cfgMaxDepth: -1}, // largest value for jsoniter without stack overflow
+	}
+
+	for _, test := range tests {
+		t.Run(fmt.Sprintf("jsonDepth:%v_cfgMaxDepth:%v", test.jsonDepth, test.cfgMaxDepth), func(t *testing.T) {
+			if testing.Short() && test.jsonDepth >= 512000 {
+				t.Skip("skipping in -short due to large input data")
+			}
+
+			should := require.New(t)
+			cfg := jsoniter.Config{MaxDepth: test.cfgMaxDepth}.Froze()
+
+			var val interface{}
+			err := cfg.Unmarshal(deepJSON(test.jsonDepth), &val)
+			if test.expectedErr != "" {
+				should.Error(err)
+				should.Contains(err.Error(), test.expectedErr)
+			} else {
+				should.NoError(err)
+			}
+		})
+	}
+}
+
 func Test_customize_tag_key(t *testing.T) {
 
 	type TestObject struct {
diff --git a/config.go b/config.go
index 8c58fcb..94a039f 100644
--- a/config.go
+++ b/config.go
@@ -11,6 +11,9 @@ import (
 	"github.com/modern-go/reflect2"
 )
 
+// limit maximum depth of nesting, as allowed by https://tools.ietf.org/html/rfc7159#section-9
+const defaultMaxDepth = 10000
+
 // Config customize how the API should behave.
 // The API is created from Config by Froze.
 type Config struct {
@@ -25,6 +28,7 @@ type Config struct {
 	ValidateJsonRawMessage        bool
 	ObjectFieldMustBeSimpleString bool
 	CaseSensitive                 bool
+	MaxDepth                      int
 }
 
 // API the public interface of this package.
@@ -56,6 +60,7 @@ var ConfigCompatibleWithStandardLibrary = Config{
 	EscapeHTML:             true,
 	SortMapKeys:            true,
 	ValidateJsonRawMessage: true,
+	MaxDepth:               -1, // encoding/json has no max depth (stack overflow at 2581101)
 }.Froze()
 
 // ConfigFastest marshals float with only 6 digits precision
@@ -80,6 +85,7 @@ type frozenConfig struct {
 	streamPool                    *sync.Pool
 	iteratorPool                  *sync.Pool
 	caseSensitive                 bool
+	maxDepth                      int
 }
 
 func (cfg *frozenConfig) initCache() {
@@ -127,6 +133,9 @@ func addFrozenConfigToCache(cfg Config, frozenConfig *frozenConfig) {
 
 // Froze forge API from config
 func (cfg Config) Froze() API {
+	if cfg.MaxDepth == 0 {
+		cfg.MaxDepth = defaultMaxDepth
+	}
 	api := &frozenConfig{
 		sortMapKeys:                   cfg.SortMapKeys,
 		indentionStep:                 cfg.IndentionStep,
@@ -134,6 +143,7 @@ func (cfg Config) Froze() API {
 		onlyTaggedField:               cfg.OnlyTaggedField,
 		disallowUnknownFields:         cfg.DisallowUnknownFields,
 		caseSensitive:                 cfg.CaseSensitive,
+		maxDepth:                      cfg.MaxDepth,
 	}
 	api.streamPool = &sync.Pool{
 		New: func() interface{} {
diff --git a/iter.go b/iter.go
index 29b31cf..cbd331e 100644
--- a/iter.go
+++ b/iter.go
@@ -327,12 +327,9 @@ func (iter *Iterator) Read() interface{} {
 	}
 }
 
-// limit maximum depth of nesting, as allowed by https://tools.ietf.org/html/rfc7159#section-9
-const maxDepth = 10000
-
 func (iter *Iterator) incrementDepth() (success bool) {
 	iter.depth++
-	if iter.depth <= maxDepth {
+	if iter.depth <= iter.cfg.maxDepth || iter.cfg.maxDepth < 0 {
 		return true
 	}
 	iter.ReportError("incrementDepth", "exceeded max depth")