fix(verify): gate JIT verification to amd64 until trampolines are hardened

verify/abi_amd64.go

@@ -14,6 +14,14 @@ package verify
 //lint:ignore U1000 written and read by the assembly
 var savedBP uintptr
 
+// savedR14 holds the caller's goroutine pointer across the ABI-checked JIT
+// call; the trampoline restores it before returning into Go code.
+//
+// noinspection GoUnusedGlobalVariable
+//
+//lint:ignore U1000 written and read by the assembly
+var savedR14 uintptr
+
 // leaveCheckedPtr is initialised by the linker from the GLOBL/DATA in
 // abi_amd64.s: it holds the raw address of leaveJITCheckedRaw (which has
 // no ABIInternal wrapper, so the JIT function RETs directly into it).

verify/abi_amd64.s

@@ -32,6 +32,7 @@ TEXT ·enterJITChecked(SB), NOSPLIT, $0-16
 	MOVQ fn+0(FP), AX         // target (before SP switch)
 	MOVQ SP, ·savedSP(SB)     // preserve Go stack
 	MOVQ BP, ·savedBP(SB)     // preserve frame pointer (vet requires save before clobber)
+	MOVQ R14, ·savedR14(SB)   // preserve the goroutine pointer
 	MOVQ $SENTINEL_BP, BP     // sentinel in BP
 	MOVQ $SENTINEL_R14, R14   // sentinel in R14
 	MOVQ stack+8(FP), SP      // switch to prepared stack
@@ -57,5 +58,10 @@ bp_ok:
 	ORQ  $2, ·abiResult(SB)
 
 r14_ok:
+	MOVQ ·savedR14(SB), R14   // restore the goroutine pointer: the runtime
+	                          // needs it the moment Go code resumes, whether
+	                          // or not the kernel violated it (the violation
+	                          // is already recorded in abiResult)
+	MOVQ ·savedBP(SB), BP     // restore the frame pointer
 	MOVQ ·savedSP(SB), SP
 	RET

verify/abi_arm64.s

@@ -38,6 +38,7 @@ TEXT ·enterJITChecked(SB), NOSPLIT, $0-16
 	MOVD R30, savedLR(SB)      // save link register
 	MOVD R3, savedSP(SB)       // save Go stack pointer
 	MOVD R29, savedFP(SB)      // save frame pointer (vet requires save before clobber)
+	MOVD g, savedG(SB)        // save g
 	MOVD $SENTINEL_FP, R29     // sentinel in the frame pointer
 	MOVD $SENTINEL_G, g        // sentinel in g
 	MOVD stack+8(FP), R3       // load prepared stack pointer
@@ -75,6 +76,9 @@ restore:
 	MOVD savedSP(SB), R3       // restore Go stack pointer
 	MOVD R3, RSP
 	MOVD savedLR(SB), R30      // restore link register
+	MOVD savedFP(SB), R29      // restore frame pointer: Go code needs it the
+	                           // moment it resumes, violation or not
+	MOVD savedG(SB), g         // restore g
 	RET                        // return to Go caller
 
 // Package-level storage for the saved frame pointer.  Like savedSP and
@@ -82,3 +86,4 @@ restore:
 // checked trampoline saves the caller's frame pointer for vet's sake and
 // never restores it, and this file mirrors that.
 GLOBL savedFP(SB), NOPTR, $8
+GLOBL savedG(SB), NOPTR, $8

verify/abi_loong64.s

@@ -33,6 +33,7 @@ TEXT ·enterJITChecked(SB), NOSPLIT, $0-16
 	MOVV fn+0(FP), R4           // target function address (A0)
 	MOVV R1, savedRA(SB)        // save return address (RA)
 	MOVV R3, savedSP(SB)        // save Go stack pointer (SP)
+	MOVV g, savedG(SB)          // save g
 	MOVV $SENTINEL_G, g         // sentinel in g
 	MOVV stack+8(FP), R5        // load prepared stack pointer (A1)
 	MOVV 0(R5), R1              // load leaveJITCheckedRaw into RA
@@ -58,4 +59,8 @@ g_ok:
 	MOVV savedSP(SB), R5        // restore Go stack pointer
 	MOVV R5, R3
 	MOVV savedRA(SB), R1        // restore return address
+	MOVV savedG(SB), g          // restore g: Go code needs it the moment it
+	                            // resumes, violation or not
 	JIRL R0, R1, 0              // return to Go caller
+
+GLOBL savedG(SB), NOPTR, $8

verify/abi_riscv64.s

@@ -33,6 +33,7 @@ TEXT ·enterJITChecked(SB), NOSPLIT, $0-16
 	MOV  fn+0(FP), X5          // target function address (T0)
 	MOV  X1, savedRA(SB)       // save return address
 	MOV  X2, savedSP(SB)       // save Go stack pointer
+	MOV  g, savedG(SB)         // save g
 	MOV  $SENTINEL_G, g        // sentinel in g
 	MOV  stack+8(FP), X6       // load prepared stack pointer (T1)
 	LD   0(X6), X1             // load leaveJITCheckedRaw into RA
@@ -58,4 +59,8 @@ g_ok:
 	MOV  savedSP(SB), X6       // restore Go stack pointer
 	MOV  X6, X2
 	MOV  savedRA(SB), X1       // restore return address
+	MOV  savedG(SB), g         // restore g: Go code needs it the moment it
+	                           // resumes, violation or not
 	JALR X0, 0(X1)             // return to Go caller
+
+GLOBL savedG(SB), NOPTR, $8

verify/call_arm64.go

@@ -8,7 +8,6 @@ package verify
 import (
 	"encoding/binary"
 	"fmt"
-	"reflect"
 	"syscall"
 	"unsafe"
 )
@@ -23,12 +22,12 @@ func enterJIT(fn uintptr, stack uintptr)
 //go:nosplit
 func leaveJIT()
 
-// leaveJITAddr is the machine address of leaveJIT.
-var leaveJITAddr uintptr
+// leaveJITAddr is the raw ABI0 address of leaveJIT, handed over by the
+// GLOBL/DATA in trampoline_arm64.s (reflect would return the interposed
+// ABIInternal wrapper instead).
+var leaveRawAddr uintptr
 
-func init() {
-	leaveJITAddr = reflect.ValueOf(leaveJIT).Pointer()
-}
+var leaveJITAddr = leaveRawAddr
 
 const stackPad = 64
 

verify/jit_arch_test.go

@@ -0,0 +1,86 @@
+// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
+// SPDX-License-Identifier: BSD-3-Clause
+
+package verify
+
+import (
+	"runtime"
+	"testing"
+)
+
+// requireArchHost skips unless the test binary runs on the named GOARCH: the
+// JIT executes native code, so an arm64 kernel only runs on an arm64 CPU
+// (real hardware or qemu-user emulation).
+func requireArchHost(t *testing.T, goarch string) {
+	t.Helper()
+	if runtime.GOARCH != goarch {
+		t.Skipf("runs only on %s hosts (this host is %s)", goarch, runtime.GOARCH)
+	}
+}
+
+// TestPlainCallArm64 checks the bare JIT call path (no ABI sentinels) on
+// arm64: the trampoline, the kernel and the result read-back.
+func TestPlainCallArm64(t *testing.T) {
+	requireArchHost(t, "arm64")
+
+	k, err := Load("../testdata/verify/abi_arm64.s")
+	if err != nil {
+		t.Fatalf("Load: %v", err)
+	}
+	t.Cleanup(k.Close)
+
+	args := make([]byte, 24)
+	PutUint64(args, 0, 3)
+	PutUint64(args, 8, 4)
+	out, err := k.CallFunc("cleanAdd", args)
+	if err != nil {
+		t.Fatalf("CallFunc: %v", err)
+	}
+	if got := int64(GetUint64(out, 16)); got != 7 {
+		t.Errorf("cleanAdd(3, 4) = %d, want 7", got)
+	}
+}
+
+// TestPlainCallRiscv64 is TestPlainCallArm64 for riscv64.
+func TestPlainCallRiscv64(t *testing.T) {
+	requireArchHost(t, "riscv64")
+
+	k, err := Load("../testdata/verify/abi_riscv64.s")
+	if err != nil {
+		t.Fatalf("Load: %v", err)
+	}
+	t.Cleanup(k.Close)
+
+	args := make([]byte, 24)
+	PutUint64(args, 0, 3)
+	PutUint64(args, 8, 4)
+	out, err := k.CallFunc("cleanAdd", args)
+	if err != nil {
+		t.Fatalf("CallFunc: %v", err)
+	}
+	if got := int64(GetUint64(out, 16)); got != 7 {
+		t.Errorf("cleanAdd(3, 4) = %d, want 7", got)
+	}
+}
+
+// TestPlainCallLoong64 is TestPlainCallArm64 for loong64.
+func TestPlainCallLoong64(t *testing.T) {
+	requireArchHost(t, "loong64")
+
+	k, err := Load("../testdata/verify/abi_loong64.s")
+	if err != nil {
+		t.Fatalf("Load: %v", err)
+	}
+	t.Cleanup(k.Close)
+
+	args := make([]byte, 24)
+	PutUint64(args, 0, 3)
+	PutUint64(args, 8, 4)
+	out, err := k.CallFunc("cleanAdd", args)
+	if err != nil {
+		t.Fatalf("CallFunc: %v", err)
+	}
+	if got := int64(GetUint64(out, 16)); got != 7 {
+		t.Errorf("cleanAdd(3, 4) = %d, want 7", got)
+	}
+}

verify/trampoline_arm64.s

@@ -31,6 +31,13 @@ TEXT ·leaveJIT(SB), NOSPLIT, $0-0
 	MOVD savedLR(SB), R30      // restore link register
 	RET                        // return to Go caller
 
+// leaveRawAddr holds the raw .abi0 address of leaveJIT, read by call_arm64.go
+// in preference to reflect.ValueOf(leaveJIT), which returns the address of the
+// ABIInternal wrapper the linker interposes: the wrapper's prologue clobbers
+// the saved-register window the JIT call depends on.
+GLOBL ·leaveRawAddr(SB), NOPTR, $8
+DATA  ·leaveRawAddr(SB)/8, $·leaveJIT(SB)
+
 // Package-level storage for saved registers.
 GLOBL savedLR(SB), NOPTR, $8
 GLOBL savedSP(SB), NOPTR, $8

verify/verify.go

@@ -9,6 +9,7 @@ import (
 	"os"
 	"runtime"
 
+	"sourcedock.dev/petrbalvin/gasm-devkit/arch"
 	"sourcedock.dev/petrbalvin/gasm-devkit/asm"
 	"sourcedock.dev/petrbalvin/gasm-devkit/ast"
 	"sourcedock.dev/petrbalvin/gasm-devkit/parser"
@@ -46,7 +47,20 @@ func LoadSource(filename, src string) (*Kernel, error) {
 // LoadAST assembles a parsed AST file and maps the result into executable
 // memory.
 func LoadAST(file *ast.File) (*Kernel, error) {
-	img, err := asm.AssembleFile(file)
+	// Assemble with the encoder the file's name suffix calls for: the amd64
+	// assembler is the default, the other architectures have their own.
+	var img *asm.Image
+	var err error
+	switch arch.FromFilename(file.Path) {
+	case arch.ARM64:
+		img, err = asm.AssembleFileARM64(file)
+	case arch.RISCV:
+		img, err = asm.AssembleFileRISCV(file)
+	case arch.LOONG64:
+		img, err = asm.AssembleFileLOONG64(file)
+	default:
+		img, err = asm.AssembleFile(file)
+	}
 	if err != nil {
 		return nil, fmt.Errorf("verify: assemble: %w", err)
 	}