[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

22

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

22

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

114

Any reason for x86_64-only?

[Amir]

bolt/unittests/Core/MCPlusBuilder.cpp

114

createCall is not defined in AArch64MCPlusBuilder

[yota9]

bolt/unittests/Core/MCPlusBuilder.cpp

114

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

[Amir]

bolt/unittests/Core/MCPlusBuilder.cpp

114

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

114

Oh yes, I forgot about that, thanks.

[maksfb]

bolt/unittests/Core/MCPlusBuilder.cpp

114

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

113–128

Still under X86 :)

[Amir]

bolt/unittests/Core/MCPlusBuilder.cpp

113–128

Good catch!

[Amir]

Move out from X86_AVAILABLE

[yota9]

Thank you @Amir!

[Amir]

Thanks @yota9, @maksfb!

bolt/include/bolt/Core/MCPlusAnnotation.h

+//===- bolt/Core/MCPlusAnnotation.h - Interface for MCPlus ------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains the definitions of annotation encoding/decoding functions.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef BOLT_CORE_MCPLUSANNOTATION_H
+#define BOLT_CORE_MCPLUSANNOTATION_H
+
+#include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/MathExtras.h"
+
+/// We encode Index and Value into a 64-bit immediate operand value.
+inline int64_t encodeAnnotationImm(uint8_t Index, int64_t Value) {
+  if (LLVM_UNLIKELY(Value !=
+                    llvm::SignExtend64<56>(Value & 0xff'ffff'ffff'ffff)))
+    llvm::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.
+inline uint8_t extractAnnotationIndex(int64_t ImmValue) {
+  return ImmValue >> 56;
+}
+
+/// Extract annotation value from immediate operand value.
+inline int64_t extractAnnotationValue(int64_t ImmValue) {
+  return llvm::SignExtend64<56>(ImmValue & 0xff'ffff'ffff'ffffULL);
+}
+
+#endif

bolt/include/bolt/Core/MCPlusBuilder.h

 // means to create/analyze/modify instructions at MCPlus level.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef BOLT_CORE_MCPLUSBUILDER_H
 #define BOLT_CORE_MCPLUSBUILDER_H
 
 #include "bolt/Core/MCPlus.h"
+#include "bolt/Core/MCPlusAnnotation.h"
 #include "bolt/Core/Relocation.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/MC/MCAsmBackend.h"
 #include "llvm/MC/MCExpr.h"
 #include "llvm/MC/MCInst.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");
-
-    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) {
-    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;
-  }
-
   MCInst *getAnnotationInst(const MCInst &Inst) const {
     if (Inst.getNumOperands() == 0)
       return nullptr;
 
     const MCOperand &LastOp = Inst.getOperand(Inst.getNumOperands() - 1);
     if (!LastOp.isInst())
       return nullptr;
 

bolt/unittests/Core/CMakeLists.txt

 set(LLVM_LINK_COMPONENTS
   BOLTCore
   BOLTRewrite
+  BOLTUtils
   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

 #ifdef AARCH64_AVAILABLE
 #include "AArch64Subtarget.h"
 #endif // AARCH64_AVAILABLE
 
 #ifdef X86_AVAILABLE
 #include "X86Subtarget.h"
 #endif // X86_AVAILABLE
 
+#include "bolt/Core/MCPlusAnnotation.h"
 #include "bolt/Rewrite/RewriteInstance.h"
 #include "llvm/BinaryFormat/ELF.h"
 #include "llvm/DebugInfo/DWARF/DWARFContext.h"
 #include "llvm/Object/ELFObjectFile.h"
 #include "llvm/Support/TargetSelect.h"
 #include "gtest/gtest.h"
 
 using namespace llvm;

 }
 
 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) {
+  // Round-trip encoding-decoding check for negative values
+  int64_t Value = INT32_MIN;
+  uint8_t Index = 123;
+  int64_t EncodedAnnotation = encodeAnnotationImm(Index, Value);
+  int64_t DecodedValue = extractAnnotationValue(EncodedAnnotation);
+  uint8_t DecodedIndex = extractAnnotationIndex(EncodedAnnotation);
+  ASSERT_EQ(Value, DecodedValue);
+  ASSERT_EQ(Index, DecodedIndex);
+
+  // Large int64 should trigger an out of range assertion
+  Value = 0x1FF'FFFF'FFFF'FFFFULL;
+  ASSERT_DEATH(encodeAnnotationImm(0, Value), "annotation value out of range");
+}