[Dwarf/AArch64] Return address signing B key dwarf support

[LukeCheeseman]

Summary

• When signing return addresses with -msign-return-address=<scope>{+<key>}, either the A key instructions or the B key instructions can be used. To correctly authenticate the return address, the unwinder/debugger must know which key was used to sign the return address.
• When and exception is thrown or a break point reached, it may be necessary to unwind the stack. To accomplish this, the unwinder/debugger must be able to first authenticate an the return address if it has been signed.
• To enable this, the augmentation string of CIEs has been extended to allow inclusion of a 'B' character. Functions that are signed using the B key variant of the instructions should have and FDE whose associated CIE has a 'B' in the augmentation string.
• One must also be able to preserve these semantics when first stepping from a high level language into assembly and then, as a second step, into an object file. To achieve this, I have introduced a new assembly directive '.cfi_b_key_frame ', that tells the assembler the current frame uses return address signing with the B key.
• This ensures that the FDE is associated with a CIE that has 'B' in the augmentation string.

[LukeCheeseman]

• std::move ExtraAugementation string in constructor
• tidy up MIR test

[olista01]

Before I review the code itself, have you agreed the .append_augmentation directive with the GCC/GAS developers to make sure we implement it the same way?

If the design isn't yet set in stone, I have a few comments on it:

• Why doesn't it begin with ".cfi_", while every other related directive does? There's a complete list of them at https://sourceware.org/binutils/docs/as/CFI-directives.html
• Why are we adding a generic directive to append to the augmentation string, rather than a directive to select the B key? This looks very similar to .cfi_signal_frame, which also adds one character to the augmentation string.

[LukeCheeseman]

• Replace .append_augmentation <string> with .cfi_b_key_frame

[olista01]

lib/MC/MCAsmStreamer.cpp

1591

This emits a pair of double quotes, which that parser doesn't expect. Also, there is no test for this.

lib/MC/MCDwarf.cpp

1735

80 columns (I'd suggest using git-clang-format on the patch).

lib/MC/MCParser/AsmParser.cpp

4167

Outdated comment.

4168

This should probably be restricted to AArch64, ideally by putting it in the AArch64 asm parser.

[olista01]

lib/MC/MCAsmStreamer.cpp

1588

This will cause an assertion if we are not emitting unwind tables, so we need to skip this this as is currently done in AsmPrinter::emitCFIInstruction.

[LukeCheeseman]

• Addressed @olista01 comment on triggering an assertion when we don't emit unwind info
• The emit b key assembly directive isn't emitted when we don't emit unwind info

[olista01]

LGTM, thanks!

[davezarzycki]

FYI – The following test fails if the ARM 32-bit target is NOT built alongside the 64-bit ARM target: MC/ELF/cfi-b-key-frame.s

[LukeCheeseman]

In D51798#1338947, @davezarzycki wrote:

FYI – The following test fails if the ARM 32-bit target is NOT built alongside the 64-bit ARM target: MC/ELF/cfi-b-key-frame.s

Ah, thanks! I was just investigating to see what was going on here.

[davezarzycki]

My mistake. Something else is going on. I assumed based on past precedent that this was what was going on. (I'm not sure how many people at ARM are verifying that their AArch64 commits work when the 32-bit target is disabled).

[davezarzycki]

Turns out my setup was busted and I thought I was building/testing AArch64 when I wasn't. I think you need to add // REQUIRES: aarch64-registered-target to the file.

include/llvm/MC/MCDwarf.h

struct MCDwarfFrameInfo {
   std::vector<MCCFIInstruction> Instructions;
   unsigned CurrentCfaRegister = 0;
   unsigned PersonalityEncoding = 0;
   unsigned LsdaEncoding = 0;
   uint32_t CompactUnwindEncoding = 0;
   bool IsSignalFrame = false;
   bool IsSimple = false;
   unsigned RAReg = static_cast<unsigned>(INT_MAX);
+  std::string ExtraAugmentation;
 };
 
 class MCDwarfFrameEmitter {
 public:
   //
   // This emits the frame info section.
   //
   static void Emit(MCObjectStreamer &streamer, MCAsmBackend *MAB, bool isEH);
   static void EmitAdvanceLoc(MCObjectStreamer &Streamer, uint64_t AddrDelta);
   static void EncodeAdvanceLoc(MCContext &Context, uint64_t AddrDelta,
                                raw_ostream &OS);
 };
 
 } // end namespace llvm
 
 #endif // LLVM_MC_MCDWARF_H

include/llvm/MC/MCStreamer.h

virtual Expected<unsigned> tryEmitDwarfFileDirective(
       unsigned CUID = 0);
 
   /// Specify the "root" file of the compilation, using the ".file 0" extension.
   virtual void emitDwarfFile0Directive(StringRef Directory, StringRef Filename,
                                        MD5::MD5Result *Checksum,
                                        Optional<StringRef> Source,
                                        unsigned CUID = 0);
 
+  virtual void EmitAppendAugmentation(StringRef String);
+
   /// This implements the DWARF2 '.loc fileno lineno ...' assembler
   /// directive.
   virtual void EmitDwarfLocDirective(unsigned FileNo, unsigned Line,
                                      unsigned Column, unsigned Flags,
                                      unsigned Isa, unsigned Discriminator,
                                      StringRef FileName);
 
   /// Associate a filename with a specified logical file number, and also

lib/DebugInfo/DWARF/DWARFDebugFrame.cpp

if (IsCIE) {
                 ReportError(StartOffset,
                             "'z' must be the first character at %lx");
               // Parse the augmentation length first.  We only parse it if
               // the string contains a 'z'.
               AugmentationLength = Data.getULEB128(&Offset);
               StartAugmentationOffset = Offset;
               EndAugmentationOffset = Offset +
                 static_cast<uint32_t>(*AugmentationLength);
+            case 'B':
+              break;
           }
         }
 
         if (AugmentationLength.hasValue()) {
           if (Offset != EndAugmentationOffset)
             ReportError(StartOffset, "Parsing augmentation data at %lx failed");
 
           AugmentationData = Data.getData().slice(StartAugmentationOffset,

lib/MC/MCAsmStreamer.cpp

void EmitCVDefRangeDirective(
       ArrayRef<std::pair<const MCSymbol *, const MCSymbol *>> Ranges,
       StringRef FixedSizePortion) override;
   void EmitCVStringTableDirective() override;
   void EmitCVFileChecksumsDirective() override;
   void EmitCVFileChecksumOffsetDirective(unsigned FileNo) override;
   void EmitCVFPOData(const MCSymbol *ProcSym, SMLoc L) override;
 
   void EmitIdent(StringRef IdentString) override;
+  void EmitAppendAugmentation(StringRef String) override;
   void EmitCFISections(bool EH, bool Debug) override;
   void EmitCFIDefCfa(int64_t Register, int64_t Offset) override;
   void EmitCFIDefCfaOffset(int64_t Offset) override;
   void EmitCFIDefCfaRegister(int64_t Register) override;
   void EmitCFIOffset(int64_t Register, int64_t Offset) override;
   void EmitCFIPersonality(const MCSymbol *Sym, unsigned Encoding) override;
   void EmitCFILsda(const MCSymbol *Sym, unsigned Encoding) override;
   void EmitCFIRememberState() override;

 
 void MCAsmStreamer::EmitIdent(StringRef IdentString) {
   assert(MAI->hasIdentDirective() && ".ident directive not supported");
   OS << "\t.ident\t";
   PrintQuotedString(IdentString, OS);
   EmitEOL();
 }
 
+void MCAsmStreamer::EmitAppendAugmentation(StringRef String)
+{
+  OS << "\t.append_augmentation \"" << String << "\"";
+  EmitEOL();
+}
+
 void MCAsmStreamer::EmitCFISections(bool EH, bool Debug) {
   MCStreamer::EmitCFISections(EH, Debug);
   OS << "\t.cfi_sections ";
   if (EH) {
     OS << ".eh_frame";
     if (Debug)
       OS << ", .debug_frame";
   } else if (Debug) {

 
 void MCAsmStreamer::EmitCFIReturnColumn(int64_t Register) {
   MCStreamer::EmitCFIReturnColumn(Register);
   OS << "\t.cfi_return_column " << Register;
   EmitEOL();
 }
 
 void MCAsmStreamer::EmitWinCFIStartProc(const MCSymbol *Symbol, SMLoc Loc) {
   MCStreamer::EmitWinCFIStartProc(Symbol, Loc);
 
   OS << ".seh_proc ";
   Symbol->print(OS, MAI);
   EmitEOL();
 }
 
 void MCAsmStreamer::EmitWinCFIEndProc(SMLoc Loc) {
   MCStreamer::EmitWinCFIEndProc(Loc);
 
   OS << "\t.seh_endproc";
   EmitEOL();

lib/MC/MCDwarf.cpp

const MCSymbol &FrameEmitterImpl::EmitCIE(const MCDwarfFrameInfo &Frame) {
   // CIE ID
   unsigned CIE_ID = IsEH ? 0 : -1;
   Streamer.EmitIntValue(CIE_ID, 4);
 
   // Version
   uint8_t CIEVersion = getCIEVersion(IsEH, context.getDwarfVersion());
   Streamer.EmitIntValue(CIEVersion, 1);
 
-  // Augmentation String
-  SmallString<8> Augmentation;
   if (IsEH) {
+    SmallString<8> Augmentation;
     Augmentation += "z";
     if (Frame.Personality)
       Augmentation += "P";
     if (Frame.Lsda)
       Augmentation += "L";
     Augmentation += "R";
     if (Frame.IsSignalFrame)
       Augmentation += "S";
+    if (!Frame.ExtraAugmentation.empty())
+      Augmentation += Frame.ExtraAugmentation;
     Streamer.EmitBytes(Augmentation);
   }
   Streamer.EmitIntValue(0, 1);
 
   if (CIEVersion >= 4) {
     // Address Size
     Streamer.EmitIntValue(context.getAsmInfo()->getCodePointerSize(), 1);
 
void FrameEmitterImpl::EmitFDE(const MCSymbol &cieStart,
 
   Streamer.EmitLabel(fdeEnd);
 }
 
 namespace {
 
 struct CIEKey {
   static const CIEKey getEmptyKey() {
-    return CIEKey(nullptr, 0, -1, false, false, static_cast<unsigned>(INT_MAX));
+    return CIEKey(nullptr, 0, -1, false, false, static_cast<unsigned>(INT_MAX), "");
   }
 
   static const CIEKey getTombstoneKey() {
-    return CIEKey(nullptr, -1, 0, false, false, static_cast<unsigned>(INT_MAX));
+    return CIEKey(nullptr, -1, 0, false, false, static_cast<unsigned>(INT_MAX), "");
   }
 
   CIEKey(const MCSymbol *Personality, unsigned PersonalityEncoding,
          unsigned LSDAEncoding, bool IsSignalFrame, bool IsSimple,
-         unsigned RAReg)
+         unsigned RAReg, std::string ExtraAugmentation)
       : Personality(Personality), PersonalityEncoding(PersonalityEncoding),
         LsdaEncoding(LSDAEncoding), IsSignalFrame(IsSignalFrame),
-        IsSimple(IsSimple), RAReg(RAReg) {}
+        IsSimple(IsSimple), RAReg(RAReg),
+        ExtraAugmentation(std::move(ExtraAugmentation)) {}
 
   explicit CIEKey(const MCDwarfFrameInfo &Frame)
       : Personality(Frame.Personality),
         PersonalityEncoding(Frame.PersonalityEncoding),
         LsdaEncoding(Frame.LsdaEncoding), IsSignalFrame(Frame.IsSignalFrame),
-        IsSimple(Frame.IsSimple), RAReg(Frame.RAReg) {}
+        IsSimple(Frame.IsSimple), RAReg(Frame.RAReg),
+        ExtraAugmentation(Frame.ExtraAugmentation) {}
 
   const MCSymbol *Personality;
   unsigned PersonalityEncoding;
   unsigned LsdaEncoding;
   bool IsSignalFrame;
   bool IsSimple;
   unsigned RAReg;
+  std::string ExtraAugmentation;
 };
 
 } // end anonymous namespace
 
 namespace llvm {
 
 template <> struct DenseMapInfo<CIEKey> {
   static CIEKey getEmptyKey() { return CIEKey::getEmptyKey(); }
   static CIEKey getTombstoneKey() { return CIEKey::getTombstoneKey(); }
 
   static unsigned getHashValue(const CIEKey &Key) {
     return static_cast<unsigned>(
         hash_combine(Key.Personality, Key.PersonalityEncoding, Key.LsdaEncoding,
-                     Key.IsSignalFrame, Key.IsSimple, Key.RAReg));
+                     Key.IsSignalFrame, Key.IsSimple, Key.RAReg,
+                     Key.ExtraAugmentation));
   }
 
   static bool isEqual(const CIEKey &LHS, const CIEKey &RHS) {
     return LHS.Personality == RHS.Personality &&
            LHS.PersonalityEncoding == RHS.PersonalityEncoding &&
            LHS.LsdaEncoding == RHS.LsdaEncoding &&
            LHS.IsSignalFrame == RHS.IsSignalFrame &&
            LHS.IsSimple == RHS.IsSimple &&
-           LHS.RAReg == RHS.RAReg;
+           LHS.RAReg == RHS.RAReg &&
+           LHS.ExtraAugmentation == RHS.ExtraAugmentation;
   }
 };
 
 } // end namespace llvm
 
 void MCDwarfFrameEmitter::Emit(MCObjectStreamer &Streamer, MCAsmBackend *MAB,
                                bool IsEH) {
   Streamer.generateCompactUnwindEncodings(MAB);

lib/MC/MCParser/AsmParser.cpp

enum DirectiveKind {
     DK_SLEB128,
     DK_ULEB128,
     DK_ERR,
     DK_ERROR,
     DK_WARNING,
     DK_PRINT,
     DK_ADDRSIG,
     DK_ADDRSIG_SYM,
+    DK_APPEND_AUGMENTATION,
     DK_END
   };
 
   /// Maps directive name --> DirectiveKind enum, for
   /// directives parsed by this class.
   StringMap<DirectiveKind> DirectiveKindMap;
 
   // ".ascii", ".asciz", ".string"
private:
   bool parseDirectiveCFIRestoreState();
   bool parseDirectiveCFISameValue(SMLoc DirectiveLoc);
   bool parseDirectiveCFIRestore(SMLoc DirectiveLoc);
   bool parseDirectiveCFIEscape();
   bool parseDirectiveCFIReturnColumn(SMLoc DirectiveLoc);
   bool parseDirectiveCFISignalFrame();
   bool parseDirectiveCFIUndefined(SMLoc DirectiveLoc);
 
+  bool parseDirectiveAppendAugmentation();
+
   // macro directives
   bool parseDirectivePurgeMacro(SMLoc DirectiveLoc);
   bool parseDirectiveExitMacro(StringRef Directive);
   bool parseDirectiveEndMacro(StringRef Directive);
   bool parseDirectiveMacro(SMLoc DirectiveLoc);
   bool parseDirectiveMacrosOnOff(StringRef Directive);
   // alternate macro mode directives
   bool parseDirectiveAltmacro(StringRef Directive);
if (IDVal.startswith(".") && IDVal != ".") {
     case DK_CFI_SIGNAL_FRAME:
       return parseDirectiveCFISignalFrame();
     case DK_CFI_UNDEFINED:
       return parseDirectiveCFIUndefined(IDLoc);
     case DK_CFI_REGISTER:
       return parseDirectiveCFIRegister(IDLoc);
     case DK_CFI_WINDOW_SAVE:
       return parseDirectiveCFIWindowSave();
+    case DK_APPEND_AUGMENTATION:
+      return parseDirectiveAppendAugmentation();
     case DK_MACROS_ON:
     case DK_MACROS_OFF:
       return parseDirectiveMacrosOnOff(IDVal);
     case DK_MACRO:
       return parseDirectiveMacro(IDLoc);
     case DK_ALTMACRO:
     case DK_NOALTMACRO:
       return parseDirectiveAltmacro(IDVal);
bool AsmParser::parseDirectiveCFIUndefined(SMLoc DirectiveLoc) {
 
   if (parseRegisterOrRegisterNumber(Register, DirectiveLoc))
     return true;
 
   getStreamer().EmitCFIUndefined(Register);
   return false;
 }
 
+
+/// parseDirectiveAppendAugmentation
+/// ::= .append_augmentation
+bool AsmParser::parseDirectiveAppendAugmentation() {
+  std::string Data;
+  if (checkForValidSection() || parseEscapedString(Data))
+    return true;
+  getStreamer().EmitAppendAugmentation(Data);;
+  return false;
+}
+
 /// parseDirectiveAltmacro
 /// ::= .altmacro
 /// ::= .noaltmacro
 bool AsmParser::parseDirectiveAltmacro(StringRef Directive) {
   if (getLexer().isNot(AsmToken::EndOfStatement))
     return TokError("unexpected token in '" + Directive + "' directive");
   if (Directive == ".altmacro")
     getLexer().SetAltMacroMode(true);
void AsmParser::initializeDirectiveKindMap() {
   DirectiveKindMap[".ds.l"] = DK_DS_L;
   DirectiveKindMap[".ds.p"] = DK_DS_P;
   DirectiveKindMap[".ds.s"] = DK_DS_S;
   DirectiveKindMap[".ds.w"] = DK_DS_W;
   DirectiveKindMap[".ds.x"] = DK_DS_X;
   DirectiveKindMap[".print"] = DK_PRINT;
   DirectiveKindMap[".addrsig"] = DK_ADDRSIG;
   DirectiveKindMap[".addrsig_sym"] = DK_ADDRSIG_SYM;
+  DirectiveKindMap[".append_augmentation"] = DK_APPEND_AUGMENTATION;
 }
 
 MCAsmMacro *AsmParser::parseMacroLikeBody(SMLoc DirectiveLoc) {
   AsmToken EndToken, StartToken = getTok();
 
   unsigned NestLevel = 0;
   while (true) {
     // Check whether we have reached the end of the file.

lib/MC/MCStreamer.cpp

void MCStreamer::emitDwarfFile0Directive(StringRef Directory,
                                          StringRef Filename,
                                          MD5::MD5Result *Checksum,
                                          Optional<StringRef> Source,
                                          unsigned CUID) {
   getContext().setMCLineTableRootFile(CUID, Directory, Filename, Checksum,
                                       Source);
 }
 
+void MCStreamer::EmitAppendAugmentation(StringRef String) {
+  getCurrentDwarfFrameInfo()->ExtraAugmentation += String.str();
+}
+
 void MCStreamer::EmitDwarfLocDirective(unsigned FileNo, unsigned Line,
                                        unsigned Column, unsigned Flags,
                                        unsigned Isa,
                                        unsigned Discriminator,
                                        StringRef FileName) {
   getContext().setCurrentDwarfLoc(FileNo, Line, Column, Flags, Isa,
                                   Discriminator);
 }

lib/Target/AArch64/AArch64AsmPrinter.cpp

void AArch64AsmPrinter::EmitInstruction(const MachineInstr *MI) {
 
   case TargetOpcode::PATCHABLE_FUNCTION_EXIT:
     LowerPATCHABLE_FUNCTION_EXIT(*MI);
     return;
 
   case TargetOpcode::PATCHABLE_TAIL_CALL:
     LowerPATCHABLE_TAIL_CALL(*MI);
     return;
+
+  case AArch64::EMITBKEY:
+    OutStreamer->EmitAppendAugmentation("B");
+    return;
   }
 
   // Finally, do the automated lowerings for everything else.
   MCInst TmpInst;
   MCInstLowering.Lower(MI, TmpInst);
   EmitToStreamer(*OutStreamer, TmpInst);
 }
 
 // Force static initialization.
 extern "C" void LLVMInitializeAArch64AsmPrinter() {
   RegisterAsmPrinter<AArch64AsmPrinter> X(getTheAArch64leTarget());
   RegisterAsmPrinter<AArch64AsmPrinter> Y(getTheAArch64beTarget());
   RegisterAsmPrinter<AArch64AsmPrinter> Z(getTheARM64Target());
 }

lib/Target/AArch64/AArch64FrameLowering.cpp

void AArch64FrameLowering::emitPrologue(MachineFunction &MF,
   // assume that's false and set it to true in the case that there's a redzone.
   AFI->setHasRedZone(false);
 
   // Debug location must be unknown since the first debug location is used
   // to determine the end of the prologue.
   DebugLoc DL;
 
   if (ShouldSignReturnAddress(MF)) {
-    BuildMI(
-        MBB, MBBI, DL,
-        TII->get(ShouldSignWithAKey(MF) ? AArch64::PACIASP : AArch64::PACIBSP))
+    if (ShouldSignWithAKey(MF))
+      BuildMI(MBB, MBBI, DL, TII->get(AArch64::PACIASP))
+          .setMIFlag(MachineInstr::FrameSetup);
+    else
+    {
+      BuildMI(MBB, MBBI, DL, TII->get(AArch64::EMITBKEY))
+          .setMIFlag(MachineInstr::FrameSetup);
+      BuildMI(MBB, MBBI, DL, TII->get(AArch64::PACIBSP))
           .setMIFlag(MachineInstr::FrameSetup);
+    }
 
     unsigned CFIIndex =
         MF.addFrameInst(MCCFIInstruction::createNegateRAState(nullptr));
     BuildMI(MBB, MBBI, DL, TII->get(TargetOpcode::CFI_INSTRUCTION))
         .addCFIIndex(CFIIndex)
         .setMIFlags(MachineInstr::FrameSetup);
   }
 

lib/Target/AArch64/AArch64InstrInfo.td

 // This is a directive-like pseudo-instruction. The purpose is to insert an
 // R_AARCH64_TLSDESC_CALL relocation at the offset of the following instruction
 // (which in the usual case is a BLR).
 let hasSideEffects = 1 in
 def TLSDESCCALL : Pseudo<(outs), (ins i64imm:$sym), []>, Sched<[]> {
   let AsmString = ".tlsdesccall $sym";
 }
 
+// Pseudo instruction to tell the streamer to emit a 'B' character into the
+// augmentation string.
+def EMITBKEY : Pseudo<(outs), (ins), []>, Sched<[]> {}
+
 // FIXME: maybe the scratch register used shouldn't be fixed to X1?
 // FIXME: can "hasSideEffects be dropped?
 let isCall = 1, Defs = [LR, X0, X1], hasSideEffects = 1,
     isCodeGenOnly = 1 in
 def TLSDESC_CALLSEQ
     : Pseudo<(outs), (ins i64imm:$sym),
              [(AArch64tlsdesc_callseq tglobaltlsaddr:$sym)]>,
       Sched<[WriteI, WriteLD, WriteI, WriteBrReg]>;

test/CodeGen/MIR/AArch64/return-address-signing.mir

frameInfo:
   hasVAStart:      false
   hasMustTailInVarArgFunc: false
   localFrameSize:  0
   savePoint:       ''
   restorePoint:    ''
 fixedStack:      
 stack:           
 constants:       
+#CHECK:    frame-setup EMITBKEY
 #CHECK:    frame-setup PACIBSP implicit-def $lr, implicit $lr, implicit $sp
 #CHECK:    frame-setup CFI_INSTRUCTION negate_ra_sign_state 
 #CHECK:    RETAB
 body:             |
   bb.0.entry:
     $w0 = MOVi32imm 2
     RET_ReallyLR implicit killed $w0
 
 ...
 

test/DebugInfo/AArch64/return-address-signing.ll

+; RUN: llc -mtriple=aarch64-arm-none-eabi < %s -filetype=obj -o - \
+; RUN:    | llvm-dwarfdump -v - | FileCheck -check-prefix=CHECK %s
+
+;CHECK: CIE
+;CHECK: Augmentation:          "zR"
+define i32 @foo()  "sign-return-address"="all" {
+  ret i32 0
+}
+
+;CHECK: CIE
+;CHECK: Augmentation:          "zRB"
+
+define i32 @bar()  "sign-return-address"="all" "sign-return-address-key"="b_key" {
+  ret i32 0
+}

test/MC/ELF/appendaugmentation.s

+// RUN: llvm-mc -filetype=obj -triple aarch64-arm-none-eabi %s -o - | llvm-dwarfdump - -v | FileCheck %s
+#CHECK:   Augmentation:          "zRB"
+f1:
+        .cfi_startproc
+        .append_augmentation "B"
+        .cfi_endproc