From bccc7bda29b66a5ed35dee1855c4e76bca7607c0 Mon Sep 17 00:00:00 2001
From: Malte Isberner <malte@redhat.com>
Date: Sun, 7 Nov 2021 17:21:24 +0100
Subject: [PATCH] allow erroring out on invalid raw message

---
 config.go                               | 23 +++++++++++++-----
 misc_tests/jsoniter_raw_message_test.go | 32 +++++++++++++++++++++----
 2 files changed, 45 insertions(+), 10 deletions(-)

diff --git a/config.go b/config.go
index 2adcdc3b..2e2ae6c3 100644
--- a/config.go
+++ b/config.go
@@ -23,6 +23,7 @@ type Config struct {
 	TagKey                        string
 	OnlyTaggedField               bool
 	ValidateJsonRawMessage        bool
+	ErrorOnInvalidJsonRawMessage  bool
 	ObjectFieldMustBeSimpleString bool
 	CaseSensitive                 bool
 }
@@ -53,9 +54,10 @@ var ConfigDefault = Config{
 
 // ConfigCompatibleWithStandardLibrary tries to be 100% compatible with standard library behavior
 var ConfigCompatibleWithStandardLibrary = Config{
-	EscapeHTML:             true,
-	SortMapKeys:            true,
-	ValidateJsonRawMessage: true,
+	EscapeHTML:                   true,
+	SortMapKeys:                  true,
+	ValidateJsonRawMessage:       true,
+	ErrorOnInvalidJsonRawMessage: true,
 }.Froze()
 
 // ConfigFastest marshals float with only 6 digits precision
@@ -158,7 +160,7 @@ func (cfg Config) Froze() API {
 		api.useNumber(decoderExtension)
 	}
 	if cfg.ValidateJsonRawMessage {
-		api.validateJsonRawMessage(encoderExtension)
+		api.validateJsonRawMessage(encoderExtension, cfg.ErrorOnInvalidJsonRawMessage)
 	}
 	api.encoderExtension = encoderExtension
 	api.decoderExtension = decoderExtension
@@ -179,14 +181,23 @@ func (cfg Config) frozeWithCacheReuse(extraExtensions []Extension) *frozenConfig
 	return api
 }
 
-func (cfg *frozenConfig) validateJsonRawMessage(extension EncoderExtension) {
+func (cfg *frozenConfig) validateJsonRawMessage(extension EncoderExtension, errorOnInvalid bool) {
 	encoder := &funcEncoder{func(ptr unsafe.Pointer, stream *Stream) {
 		rawMessage := *(*json.RawMessage)(ptr)
+		if rawMessage == nil {
+			stream.WriteRaw("null")
+			return
+		}
+
 		iter := cfg.BorrowIterator([]byte(rawMessage))
 		defer cfg.ReturnIterator(iter)
 		iter.Read()
 		if iter.Error != nil && iter.Error != io.EOF {
-			stream.WriteRaw("null")
+			if errorOnInvalid {
+				stream.Error = iter.Error
+			} else {
+				stream.WriteRaw("null")
+			}
 		} else {
 			stream.WriteRaw(string(rawMessage))
 		}
diff --git a/misc_tests/jsoniter_raw_message_test.go b/misc_tests/jsoniter_raw_message_test.go
index 86af5efe..5abfff43 100644
--- a/misc_tests/jsoniter_raw_message_test.go
+++ b/misc_tests/jsoniter_raw_message_test.go
@@ -2,10 +2,11 @@ package misc_tests
 
 import (
 	"encoding/json"
-	"github.com/json-iterator/go"
-	"github.com/stretchr/testify/require"
 	"strings"
 	"testing"
+
+	jsoniter "github.com/json-iterator/go"
+	"github.com/stretchr/testify/require"
 )
 
 func Test_jsoniter_RawMessage(t *testing.T) {
@@ -28,7 +29,7 @@ func Test_encode_map_of_jsoniter_raw_message(t *testing.T) {
 	should.Equal(`{"hello":[]}`, output)
 }
 
-func Test_marshal_invalid_json_raw_message(t *testing.T) {
+func Test_marshal_nil_json_raw_message(t *testing.T) {
 	type A struct {
 		Raw json.RawMessage `json:"raw"`
 	}
@@ -42,7 +43,7 @@ func Test_marshal_invalid_json_raw_message(t *testing.T) {
 	should.Nil(aouterr)
 }
 
-func Test_marshal_nil_json_raw_message(t *testing.T) {
+func Test_marshal_invalid_json_raw_message_default(t *testing.T) {
 	type A struct {
 		Nil1 jsoniter.RawMessage `json:"raw1"`
 		Nil2 json.RawMessage     `json:"raw2"`
@@ -61,6 +62,29 @@ func Test_marshal_nil_json_raw_message(t *testing.T) {
 	should.Nil(a.Nil2)
 }
 
+func Test_marshal_invalid_json_raw_message_stdcompat(t *testing.T) {
+	type A struct {
+		Nil1 jsoniter.RawMessage `json:"raw1"`
+		Nil2 json.RawMessage     `json:"raw2"`
+	}
+
+	a := A{}
+	should := require.New(t)
+	aout, aouterr := jsoniter.ConfigCompatibleWithStandardLibrary.Marshal(&a)
+	should.Equal(`{"raw1":null,"raw2":null}`, string(aout))
+	should.Nil(aouterr)
+
+	a.Nil1 = []byte(`Any`)
+	a.Nil2 = []byte(`Any`)
+	aout, aouterr = jsoniter.ConfigCompatibleWithStandardLibrary.Marshal(&a)
+	should.Error(aouterr)
+
+	aout = []byte(`{"raw1":null,"raw2":null}`)
+	should.NoError(jsoniter.ConfigCompatibleWithStandardLibrary.Unmarshal(aout, &a))
+	should.Nil(a.Nil1)
+	should.Nil(a.Nil2)
+}
+
 func Test_raw_message_memory_not_copied_issue(t *testing.T) {
 	jsonStream := `{"name":"xxxxx","bundle_id":"com.zonst.majiang","app_platform":"ios","app_category":"100103", "budget_day":1000,"bidding_min":1,"bidding_max":2,"bidding_type":"CPM", "freq":{"open":true,"type":"day","num":100},"speed":1, "targeting":{"vendor":{"open":true,"list":["zonst"]}, "geo_code":{"open":true,"list":["156110100"]},"app_category":{"open":true,"list":["100101"]}, "day_parting":{"open":true,"list":["100409","100410"]},"device_type":{"open":true,"list":["ipad"]}, "os_version":{"open":true,"list":[10]},"carrier":{"open":true,"list":["mobile"]}, "network":{"open":true,"list":["4G"]}},"url":{"tracking_imp_url":"http://www.baidu.com", "tracking_clk_url":"http://www.baidu.com","jump_url":"http://www.baidu.com","deep_link_url":"http://www.baidu.com"}}`
 	type IteratorObject struct {