Arm64: [PAC-RET] Jit Support

src/coreclr/inc/cfi.h

@@ -9,7 +9,8 @@ enum CFI_OPCODE
 {
    CFI_ADJUST_CFA_OFFSET,    // Offset is adjusted relative to the current one.
    CFI_DEF_CFA_REGISTER,     // New register is used to compute CFA
-   CFI_REL_OFFSET            // Register is saved at offset from the current CFA
+   CFI_REL_OFFSET,           // Register is saved at offset from the current CFA
+   CFI_NEGATE_RA_STATE,      // Sign the return address in lr with paciasp
 };
 
 struct CFI_CODE

src/coreclr/jit/codegenarm64.cpp

@@ -251,6 +251,11 @@ void CodeGen::genPopCalleeSavedRegistersAndFreeLclFrame(bool jmpEpilog)
         }
     }
 
+    if (JitConfig.JitPacEnabled() != 0)
+    {
+        GetEmitter()->emitPacInEpilog();
+    }
+
     // For OSR, we must also adjust the SP to remove the Tier0 frame.
     //
     if (m_compiler->opts.IsOSR())
@@ -487,9 +492,10 @@ void CodeGen::genPrologSaveRegPair(regNumber reg1,
 
         if ((spOffset == 0) && (spDelta >= -512))
         {
-            // We can use pre-indexed addressing.
+            // We can use pre-indexed addressing when the stack adjustment fits in the instruction.
             // stp REG, REG + 1, [SP, #spDelta]!
             // 64-bit STP offset range: -512 to 504, multiple of 8.
+            assert(reg1 != REG_LR);
             GetEmitter()->emitIns_R_R_R_I(INS_stp, EA_PTRSIZE, reg1, reg2, REG_SPBASE, spDelta, INS_OPTS_PRE_INDEX);
             m_compiler->unwindSaveRegPairPreindexed(reg1, reg2, spDelta);
 
@@ -511,6 +517,8 @@ void CodeGen::genPrologSaveRegPair(regNumber reg1,
         // 64-bit STP offset range: -512 to 504, multiple of 8.
         assert(spOffset <= 504);
         assert((spOffset % 8) == 0);
+        assert(reg1 != REG_LR);
+
         GetEmitter()->emitIns_R_R_R_I(INS_stp, EA_PTRSIZE, reg1, reg2, REG_SPBASE, spOffset);
 
         if (TargetOS::IsUnix && m_compiler->generateCFIUnwindCodes())
@@ -622,6 +630,7 @@ void CodeGen::genRestoreRegPair(regNumber reg1,
     assert((spDelta % 16) == 0);                                  // SP changes must be 16-byte aligned
     assert(genIsValidFloatReg(reg1) == genIsValidFloatReg(reg2)); // registers must be both general-purpose, or both
                                                                   // FP/SIMD
+    assert(reg1 != REG_LR);
 
     if (spDelta != 0)
     {
@@ -1384,6 +1393,11 @@ void CodeGen::genFuncletProlog(BasicBlock* block)
 
     m_compiler->unwindBegProlog();
 
+    if (JitConfig.JitPacEnabled() != 0)
+    {
+        GetEmitter()->emitPacInProlog();
+    }
+
     regMaskTP maskSaveRegsFloat = genFuncletInfo.fiSaveRegs & RBM_ALLFLOAT;
     regMaskTP maskSaveRegsInt   = genFuncletInfo.fiSaveRegs & ~maskSaveRegsFloat;
 
@@ -1669,6 +1683,11 @@ void CodeGen::genFuncletEpilog(BasicBlock* /* block */)
         }
     }
 
+    if (JitConfig.JitPacEnabled() != 0)
+    {
+        GetEmitter()->emitPacInEpilog();
+    }
+
     inst_RV(INS_ret, REG_LR, TYP_I_IMPL);
     m_compiler->unwindReturn(REG_LR);
 
@@ -5675,6 +5694,20 @@ void CodeGen::genOSRHandleTier0CalleeSavedRegistersAndFrame()
     genRestoreRegPair(REG_FP, REG_LR, REG_FP, 0, 0, false, REG_IP1, nullptr,
                       /* reportUnwindData */ false);
 
+    if (JitConfig.JitPacEnabled() != 0)
+    {
+        // Tier0 signed LR with the Tier0 caller SP before allocating its frame.
+        // Recreate that SP from the current Tier0 body SP so we can authenticate
+        // LR before the OSR prolog later re-signs it with the OSR SP.
+        // TODO-PAC: Avoid authenticating and re-signing so the signing SP will point to the start of the frame. It may
+        // require adding a phantom pac_sign_lr unwind code.
+        genInstrWithConstant(INS_add, EA_PTRSIZE, REG_IP0, REG_SPBASE, patchpointInfo->TotalFrameSize(), REG_IP0,
+                             /* inUnwindRegion */ false);
+        GetEmitter()->emitIns_Mov(INS_mov, EA_PTRSIZE, REG_IP1, REG_LR, /* canSkip */ false);
+        GetEmitter()->emitIns(TargetOS::IsWindows ? INS_autib1716 : INS_autia1716);
+        GetEmitter()->emitIns_Mov(INS_mov, EA_PTRSIZE, REG_LR, REG_IP1, /* canSkip */ false);
+    }
+
     // Emit phantom unwind data for the tier0 frame.
     m_compiler->unwindAllocStack(patchpointInfo->TotalFrameSize());
     // Emit nops to make the prolog 1:1 in unwind codes to instructions. This

src/coreclr/jit/codegenarmarch.cpp

@@ -4484,6 +4484,13 @@ void CodeGen::genPushCalleeSavedRegisters(regNumber initReg, bool* pInitRegZeroe
     }
 #endif // DEBUG
 
+#if defined(TARGET_ARM64)
+    if (JitConfig.JitPacEnabled() != 0)
+    {
+        GetEmitter()->emitPacInProlog();
+    }
+#endif // TARGET_ARM64
+
     // The frameType number is arbitrary, is defined below, and corresponds to one of the frame styles we
     // generate based on various sizes.
     int frameType = 0;

src/coreclr/jit/compiler.h

@@ -9796,6 +9796,7 @@ class Compiler
     void unwindSaveRegPair(regNumber reg1, regNumber reg2, int offset);           // stp reg1, reg2, [sp, #offset]
     void unwindSaveRegPairPreindexed(regNumber reg1, regNumber reg2, int offset); // stp reg1, reg2, [sp, #offset]!
     void unwindSaveNext();                                                        // unwind code: save_next
+    void unwindPacSignLR();                                                       // unwind code: pac_sign_lr
     void unwindReturn(regNumber reg);                                             // ret lr
 #endif                                                                            // defined(TARGET_ARM64)
 

src/coreclr/jit/emit.h

@@ -3347,6 +3347,11 @@ class emitter
     instrDescAlign* emitNewInstrAlign();
 #endif
 
+#if defined(TARGET_ARM64)
+    void emitPacInProlog();
+    void emitPacInEpilog();
+#endif
+
     instrDesc* emitNewInstrSmall(emitAttr attr);
     instrDesc* emitNewInstr(emitAttr attr = EA_4BYTE);
     instrDesc* emitNewInstrSC(emitAttr attr, cnsval_ssize_t cns);

src/coreclr/jit/emitarm64.cpp

@@ -1435,6 +1435,32 @@ static const char * const  bRegNames[] =
 
 // clang-format on
 
+//------------------------------------------------------------------------
+// emitPacInProlog: Sign LR as part of Pointer Authentication (PAC) support
+//
+void emitter::emitPacInProlog()
+{
+    if (JitConfig.JitPacEnabled() == 0)
+    {
+        return;
+    }
+    emitIns(TargetOS::IsWindows ? INS_pacibsp : INS_paciasp);
+    m_compiler->unwindPacSignLR();
+}
+
+//------------------------------------------------------------------------
+// emitPacInEpilog: unsign LR as part of Pointer Authentication (PAC) support
+//
+void emitter::emitPacInEpilog()
+{
+    if (JitConfig.JitPacEnabled() == 0)
+    {
+        return;
+    }
+    emitIns(TargetOS::IsWindows ? INS_autibsp : INS_autiasp);
+    m_compiler->unwindPacSignLR();
+}
+
 //------------------------------------------------------------------------
 // emitRegName: Returns a general-purpose register name or SIMD and floating-point scalar register name.
 //
@@ -16250,9 +16276,10 @@ emitter::insExecutionCharacteristics emitter::getInsExecutionCharacteristics(ins
             }
             break;
 
-        case IF_PC_0A:
-        case IF_PC_1A:
-        case IF_PC_2A:
+        case IF_PC_0A: // autia1716, autiasp, autib1716, autibsp, autibz, autiaz, pacia1716, paciasp, pacib1716,
+                       // pacibsp, pacibz, paciaz, xpaclri
+        case IF_PC_1A: // autiza, autizb, paciza, pacizb, xpacd, xpaci
+        case IF_PC_2A: // autia, autib, pacia, pacib
             switch (ins)
             {
                 case INS_xpacd:

src/coreclr/jit/jitconfigvalues.h

@@ -126,6 +126,9 @@ CONFIG_STRING(JitInlineMethodsWithEHRange, "JitInlineMethodsWithEHRange")
 
 CONFIG_INTEGER(JitLongAddress, "JitLongAddress", 0) // Force using the large pseudo instruction form for long address
 CONFIG_INTEGER(JitMaxUncheckedOffset, "JitMaxUncheckedOffset", 8)
+#if defined(TARGET_ARM64)
+RELEASE_CONFIG_INTEGER(JitPacEnabled, "JitPacEnabled", 0)
+#endif
 
 // Enable devirtualization for generic virtual methods
 RELEASE_CONFIG_INTEGER(JitEnableGenericVirtualDevirtualization, "JitEnableGenericVirtualDevirtualization", 1)

src/coreclr/jit/unwind.cpp

@@ -398,6 +398,11 @@ void Compiler::DumpCfiInfo(bool                  isHotCode,
                 assert(dwarfReg == DWARF_REG_ILLEGAL);
                 printf("    CodeOffset: 0x%02X Op: AdjustCfaOffset Offset:0x%X\n", codeOffset, offset);
                 break;
+            case CFI_NEGATE_RA_STATE:
+                assert(dwarfReg == DWARF_REG_ILLEGAL);
+                assert(offset == 0);
+                printf("    CodeOffset: 0x%02X Op: NegateRAState\n", codeOffset);
+                break;
             default:
                 printf("    Unrecognized CFI_CODE: 0x%llX\n", *(UINT64*)pCode);
                 break;

src/coreclr/jit/unwindarm64.cpp

@@ -635,6 +635,33 @@ void Compiler::unwindSaveNext()
     pu->AddCode(0xE6);
 }
 
+void Compiler::unwindPacSignLR()
+{
+    if (JitConfig.JitPacEnabled() == 0)
+    {
+        return;
+    }
+#if defined(FEATURE_CFI_SUPPORT)
+    if (generateCFIUnwindCodes())
+    {
+        // Emit NEGATE_RA_STATE opcode in prologs.
+        if (!compGeneratingProlog)
+        {
+            return;
+        }
+        FuncInfoDsc*   func     = funCurrentFunc();
+        UNATIVE_OFFSET cbProlog = unwindGetCurrentOffset(func);
+        // Maps to DW_CFA_AARCH64_negate_ra_state
+        createCfiCode(func, cbProlog, CFI_NEGATE_RA_STATE, DWARF_REG_ILLEGAL);
+
+        return;
+    }
+#endif // FEATURE_CFI_SUPPORT
+
+    // pac_sign_lr: 11111100: sign the return address in lr with the platform PAC key
+    funCurrentFunc()->uwi.AddCode(0xFC);
+}
+
 void Compiler::unwindReturn(regNumber reg)
 {
     // Nothing to do; we will always have at least one trailing "end" opcode in our padding.
@@ -1081,6 +1108,12 @@ void DumpUnwindInfo(Compiler*         comp,
 
             printf("    %02X          save_next\n", b1);
         }
+        else if (b1 == 0xFC)
+        {
+            // pac_sign_lr: 11111100 : sign the return address in lr with the platform PAC key.
+
+            printf("    %02X          pac_sign_lr\n", b1);
+        }
         else
         {
             // Unknown / reserved unwind code