Use RAII for ensuring that mprotect calls are undone

[dberris]

Summary

This fixes an mprotect leak identified in D21612.

[echristo]

lib/xray/xray_interface.cc

29 ↗

(On Diff #65346)

Comment describing the class.

[dberris]

• Document the MProtectHelper class

[rSerge]

Looks good.

compiler-rt/trunk/lib/xray/xray_interface.cc

 #include <limits>
 #include <sys/mman.h>
 
 namespace __xray {
 
 // This is the function to call when we encounter the entry or exit sleds.
 std::atomic<void (*)(int32_t, XRayEntryType)> XRayPatchedFunction{nullptr};
 
+// MProtectHelper is an RAII wrapper for calls to mprotect(...) that will undo
+// any successful mprotect(...) changes. This is used to make a page writeable
+// and executable, and upon destruction if it was successful in doing so returns
+// the page into a read-only and executable page.
+//
+// This is only used specifically for runtime-patching of the XRay
+// instrumentation points. This assumes that the executable pages are originally
+// read-and-execute only.
+class MProtectHelper {
+  void *PageAlignedAddr;
+  std::size_t MProtectLen;
+  bool MustCleanup;
+
+public:
+  explicit MProtectHelper(void *PageAlignedAddr, std::size_t MProtectLen)
+      : PageAlignedAddr(PageAlignedAddr), MProtectLen(MProtectLen),
+        MustCleanup(false) {}
+
+  int MakeWriteable() {
+    auto R = mprotect(PageAlignedAddr, MProtectLen,
+                      PROT_READ | PROT_WRITE | PROT_EXEC);
+    if (R != -1)
+      MustCleanup = true;
+    return R;
+  }
+
+  ~MProtectHelper() {
+    if (MustCleanup) {
+      mprotect(PageAlignedAddr, MProtectLen, PROT_READ | PROT_EXEC);
+    }
+  }
+};
+
 } // namespace __xray
 
 extern "C" {
 // The following functions have to be defined in assembler, on a per-platform
 // basis. See xray_trampoline_*.s files for implementations.
 extern void __xray_FunctionEntry();
 extern void __xray_FunctionExit();
 }
 
 extern std::atomic<bool> XRayInitialized;
 extern std::atomic<__xray::XRaySledMap> XRayInstrMap;
 
 int __xray_set_handler(void (*entry)(int32_t, XRayEntryType)) {
   if (XRayInitialized.load(std::memory_order_acquire)) {
     __xray::XRayPatchedFunction.store(entry, std::memory_order_release);
     return 1;
   }
   return 0;
 }
 
 std::atomic<bool> XRayPatching{false};
 
+using namespace __xray;
+
 XRayPatchingStatus __xray_patch() {
   // FIXME: Make this happen asynchronously. For now just do this sequentially.
   if (!XRayInitialized.load(std::memory_order_acquire))
     return XRayPatchingStatus::NOT_INITIALIZED; // Not initialized.
 
   static bool NotPatching = false;
   if (!XRayPatching.compare_exchange_strong(NotPatching, true,
                                             std::memory_order_acq_rel,
                                             std::memory_order_acquire)) {
     return XRayPatchingStatus::ONGOING; // Already patching.
   }
 
   // Step 1: Compute the function id, as a unique identifier per function in the
   // instrumentation map.
-  __xray::XRaySledMap InstrMap = XRayInstrMap.load(std::memory_order_acquire);
+  XRaySledMap InstrMap = XRayInstrMap.load(std::memory_order_acquire);
   if (InstrMap.Entries == 0)
     return XRayPatchingStatus::NOT_INITIALIZED;
 
   int32_t FuncId = 1;
   static constexpr uint8_t CallOpCode = 0xe8;
   static constexpr uint16_t MovR10Seq = 0xba41;
   static constexpr uint8_t JmpOpCode = 0xe9;
   uint64_t CurFun = 0;
   for (std::size_t I = 0; I < InstrMap.Entries; I++) {
     auto Sled = InstrMap.Sleds[I];
     auto F = Sled.Function;
     if (CurFun == 0)
       CurFun = F;
     if (F != CurFun) {
       ++FuncId;
       CurFun = F;
     }
 
     // While we're here, we should patch the nop sled. To do that we mprotect
     // the page containing the function to be writeable.
     void *PageAlignedAddr =
         reinterpret_cast<void *>(Sled.Address & ~((2 << 16) - 1));
     std::size_t MProtectLen =
         (Sled.Address + 12) - reinterpret_cast<uint64_t>(PageAlignedAddr);
-    if (mprotect(PageAlignedAddr, MProtectLen,
-                 PROT_READ | PROT_WRITE | PROT_EXEC) == -1) {
+    MProtectHelper Protector(PageAlignedAddr, MProtectLen);
+    if (Protector.MakeWriteable() == -1) {
       printf("Failed mprotect: %d\n", errno);
       return XRayPatchingStatus::FAILED;
     }
 
     static constexpr int64_t MinOffset{std::numeric_limits<int32_t>::min()};
     static constexpr int64_t MaxOffset{std::numeric_limits<int32_t>::max()};
     if (Sled.Kind == XRayEntryType::ENTRY) {
       // Here we do the dance of replacing the following sled:
if (Sled.Kind == XRayEntryType::EXIT) {
       }
       *reinterpret_cast<uint32_t *>(Sled.Address + 2) = FuncId;
       *reinterpret_cast<uint8_t *>(Sled.Address + 6) = JmpOpCode;
       *reinterpret_cast<uint32_t *>(Sled.Address + 7) = TrampolineOffset;
       std::atomic_store_explicit(
           reinterpret_cast<std::atomic<uint16_t> *>(Sled.Address), MovR10Seq,
           std::memory_order_release);
     }
-
-    if (mprotect(PageAlignedAddr, MProtectLen, PROT_READ | PROT_EXEC) == -1) {
-      printf("Failed mprotect: %d\n", errno);
-      return XRayPatchingStatus::FAILED;
-    }
   }
   XRayPatching.store(false, std::memory_order_release);
   return XRayPatchingStatus::NOTIFIED;
 }