Subzero: Replace global pointers to allocation functions

src/IceASanInstrumentation.cpp

@@ -99,9 +99,12 @@ void ASanInstrumentation::instrumentGlobals(VariableDeclarationList &Globals) {
   NewGlobals.push_back(RzArray);
   NewGlobals.push_back(RzSizes);
 
+  using PrototypeMap = std::unordered_map<std::string, FunctionDeclaration *>;
+  PrototypeMap ProtoSubstitutions;
   for (VariableDeclaration *Global : Globals) {
     assert(Global->getAlignment() <= RzSize);
     VariableDeclaration *RzLeft = VariableDeclaration::create(&NewGlobals);
+    VariableDeclaration *NewGlobal = Global;
     VariableDeclaration *RzRight = VariableDeclaration::create(&NewGlobals);
     RzLeft->setName(Ctx, nextRzName());
     RzRight->setName(Ctx, nextRzName());
@@ -109,22 +112,68 @@ void ASanInstrumentation::instrumentGlobals(VariableDeclarationList &Globals) {
     SizeT RzLeftSize = Alignment;
     SizeT RzRightSize =
         RzSize + Utils::OffsetToAlignment(Global->getNumBytes(), Alignment);
-    if (Global->hasNonzeroInitializer()) {
+    if (!Global->hasNonzeroInitializer()) {
+      RzLeft->addInitializer(VariableDeclaration::ZeroInitializer::create(
+          &NewGlobals, RzLeftSize));
+      RzRight->addInitializer(VariableDeclaration::ZeroInitializer::create(
+          &NewGlobals, RzRightSize));
+    } else {
       RzLeft->addInitializer(VariableDeclaration::DataInitializer::create(
           &NewGlobals, llvm::NaClBitcodeRecord::RecordVector(RzLeftSize, 'R')));
       RzRight->addInitializer(VariableDeclaration::DataInitializer::create(
           &NewGlobals,
           llvm::NaClBitcodeRecord::RecordVector(RzRightSize, 'R')));
-    } else {
-      RzLeft->addInitializer(VariableDeclaration::ZeroInitializer::create(
-          &NewGlobals, RzLeftSize));
-      RzRight->addInitializer(VariableDeclaration::ZeroInitializer::create(
-          &NewGlobals, RzRightSize));
+
+      // replace any pointers to allocator functions
+      NewGlobal = VariableDeclaration::create(&NewGlobals);
+      NewGlobal->setName(Global->getName());
+      std::vector<VariableDeclaration::Initializer *> GlobalInits =
+          Global->getInitializers();
+      for (VariableDeclaration::Initializer *Init : GlobalInits) {
+        auto *RelocInit =
+            llvm::dyn_cast<VariableDeclaration::RelocInitializer>(Init);
+        if (RelocInit == nullptr) {
+          NewGlobal->addInitializer(Init);
+          continue;
+        }
+        const GlobalDeclaration *TargetDecl = RelocInit->getDeclaration();
+        const auto *TargetFunc =
+            llvm::dyn_cast<FunctionDeclaration>(TargetDecl);
+        if (TargetFunc == nullptr) {
+          NewGlobal->addInitializer(Init);
+          continue;
+        }
+        std::string TargetName = TargetDecl->getName().toStringOrEmpty();
+        StringMap::const_iterator Subst = FuncSubstitutions.find(TargetName);
+        if (Subst == FuncSubstitutions.end()) {
+          NewGlobal->addInitializer(Init);
+          continue;
+        }
+        std::string SubstName = Subst->second;
+        PrototypeMap::iterator SubstProtoEntry =
+            ProtoSubstitutions.find(SubstName);
+        FunctionDeclaration *SubstProto;
+        if (SubstProtoEntry != ProtoSubstitutions.end())
+          SubstProto = SubstProtoEntry->second;
+        else {
+          constexpr bool IsProto = true;
+          SubstProto = FunctionDeclaration::create(
+              Ctx, TargetFunc->getSignature(), TargetFunc->getCallingConv(),
+              llvm::GlobalValue::ExternalLinkage, IsProto);
+          SubstProto->setName(Ctx, SubstName);
+          ProtoSubstitutions.insert({SubstName, SubstProto});
+        }
+
+        NewGlobal->addInitializer(VariableDeclaration::RelocInitializer::create(
+            &NewGlobals, SubstProto, RelocOffsetArray(0)));
+      }
     }
+
     RzLeft->setIsConstant(Global->getIsConstant());
+    NewGlobal->setIsConstant(Global->getIsConstant());
     RzRight->setIsConstant(Global->getIsConstant());
     RzLeft->setAlignment(Alignment);
-    Global->setAlignment(Alignment);
+    NewGlobal->setAlignment(Alignment);
     RzRight->setAlignment(1);
     RzArray->addInitializer(VariableDeclaration::RelocInitializer::create(
         &NewGlobals, RzLeft, RelocOffsetArray(0)));
@@ -136,11 +185,11 @@ void ASanInstrumentation::instrumentGlobals(VariableDeclarationList &Globals) {
         &NewGlobals, sizeToByteVec(RzRightSize)));
 
     NewGlobals.push_back(RzLeft);
-    NewGlobals.push_back(Global);
+    NewGlobals.push_back(NewGlobal);
     NewGlobals.push_back(RzRight);
     RzGlobalsNum += 2;
 
-    GlobalSizes.insert({Global->getName(), Global->getNumBytes()});
+    GlobalSizes.insert({NewGlobal->getName(), NewGlobal->getNumBytes()});
   }
 
   // Replace old list of globals, without messing up arena allocators

tests_lit/asan_tests/globalreplacement.ll

+; Test that global pointers to allocation functions are replaced
+
+; REQUIRES: allow_dump
+
+; RUN: %p2i -i %s --args -verbose=global_init -threads=0 -fsanitize-address \
+; RUN:     -allow-externally-defined-symbols | FileCheck --check-prefix=DUMP %s
+
+declare external i32 @malloc(i32)
+declare external i32 @realloc(i32, i32)
+declare external i32 @calloc(i32, i32)
+declare external void @free(i32)
+declare external void @foo()
+
+@global_malloc = internal global i32 ptrtoint (i32 (i32)* @malloc to i32)
+@global_realloc = internal global i32 ptrtoint (i32 (i32, i32)* @realloc to i32)
+@global_calloc = internal global i32 ptrtoint (i32 (i32, i32)* @calloc to i32)
+@global_free = internal global i32 ptrtoint (void (i32)* @free to i32)
+@global_foo = internal global i32 ptrtoint (void ()* @foo to i32)
+
+@constant_malloc = internal constant i32 ptrtoint (i32 (i32)* @malloc to i32)
+@constant_realloc = internal constant i32 ptrtoint (i32 (i32, i32)* @realloc to i32)
+@constant_calloc = internal constant i32 ptrtoint (i32 (i32, i32)* @calloc to i32)
+@constant_free = internal constant i32 ptrtoint (void (i32)* @free to i32)
+@constant_foo = internal constant i32 ptrtoint (void ()* @foo to i32)
+
+@multiple_initializers = internal global <{i32, i32}> <{i32 ptrtoint (i32 (i32)* @malloc to i32), i32 ptrtoint (void (i32)* @free to i32)}>
+
+define void @func() {
+  ret void
+}
+
+; DUMP: Instrumented Globals
+; DUMP-NEXT: @__$rz_array
+; DUMP-NEXT: @__$rz_sizes
+; DUMP-NEXT: @__$rz0
+; DUMP-NEXT: @global_malloc = internal global i32
+; DUMP-SAME:   ptrtoint (i32 (i32)* @__asan_malloc to i32)
+; DUMP-NEXT: @__$rz1
+; DUMP-NEXT: @__$rz2
+; DUMP-NEXT: @global_realloc = internal global i32
+; DUMP-SAME:   ptrtoint (i32 (i32, i32)* @__asan_realloc to i32)
+; DUMP-NEXT: @__$rz3
+; DUMP-NEXT: @__$rz4
+; DUMP-NEXT: @global_calloc = internal global i32
+; DUMP-SAME:   ptrtoint (i32 (i32, i32)* @__asan_calloc to i32)
+; DUMP-NEXT: @__$rz5
+; DUMP-NEXT: @__$rz6
+; DUMP-NEXT: @global_free = internal global i32
+; DUMP-SAME:   ptrtoint (void (i32)* @__asan_free to i32)
+; DUMP-NEXT: @__$rz7
+; DUMP-NEXT: @__$rz8
+; DUMP-NEXT: @global_foo = internal global i32
+; DUMP-SAME:   ptrtoint (void ()* @foo to i32)
+; DUMP-NEXT: @__$rz9
+; DUMP-NEXT: @__$rz10
+; DUMP-NEXT: @constant_malloc = internal constant i32
+; DUMP-SAME:   ptrtoint (i32 (i32)* @__asan_malloc to i32)
+; DUMP-NEXT: @__$rz11
+; DUMP-NEXT: @__$rz12
+; DUMP-NEXT: @constant_realloc = internal constant i32
+; DUMP-SAME:   ptrtoint (i32 (i32, i32)* @__asan_realloc to i32)
+; DUMP-NEXT: @__$rz13
+; DUMP-NEXT: @__$rz14
+; DUMP-NEXT: @constant_calloc = internal constant i32
+; DUMP-SAME:   ptrtoint (i32 (i32, i32)* @__asan_calloc to i32)
+; DUMP-NEXT: @__$rz15
+; DUMP-NEXT: @__$rz16
+; DUMP-NEXT: @constant_free = internal constant i32
+; DUMP-SAME:   ptrtoint (void (i32)* @__asan_free to i32)
+; DUMP-NEXT: @__$rz17
+; DUMP-NEXT: @__$rz18
+; DUMP-NEXT: @constant_foo = internal constant i32
+; DUMP-SAME:   ptrtoint (void ()* @foo to i32)
+; DUMP-NEXT: @__$rz19
+; DUMP-NEXT: @__$rz20
+; DUMP-NEXT: @multiple_initializers = internal global <{ i32, i32 }>
+; DUMP-SAME:   <{ i32 ptrtoint (i32 (i32)* @__asan_malloc to i32),
+; DUMP-SAME:      i32 ptrtoint (void (i32)* @__asan_free to i32) }>
+; DUMP-NEXT: @__$rz21
\ No newline at end of file