feat: implement function cloning for arm64

Author: pboyd

README.md

@@ -36,11 +36,11 @@ It's 5:00 PM somewhere
 | Linux | amd64 | Full | |
 | Windows | amd64 | Full | |
 | Darwin (macOS) | amd64 | Full | |
-| Linux | arm64 | Partial | `redefine.Func` and `redefine.Restore` work |
+| Linux | arm64 | Full | |
 | FreeBSD | amd64 | Untested | Compiles but untested |
 | OpenBSD | amd64 | Untested | Compiles but untested |
 | NetBSD | amd64 | Untested | Compiles but untested |
-| Windows | arm64 | Untested | No build environment available |
+| Windows | arm64 | Untested | |
 | Darwin (macOS) | arm64 | Broken | `mprotect` returns EACCES |
 
 ## FAQ

asm_amd64.go

@@ -19,6 +19,9 @@ const (
 	opcodeJMP     = 0xe9 // JMP rel32
 )
 
+// The maximum acceptable distance from the text and data segments.
+const maxCloneDistance = math.MaxInt32
+
 func insertJump(buf []byte, dest uintptr) error {
 	const instructionSize = 5 // 1 byte opcode + 4 byte address
 
@@ -45,12 +48,10 @@ func insertJump(buf []byte, dest uintptr) error {
 }
 
 // relocateFunc copies machine instructions from src into dest translating
-// relative instructions as it goes. dest must be larger than src.
+// relative instructions as it goes. dest must be at least as large as src.
 //
 // The data underlying the slices is assumed to be the same address the code
 // would execute from.
-//
-// The dest slice is returned after being resized.
 func relocateFunc(src, dest []byte) ([]byte, error) {
 	dest = dest[:len(src)]
 

asm_arm64.go

@@ -1,11 +1,38 @@
 package redefine
 
 import (
+	"bytes"
 	"encoding/binary"
+	"encoding/hex"
 	"errors"
+	"fmt"
 	"unsafe"
+
+	"golang.org/x/arch/arm64/arm64asm"
+)
+
+const (
+	// -----------------------------------
+	// | 000101 | ... 26 bit address ... |
+	// -----------------------------------
+	_B = uint32(5 << 26)
+
+	// -----------------------------------
+	// | 100101 | ... 26 bit address ... |
+	// -----------------------------------
+	_BL = uint32(1<<31 | _B)
+
+	// ADR/ADRP is encoded as:
+	// --------------------------------------------------
+	// | P | lo 2 bits | 10000 | hi 19 bits | 5-bit reg |
+	// --------------------------------------------------
+	// Mask for the address:
+	adrAddressMask = uint32(3<<29 | 0x7ffff<<5)
 )
 
+// The maximum acceptable distance from the text and data segments.
+const maxCloneDistance = 128 * 1024 * 1024
+
 func insertJump(buf []byte, dest uintptr) error {
 	if len(buf) < 4 {
 		return errors.New("buffer too small")
@@ -14,11 +41,7 @@ func insertJump(buf []byte, dest uintptr) error {
 	addr := uintptr(unsafe.Pointer(unsafe.SliceData(buf)))
 	offset := int32(dest - addr)
 
-	// Encode the instruction:
-	// -----------------------------------
-	// | 000101 | ... 26 bit address ... |
-	// -----------------------------------
-	inst := (5 << 26) | (uint32(offset>>2) & (1<<26 - 1))
+	inst := _B | (uint32(offset>>2) & (1<<26 - 1))
 	binary.LittleEndian.PutUint32(buf, inst)
 
 	// Pad the rest of the buffer with nulls
@@ -28,3 +51,88 @@ func insertJump(buf []byte, dest uintptr) error {
 
 	return nil
 }
+
+// relocateFunc copies machine instructions from src into dest translating
+// relative instructions as it goes. dest must be at least as large as src.
+//
+// The data underlying the slices is assumed to be the same address the code
+// would execute from.
+func relocateFunc(src, dest []byte) ([]byte, error) {
+	dest = dest[:len(src)]
+	copy(dest, src)
+
+	srcPC := uintptr(unsafe.Pointer(unsafe.SliceData(src)))
+
+	for i := 0; i < len(src); i += 4 {
+		srcPC += 4
+		raw := dest[i : i+4]
+
+		instruction, err := arm64asm.Decode(raw)
+		if err != nil {
+			// Stop if the bad instruction was padding
+			if bytes.Equal(raw, []byte{0, 0, 0, 0}) {
+				break
+			}
+			return nil, fmt.Errorf("decode error at offset %d %v: %w", i, raw, err)
+		}
+
+		for _, arg := range instruction.Args {
+			if _, ok := arg.(arm64asm.PCRel); ok {
+				err = fixPCRelAddress(instruction, srcPC, raw)
+				if err != nil {
+					return nil, err
+				}
+			}
+		}
+	}
+
+	return dest, nil
+}
+
+func fixPCRelAddress(inst arm64asm.Inst, srcPC uintptr, dest []byte) error {
+	destPC := uintptr(unsafe.Pointer(unsafe.SliceData(dest)))
+
+	switch inst.Op {
+	case arm64asm.ADRP:
+		// Get the offset (arm64asm converts it to bytes)
+		oldOffset := int64(inst.Args[1].(arm64asm.PCRel))
+
+		// Calculate the offset in bytes, then divide it by 4096 to get it in pages
+		newOffsetPages := uint32((int64(srcPC)+oldOffset-int64(destPC))>>12) + 1
+
+		encoded := binary.LittleEndian.Uint32(dest) &^ adrAddressMask
+		encoded |= (newOffsetPages & 3) << 29 // Lowest 2 bits to bits 30 and 29
+		encoded |= (newOffsetPages >> 2) << 5 // Highest 19 bits to bits 23 to 5
+		binary.LittleEndian.PutUint32(dest, encoded)
+
+	case arm64asm.BL:
+		oldOffset := int64(inst.Args[0].(arm64asm.PCRel))
+		offset := int64(srcPC) + oldOffset - int64(destPC)
+		binary.LittleEndian.PutUint32(dest, _BL|(uint32(offset>>2)&(1<<26-1)))
+
+	default:
+		// Most PC-relative addresses are local. Go only seems to
+		// generate ADRP and BL that are external to the function.
+	}
+
+	return nil
+}
+
+func disassemble(code []byte) (string, error) {
+	var buf bytes.Buffer
+
+	baseAddr := uintptr(unsafe.Pointer(unsafe.SliceData(code)))
+
+	for i := 0; i < len(code); i += 4 {
+		var asm string
+		instruction, err := arm64asm.Decode(code[i:])
+		if err == nil {
+			asm = instruction.String()
+		} else {
+			asm = "?"
+		}
+		fmt.Fprintf(&buf, "0x%08x\t%-20s\t%s\n", baseAddr+uintptr(i), hex.EncodeToString(code[i:i+4]), asm)
+	}
+
+	return buf.String(), nil
+}

clone.go

@@ -72,10 +72,6 @@ func (a *allocator) init(startSize int) error {
 	return err
 }
 
-// The maximum acceptable distance from the text and data segments.
-// This is the value from amd64, arm64 will be less.
-const maxCloneDistance = math.MaxInt32
-
 // The lowest address to consider for our cloned functions.
 const absMinAddress = 0x100000
 

clone_fallback.go

@@ -1,4 +1,4 @@
-//go:build !amd64
+//go:build !amd64 && !arm64
 
 package redefine
 

clone_amd64.go clone_full.goclone_amd64.go renamed to clone_full.go

@@ -1,4 +1,4 @@
-//go:build amd64
+//go:build amd64 || arm64
 
 package redefine
 
@@ -25,6 +25,8 @@ func _cloneFunc[T any](fn T, originalCode []byte) (*clonedFunc[T], error) {
 		return nil, err
 	}
 
+	cacheflush(newCode)
+
 	//fmt.Println(disassemble(newCode))
 
 	// This seems too complicated. The idea is to take our newly allocated

clone_test.go

@@ -1,4 +1,4 @@
-//go:build amd64
+//go:build amd64 || arm64
 
 package redefine
 

doc.go

@@ -5,8 +5,7 @@
 // is a fun experiment, but do not use it for production code.
 //
 // This project is fundamentally non-portable. OS/Arch support:
-//   - Full support: Linux/amd64, Windows/amd64, Darwin/amd64
-//   - Partial support: Linux/arm64 (redefine.Func and redefine.Restore work)
+//   - Full support: Linux/amd64, Windows/amd64, Darwin/amd64, Linux/arm64
 //   - Might work (untested, but it compiles): FreeBSD/amd64, OpenBSD/amd64, NetBSD/amd64
 //   - Also might work: Windows/arm64 (I lack a working build environment)
 //   - Known broken: Darwin/arm64 (EACCES errors from mprotect)