fix: function cloning on Darwin

- On Darwin, use MAP_JIT when allocating memory for cloned functions and
  call `pthread_jit_write_protect_np()` as needed.
- On arm64, try to reserve memory within 128 MiB of the text/data
  segments, but allow addresses up to 4 GiB away. When a BL address is
  more than 128 MiB away insert a trampoline at the end of the function
  to branch to the fixed address.
- On amd64, the search code still still applies, but we only need to be
  within range of a 32-bit signed int.

Author: pboyd

.github/workflows/test-darwin-arm64.yaml

@@ -0,0 +1,37 @@
+name: Test
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+
+jobs:
+  test:
+    #if: ${{ 0 == 1}} # Disable for now since Darwin on arm64 is still broken
+    name: Test Darwin (go ${{ matrix.go }}, ${{ matrix.host }})
+    runs-on: ${{ matrix.host }}
+    timeout-minutes: 10
+    strategy:
+      matrix:
+        #go: ['1.25', '1.26']
+        go: ['1.26']
+        host: ['macos-latest']
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Setup Go
+        uses: actions/setup-go@v5
+        with:
+          go-version: ${{ matrix.go }}
+          cache: true
+
+      - name: Run tests (buildmode=exe)
+        run: |
+          go test -buildmode=exe -run 'TestCloneFunc'
+
+      - name: Run tests (buildmode=pie)
+        run: |
+          go test -buildmode=pie -run 'TestCloneFunc'

.github/workflows/test-windows-arm.yaml .github/workflows/test-windows-arm64.yaml.github/workflows/test-windows-arm.yaml renamed to .github/workflows/test-windows-arm64.yaml

@@ -19,7 +19,8 @@ const (
 	opcodeJMP     = 0xe9 // JMP rel32
 )
 
-// The maximum acceptable distance from the text and data segments.
+// Cloned functions need to be within range of a signed 32-bit JMP.
+const idealCloneDistance = 0
 const maxCloneDistance = math.MaxInt32
 
 func insertJump(buf []byte, dest uintptr) error {

asm_amd64.go

@@ -22,6 +22,21 @@ const (
 	// -----------------------------------
 	_BL = uint32(1<<31 | _B)
 
+	// ----------------------------------------------
+	// | 1101011000111111000000 | 5-bit reg | 00000 |
+	// ----------------------------------------------
+	_BLR = uint32(0xd63f0000)
+
+	// -----------------------------------------------------------
+	// | 1-bit sf | 10100101 | 2-bit hw | 16-bit imm | 5-bit reg |
+	// -----------------------------------------------------------
+	_MOVZ = uint32(0xd2800000) // sf is 1
+
+	// -----------------------------------------------------------
+	// | 1-bit sf | 11100101 | 2-bit hw | 16-bit imm | 5-bit reg |
+	// -----------------------------------------------------------
+	_MOVK = uint32(0xf2800000) // sf is 1
+
 	// ADR/ADRP is encoded as:
 	// --------------------------------------------------
 	// | P | lo 2 bits | 10000 | hi 19 bits | 5-bit reg |
@@ -30,8 +45,15 @@ const (
 	adrAddressMask = uint32(3<<29 | 0x7ffff<<5)
 )
 
-// The maximum acceptable distance from the text and data segments.
-const maxCloneDistance = 128 * 1024 * 1024
+const scratchRegister = 16
+
+// Ideally, cloned functions will be within 128 MiB of the original function.
+// But it's acceptable to be within the 4 GiB range for ADRP because there's code
+// to generate trampolines for BLs.
+const idealCloneDistance = 0 //128 * 1024 * 1024
+const maxCloneDistance = 4 * 1024 * 1024 * 1024
+
+var errAddressOutOfRange = errors.New("address out of range")
 
 func insertJump(buf []byte, dest uintptr) error {
 	if len(buf) < 4 {
@@ -45,8 +67,7 @@ func insertJump(buf []byte, dest uintptr) error {
 		return fmt.Errorf("B target out of range: %d bytes exceeds 128MiB", offset)
 	}
 
-	inst := _B | (uint32(offset>>2) & (1<<26 - 1))
-	binary.LittleEndian.PutUint32(buf, inst)
+	encodeB(buf, uint32(offset))
 
 	// Pad the rest of the buffer with nulls
 	for i := 4; i < len(buf); i++ {
@@ -62,6 +83,7 @@ func insertJump(buf []byte, dest uintptr) error {
 // The data underlying the slices is assumed to be the same address the code
 // would execute from.
 func relocateFunc(src, dest []byte) ([]byte, error) {
+	src = trimPadding(src)
 	dest = dest[:len(src)]
 	copy(dest, src)
 
@@ -83,7 +105,15 @@ func relocateFunc(src, dest []byte) ([]byte, error) {
 			if _, ok := arg.(arm64asm.PCRel); ok {
 				err = fixPCRelAddress(instruction, srcPC, raw)
 				if err != nil {
-					return nil, err
+					if errors.Is(err, errAddressOutOfRange) && instruction.Op == arm64asm.BL {
+						var trErr error
+						dest, trErr = makeBLTrampoline(instruction, srcPC, dest, i)
+						if trErr != nil {
+							return nil, fmt.Errorf("unable to make trampoline: %w (original error: %w)", trErr, err)
+						}
+					} else {
+						return nil, err
+					}
 				}
 			}
 		}
@@ -93,6 +123,17 @@ func relocateFunc(src, dest []byte) ([]byte, error) {
 	return dest, nil
 }
 
+func trimPadding(buf []byte) []byte {
+	newLen := len(buf)
+	for i := len(buf) - 4; i >= 0; i -= 4 {
+		if bytes.Equal(buf[i:i+4], []byte{0, 0, 0, 0}) {
+			newLen = i
+		}
+	}
+
+	return buf[:newLen]
+}
+
 func fixPCRelAddress(inst arm64asm.Inst, srcPC uintptr, dest []byte) error {
 	destPC := uintptr(unsafe.Pointer(unsafe.SliceData(dest)))
 
@@ -105,12 +146,12 @@ func fixPCRelAddress(inst arm64asm.Inst, srcPC uintptr, dest []byte) error {
 		newOffsetPages := (int64(srcPC&^uintptr(0xfff)) + oldOffset - int64(destPC&^uintptr(0xfff))) >> 12
 
 		if newOffsetPages < -(1<<20) || newOffsetPages >= (1<<20) {
-			return fmt.Errorf("ADRP target out of range: %d pages exceeds 4GiB", newOffsetPages)
+			return fmt.Errorf("%w: ADRP target out of range: %d pages exceeds 4GiB", errAddressOutOfRange, newOffsetPages)
 		}
 
 		p := uint32(newOffsetPages)
 		encoded := binary.LittleEndian.Uint32(dest) &^ adrAddressMask
-		encoded |= (p & 3) << 29 // Lowest 2 bits to bits 30 and 29
+		encoded |= (p & 3) << 29             // Lowest 2 bits to bits 30 and 29
 		encoded |= ((p >> 2) & 0x7ffff) << 5 // Highest 19 bits to bits 23 to 5
 		binary.LittleEndian.PutUint32(dest, encoded)
 
@@ -120,7 +161,7 @@ func fixPCRelAddress(inst arm64asm.Inst, srcPC uintptr, dest []byte) error {
 
 		// BL encodes a 26-bit signed instruction offset.
 		if offset < -(1<<27) || offset >= (1<<27) {
-			return fmt.Errorf("BL target out of range: %d bytes exceeds 128MiB", offset)
+			return fmt.Errorf("%w: BL target out of range: %d bytes exceeds 128MiB", errAddressOutOfRange, offset)
 		}
 
 		binary.LittleEndian.PutUint32(dest, _BL|(uint32(offset>>2)&(1<<26-1)))
@@ -133,6 +174,57 @@ func fixPCRelAddress(inst arm64asm.Inst, srcPC uintptr, dest []byte) error {
 	return nil
 }
 
+func makeBLTrampoline(inst arm64asm.Inst, srcPC uintptr, dest []byte, blOffset int) ([]byte, error) {
+	if cap(dest)-len(dest) < 24 {
+		return nil, errors.New("destination is too small for BL trampoline")
+	}
+	origLen := len(dest)
+	dest = dest[:len(dest)+24]
+
+	//destPC := uintptr(unsafe.Pointer(unsafe.SliceData(dest))) + uintptr(blOffset)
+	blrTarget := uintptr(int64(srcPC) + int64(inst.Args[0].(arm64asm.PCRel)))
+
+	// Encode the trampoline itself.
+	// It uses 4 instructions to store a 64-bit number in x16, then calls BLR x16.
+	trampoline := dest[origLen:]
+	encodeMov(trampoline, true, 0, uint16(blrTarget), scratchRegister)
+	encodeMov(trampoline[4:], false, 16, uint16(blrTarget>>16), scratchRegister)
+	encodeMov(trampoline[8:], false, 32, uint16(blrTarget>>32), scratchRegister)
+	encodeMov(trampoline[12:], false, 48, uint16(blrTarget>>48), scratchRegister)
+	binary.LittleEndian.PutUint32(trampoline[16:], uint32(_BLR|uint32(scratchRegister<<5)))
+
+	// Replace the original BL with a B to the beginning of the trampoline
+	blAddr := uintptr(unsafe.Pointer(unsafe.SliceData(dest))) + uintptr(blOffset)
+	trampolineAddr := uintptr(unsafe.Pointer(unsafe.SliceData(trampoline)))
+	encodeB(dest[blOffset:], uint32(int32(trampolineAddr)-int32(blAddr)))
+
+	// The last instruction in the trampoline needs to jump back to the
+	// instruction after the original BL
+	encodeB(trampoline[20:], uint32(int32(blAddr+4)-int32(trampolineAddr+20)))
+
+	return dest, nil
+}
+
+func encodeB(dest []byte, offset uint32) {
+	inst := _B | (uint32(offset>>2) & 0x3ffffff)
+	binary.LittleEndian.PutUint32(dest, inst)
+}
+
+func encodeMov(dest []byte, zero bool, lsl uint8, imm uint16, register uint8) {
+	var mov uint32
+	if zero {
+		mov = _MOVZ
+	} else {
+		mov = _MOVK
+	}
+
+	mov |= (uint32(lsl>>4) & 3) << 21
+	mov |= uint32(imm) << 5
+	mov |= uint32(register & 0x1f)
+
+	binary.LittleEndian.PutUint32(dest, mov)
+}
+
 func disassemble(code []byte) (string, error) {
 	var buf bytes.Buffer
 

asm_arm64.go

@@ -29,7 +29,7 @@ func cloneFunc[T any](fn T) (*clonedFunc[T], error) {
 	cloneAllocator.BeginMutate()
 	defer cloneAllocator.EndMutate()
 
-	newCode, err := cloneAllocator.Allocate(len(originalCode))
+	newCode, err := cloneAllocator.Allocate(len(originalCode) * 2)
 	if err != nil {
 		return nil, err
 	}
@@ -83,7 +83,7 @@ func (a *allocator) init(startSize int) error {
 		}
 
 		if protBE, ok := be.(malloc.ProtectedArenaBackend); ok {
-			a.mprotect = protBE.Protect
+			a.mprotect = mprotectHook(protBE.Protect)
 		} else {
 			// No real mprotect for some reason. This shouldn't
 			// really happen, but continue with a no-op mprotect.
@@ -126,35 +126,71 @@ func initMallocBackend() (malloc.ArenaBackend, error) {
 	// instructions that the original function used. There's often enough
 	// space right before the text segment but that's not guaranteed
 	// (particularly when buildmode=pie).
+	// The distances are calculated from the furthest address in the
+	// text/data segment so that every address in the reserved block can
+	// reach every address in the target.
+
+	// If there's an ideal range for the architecture, try that first.
+	if idealCloneDistance > 0 {
+		// Search before text
+		minAddress := end - idealCloneDistance
+		if minAddress > end || minAddress < absMinAddress {
+			minAddress = absMinAddress
+		}
+		be := tryBackendRange(size, minAddress, text-pageSize-size)
+		if be != nil {
+			return be, nil
+		}
+
+		// Search after end
+		maxAddress := text + idealCloneDistance - size
+		if maxAddress < text {
+			maxAddress = math.MaxUint
+		}
+		be = tryBackendRange(size, end, maxAddress)
+		if be != nil {
+			return be, nil
+		}
+	}
 
-	// The minimum acceptable address is where the first
-	// instruction in the code segment can still reach the final
-	// address before end. These are unsigned so watch for wrap-around.
+	// Nothing in the ideal range, so search within the acceptable range
 	minAddress := end - maxCloneDistance
 	if minAddress > end || minAddress < absMinAddress {
 		minAddress = absMinAddress
 	}
-	for addr := text - pageSize - size; addr >= minAddress; addr -= 0x100000 {
-		be, err := malloc.VirtBackend(size, malloc.MmapAddr(addr), malloc.MmapProt(mprotectExec), malloc.MmapFlags(_MAP_FIXED_NOREPLACE))
-		if err == nil {
-			return be, nil
-		}
+	be := tryBackendRange(size, minAddress, text-pageSize-size)
+	if be != nil {
+		return be, nil
 	}
 
-	// Nothing was found before the text segment, repeat the process for
-	// the space after end.
 	maxAddress := text + maxCloneDistance - size
 	if maxAddress < text {
 		maxAddress = math.MaxUint
 	}
-	for addr := end; addr <= maxAddress; addr += 0x100000 {
-		be, err := malloc.VirtBackend(size, malloc.MmapAddr(addr), malloc.MmapProt(mprotectExec), malloc.MmapFlags(_MAP_FIXED_NOREPLACE))
+	be = tryBackendRange(size, end, maxAddress)
+	if be != nil {
+		return be, nil
+	}
+
+	// Well, we tried. We tried really hard. There's nothing left to do but
+	// take whatever address the OS gives us.
+	return malloc.VirtBackend(size, malloc.MmapAddr(minAddress), malloc.MmapProt(mprotectExec), malloc.MmapFlags(_MMAP_FLAGS))
+}
+
+func tryBackendRange(size, minAddress, maxAddress uintptr) malloc.ArenaBackend {
+	for addr := minAddress; addr <= maxAddress; addr += 0x100000 {
+		be, err := malloc.VirtBackend(size, malloc.MmapAddr(addr), malloc.MmapProt(mprotectExec), malloc.MmapFlags(_MMAP_FLAGS))
 		if err == nil {
-			return be, nil
+			if be.Addr() < minAddress || be.Addr() > maxAddress {
+				// No good, try again.
+				be.Release()
+			} else {
+				return be
+			}
 		}
 	}
 
-	return nil, errors.New("no suitable virtual memory space found")
+	return nil
 }
 
 func (a *allocator) BeginMutate() error {

clone.go

@@ -0,0 +1,7 @@
+//go:build !(darwin && arm64)
+
+package redefine
+
+func mprotectHook(inner func(int) error) func(int) error {
+	return inner
+}

clone_mprotect.go

@@ -0,0 +1,29 @@
+//go:build darwin && arm64
+
+package redefine
+
+import "golang.org/x/sys/unix"
+
+/*
+	#include <pthread.h>
+*/
+import "C"
+
+func mprotectHook(inner func(int) error) func(int) error {
+	return func(prot int) error {
+		if prot&unix.PROT_WRITE != 0 {
+			C.pthread_jit_write_protect_np(0)
+		} else {
+			C.pthread_jit_write_protect_np(1)
+		}
+
+		err := inner(prot)
+
+		// Restore write protection after an error
+		if err != nil && prot&unix.PROT_WRITE != 0 {
+			C.pthread_jit_write_protect_np(1)
+		}
+
+		return err
+	}
+}

clone_mprotect_darwin.go

@@ -3,15 +3,15 @@ module github.com/pboyd/redefine
 go 1.25.0
 
 require (
-	github.com/pboyd/malloc v1.2.0
+	github.com/pboyd/malloc v1.2.1
 	github.com/stretchr/testify v1.11.1
 	golang.org/x/arch v0.23.0
+	golang.org/x/sys v0.41.0
 )
 
 require (
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	golang.org/x/exp v0.0.0-20260218203240-3dfff04db8fa // indirect
-	golang.org/x/sys v0.41.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )