fix(jit): no panic when I2I failure of NocopyWriter (#63)

cloudwego · Oct 14, 2024 · 4378967 · 4378967
1 parent a1e494c
commit 4378967
Show file tree

Hide file tree

Showing 6 changed files with 35 additions and 11 deletions.
diff --git a/internal/jit/encoder/bucket.go b/internal/jit/encoder/bucket.go
@@ -33,9 +33,7 @@ func (self _Bucket) String() string {
 	return rt.StringFrom(self.p, int(self.v))
 }
 
-var (
-	bucketPool sync.Pool
-)
+var bucketPool sync.Pool
 
 const (
 	_BucketSize = unsafe.Sizeof(_Bucket{})

diff --git a/internal/jit/encoder/compiler.go b/internal/jit/encoder/compiler.go
@@ -136,6 +136,7 @@ func (self *Program) i64(op OpCode, iv int64) { self.ins(Instr{Op: op, Iv: iv})
 func (self *Program) str(op OpCode, sv string) {
 	self.ins(Instr{Op: op, Iv: int64(len(sv)), Pr: rt.StringPtr(sv)})
 }
+
 func (self *Program) rtt(op OpCode, vt reflect.Type) {
 	self.ins(Instr{Op: op, Pr: unsafe.Pointer(rt.UnpackType(vt))})
 }

diff --git a/internal/jit/encoder/encoder.go b/internal/jit/encoder/encoder.go
@@ -41,9 +41,7 @@ var (
 	TypeCount uint64 = 0
 )
 
-var (
-	programCache = utils.CreateProgramCache()
-)
+var programCache = utils.CreateProgramCache()
 
 func encode(vt *rt.GoType, buf unsafe.Pointer, len int, mem thrift.NocopyWriter, p unsafe.Pointer, rs *RuntimeState, st int) (int, error) {
 	if enc, err := resolve(vt); err != nil {
@@ -65,9 +63,7 @@ func resolve(vt *rt.GoType) (Encoder, error) {
 	/* record the cache miss, and compile the type */
 	atomic.AddUint64(&MissCount, 1)
 	val, err = programCache.Compute(vt, compile)
-
-	/* check for errors */
-	if err != nil {
+	/* check for errors */ if err != nil {
 		return nil, err
 	}
 
@@ -118,6 +114,15 @@ func EncodeObject(buf []byte, mem thrift.NocopyWriter, val interface{}) (ret int
 	efv := rt.UnpackEface(val)
 	out := (*rt.GoSlice)(unsafe.Pointer(&buf))
 
+	if mem == nil {
+		// starting from go1.22,
+		// even though `mem`==nil, it may equal to (0x0, 0xc0000fe1e0).
+		// it keeps original data pointer with itab = nil.
+		// this would cause JIT panic when we only use data pointer to call its methods.
+		// updating `mem` to nil, (0x0, 0xc0000fe1e0) -> (0x0, 0x0), is a quick fix for this case.
+		mem = nil
+	}
+
 	/* check for indirect types */
 	if efv.Type.IsIndirect() {
 		ret, err = encode(efv.Type, out.Ptr, out.Len, mem, efv.Value, rst, 0)

diff --git a/internal/jit/encoder/encoder_test.go b/internal/jit/encoder/encoder_test.go
@@ -21,6 +21,7 @@ import (
 	"encoding/base64"
 	"testing"
 
+	"github.com/cloudwego/gopkg/protocol/thrift"
 	"github.com/davecgh/go-spew/spew"
 	"github.com/stretchr/testify/require"
 )
@@ -64,6 +65,25 @@ func TestEncoder_Encode(t *testing.T) {
 	println("Base64 Encoded Message:", base64.StdEncoding.EncodeToString(mm.Bytes()))
 }
 
+func TestEncoder_Encode_NocopyWriter(t *testing.T) {
+	p := &TranslatorTestStruct{
+		H: make([]byte, 16<<10), // it should trigger NocopyWriter
+	}
+	sz := EncodedSize(p)
+	require.Less(t, 16<<10, sz)
+
+	b := make([]byte, sz)
+
+	var w interface{} = &bytes.Buffer{}
+	// for >=go1.22, nw != (0x0, 0x0)
+	// it keeps data pointer with nil itab like (0x0,0xc0000fe1e0)
+	nw, _ := w.(thrift.NocopyWriter)
+
+	ret, err := EncodeObject(b, nw, p)
+	require.Equal(t, sz, ret)
+	require.NoError(t, err)
+}
+
 type StructSeekTest struct {
 	H StructSeekTestSubStruct `frugal:"0,default,StructSeekTestSubStruct"`
 	O []int8                  `frugal:"1,default,list<i8>"`

diff --git a/internal/jit/encoder/optimizer.go b/internal/jit/encoder/optimizer.go
@@ -216,7 +216,7 @@ func Optimize(p Program) Program {
 	}
 
 	/* sort the blocks by entry point */
-	sort.Slice(ctx.buf, func(i int, j int) bool {
+	sort.Slice(ctx.buf, func(i, j int) bool {
 		return ctx.buf[i].Src < ctx.buf[j].Src
 	})
 

diff --git a/internal/jit/encoder/translator.go b/internal/jit/encoder/translator.go
@@ -539,7 +539,7 @@ func translate_OP_unique_int(p *hir.Builder) {
 	}
 }
 
-func translate_OP_unique_small(p *hir.Builder, nb int64, dv int64, ld func(hir.PointerRegister, int64, hir.GenericRegister) *hir.Ir) {
+func translate_OP_unique_small(p *hir.Builder, nb, dv int64, ld func(hir.PointerRegister, int64, hir.GenericRegister) *hir.Ir) {
 	p.ADDPI(RS, BmOffset, ET)
 	p.BZERO(nb, ET)
 	p.LP(WP, 0, EP)