[BOLT][NFC] Fix undefined behavior in encodeAnnotationImm

[Amir]

Summary

Fix UBSan-reported issue in MCPlusBuilder::encodeAnnotationImm (left shift of a
negative value).

Test Plan:

ninja check-bolt
...
PASS: BOLT-Unit :: Core/./CoreTests/AArch64/MCPlusBuilderTester.Annotation/0 (1 of 140)
PASS: BOLT-Unit :: Core/./CoreTests/X86/MCPlusBuilderTester.Annotation/0 (131 of 134)

[yota9]

bolt/include/bolt/Core/MCPlusAnnotation.h

21 ↗

(On Diff #410328)

May I suggest with replacing assert with if condition? I think it is the important thing that we should check with disabled assertions too.

[Amir]

bolt/include/bolt/Core/MCPlusAnnotation.h

21 ↗

(On Diff #410328)

Good point.

[Amir]

Check annotation value in release builds

[yota9]

Sorry, why the index assert has gone? It is unlikely to reach this limit, but why not to check it anyway :)

[Amir]

In D120260#3338530, @yota9 wrote:

Sorry, why the index assert has gone? It is unlikely to reach this limit, but why not to check it anyway :)

It's gone because I changed Index type from unsigned to uint8_t. We're only using 1 byte of data, no need to pass larger type.

[yota9]

@Amir Sorry, missed that since the extractAnnotationIndex and the test are still using unsigned values type. Could you please fix that?

[Amir]

I wanted to test MCPlusBuilder’s private methods (encodeAnnotationImm and friends) that are not directly exposed through public methods. The options were:

1. add test class as a friend of MCPlusBuilder (less invasive, future unit tests may follow)
2. move these methods into another file (eg MCPlusAnnotation.h) and include it into MCPlusBuilder.h and test source, the downside is code churn and pollution of the global namespace with these new functions,
3. test using hackish indirect way (through GnuArgsSize annotation) that doesn’t require any changes to MCPlusBuilder source, does the job but test code would be ugly.

This diff goes with option 2, but I'm open to the feedback.

[Amir]

Change extractAnnotationIndex return type from unsigned to uint8_t

[yota9]

LGTM

[Amir]

Use addEHInfo to test encodeAnnotationImm without exposing it as public method/function

[Amir]

Fix shared build

[Amir]

Use extractAnnotationValue in the check

[maksfb]

bolt/unittests/Core/MCPlusBuilder.cpp

113

Any reason for x86_64-only?

[Amir]

bolt/unittests/Core/MCPlusBuilder.cpp

113

createCall is not defined in AArch64MCPlusBuilder

[yota9]

bolt/unittests/Core/MCPlusBuilder.cpp

113

Sorry, but why not just use BC->MIB->addAnnotation directly?

[Amir]

bolt/unittests/Core/MCPlusBuilder.cpp

113

There's no way to put a predefined Value into an encoded annotation immediate using addAnnotation.
addAnnotation allocates a MCSimpleAnnotation<ValueType> using placement new, and passes the pointer to allocated object to
setAnnotationOpValue.

https://github.com/llvm/llvm-project/blob/main/bolt/include/bolt/Core/MCPlusBuilder.h#L1639

[yota9]

bolt/unittests/Core/MCPlusBuilder.cpp

113

Oh yes, I forgot about that, thanks.

[maksfb]

bolt/unittests/Core/MCPlusBuilder.cpp

113

createCall is not defined in AArch64MCPlusBuilder

Interesting. But it has `createTailCall(). Can we use that for testing?

[Amir]

Use createTailCall

[maksfb]

LGTM. Thanks for the fix.

[yota9]

bolt/unittests/Core/MCPlusBuilder.cpp

112–138

Still under X86 :)

[Amir]

bolt/unittests/Core/MCPlusBuilder.cpp

112–138

Good catch!

[Amir]

Move out from X86_AVAILABLE

[yota9]

Thank you @Amir!

[Amir]

Thanks @yota9, @maksfb!

bolt/include/bolt/Core/MCPlusBuilder.h

private:
 
   /// A set of annotation allocators
   std::unordered_map<AllocatorIdTy, AnnotationAllocator> AnnotationAllocators;
 
   /// A variable that is used to generate unique ids for annotation allocators
   AllocatorIdTy MaxAllocatorId = 0;
 
   /// We encode Index and Value into a 64-bit immediate operand value.
-  static int64_t encodeAnnotationImm(unsigned Index, int64_t Value) {
-    assert(Index < 256 && "annotation index max value exceeded");
-    assert((Value == (Value << 8) >> 8) && "annotation value out of range");
+  static int64_t encodeAnnotationImm(uint8_t Index, int64_t Value) {
+    if (LLVM_UNLIKELY(Value != extractAnnotationValue(Value)))
+      report_fatal_error("annotation value out of range");
 
     Value &= 0xff'ffff'ffff'ffff;
     Value |= (int64_t)Index << 56;
 
     return Value;
   }
 
   /// Extract annotation index from immediate operand value.
-  static unsigned extractAnnotationIndex(int64_t ImmValue) {
+  static uint8_t extractAnnotationIndex(int64_t ImmValue) {
     return ImmValue >> 56;
   }
 
   /// Extract annotation value from immediate operand value.
   static int64_t extractAnnotationValue(int64_t ImmValue) {
-    ImmValue &= 0xff'ffff'ffff'ffff;
-    return (ImmValue << 8) >> 8;
+    return SignExtend64<56>(ImmValue & 0xff'ffff'ffff'ffffULL);
   }
 
   MCInst *getAnnotationInst(const MCInst &Inst) const {
     if (Inst.getNumOperands() == 0)
       return nullptr;
 
     const MCOperand &LastOp = Inst.getOperand(Inst.getNumOperands() - 1);
     if (!LastOp.isInst())

bolt/unittests/Core/CMakeLists.txt

 set(LLVM_LINK_COMPONENTS
   BOLTCore
   BOLTRewrite
   DebugInfoDWARF
   Object
+  MC
   ${LLVM_TARGETS_TO_BUILD}
   )
 
 add_bolt_unittest(CoreTests
   MCPlusBuilder.cpp
   )
 
 if ("AArch64" IN_LIST LLVM_TARGETS_TO_BUILD)

bolt/unittests/Core/MCPlusBuilder.cpp

 }
 
 TEST_P(MCPlusBuilderTester, AliasSmallerAX) {
   uint64_t AliasesAX[] = {X86::AX, X86::AL, X86::AH};
   size_t AliasesAXCount = sizeof(AliasesAX) / sizeof(*AliasesAX);
   testRegAliases(Triple::x86_64, X86::AX, AliasesAX, AliasesAXCount, true);
 }
 
 #endif // X86_AVAILABLE
+
+TEST_P(MCPlusBuilderTester, Annotation) {
+  MCInst Inst;
+  bool Success = BC->MIB->createTailCall(Inst, BC->Ctx->createNamedTempSymbol(),
+                                         BC->Ctx.get());
+  ASSERT_TRUE(Success);
+  MCSymbol *LPSymbol = BC->Ctx->createNamedTempSymbol("LP");
+  uint64_t Value = INT32_MIN;
+  // Test encodeAnnotationImm using this indirect way
+  BC->MIB->addEHInfo(Inst, MCPlus::MCLandingPad(LPSymbol, Value));
+  // Round-trip encoding-decoding check for negative values
+  Optional<MCPlus::MCLandingPad> EHInfo = BC->MIB->getEHInfo(Inst);
+  ASSERT_TRUE(EHInfo.hasValue());
+  MCPlus::MCLandingPad LP = EHInfo.getValue();
+  uint64_t DecodedValue = LP.second;
+  ASSERT_EQ(Value, DecodedValue);
+
+  // Large int64 should trigger an out of range assertion
+  Value = 0x1FF'FFFF'FFFF'FFFFULL;
+  Inst.clear();
+  Success = BC->MIB->createTailCall(Inst, BC->Ctx->createNamedTempSymbol(),
+                                    BC->Ctx.get());
+  ASSERT_TRUE(Success);
+  ASSERT_DEATH(BC->MIB->addEHInfo(Inst, MCPlus::MCLandingPad(LPSymbol, Value)),
+               "annotation value out of range");
+}