fix(loader): invalid binary length in stackmap (#754)

bytedance · Feb 26, 2025 · 86fae91 · 86fae91
1 parent c5a58f9
commit 86fae91
Show file tree

Hide file tree

Showing 2 changed files with 67 additions and 201 deletions.
diff --git a/internal/rt/stackmap.go b/internal/rt/stackmap.go
diff --git a/loader/internal/rt/stackmap.go b/loader/internal/rt/stackmap.go
@@ -62,10 +62,13 @@ func (self *Bitmap) AppendMany(n int, bv int) {
     }
 }
 
-// var (
-//     _stackMapLock  = sync.Mutex{}
-//     _stackMapCache = make(map[*StackMap]struct{})
-// )
+func (b *Bitmap) String() string {
+    var buf strings.Builder
+    for _, byteVal := range b.B {
+        fmt.Fprintf(&buf, "%08b ", byteVal)
+    }
+    return fmt.Sprintf("Bits: %s(total %d bits, %d bytes)", buf.String(), b.N, len(b.B))
+}
 
 type BitVec struct {
     N uintptr
@@ -92,27 +95,44 @@ func (self BitVec) String() string {
     )
 }
 
+/*
+reference Golang 1.22.0 code:
+
+```
+args := bitvec.New(int32(maxArgs / int64(types.PtrSize)))
+aoff := objw.Uint32(&argsSymTmp, 0, uint32(len(lv.stackMaps))) // number of bitmaps
+aoff = objw.Uint32(&argsSymTmp, aoff, uint32(args.N))          // number of bits in each bitmap
+
+locals := bitvec.New(int32(maxLocals / int64(types.PtrSize)))
+loff := objw.Uint32(&liveSymTmp, 0, uint32(len(lv.stackMaps))) // number of bitmaps
+loff = objw.Uint32(&liveSymTmp, loff, uint32(locals.N))        // number of bits in each bitmap
+
+for _, live := range lv.stackMaps {
+    args.Clear()
+    locals.Clear()
+
+    lv.pointerMap(live, lv.vars, args, locals)
+
+    aoff = objw.BitVec(&argsSymTmp, aoff, args)
+    loff = objw.BitVec(&liveSymTmp, loff, locals)
+}
+```
+
+*/
+
 type StackMap struct {
+    // number of bitmaps
     N int32
+    // number of bits of each bitmap
     L int32
+    // bitmap1, bitmap2, ... bitmapN
     B [1]byte
 }
 
-// func (self *StackMap) add() {
-//     _stackMapLock.Lock()
-//     _stackMapCache[self] = struct{}{}
-//     _stackMapLock.Unlock()
-// }
-
-func (self *StackMap) Pin() uintptr {
-    // self.add()
-    return uintptr(unsafe.Pointer(self))
-}
-
-func (self *StackMap) Get(i int32) BitVec {
+func (self *StackMap) Get(i int) BitVec {
     return BitVec {
         N: uintptr(self.L),
-        B: unsafe.Pointer(uintptr(unsafe.Pointer(&self.B)) + uintptr(i * ((self.L + 7) >> 3))),
+        B: unsafe.Pointer(uintptr(unsafe.Pointer(&self.B)) + uintptr(i * self.BitmapBytes())),
     }
 }
 
@@ -121,7 +141,7 @@ func (self *StackMap) String() string {
     sb.WriteString("StackMap {")
 
     /* dump every stack map */
-    for i := int32(0); i < self.N; i++ {
+    for i := 0; i < int(self.N); i++ {
         sb.WriteRune('\n')
         sb.WriteString("    " + self.Get(i).String())
     }
@@ -131,11 +151,31 @@ func (self *StackMap) String() string {
     return sb.String()
 }
 
+func (self *StackMap) BitmapLen() int {
+    return int(self.L)
+}
+
+func (self *StackMap) BitmapBytes() int {
+    return int(self.L + 7) >> 3
+}
+
+func (self *StackMap) BitmapNums() int {
+    return int(self.N)
+}
+
+func (self *StackMap) StackMapHeaderSize() int {
+    return int(unsafe.Sizeof(self.L)) + int(unsafe.Sizeof(self.N)) 
+}
+
 func (self *StackMap) MarshalBinary() ([]byte, error) {
-    size := int(self.N) * int(self.L) + int(unsafe.Sizeof(self.L)) + int(unsafe.Sizeof(self.N))
+    size := self.BinaryLen()
     return BytesFrom(unsafe.Pointer(self), size, size), nil
 }
 
+func (self *StackMap) BinaryLen() int {
+    return self.StackMapHeaderSize() + self.BitmapBytes() * self.BitmapNums()
+}
+
 var (
     byteType = UnpackEface(byte(0)).Type
 )
@@ -152,15 +192,22 @@ type StackMapBuilder struct {
     b Bitmap
 }
 
+
 //go:nocheckptr
 func (self *StackMapBuilder) Build() (p *StackMap) {
     nb := len(self.b.B)
-    bm := mallocgc(_StackMapSize + uintptr(nb) - 1, byteType, false)
+    allocatedSize := _StackMapSize + uintptr(nb) - 1
+    bm := mallocgc(allocatedSize, byteType, false)
 
     /* initialize as 1 bitmap of N bits */
     p = (*StackMap)(bm)
     p.N, p.L = 1, int32(self.b.N)
     copy(BytesFrom(unsafe.Pointer(&p.B), nb, nb), self.b.B)
+
+    /* assert length */
+    if allocatedSize < uintptr(p.BinaryLen()) {
+        panic(fmt.Sprintf("stackmap allocation too small: allocated %d, required %d", allocatedSize, p.BinaryLen()))
+    }
     return
 }