[cfi] Support for dlopen and dlclose

[eugenis]

Summary

Add dlopen/dlclose interceptors to update CFI shadow for loaded/unloaded libraries.

[pcc]

lib/cfi/cfi.cc

331

As you mentioned offline, this won't correctly handle dependencies.

Perhaps the code should be building the shadow from scratch using dl_iterate_phdr every time a rebuild is needed? That way, you save needing to worry about dependencies. This may require a redesign of D16098 to avoid the copy.

[eugenis]

A new approach.
dlclose may unmap (and dlopen may map) multiple libraries, so we actually need to do the entire dl_iterate_phdr run and update all libraries. For dlclose, we also need to detect libraries that have disappeared from the address space, so we do dl_iterate_phdr before and after dlclose and compare the results.

[pcc]

lib/cfi/cfi.cc

40

Looks like we already have a kMaxNumberOfModules constant defined.

172

I changed this interface in r257858.

235–261

Can you avoid needing to take the set difference by rebuilding from scratch?

lib/sanitizer_common/sanitizer_common.cc

344

Should there be a single init() function that combines set and set_phdr?

lib/sanitizer_common/sanitizer_common.h

635

Is this function needed? Looks like you're passing a comparator to your sort function above.

[eugenis]

lib/cfi/cfi.cc

235–261

Do you mean unmap or clean the entire shadow and then add the existing modules? That would be racy, we can not ensure that there are no concurrent virtual/indirect calls.

lib/sanitizer_common/sanitizer_common.h

635

removed

[pcc]

lib/cfi/cfi.cc

235–261

Yes, that's why I was thinking this would need to be done after D16098. If it's too painful to re-order the patches I'd be fine with combining them.

[eugenis]

Yet another new approach, way simpler than before.
Deprecates D16098.

[eugenis]

The previous approach contained a data race on the shadow pointer, which can not be updated without a lock in __cfi_slowpath. That would not be acceptable performance-wise. Probably.

This new version is linux-only as it requires mremap().

[eugenis]

how does it look?

[pcc]

lib/cfi/cfi.cc

34

Maybe assert somewhere that this is equal to GetPageSize()?

38

This name is unclear. It should reflect that the shadow pointer is stored here, not the shadow itself.

48

Please inline into ShadowBuilder::Install, as the pre-conditions of this function aren't clear without context.

101

Do you need to set the memory protection for shadow_ first?

327

You can use it now I guess.

lib/sanitizer_common/sanitizer_posix.cc

178 ↗

(On Diff #45288)

This function is unused.

[eugenis]

lib/cfi/cfi.cc

101

Of course.
Good catch.

[pcc]

lib/cfi/cfi.cc

17

This is done now, right?

101

Thanks. Is it possible to construct a test case that makes sure we get the memory protection right? I was thinking that the test case would try to write to the shadow at various points (e.g. start of main, after dlopen, after dlclose), and use %expect_crash to test that the write fails.

[eugenis]

lib/cfi/cfi.cc

101

But we don't really want there to be an easy way to find the cfi shadow.

[eugenis]

I'll move most of the implementation under a namespace in a separate commit.

[pcc]

LGTM

[eugenis]

Thanks, http://reviews.llvm.org/rL258857

lib/cfi/cfi.cc

 //===-------- cfi.cc ------------------------------------------------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the runtime support for the cross-DSO CFI.
 //
 //===----------------------------------------------------------------------===//
 
 // FIXME: Intercept dlopen/dlclose.
 // FIXME: Support diagnostic mode.
 // FIXME: Harden:
 //  * mprotect shadow, use mremap for updates
 //  * something else equally important
 
 #include <assert.h>
 #include <elf.h>
 #include <link.h>
 #include <string.h>
 
 typedef ElfW(Phdr) Elf_Phdr;
 typedef ElfW(Ehdr) Elf_Ehdr;
 
 #include "interception/interception.h"
 #include "sanitizer_common/sanitizer_common.h"
 #include "sanitizer_common/sanitizer_flag_parser.h"
 #include "ubsan/ubsan_init.h"
 #include "ubsan/ubsan_flags.h"
 
 static uptr __cfi_shadow;
 static constexpr uptr kShadowGranularity = 12;
 static constexpr uptr kShadowAlign = 1UL << kShadowGranularity; // 4096
 
 static constexpr uint16_t kInvalidShadow = 0;
 static constexpr uint16_t kUncheckedShadow = 0xFFFFU;
 
 static uint16_t *mem_to_shadow(uptr x) {
   return (uint16_t *)(__cfi_shadow + ((x >> kShadowGranularity) << 1));
 }
 
 typedef int (*CFICheckFn)(u64, void *);
 
 class ShadowValue {
   uptr addr;
   uint16_t v;
   explicit ShadowValue(uptr addr, uint16_t v) : addr(addr), v(v) {}
 
 public:
   bool is_invalid() const { return v == kInvalidShadow; }
 
   bool is_unchecked() const { return v == kUncheckedShadow; }
 
for (; q < end; q += kShadowAlign) {
     assert((uptr)ShadowValue::load(q).get_cfi_check() == cfi_check);
     assert((uptr)ShadowValue::load(q + kShadowAlign / 2).get_cfi_check() ==
            cfi_check);
     assert((uptr)ShadowValue::load(q + kShadowAlign - 1).get_cfi_check() ==
            cfi_check);
   }
 }
 
 // This is a workaround for a glibc bug:
 // https://sourceware.org/bugzilla/show_bug.cgi?id=15199
 // Other platforms can, hopefully, just do
 //    dlopen(RTLD_NOLOAD | RTLD_LAZY)
 //    dlsym("__cfi_check").
-static uptr find_cfi_check_in_dso(dl_phdr_info *info) {
+static uptr find_cfi_check_in_dso(LoadedModule *module) {
+  const ElfW(Phdr) *phdr = (const ElfW(Phdr) *)module->phdr();
+  if (!phdr) {
+    Printf("%s %zx\n", module->full_name(), module->base_address());
+  }
+  CHECK(phdr);
+  int phnum = module->phnum();
+  uptr base = module->base_address();
   const ElfW(Dyn) *dynamic = nullptr;
-  for (int i = 0; i < info->dlpi_phnum; ++i) {
-    if (info->dlpi_phdr[i].p_type == PT_DYNAMIC) {
-      dynamic =
-          (const ElfW(Dyn) *)(info->dlpi_addr + info->dlpi_phdr[i].p_vaddr);
+  for (int i = 0; i < phnum; ++i) {
+    if (phdr[i].p_type == PT_DYNAMIC) {
+      dynamic = (const ElfW(Dyn) *)(base + phdr[i].p_vaddr);
       break;
     }
   }
-  if (!dynamic) return 0;
+  if (!dynamic)
+    return 0;
   uptr strtab = 0, symtab = 0;
   for (const ElfW(Dyn) *p = dynamic; p->d_tag != PT_NULL; ++p) {
     if (p->d_tag == DT_SYMTAB)
       symtab = p->d_un.d_ptr;
     else if (p->d_tag == DT_STRTAB)
       strtab = p->d_un.d_ptr;
   }
 
   if (symtab > strtab) {
     VReport(1, "Can not handle: symtab > strtab (%p > %zx)\n", symtab, strtab);
     return 0;
   }
 
   // Verify that strtab and symtab are inside of the same LOAD segment.
   // This excludes VDSO, which has (very high) bogus strtab and symtab pointers.
   int phdr_idx;
-  for (phdr_idx = 0; phdr_idx < info->dlpi_phnum; phdr_idx++) {
-    const Elf_Phdr *phdr = &info->dlpi_phdr[phdr_idx];
-    if (phdr->p_type == PT_LOAD) {
-      uptr beg = info->dlpi_addr + phdr->p_vaddr;
-      uptr end = beg + phdr->p_memsz;
+  for (phdr_idx = 0; phdr_idx < phnum; phdr_idx++) {
+    const Elf_Phdr *ph = &phdr[phdr_idx];
+    if (ph->p_type == PT_LOAD) {
+      uptr beg = base + ph->p_vaddr;
+      uptr end = beg + ph->p_memsz;
       if (strtab >= beg && strtab < end && symtab >= beg && symtab < end)
         break;
     }
   }
-  if (phdr_idx == info->dlpi_phnum) {
+  if (phdr_idx == phnum) {
     // Nope, either different segments or just bogus pointers.
     // Can not handle this.
     VReport(1, "Can not handle: symtab %p, strtab %zx\n", symtab, strtab);
     return 0;
   }
 
   for (const ElfW(Sym) *p = (const ElfW(Sym) *)symtab; (ElfW(Addr))p < strtab;
        ++p) {
-    char *name = (char*)(strtab + p->st_name);
+    char *name = (char *)(strtab + p->st_name);
     if (strcmp(name, "__cfi_check") == 0) {
       assert(p->st_info == ELF32_ST_INFO(STB_GLOBAL, STT_FUNC));
-      uptr addr = info->dlpi_addr + p->st_value;
+      uptr addr = base + p->st_value;
       return addr;
     }
   }
   return 0;
 }
 
-static int dl_iterate_phdr_cb(dl_phdr_info *info, size_t size, void *data) {
-  uptr cfi_check = find_cfi_check_in_dso(info);
+void add_module(LoadedModule *module) {
+  uptr cfi_check = find_cfi_check_in_dso(module);
   if (cfi_check)
-    VReport(1, "Module '%s' __cfi_check %zx\n", info->dlpi_name, cfi_check);
+    VReport(1, "Module '%s' __cfi_check %zx\n", module->full_name(), cfi_check);
 
-  for (int i = 0; i < info->dlpi_phnum; i++) {
-    const Elf_Phdr *phdr = &info->dlpi_phdr[i];
-    if (phdr->p_type == PT_LOAD) {
-      // Jump tables are in the executable segment.
-      // VTables are in the non-executable one.
-      // Need to fill shadow for both.
-      // FIXME: reject writable if vtables are in the r/o segment. Depend on
-      // PT_RELRO?
-      uptr cur_beg = info->dlpi_addr + phdr->p_vaddr;
-      uptr cur_end = cur_beg + phdr->p_memsz;
-      if (cfi_check) {
-        VReport(1, "   %zx .. %zx\n", cur_beg, cur_end);
-        fill_shadow(cur_beg, cur_end, cfi_check ? cfi_check : (uptr)(-1));
-      } else {
-        fill_shadow_constant(cur_beg, cur_end, kUncheckedShadow);
-      }
-    }
-  }
-  return 0;
+  for (const auto &r : module->ranges()) {
+      VReport(1, "   %zx .. %zx\n", r.beg, r.end);
+      fill_shadow(r.beg, r.end, cfi_check);
+      fill_shadow_constant(r.beg, r.end, kUncheckedShadow);
+
+void remove_module(LoadedModule *module) {
+  for (const auto &r : module->ranges())
+    fill_shadow_constant(r.beg, r.end, kInvalidShadow);
+}
+
+// Update CFI shadow on dlopen/dlclose.
+// Records the current set of loaded DSOs in start(),
+// then each call to update() detects changes to that set and updates shadow
+// appropriately.
+// stop() releases all allocated memory.
+
+// This is a very approximate way of handling dlopen/dlclose and it has possible
+// false positives around module constructors/destructors. Ex., the shadow for a
+// loaded library is not updated until after dlopen() returns, which causes the
+// library to be invalid target for indirect/virtual calls until the end of its
+// constructors.
+// Ideally, the shadow must be updated inside dlopen (after the library has been
+// relocated but before any constructors are run) and dlclose (immediately
+// before unmapping the library pages).
+class ShadowUpdater {
+  LoadedModule *modules_before;
+  uptr modules_before_cnt;
+
+public:
+  void start();
+  void update();
+  void stop();
+};
+
+void ShadowUpdater::start() {
+  CHECK_EQ(nullptr, modules_before);
+  modules_before = (LoadedModule *)MmapOrDie(
+      kMaxNumberOfModules * sizeof(LoadedModule), "ShadowUpdater");
+  modules_before_cnt =
+      GetListOfModules(modules_before, kMaxNumberOfModules, nullptr);
+}
+
+void ShadowUpdater::stop() {
+  UnmapOrDie(modules_before, kMaxNumberOfModules * sizeof(LoadedModule));
+  modules_before = nullptr;
+}
+
+bool LoadedModuleCompare(const LoadedModule &a, const LoadedModule &b) {
+  return a.base_address() < b.base_address();
+}
+
+void ShadowUpdater::update() {
+  LoadedModule *modules_after = (LoadedModule *)MmapOrDie(
+      kMaxNumberOfModules * sizeof(LoadedModule), "ShadowUpdater");
+  uptr modules_after_cnt =
+      GetListOfModules(modules_after, kMaxNumberOfModules, nullptr);
+
+  InternalSort(&modules_before, modules_before_cnt, LoadedModuleCompare);
+  InternalSort(&modules_after, modules_after_cnt, LoadedModuleCompare);
+
+  LoadedModule *p = modules_before;
+  LoadedModule *p_end = p + modules_before_cnt;
+  LoadedModule *q = modules_after;
+  LoadedModule *q_end = q + modules_after_cnt;
+  while (p < p_end && q < q_end) {
+    if (p->base_address() < q->base_address()) {
+      remove_module(p++);
+    } else if (p->base_address() > q->base_address()) {
+      add_module(q++);
+    } else {
+      p++;
+      q++;
+    }
+  }
+  while (p < p_end)
+    remove_module(p++);
+  while (q < q_end)
+    add_module(q++);
+
+  UnmapOrDie(modules_before, kMaxNumberOfModules * sizeof(LoadedModule));
+  modules_before = modules_after;
+  modules_before_cnt = modules_after_cnt;
 }
 
 // Fill shadow for the initial libraries.
 static void init_shadow() {
-  dl_iterate_phdr(dl_iterate_phdr_cb, nullptr);
+  LoadedModule *modules = (LoadedModule *)MmapOrDie(
+      kMaxNumberOfModules * sizeof(LoadedModule), "init_shadow");
+  uptr modules_cnt = GetListOfModules(modules, kMaxNumberOfModules, nullptr);
+  for (uptr i = 0; i < modules_cnt; ++i)
+    add_module(&modules[i]);
+  UnmapOrDie(modules, kMaxNumberOfModules * sizeof(LoadedModule));
+}
+
+static THREADLOCAL int in_loader;
+static ShadowUpdater shadow_update;
+static BlockingMutex shadow_update_lock(LINKER_INITIALIZED);
+
+void EnterLoader() {
+  if (in_loader == 0) {
+    shadow_update_lock.Lock();
+    shadow_update.start();
+  }
+  ++in_loader;
+}
+
+void ExitLoader() {
+  CHECK(in_loader > 0);
+  --in_loader;
+  shadow_update.update();
+  if (in_loader == 0) {
+    shadow_update.stop();
+    shadow_update_lock.Unlock();
+  }
+}
+
+// Setup shadow for dlopen()ed libraries.
+// The actual shadow setup happens after dlopen() returns, which means that
+// a library can not be a target of any CFI checks while its constructors are
+// running. It's unclear how to fix this without some extra help from libc.
+// In glibc, mmap inside dlopen is not interceptable.
+// Maybe a seccomp-bpf filter?
+// We could insert a high-priority constructor into the library, but that would
+// not help with the uninstrumented libraries.
+INTERCEPTOR(void*, dlopen, const char *filename, int flag) {
+  EnterLoader();
+  void *handle = REAL(dlopen)(filename, flag);
+  ExitLoader();
+  return handle;
+}
+
+INTERCEPTOR(int, dlclose, void *handle) {
+  EnterLoader();
+  int res = REAL(dlclose)(handle);
+  ExitLoader();
+  return res;
 }
 
 extern "C" SANITIZER_INTERFACE_ATTRIBUTE
 void __cfi_slowpath(u64 CallSiteTypeId, void *Ptr) {
   uptr Addr = (uptr)Ptr;
   VReport(3, "__cfi_slowpath: %llx, %p\n", CallSiteTypeId, Ptr);
   ShadowValue sv = ShadowValue::load(Addr);
   if (sv.is_invalid()) {
     VReport(2, "CFI: invalid memory region for a function pointer (shadow==0): %p\n", Ptr);
-    Die();
+    __builtin_trap();
   }
   if (sv.is_unchecked()) {
     VReport(2, "CFI: unchecked call (shadow=FFFF): %p\n", Ptr);
     return;
   }
   CFICheckFn cfi_check = sv.get_cfi_check();
   VReport(2, "__cfi_check at %p\n", cfi_check);
   cfi_check(CallSiteTypeId, Ptr);
 }
 
 static void InitializeFlags() {
   SetCommonFlagsDefaults();
void __cfi_init() {
   uptr shadow_size = (vma >> (kShadowGranularity - 1)) + 1;
   VReport(1, "CFI: VMA size %zx, shadow size %zx\n", vma, shadow_size);
   void *shadow = MmapNoReserveOrDie(shadow_size, "CFI shadow");
   VReport(1, "CFI: shadow at %zx .. %zx\n", shadow,
           reinterpret_cast<uptr>(shadow) + shadow_size);
   __cfi_shadow = (uptr)shadow;
   init_shadow();
 
+  INTERCEPT_FUNCTION(dlopen);
+  INTERCEPT_FUNCTION(dlclose);
+
 #ifdef CFI_ENABLE_DIAG
   __ubsan::InitAsPlugin();
 #endif
 }
 
 #if SANITIZER_CAN_USE_PREINIT_ARRAY
 // On ELF platforms, run cfi initialization before any other constructors.
 // On other platforms we use the constructor attribute to arrange to run our
 // initialization early.
 extern "C" {
 __attribute__((section(".preinit_array"),
                used)) void (*__cfi_preinit)(void) = __cfi_init;
 }
 #endif

lib/sanitizer_common/sanitizer_common.h

uptr InternalBinarySearch(const Container &v, uptr first, uptr last,
   return not_found;
 }
 
 // Represents a binary loaded into virtual memory (e.g. this can be an
 // executable or a shared object).
 class LoadedModule {
  public:
   LoadedModule() : full_name_(nullptr), base_address_(0) { ranges_.clear(); }
-  void set(const char *module_name, uptr base_address);
+  void init(const char *module_name, uptr base_address, uptr phdr, int phnum);
   void clear();
   void addAddressRange(uptr beg, uptr end, bool executable);
   bool containsAddress(uptr address) const;
 
   const char *full_name() const { return full_name_; }
   uptr base_address() const { return base_address_; }
+  uptr phdr() { return phdr_; }
+  int phnum() { return phnum_; }
 
   struct AddressRange {
     AddressRange *next;
     uptr beg;
     uptr end;
     bool executable;
 
     AddressRange(uptr beg, uptr end, bool executable)
         : next(nullptr), beg(beg), end(end), executable(executable) {}
   };
 
   const IntrusiveList<AddressRange> &ranges() const { return ranges_; }
 
  private:
   char *full_name_;  // Owned.
   uptr base_address_;
+  uptr phdr_;
+  int phnum_;
   IntrusiveList<AddressRange> ranges_;
 };
 
 // OS-dependent function that fills array with descriptions of at most
 // "max_modules" currently loaded modules. Returns the number of
 // initialized modules. If filter is nonzero, ignores modules for which
 // filter(full_name) is false.
 typedef bool (*string_predicate_t)(const char *);

lib/sanitizer_common/sanitizer_common.cc

while (*s != '\0') {
     z++;
     s++;
   }
 
   // Null terminate the string.
   *z = '\0';
 }
 
-void LoadedModule::set(const char *module_name, uptr base_address) {
+void LoadedModule::init(const char *module_name, uptr base_address, uptr phdr,
+                        int phnum) {
   clear();
   full_name_ = internal_strdup(module_name);
   base_address_ = base_address;
+  phdr_ = phdr;
+  phnum_ = phnum;
 }
 
 void LoadedModule::clear() {
   InternalFree(full_name_);
   full_name_ = nullptr;
+  phdr_ = 0;
+  phnum_ = 0;
   while (!ranges_.empty()) {
     AddressRange *r = ranges_.front();
     ranges_.pop_front();
     InternalFree(r);
   }
 }
 
 void LoadedModule::addAddressRange(uptr beg, uptr end, bool executable) {

lib/sanitizer_common/sanitizer_linux_libcdep.cc

static int dl_iterate_phdr_cb(dl_phdr_info *info, size_t size, void *arg) {
   } else if (info->dlpi_name) {
     module_name.append("%s", info->dlpi_name);
   }
   if (module_name[0] == '\0')
     return 0;
   if (data->filter && !data->filter(module_name.data()))
     return 0;
   LoadedModule *cur_module = &data->modules[data->current_n];
-  cur_module->set(module_name.data(), info->dlpi_addr);
+  cur_module->init(module_name.data(), info->dlpi_addr, (uptr)info->dlpi_phdr,
+                   info->dlpi_phnum);
   data->current_n++;
   for (int i = 0; i < info->dlpi_phnum; i++) {
     const Elf_Phdr *phdr = &info->dlpi_phdr[i];
     if (phdr->p_type == PT_LOAD) {
       uptr cur_beg = info->dlpi_addr + phdr->p_vaddr;
       uptr cur_end = cur_beg + phdr->p_memsz;
       bool executable = phdr->p_flags & PF_X;
       cur_module->addAddressRange(cur_beg, cur_end, executable);

lib/sanitizer_common/sanitizer_procmaps_common.cc

for (uptr i = 0; n_modules < max_modules &&
     //   instruction address in virtual memory (as code section
     //   is mapped to a fixed memory range).
     // * If a binary is compiled with -pie, all the modules are
     //   mapped high at address space (in particular, higher than
     //   shadow memory of the tool), so the module can't be the
     //   first entry.
     uptr base_address = (i ? cur_beg : 0) - cur_offset;
     LoadedModule *cur_module = &modules[n_modules];
-    cur_module->set(cur_name, base_address);
+    cur_module->init(cur_name, base_address, 0, 0);
     cur_module->addAddressRange(cur_beg, cur_end, prot & kProtectionExecute);
     n_modules++;
   }
   return n_modules;
 }
 
 void GetMemoryProfile(fill_profile_f cb, uptr *stats, uptr stats_size) {
   char *smaps = nullptr;

lib/sanitizer_common/sanitizer_procmaps_mac.cc

for (uptr i = 0; n_modules < max_modules &&
     if (filter && !filter(cur_name))
       continue;
     LoadedModule *cur_module = nullptr;
     if (n_modules > 0 &&
         0 == internal_strcmp(cur_name, modules[n_modules - 1].full_name())) {
       cur_module = &modules[n_modules - 1];
     } else {
       cur_module = &modules[n_modules];
-      cur_module->set(cur_name, cur_beg);
+      cur_module->init(cur_name, cur_beg, 0, 0);
       n_modules++;
     }
     cur_module->addAddressRange(cur_beg, cur_end, prot & kProtectionExecute);
   }
   return n_modules;
 }
 
 }  // namespace __sanitizer
 
 #endif  // SANITIZER_MAC

lib/sanitizer_common/sanitizer_win.cc

for (size_t i = 0; i < nun_modules && count < max_modules; ++i) {
     // Adjust the base address of the module so that we get a VA instead of an
     // RVA when computing the module offset. This helps llvm-symbolizer find the
     // right DWARF CU. In the common case that the image is loaded at it's
     // preferred address, we will now print normal virtual addresses.
     uptr preferred_base = GetPreferredBase(&module_name[0]);
     uptr adjusted_base = base_address - preferred_base;
 
     LoadedModule *cur_module = &modules[count];
-    cur_module->set(module_name, adjusted_base);
+    cur_module->init(module_name, adjusted_base, 0, 0);
     // We add the whole module as one single address range.
     cur_module->addAddressRange(base_address, end_address, /*executable*/ true);
     count++;
   }
   UnmapOrDie(hmodules, modules_buffer_size);
 
   return count;
 };

test/cfi/cross-dso/dlopen.cpp

+// RUN: %clangxx_cfi_dso -DSHARED_LIB %s -fPIC -shared -o %t1-so.so
+// RUN: %clangxx_cfi_dso %s -o %t1
+// RUN: %expect_crash %t1 2>&1 | FileCheck --check-prefix=CFI %s
+// RUN: %expect_crash %t1 cast 2>&1 | FileCheck --check-prefix=CFI-CAST %s
+// RUN: %expect_crash %t1 dlclose 2>&1 | FileCheck --check-prefix=CFI %s
+
+// RUN: %clangxx_cfi_dso -DB32 -DSHARED_LIB %s -fPIC -shared -o %t2-so.so
+// RUN: %clangxx_cfi_dso -DB32 %s -o %t2
+// RUN: %expect_crash %t2 2>&1 | FileCheck --check-prefix=CFI %s
+// RUN: %expect_crash %t2 cast 2>&1 | FileCheck --check-prefix=CFI-CAST %s
+// RUN: %expect_crash %t2 dlclose 2>&1 | FileCheck --check-prefix=CFI %s
+
+// RUN: %clangxx_cfi_dso -DB64 -DSHARED_LIB %s -fPIC -shared -o %t3-so.so
+// RUN: %clangxx_cfi_dso -DB64 %s -o %t3
+// RUN: %expect_crash %t3 2>&1 | FileCheck --check-prefix=CFI %s
+// RUN: %expect_crash %t3 cast 2>&1 | FileCheck --check-prefix=CFI-CAST %s
+// RUN: %expect_crash %t3 dlclose 2>&1 | FileCheck --check-prefix=CFI %s
+
+// RUN: %clangxx_cfi_dso -DBM -DSHARED_LIB %s -fPIC -shared -o %t4-so.so
+// RUN: %clangxx_cfi_dso -DBM %s -o %t4
+// RUN: %expect_crash %t4 2>&1 | FileCheck --check-prefix=CFI %s
+// RUN: %expect_crash %t4 cast 2>&1 | FileCheck --check-prefix=CFI-CAST %s
+// RUN: %expect_crash %t4 dlclose 2>&1 | FileCheck --check-prefix=CFI %s
+
+// RUN: %clangxx -g -DBM -DSHARED_LIB -DNOCFI %s -fPIC -shared -o %t5-so.so
+// RUN: %clangxx -g -DBM -DNOCFI %s -ldl -o %t5
+// RUN: %t5 2>&1 | FileCheck --check-prefix=NCFI %s
+// RUN: %t5 cast 2>&1 | FileCheck --check-prefix=NCFI %s
+// RUN: %t5 dlclose 2>&1 | FileCheck --check-prefix=NCFI %s
+
+// Test that calls to uninstrumented library are unchecked.
+// RUN: %clangxx -DBM -DSHARED_LIB %s -fPIC -shared -o %t6-so.so
+// RUN: %clangxx_cfi_dso -DBM %s -o %t6
+// RUN: %t6 2>&1 | FileCheck --check-prefix=NCFI %s
+// RUN: %t6 cast 2>&1 | FileCheck --check-prefix=NCFI %s
+
+// Call-after-dlclose is checked on the caller side.
+// RUN: %expect_crash %t6 dlclose 2>&1 | FileCheck --check-prefix=CFI %s
+
+// Tests calls into dlopen-ed library.
+// REQUIRES: cxxabi
+
+#include <assert.h>
+#include <dlfcn.h>
+#include <stdio.h>
+#include <stdint.h>
+#include <string.h>
+#include <sys/mman.h>
+
+struct A {
+  virtual void f();
+};
+
+#ifdef SHARED_LIB
+
+#include "../utils.h"
+struct B {
+  virtual void f();
+};
+void B::f() {}
+
+extern "C" void *create_B() {
+  create_derivers<B>();
+  return (void *)(new B());
+}
+
+extern "C" void do_nothing() __attribute__((aligned(4096))) {}
+
+#else
+
+void A::f() {}
+
+static const int kCodeAlign = 4096;
+static const int kCodeSize = 4096;
+static char saved_code[kCodeSize];
+static char *real_start;
+
+static void save_code(char *p) {
+  real_start = (char *)(((uintptr_t)p) & ~(kCodeAlign - 1));
+  memcpy(saved_code, real_start, kCodeSize);
+}
+
+static void restore_code() {
+  char *code = (char *)mmap(real_start, kCodeSize, PROT_WRITE | PROT_EXEC,
+                            MAP_PRIVATE | MAP_ANONYMOUS | MAP_FIXED, 0, 0);
+  assert(code == real_start);
+  memcpy(code, saved_code, kCodeSize);
+}
+
+int main(int argc, char *argv[]) {
+  const bool test_cast = argc > 1 && strcmp(argv[1], "cast") == 0;
+  const bool test_dlclose = argc > 1 && strcmp(argv[1], "dlclose") == 0;
+
+  char name[100];
+  snprintf(name, sizeof(name), "%s-so.so", argv[0]);
+  void *handle = dlopen(name, RTLD_NOW);
+  assert(handle);
+  void *(*create_B)() = (void *(*)())dlsym(handle, "create_B");
+  assert(create_B);
+
+  void *p = create_B();
+  A *a;
+
+  // CFI: =0=
+  // CFI-CAST: =0=
+  // NCFI: =0=
+  fprintf(stderr, "=0=\n");
+
+  if (test_cast) {
+    // Test cast. BOOM.
+    a = (A*)p;
+  } else {
+    // Invisible to CFI. Test virtual call later.
+    memcpy(&a, &p, sizeof(a));
+  }
+
+  // CFI: =1=
+  // CFI-CAST-NOT: =1=
+  // NCFI: =1=
+  fprintf(stderr, "=1=\n");
+
+  if (test_dlclose) {
+    // Imitate an attacker sneaking in an executable page where a dlclose()d
+    // library was loaded. This needs to pass w/o CFI, so for the testing
+    // purpose, we just copy the bytes of a "void f() {}" function back and
+    // forth.
+    void (*do_nothing)() = (void (*)())dlsym(handle, "do_nothing");
+    assert(do_nothing);
+    save_code((char *)do_nothing);
+
+    int res = dlclose(handle);
+    assert(res == 0);
+
+    restore_code();
+
+    do_nothing(); // UB here
+  } else {
+    a->f(); // UB here
+  }
+
+  // CFI-NOT: =2=
+  // CFI-CAST-NOT: =2=
+  // NCFI: =2=
+  fprintf(stderr, "=2=\n");
+}
+#endif