fix(arm): fix out of bound in native decoder

bytedance · Jul 24, 2024 · 2a6faf2 · 2a6faf2
1 parent 15dff36
commit 2a6faf2
Show file tree

Hide file tree

Showing 8 changed files with 29,045 additions and 29,118 deletions.
diff --git a/internal/decoder/optdec/native.go b/internal/decoder/optdec/native.go
@@ -1,6 +1,7 @@
 package optdec
 
 import (
+	"fmt"
 	"reflect"
 	"unsafe"
 
@@ -174,8 +175,12 @@ func (p *Parser) JsonBytes() []byte {
 
 var nodeType = rt.UnpackType(reflect.TypeOf(node{}))
 
+//go:inline
+func calMaxNodeCap(jsonSize int) int {
+	return jsonSize / 2 + 2
+}
+
 func (p *Parser) parse() ErrorCode {
-	var offset uintptr
 	// when decode into struct, we should decode number as possible
 	old := p.options
 	if !p.isEface {
@@ -189,15 +194,22 @@ func (p *Parser) parse() ErrorCode {
 		return err
 	}
 
+	// check OoB here
+	println("retyru origin json is ", p.Json)
+	offset := p.nbuf.ncur - p.nbuf.nstart
+	curLen :=  offset / unsafe.Sizeof(node{})
+	if curLen != uintptr(len(p.nodes)) {
+		panic(fmt.Sprintf("current len: %d, real len: %d cap: %d", curLen, len(p.nodes), cap(p.nodes)))
+	}
+
 	// node buf is not enough, continue parse
 	// the maxCap is always meet all valid JSON
-	maxCap := len(p.Json) / 2 + 2 
+	maxCap := calMaxNodeCap(len(p.Json))
 	slice := rt.GoSlice{
 		Ptr: rt.Mallocgc(uintptr(maxCap) * nodeType.Size, nodeType, false),
 		Len: maxCap,
 		Cap: maxCap,
 	}
-	offset = p.nbuf.ncur - p.nbuf.nstart
 	rt.Memmove(unsafe.Pointer(slice.Ptr), unsafe.Pointer(&p.nodes[0]), offset)
 	p.backup = p.nodes
 	p.nodes = *(*[]node)(unsafe.Pointer(&slice))

diff --git a/internal/decoder/optdec/native_test.go b/internal/decoder/optdec/native_test.go
@@ -32,6 +32,18 @@ func TestParseNativeRetryLargeJson(t *testing.T) {
 		assert.Equal(t, int(p.nbuf.stat.number), 1 << 20 + 1)
 	})
 
+	t.Run("ObjectNull", func (t *testing.T)  {
+		data := "{" + strings.Repeat("\"a\":null,", 1 << 20) + "\"a\":null}"
+		p := newParser(data, 0, 0)
+		ecode := p.parse()
+		defer p.free()
+		assert.Equal(t, int(ecode), 0)
+		assert.Equal(t, int(p.Pos()), len(data))
+		assert.Equal(t, int(p.nbuf.stat.object),  1)
+		assert.Equal(t, int(p.nbuf.stat.object_keys), 1 << 20 + 1)
+		assert.Equal(t, int(p.nbuf.stat.max_depth), 1)
+	})
+
 	t.Run("Object2", func (t *testing.T)  {
 		data := "{\"top\": {" + strings.Repeat("\"a\":1,", 1 << 20) + "\"a\":1}, \"final\": true}"
 		p := newParser(data, 0, 0)

diff --git a/internal/native/avx2/parse_with_padding_subr.go b/internal/native/avx2/parse_with_padding_subr.go
diff --git a/internal/native/avx2/parse_with_padding_text_amd64.go b/internal/native/avx2/parse_with_padding_text_amd64.go
diff --git a/internal/native/sse/parse_with_padding_subr.go b/internal/native/sse/parse_with_padding_subr.go
diff --git a/internal/native/sse/parse_with_padding_text_amd64.go b/internal/native/sse/parse_with_padding_text_amd64.go
diff --git a/internal/native/traceback_test.mock_tmpl b/internal/native/traceback_test.mock_tmpl
@@ -299,6 +299,22 @@ func mock_value() {
     F_mock_value()
 }
 
+var F_mock_parse_with_padding func()
+var S_mock_parse_with_padding uintptr
+
+//go:nosplit
+func mock_parse_with_padding() {
+    F_mock_parse_with_padding()
+}
+
+var F_mock_lookup_small_key func()
+var S_mock_lookup_small_key uintptr
+
+//go:nosplit
+func mock_lookup_small_key() {
+    F_mock_lookup_small_key()
+}
+
 func init() {
     Mock()
 }
@@ -326,6 +342,8 @@ func Mock() {
     loader.WrapGoC(_mock_text_vunsigned, _cfunc_vunsigned, []loader.GoC{{"_vunsigned", &S_mock_vunsigned, &F_mock_vunsigned}}, "{{PACKAGE}}", "{{PACKAGE}}/vunsigned.c")
     loader.WrapGoC(_mock_text_vstring, _cfunc_vstring, []loader.GoC{{"_vstring", &S_mock_vstring, &F_mock_vstring}}, "{{PACKAGE}}", "{{PACKAGE}}/vstring.c")
     loader.WrapGoC(_mock_text_value, _cfunc_value, []loader.GoC{{"_value", &S_mock_value, &F_mock_value}}, "{{PACKAGE}}", "{{PACKAGE}}/value.c")
+    loader.WrapGoC(_mock_text_parse_with_padding, _cfunc_parse_with_padding, []loader.GoC{{"_parse_with_padding", &S_mock_parse_with_padding, &F_mock_parse_with_padding}}, "{{PACKAGE}}", "{{PACKAGE}}/parse_with_padding.c")
+    loader.WrapGoC(_mock_text_lookup_small_key, _cfunc_lookup_small_key, []loader.GoC{{"_lookup_small_key", &S_mock_lookup_small_key, &F_mock_lookup_small_key}}, "{{PACKAGE}}", "{{PACKAGE}}/lookup_small_key.c")
 }
 
 func TestTraceback(t *testing.T) {
@@ -352,6 +370,8 @@ func TestTraceback(t *testing.T) {
         mock_vunsigned,
         mock_vstring,
         mock_value,
+        mock_parse_with_padding,
+        mock_lookup_small_key,
     }
     for _, f := range tests {
         f()

diff --git a/native/parse_with_padding.c b/native/parse_with_padding.c
@@ -244,9 +244,9 @@ static always_inline bool top_is_key(node_buf* buf) {
 
 static always_inline bool node_on_null(void* ctx, size_t pos) {
     node_buf* buf = (node_buf*)ctx;
-    node_buf_grow(buf, 1);
     buf->cur->typ = node_pack_type(KNULL, pos);
     buf->cur++;
+    node_buf_grow(buf, 1);
     return true;
 }
 
@@ -1143,6 +1143,8 @@ static always_inline error_code parse(GoParser* slf, reader* rdr, visitor* vis)
             break;
         default: err =  SONIC_INVALID_CHAR;
     }
+
+    check_error();
     check_visit();
     return err;