ARM64: Add support for ARM64e pointer authentication codes (PAC).

Contributed by Peter Cawley. #559

14 files changed, 115 insertions, 32 deletions

diff --git a/src/Makefile b/src/Makefile
index 30d64be2..f6d093bb 100644
--- a/src/Makefile
+++ b/src/Makefile
@@ -433,6 +433,10 @@ ifneq (,$(findstring LJ_NO_UNWIND 1,$(TARGET_TESTARCH)))
   DASM_AFLAGS+= -D NO_UNWIND
   TARGET_ARCH+= -DLUAJIT_NO_UNWIND
 endif
+ifneq (,$(findstring LJ_ABI_PAUTH 1,$(TARGET_TESTARCH)))
+  DASM_AFLAGS+= -D PAUTH
+  TARGET_ARCH+= -DLJ_ABI_PAUTH=1
+endif
 DASM_AFLAGS+= -D VER=$(subst LJ_ARCH_VERSION_,,$(filter LJ_ARCH_VERSION_%,$(subst LJ_ARCH_VERSION ,LJ_ARCH_VERSION_,$(TARGET_TESTARCH))))
 ifeq (Windows,$(TARGET_SYS))
   DASM_AFLAGS+= -D WIN
diff --git a/src/lib_ffi.c b/src/lib_ffi.c
index 1b1fa389..3133cab2 100644
--- a/src/lib_ffi.c
+++ b/src/lib_ffi.c
@@ -745,6 +745,9 @@ LJLIB_CF(ffi_abi)	LJLIB_REC(.)
 #if LJ_ABI_WIN
     "\003win"
 #endif
+#if LJ_ABI_PAUTH
+    "\007pauth"
+#endif
 #if LJ_TARGET_UWP
     "\003uwp"
 #endif
diff --git a/src/lib_jit.c b/src/lib_jit.c
index 2867d420..2300f1da 100644
--- a/src/lib_jit.c
+++ b/src/lib_jit.c
@@ -422,7 +422,8 @@ LJLIB_CF(jit_util_ircalladdr)
 {
   uint32_t idx = (uint32_t)lj_lib_checkint(L, 1);
   if (idx < IRCALL__MAX) {
-    setintptrV(L->top-1, (intptr_t)(void *)lj_ir_callinfo[idx].func);
+    ASMFunction func = lj_ir_callinfo[idx].func;
+    setintptrV(L->top-1, (intptr_t)(void *)lj_ptr_strip(func));
     return 1;
   }
   return 0;
diff --git a/src/lj_arch.h b/src/lj_arch.h
index bddd757d..d354fc69 100644
--- a/src/lj_arch.h
+++ b/src/lj_arch.h
@@ -259,6 +259,9 @@
 #define LJ_ARCH_NAME            "arm64"
 #define LJ_ARCH_ENDIAN          LUAJIT_LE
 #endif
+#if !defined(LJ_ABI_PAUTH) && defined(__arm64e__)
+#define LJ_ABI_PAUTH            1
+#endif
 #define LJ_TARGET_ARM64         1
 #define LJ_TARGET_EHRETREG      0
 #define LJ_TARGET_EHRAREG       30
@@ -603,6 +606,10 @@
 #define LJ_SOFTFP               (!LJ_ARCH_HASFPU)
 #define LJ_SOFTFP32             (LJ_SOFTFP && LJ_32)
 
+#ifndef LJ_ABI_PAUTH
+#define LJ_ABI_PAUTH            0
+#endif
+
 #if LJ_ARCH_ENDIAN == LUAJIT_BE
 #define LJ_LE                   0
 #define LJ_BE                   1
diff --git a/src/lj_asm_arm64.h b/src/lj_asm_arm64.h
index f761525f..c537c514 100644
--- a/src/lj_asm_arm64.h
+++ b/src/lj_asm_arm64.h
@@ -421,8 +421,8 @@ static void asm_gencall(ASMState *as, const CCallInfo *ci, IRRef *args)
   uint32_t n, nargs = CCI_XNARGS(ci);
   int32_t ofs = 0;
   Reg gpr, fpr = REGARG_FIRSTFPR;
-  if ((void *)ci->func)
-    emit_call(as, (void *)ci->func);
+  if (ci->func)
+    emit_call(as, ci->func);
   for (gpr = REGARG_FIRSTGPR; gpr <= REGARG_LASTGPR; gpr++)
     as->cost[gpr] = REGCOST(~0u, ASMREF_L);
   gpr = REGARG_FIRSTGPR;
@@ -501,7 +501,7 @@ static void asm_callx(ASMState *as, IRIns *ir)
     ci.func = (ASMFunction)(ir_k64(irf)->u64);
   } else {  /* Need a non-argument register for indirect calls. */
     Reg freg = ra_alloc1(as, func, RSET_RANGE(RID_X8, RID_MAX_GPR)-RSET_FIXED);
-    emit_n(as, A64I_BLR, freg);
+    emit_n(as, A64I_BLR_AUTH, freg);
     ci.func = (ASMFunction)(void *)0;
   }
   asm_gencall(as, &ci, args);
diff --git a/src/lj_ccallback.c b/src/lj_ccallback.c
index 43e44305..98e9e02b 100644
--- a/src/lj_ccallback.c
+++ b/src/lj_ccallback.c
@@ -171,13 +171,13 @@ static void *callback_mcode_init(global_State *g, uint32_t *page)
 static void *callback_mcode_init(global_State *g, uint32_t *page)
 {
   uint32_t *p = page;
-  void *target = (void *)lj_vm_ffi_callback;
+  ASMFunction target = lj_vm_ffi_callback;
   MSize slot;
   *p++ = A64I_LE(A64I_LDRLx | A64F_D(RID_X11) | A64F_S19(4));
   *p++ = A64I_LE(A64I_LDRLx | A64F_D(RID_X10) | A64F_S19(5));
-  *p++ = A64I_LE(A64I_BR | A64F_N(RID_X11));
+  *p++ = A64I_LE(A64I_BR_AUTH | A64F_N(RID_X11));
   *p++ = A64I_LE(A64I_NOP);
-  ((void **)p)[0] = target;
+  ((ASMFunction *)p)[0] = target;
   ((void **)p)[1] = g;
   p += 4;
   for (slot = 0; slot < CALLBACK_MAX_SLOT; slot++) {
diff --git a/src/lj_emit_arm64.h b/src/lj_emit_arm64.h
index fcc9c1d8..65463a5e 100644
--- a/src/lj_emit_arm64.h
+++ b/src/lj_emit_arm64.h
@@ -348,16 +348,22 @@ static void emit_cnb(ASMState *as, A64Ins ai, Reg r, MCode *target)
 
 #define emit_jmp(as, target)    emit_branch(as, A64I_B, (target))
 
-static void emit_call(ASMState *as, void *target)
+static void emit_call(ASMState *as, ASMFunction target)
 {
   MCode *p = --as->mcp;
-  ptrdiff_t delta = (char *)target - (char *)p;
+#if LJ_ABI_PAUTH
+  char *targetp = ptrauth_auth_data((char *)target,
+                                    ptrauth_key_function_pointer, 0);
+#else
+  char *targetp = (char *)target;
+#endif
+  ptrdiff_t delta = targetp - (char *)p;
   if (A64F_S_OK(delta>>2, 26)) {
     *p = A64I_BL | A64F_S26(delta>>2);
   } else {  /* Target out of range: need indirect call. But don't use R0-R7. */
     Reg r = ra_allock(as, i64ptr(target),
                       RSET_RANGE(RID_X8, RID_MAX_GPR)-RSET_FIXED);
-    *p = A64I_BLR | A64F_N(r);
+    *p = A64I_BLR_AUTH | A64F_N(r);
   }
 }
 
diff --git a/src/lj_err.c b/src/lj_err.c
index 9652ef35..a0a28692 100644
--- a/src/lj_err.c
+++ b/src/lj_err.c
@@ -444,10 +444,10 @@ LJ_FUNCA int lj_err_unwind_dwarf(int version, int actions,
     if ((actions & _UA_FORCE_UNWIND)) {
       return _URC_CONTINUE_UNWIND;
     } else if (cf) {
+      ASMFunction ip;
       _Unwind_SetGR(ctx, LJ_TARGET_EHRETREG, errcode);
-      _Unwind_SetIP(ctx, (uintptr_t)(cframe_unwind_ff(cf) ?
-                                     lj_vm_unwind_ff_eh :
-                                     lj_vm_unwind_c_eh));
+      ip = cframe_unwind_ff(cf) ? lj_vm_unwind_ff_eh : lj_vm_unwind_c_eh;
+      _Unwind_SetIP(ctx, (uintptr_t)lj_ptr_strip(ip));
       return _URC_INSTALL_CONTEXT;
     }
 #if LJ_TARGET_X86ORX64
@@ -580,9 +580,17 @@ extern void __deregister_frame(const void *);
 
 uint8_t *lj_err_register_mcode(void *base, size_t sz, uint8_t *info)
 {
-  void **handler;
+  ASMFunction handler = (ASMFunction)err_unwind_jit;
   memcpy(info, err_frame_jit_template, sizeof(err_frame_jit_template));
-  handler = (void *)err_unwind_jit;
+#if LJ_ABI_PAUTH
+#if LJ_TARGET_ARM64
+  handler = ptrauth_auth_and_resign(handler,
+    ptrauth_key_function_pointer, 0,
+    ptrauth_key_process_independent_code, info + ERR_FRAME_JIT_OFS_HANDLER);
+#else
+#error "missing pointer authentication support for this architecture"
+#endif
+#endif
   memcpy(info + ERR_FRAME_JIT_OFS_HANDLER, &handler, sizeof(handler));
   *(uint32_t *)(info + ERR_FRAME_JIT_OFS_CODE_SIZE) =
     (uint32_t)(sz - sizeof(err_frame_jit_template) - (info - (uint8_t *)base));
diff --git a/src/lj_jit.h b/src/lj_jit.h
index 7f081730..0fae60ad 100644
--- a/src/lj_jit.h
+++ b/src/lj_jit.h
@@ -273,6 +273,9 @@ typedef struct GCtrace {
   BCIns startins;       /* Original bytecode of starting instruction. */
   MSize szmcode;        /* Size of machine code. */
   MCode *mcode;         /* Start of machine code. */
+#if LJ_ABI_PAUTH
+  ASMFunction mcauth;   /* Start of machine code, with ptr auth applied. */
+#endif
   MSize mcloop;         /* Offset of loop start in machine code. */
   uint16_t nchild;      /* Number of child traces (root trace only). */
   uint16_t spadjust;    /* Stack pointer adjustment (offset in bytes). */
diff --git a/src/lj_obj.h b/src/lj_obj.h
index e541387f..52c7bc03 100644
--- a/src/lj_obj.h
+++ b/src/lj_obj.h
@@ -1042,4 +1042,18 @@ LJ_DATA const char *const lj_obj_itypename[~LJ_TNUMX+1];
 LJ_FUNC int LJ_FASTCALL lj_obj_equal(cTValue *o1, cTValue *o2);
 LJ_FUNC const void * LJ_FASTCALL lj_obj_ptr(global_State *g, cTValue *o);
 
+#if LJ_ABI_PAUTH
+#if LJ_TARGET_ARM64
+#include <ptrauth.h>
+#define lj_ptr_sign(ptr, ctx) \
+  ptrauth_sign_unauthenticated((ptr), ptrauth_key_function_pointer, (ctx))
+#define lj_ptr_strip(ptr) ptrauth_strip((ptr), ptrauth_key_function_pointer)
+#else
+#error "No support for pointer authentication for this architecture"
+#endif
+#else
+#define lj_ptr_sign(ptr, ctx) (ptr)
+#define lj_ptr_strip(ptr) (ptr)
+#endif
+
 #endif
diff --git a/src/lj_target_arm64.h b/src/lj_target_arm64.h
index d45af2e4..c9c6b80f 100644
--- a/src/lj_target_arm64.h
+++ b/src/lj_target_arm64.h
@@ -260,6 +260,9 @@ typedef enum A64Ins {
   A64I_CBZ = 0x34000000,
   A64I_CBNZ = 0x35000000,
 
+  A64I_BRAAZ = 0xd61f081f,
+  A64I_BLRAAZ = 0xd63f081f,
+
   A64I_NOP = 0xd503201f,
 
   /* FP */
@@ -317,6 +320,9 @@ typedef enum A64Ins {
   A64I_FMOV_DI = 0x1e601000,
 } A64Ins;
 
+#define A64I_BR_AUTH    (LJ_ABI_PAUTH ? A64I_BRAAZ : A64I_BR)
+#define A64I_BLR_AUTH   (LJ_ABI_PAUTH ? A64I_BLRAAZ : A64I_BLR)
+
 typedef enum A64Shift {
   A64SH_LSL, A64SH_LSR, A64SH_ASR, A64SH_ROR
 } A64Shift;
diff --git a/src/lj_trace.c b/src/lj_trace.c
index c2329394..03c8d1d0 100644
--- a/src/lj_trace.c
+++ b/src/lj_trace.c
@@ -153,6 +153,9 @@ static void trace_save(jit_State *J, GCtrace *T)
   newwhite(J2G(J), T);
   T->gct = ~LJ_TTRACE;
   T->ir = (IRIns *)p - J->cur.nk;  /* The IR has already been copied above. */
+#if LJ_ABI_PAUTH
+  T->mcauth = lj_ptr_sign((ASMFunction)T->mcode, T);
+#endif
   p += szins;
   TRACE_APPENDVEC(snap, nsnap, SnapShot)
   TRACE_APPENDVEC(snapmap, nsnapmap, SnapEntry)
diff --git a/src/lj_vm.h b/src/lj_vm.h
index c66db004..c7095941 100644
--- a/src/lj_vm.h
+++ b/src/lj_vm.h
@@ -54,8 +54,8 @@ LJ_ASMF void lj_vm_profhook(void);
 LJ_ASMF void lj_vm_IITERN(void);
 
 /* Trace exit handling. */
-LJ_ASMF void lj_vm_exit_handler(void);
-LJ_ASMF void lj_vm_exit_interp(void);
+LJ_ASMF char lj_vm_exit_handler[];
+LJ_ASMF char lj_vm_exit_interp[];
 
 /* Internal math helper functions. */
 #if LJ_TARGET_PPC || LJ_TARGET_ARM64 || (LJ_TARGET_MIPS && LJ_ABI_SOFTFP)
@@ -111,6 +111,6 @@ LJ_ASMF void lj_cont_stitch(void);  /* Trace stitching. */
 LJ_ASMF char lj_vm_asm_begin[];
 
 /* Bytecode offsets are relative to lj_vm_asm_begin. */
-#define makeasmfunc(ofs)        ((ASMFunction)(lj_vm_asm_begin + (ofs)))
+#define makeasmfunc(ofs) lj_ptr_sign((ASMFunction)(lj_vm_asm_begin + (ofs)), 0)
 
 #endif
diff --git a/src/vm_arm64.dasc b/src/vm_arm64.dasc
index 36a036ae..d45cc86b 100644
--- a/src/vm_arm64.dasc
+++ b/src/vm_arm64.dasc
@@ -77,6 +77,23 @@
 |.define CRET1,         x0
 |.define CRET1w,        w0
 |
+|//-----------------------------------------------------------------------
+|
+|// ARM64e pointer authentication codes (PAC).
+|.if PAUTH
+|.macro sp_auth; pacibsp; .endmacro
+|.macro br_auth, reg; braaz reg; .endmacro
+|.macro blr_auth, reg; blraaz reg; .endmacro
+|.macro ret_auth; retab; .endmacro
+|.else
+|.macro sp_auth; .endmacro
+|.macro br_auth, reg; br reg; .endmacro
+|.macro blr_auth, reg; blr reg; .endmacro
+|.macro ret_auth; ret; .endmacro
+|.endif
+|
+|//-----------------------------------------------------------------------
+|
 |// Stack layout while in interpreter. Must match with lj_frame.h.
 |
 |.define CFRAME_SPACE,  208
@@ -106,6 +123,7 @@
 |.endmacro
 |
 |.macro saveregs
+|  sp_auth
 |  sub sp, sp, # CFRAME_SPACE
 |  stp fp, lr, [sp, # SAVE_FP_LR_]
 |  add fp, sp, # SAVE_FP_LR_
@@ -180,7 +198,7 @@
 |   decode_RA RA, INS
 |  ldr TMP0, [TMP1, #GG_G2DISP]
 |   decode_RD RC, INS
-|  br TMP0
+|  br_auth TMP0
 |.endmacro
 |
 |// Instruction footer.
@@ -209,7 +227,7 @@
 |   decode_RA RA, INS
 |  ldr TMP0, [TMP1, #GG_G2DISP]
 |   add RA, BASE, RA, lsl #3
-|  br TMP0
+|  br_auth TMP0
 |.endmacro
 |
 |.macro ins_call
@@ -356,7 +374,7 @@ static void build_subroutines(BuildCtx *ctx)
   |
   |->vm_leave_unw:
   |  restoreregs
-  |  ret
+  |  ret_auth
   |
   |6:
   |  bgt >7                             // Less results wanted?
@@ -542,7 +560,7 @@ static void build_subroutines(BuildCtx *ctx)
   |   str RC, SAVE_CFRAME
   |  str TMP0, L->cframe                // Add our C frame to cframe chain.
   |    str L, GL->cur_L
-  |  blr CARG4                  // (lua_State *L, lua_CFunction func, void *ud)
+  |  blr_auth CARG4             // (lua_State *L, lua_CFunction func, void *ud)
   |  mov BASE, CRET1
   |   mov PC, #FRAME_CP
   |  cbnz BASE, <3                      // Else continue with the call.
@@ -573,7 +591,7 @@ static void build_subroutines(BuildCtx *ctx)
   |  ldr CARG3, LFUNC:CARG3->pc
   |  ldr KBASE, [CARG3, #PC2PROTO(k)]
   |  // BASE = base, RA = resultptr, CARG4 = meta base
-  |    br CARG1
+  |    br_auth CARG1
   |
   |.if FFI
   |1:
@@ -1707,7 +1725,7 @@ static void build_subroutines(BuildCtx *ctx)
   |  cmp TMP1, TMP2
   |   mov CARG1, L
   |  bhi >5                             // Need to grow stack.
-  |   blr CARG3                         // (lua_State *L)
+  |   blr_auth CARG3                    // (lua_State *L)
   |  // Either throws an error, or recovers and returns -1, 0 or nresults+1.
   |   ldr BASE, L->base
   |  cmp CRET1w, #0
@@ -1743,6 +1761,7 @@ static void build_subroutines(BuildCtx *ctx)
   |
   |->fff_gcstep:                        // Call GC step function.
   |  // BASE = new base, RC = nargs*8
+  |  sp_auth
   |   add CARG2, BASE, NARGS8:RC        // Calculate L->top.
   |  mov RA, lr
   |   stp BASE, CARG2, L->base
@@ -1754,7 +1773,7 @@ static void build_subroutines(BuildCtx *ctx)
   |  mov lr, RA                         // Help return address predictor.
   |  sub NARGS8:RC, CARG2, BASE         // Calculate nargs*8.
   |   and CFUNC:CARG3, CARG3, #LJ_GCVMASK
-  |  ret
+  |  ret_auth
   |
   |//-----------------------------------------------------------------------
   |//-- Special dispatch targets -------------------------------------------
@@ -1781,7 +1800,7 @@ static void build_subroutines(BuildCtx *ctx)
   |  tbz TMP2w, #HOOK_ACTIVE_SHIFT, >1  // Hook already active?
   |5:  // Re-dispatch to static ins.
   |  ldr TMP0, [TMP1, #GG_G2DISP+GG_DISP2STATIC]
-  |  br TMP0
+  |  br_auth TMP0
   |
   |->vm_inshook:                        // Dispatch target for instr/line hooks.
   |  ldrb TMP2w, GL->hookmask
@@ -1807,7 +1826,7 @@ static void build_subroutines(BuildCtx *ctx)
   |   decode_RA RA, INS
   |  ldr TMP0, [TMP1, #GG_G2DISP+GG_DISP2STATIC]
   |   decode_RD RC, INS
-  |  br TMP0
+  |  br_auth TMP0
   |
   |->cont_hook:                         // Continue from hook yield.
   |  ldr CARG1, [CARG4, #-40]
@@ -1857,7 +1876,7 @@ static void build_subroutines(BuildCtx *ctx)
   |  sub NARGS8:RC, TMP1, BASE
   |   ldr INSw, [PC, #-4]
   |  and LFUNC:CARG3, CARG3, #LJ_GCVMASK
-  |  br CRET1
+  |  br_auth CRET1
   |
   |->cont_stitch:                       // Trace stitching.
   |.if JIT
@@ -2020,7 +2039,7 @@ static void build_subroutines(BuildCtx *ctx)
   |   add RA, BASE, RA, lsl #3  // Yes: RA = BASE+framesize*8, RC = nargs*8
   |   and LFUNC:CARG3, CARG3, #LJ_GCVMASK
   |5:
-  |  br RB
+  |  br_auth RB
   |
   |4:  // Check frame below fast function.
   |  ldr CARG1, [BASE, FRAME_PC]
@@ -2182,6 +2201,7 @@ static void build_subroutines(BuildCtx *ctx)
   |  // Caveat: needs special frame unwinding, see below.
   |.if FFI
   |  .type CCSTATE, CCallState, x19
+  |  sp_auth
   |  stp x20, CCSTATE, [sp, #-32]!
   |  stp fp, lr, [sp, #16]
   |  add fp, sp, #16
@@ -2208,14 +2228,14 @@ static void build_subroutines(BuildCtx *ctx)
   |  ldp x6, x7, CCSTATE->gpr[6]
   |   ldp d6, d7, CCSTATE->fpr[6]
   |  ldr x8, CCSTATE->retp
-  |  blr TMP3
+  |  blr_auth TMP3
   |  sub sp, fp, #16
   |  stp x0, x1, CCSTATE->gpr[0]
   |   stp d0, d1, CCSTATE->fpr[0]
   |   stp d2, d3, CCSTATE->fpr[2]
   |  ldp fp, lr, [sp, #16]
   |  ldp x20, CCSTATE, [sp], #32
-  |  ret
+  |  ret_auth
   |.endif
   |// Note: vm_ffi_call must be the last function in this object file!
   |
@@ -3786,12 +3806,20 @@ static void build_ins(BuildCtx *ctx, BCOp op, int defop)
     |   mov CARG2w, #0  // Traces on ARM64 don't store the trace #, so use 0.
     |  ldr TRACE:RC, [CARG1, RC, lsl #3]
     |   st_vmstate CARG2w
+    |.if PAUTH
+    |  ldr RA, TRACE:RC->mcauth
+    |.else
     |  ldr RA, TRACE:RC->mcode
+    |.endif
     |   str BASE, GL->jit_base
     |   str L, GL->tmpbuf.L
     |  sub sp, sp, #16  // See SPS_FIXED. Avoids sp adjust in every root trace.
+    |.if PAUTH
+    |  braa RA, RC
+    |.else
     |  br RA
     |.endif
+    |.endif
     break;
 
   case BC_JMP:
@@ -3901,7 +3929,7 @@ static void build_ins(BuildCtx *ctx, BCOp op, int defop)
     |  mov CARG1, L
     |   bhi ->vm_growstack_c            // Need to grow stack.
     |    st_vmstate TMP0w
-    |  blr CARG4                        // (lua_State *L [, lua_CFunction f])
+    |  blr_auth CARG4                   // (lua_State *L [, lua_CFunction f])
     |  // Returns nresults.
     |  ldp BASE, TMP1, L->base
     |    str L, GL->cur_L