Add support for enforcing a limit on the size of the marshalled output

config.go

@@ -25,6 +25,13 @@ type Config struct {
 	ValidateJsonRawMessage        bool
 	ObjectFieldMustBeSimpleString bool
 	CaseSensitive                 bool
+
+	// MaxMarshalledBytes limits the maximum size of the output.
+	//
+	// While it guarantees not to return more bytes than MaxMarshalledBytes,
+	// it does not guarantee that the internal buffer will be smaller than MaxMarshalledBytes.
+	// In most cases, the internal buffer may be larger by only a few bytes.
+	MaxMarshalledBytes uint64
 }
 
 // API the public interface of this package.
@@ -80,6 +87,7 @@ type frozenConfig struct {
 	streamPool                    *sync.Pool
 	iteratorPool                  *sync.Pool
 	caseSensitive                 bool
+	maxMarshalledBytes            uint64
 }
 
 func (cfg *frozenConfig) initCache() {
@@ -134,6 +142,7 @@ func (cfg Config) Froze() API {
 		onlyTaggedField:               cfg.OnlyTaggedField,
 		disallowUnknownFields:         cfg.DisallowUnknownFields,
 		caseSensitive:                 cfg.CaseSensitive,
+		maxMarshalledBytes:            cfg.MaxMarshalledBytes,
 	}
 	api.streamPool = &sync.Pool{
 		New: func() interface{} {
@@ -293,9 +302,22 @@ func (cfg *frozenConfig) MarshalToString(v interface{}) (string, error) {
 	return string(stream.Buffer()), nil
 }
 
-func (cfg *frozenConfig) Marshal(v interface{}) ([]byte, error) {
+func (cfg *frozenConfig) Marshal(v interface{}) (_ []byte, err error) {
 	stream := cfg.BorrowStream(nil)
 	defer cfg.ReturnStream(stream)
+
+	defer func() {
+		// See Stream.enforceMaxBytes() for an explanation of this.
+		if r := recover(); r != nil {
+			if limitError, ok := r.(ExceededMaxMarshalledBytesError); ok {
+				err = limitError
+				return
+			}
+
+			panic(r)
+		}
+	}()
+
 	stream.WriteVal(v)
 	if stream.Error != nil {
 		return nil, stream.Error

misc_tests/max_marshalled_size_test.go

@@ -0,0 +1,77 @@
+package misc_tests
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	jsoniter "github.com/json-iterator/go"
+)
+
+func TestMaxMarshalledSize(t *testing.T) {
+	testCases := []interface{}{
+		nil,
+		"",
+		false,
+		123,
+		123.123,
+		[]string{"foo", "bar"},
+		map[string]int{"foo": 123},
+	}
+
+	for _, testCase := range testCases {
+		t.Run(fmt.Sprintf("%#v", testCase), func(t *testing.T) {
+			expectedBytes, err := jsoniter.Marshal(testCase)
+			require.NoError(t, err)
+
+			expectedLength := uint64(len(expectedBytes))
+
+			expectSuccessfulMarshalling := func(t *testing.T, limit uint64) {
+				cfg := jsoniter.Config{
+					MaxMarshalledBytes: limit,
+				}
+
+				api := cfg.Froze()
+				actualBytes, err := api.Marshal(testCase)
+				require.NoError(t, err)
+				require.Equal(t, expectedBytes, actualBytes)
+			}
+
+			expectFailedMarshalling := func(t *testing.T, limit uint64) {
+				cfg := jsoniter.Config{
+					MaxMarshalledBytes: limit,
+				}
+
+				api := cfg.Froze()
+				actualBytes, err := api.Marshal(testCase)
+				require.ErrorContains(t, err, fmt.Sprintf("marshalling produced a result over the configured limit of %d bytes", limit))
+				require.Nil(t, actualBytes)
+			}
+
+			t.Run("limit set to 0 (unlimited)", func(t *testing.T) {
+				expectSuccessfulMarshalling(t, 0)
+			})
+
+			t.Run("limit set to exact length of output", func(t *testing.T) {
+				expectSuccessfulMarshalling(t, expectedLength)
+			})
+
+			t.Run("limit set to just under length of output", func(t *testing.T) {
+				expectFailedMarshalling(t, expectedLength-1)
+			})
+
+			t.Run("limit set to well under length of output", func(t *testing.T) {
+				expectFailedMarshalling(t, 1)
+			})
+
+			t.Run("limit set to just over length of output", func(t *testing.T) {
+				expectSuccessfulMarshalling(t, expectedLength+1)
+			})
+
+			t.Run("limit set to well over length of output", func(t *testing.T) {
+				expectSuccessfulMarshalling(t, expectedLength+100)
+			})
+		})
+	}
+}

reflect_native.go

@@ -444,6 +444,7 @@ func (codec *base64Codec) Encode(ptr unsafe.Pointer, stream *Stream) {
 		buf := make([]byte, size)
 		encoding.Encode(buf, src)
 		stream.buf = append(stream.buf, buf...)
+		stream.enforceMaxBytes()
 	}
 	stream.writeByte('"')
 }

stream.go

@@ -1,31 +1,40 @@
 package jsoniter
 
 import (
+	"fmt"
 	"io"
 )
 
 // stream is a io.Writer like object, with JSON specific write functions.
 // Error is not returned as return value, but stored as Error member on this stream instance.
 type Stream struct {
-	cfg        *frozenConfig
-	out        io.Writer
-	buf        []byte
-	Error      error
-	indention  int
-	Attachment interface{} // open for customized encoder
+	cfg                         *frozenConfig
+	out                         io.Writer
+	buf                         []byte
+	Error                       error
+	indention                   int
+	Attachment                  interface{} // open for customized encoder
+	enforceMarshalledBytesLimit bool
+
+	// Number of bytes remaining before marshalled size exceeds cfg.maxMarshalledBytes.
+	// This is tracked as an amount remaining to account for bytes already flushed in Write().
+	marshalledBytesLimitRemaining uint64
 }
 
 // NewStream create new stream instance.
 // cfg can be jsoniter.ConfigDefault.
 // out can be nil if write to internal buffer.
 // bufSize is the initial size for the internal buffer in bytes.
 func NewStream(cfg API, out io.Writer, bufSize int) *Stream {
+	config := cfg.(*frozenConfig)
 	return &Stream{
-		cfg:       cfg.(*frozenConfig),
-		out:       out,
-		buf:       make([]byte, 0, bufSize),
-		Error:     nil,
-		indention: 0,
+		cfg:                           config,
+		out:                           out,
+		buf:                           make([]byte, 0, bufSize),
+		Error:                         nil,
+		indention:                     0,
+		enforceMarshalledBytesLimit:   config.maxMarshalledBytes > 0,
+		marshalledBytesLimitRemaining: config.maxMarshalledBytes,
 	}
 }
 
@@ -38,6 +47,7 @@ func (stream *Stream) Pool() StreamPool {
 func (stream *Stream) Reset(out io.Writer) {
 	stream.out = out
 	stream.buf = stream.buf[:0]
+	stream.marshalledBytesLimitRemaining = stream.cfg.maxMarshalledBytes
 }
 
 // Available returns how many bytes are unused in the buffer.
@@ -66,9 +76,12 @@ func (stream *Stream) SetBuffer(buf []byte) {
 // why the write is short.
 func (stream *Stream) Write(p []byte) (nn int, err error) {
 	stream.buf = append(stream.buf, p...)
+	stream.enforceMaxBytes()
+
 	if stream.out != nil {
 		nn, err = stream.out.Write(stream.buf)
 		stream.buf = stream.buf[nn:]
+		stream.marshalledBytesLimitRemaining -= uint64(nn)
 		return
 	}
 	return len(p), nil
@@ -77,22 +90,51 @@ func (stream *Stream) Write(p []byte) (nn int, err error) {
 // WriteByte writes a single byte.
 func (stream *Stream) writeByte(c byte) {
 	stream.buf = append(stream.buf, c)
+	stream.enforceMaxBytes()
 }
 
 func (stream *Stream) writeTwoBytes(c1 byte, c2 byte) {
 	stream.buf = append(stream.buf, c1, c2)
+	stream.enforceMaxBytes()
 }
 
 func (stream *Stream) writeThreeBytes(c1 byte, c2 byte, c3 byte) {
 	stream.buf = append(stream.buf, c1, c2, c3)
+	stream.enforceMaxBytes()
 }
 
 func (stream *Stream) writeFourBytes(c1 byte, c2 byte, c3 byte, c4 byte) {
 	stream.buf = append(stream.buf, c1, c2, c3, c4)
+	stream.enforceMaxBytes()
 }
 
 func (stream *Stream) writeFiveBytes(c1 byte, c2 byte, c3 byte, c4 byte, c5 byte) {
 	stream.buf = append(stream.buf, c1, c2, c3, c4, c5)
+	stream.enforceMaxBytes()
+}
+
+func (stream *Stream) enforceMaxBytes() {
+	if !stream.enforceMarshalledBytesLimit {
+		return
+	}
+
+	if uint64(len(stream.buf)) > stream.marshalledBytesLimitRemaining {
+		// Why do we do this rather than return an error?
+		// Most of the writing methods on Stream do not return an error, and introducing this would be a
+		// breaking change for custom encoders.
+		// Furthermore, nothing checks if the stream has failed until the object has been completely written
+		// so if we don't panic here, we'd continue writing the rest of the object, negating the purpose of
+		// this limit.
+		panic(ExceededMaxMarshalledBytesError{stream.cfg.maxMarshalledBytes})
+	}
+}
+
+type ExceededMaxMarshalledBytesError struct {
+	MaxMarshalledBytes uint64
+}
+
+func (err ExceededMaxMarshalledBytesError) Error() string {
+	return fmt.Sprintf("marshalling produced a result over the configured limit of %d bytes", err.MaxMarshalledBytes)
 }
 
 // Flush writes any buffered data to the underlying io.Writer.
@@ -117,6 +159,7 @@ func (stream *Stream) Flush() error {
 // WriteRaw write string out without quotes, just like []byte
 func (stream *Stream) WriteRaw(s string) {
 	stream.buf = append(stream.buf, s...)
+	stream.enforceMaxBytes()
 }
 
 // WriteNil write null to stream
@@ -207,4 +250,5 @@ func (stream *Stream) writeIndention(delta int) {
 	for i := 0; i < toWrite; i++ {
 		stream.buf = append(stream.buf, ' ')
 	}
+	stream.enforceMaxBytes()
 }

stream_float.go

@@ -35,6 +35,7 @@ func (stream *Stream) WriteFloat32(val float32) {
 			stream.buf = stream.buf[:n-1]
 		}
 	}
+	stream.enforceMaxBytes()
 }
 
 // WriteFloat32Lossy write float32 to stream with ONLY 6 digits precision although much much faster
@@ -92,6 +93,7 @@ func (stream *Stream) WriteFloat64(val float64) {
 			stream.buf = stream.buf[:n-1]
 		}
 	}
+	stream.enforceMaxBytes()
 }
 
 // WriteFloat64Lossy write float64 to stream with ONLY 6 digits precision although much much faster

stream_int.go

@@ -32,6 +32,7 @@ func writeBuf(buf []byte, v uint32) []byte {
 // WriteUint8 write uint8 to stream
 func (stream *Stream) WriteUint8(val uint8) {
 	stream.buf = writeFirstBuf(stream.buf, digits[val])
+	stream.enforceMaxBytes()
 }
 
 // WriteInt8 write int8 to stream
@@ -44,6 +45,7 @@ func (stream *Stream) WriteInt8(nval int8) {
 		val = uint8(nval)
 	}
 	stream.buf = writeFirstBuf(stream.buf, digits[val])
+	stream.enforceMaxBytes()
 }
 
 // WriteUint16 write uint16 to stream
@@ -56,6 +58,7 @@ func (stream *Stream) WriteUint16(val uint16) {
 	r1 := val - q1*1000
 	stream.buf = writeFirstBuf(stream.buf, digits[q1])
 	stream.buf = writeBuf(stream.buf, digits[r1])
+	stream.enforceMaxBytes()
 	return
 }
 
@@ -76,13 +79,15 @@ func (stream *Stream) WriteUint32(val uint32) {
 	q1 := val / 1000
 	if q1 == 0 {
 		stream.buf = writeFirstBuf(stream.buf, digits[val])
+		stream.enforceMaxBytes()
 		return
 	}
 	r1 := val - q1*1000
 	q2 := q1 / 1000
 	if q2 == 0 {
 		stream.buf = writeFirstBuf(stream.buf, digits[q1])
 		stream.buf = writeBuf(stream.buf, digits[r1])
+		stream.enforceMaxBytes()
 		return
 	}
 	r2 := q1 - q2*1000
@@ -96,6 +101,7 @@ func (stream *Stream) WriteUint32(val uint32) {
 	}
 	stream.buf = writeBuf(stream.buf, digits[r2])
 	stream.buf = writeBuf(stream.buf, digits[r1])
+	stream.enforceMaxBytes()
 }
 
 // WriteInt32 write int32 to stream
@@ -115,13 +121,15 @@ func (stream *Stream) WriteUint64(val uint64) {
 	q1 := val / 1000
 	if q1 == 0 {
 		stream.buf = writeFirstBuf(stream.buf, digits[val])
+		stream.enforceMaxBytes()
 		return
 	}
 	r1 := val - q1*1000
 	q2 := q1 / 1000
 	if q2 == 0 {
 		stream.buf = writeFirstBuf(stream.buf, digits[q1])
 		stream.buf = writeBuf(stream.buf, digits[r1])
+		stream.enforceMaxBytes()
 		return
 	}
 	r2 := q1 - q2*1000
@@ -130,6 +138,7 @@ func (stream *Stream) WriteUint64(val uint64) {
 		stream.buf = writeFirstBuf(stream.buf, digits[q2])
 		stream.buf = writeBuf(stream.buf, digits[r2])
 		stream.buf = writeBuf(stream.buf, digits[r1])
+		stream.enforceMaxBytes()
 		return
 	}
 	r3 := q2 - q3*1000
@@ -139,6 +148,7 @@ func (stream *Stream) WriteUint64(val uint64) {
 		stream.buf = writeBuf(stream.buf, digits[r3])
 		stream.buf = writeBuf(stream.buf, digits[r2])
 		stream.buf = writeBuf(stream.buf, digits[r1])
+		stream.enforceMaxBytes()
 		return
 	}
 	r4 := q3 - q4*1000
@@ -149,6 +159,7 @@ func (stream *Stream) WriteUint64(val uint64) {
 		stream.buf = writeBuf(stream.buf, digits[r3])
 		stream.buf = writeBuf(stream.buf, digits[r2])
 		stream.buf = writeBuf(stream.buf, digits[r1])
+		stream.enforceMaxBytes()
 		return
 	}
 	r5 := q4 - q5*1000
@@ -165,6 +176,7 @@ func (stream *Stream) WriteUint64(val uint64) {
 	stream.buf = writeBuf(stream.buf, digits[r3])
 	stream.buf = writeBuf(stream.buf, digits[r2])
 	stream.buf = writeBuf(stream.buf, digits[r1])
+	stream.enforceMaxBytes()
 }
 
 // WriteInt64 write int64 to stream