Fix use-after-free for ATT code objects (#35)

* Fix use-after-free for ATT code objects

* Formatting

* Temporary fix for last kernel

---------

Co-authored-by: Giovanni Baraldi <gbaraldi@amd.com>

[ROCm/rocprofiler-sdk commit: 21c577ba60]

projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/thread_trace/att_core.cpp

@@ -565,6 +565,8 @@ finalize()
     {
         if(ctx->agent_thread_trace) ctx->agent_thread_trace->resource_deinit();
     }
+
+    code_object::finalize();
 }
 
 }  // namespace thread_trace

projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/thread_trace/code_object.cpp

@@ -21,55 +21,49 @@
 // SOFTWARE.
 
 #include "lib/rocprofiler-sdk/thread_trace/code_object.hpp"
+#include "lib/common/static_object.hpp"
+#include "lib/common/synchronized.hpp"
 #include "lib/rocprofiler-sdk/code_object/code_object.hpp"
 
+#include <mutex>
+#include <set>
+
 namespace rocprofiler
 {
 namespace thread_trace
 {
 namespace code_object
 {
-std::mutex                         CodeobjCallbackRegistry::mut;
-std::set<CodeobjCallbackRegistry*> CodeobjCallbackRegistry::all_registries{};
+using set_type_t = std::set<CodeobjCallbackRegistry*>;
+
+auto*
+_static_registries()
+{
+    static auto*& reg = common::static_object<common::Synchronized<set_type_t>>::construct();
+    return reg;
+}
+
+auto&
+get_registries()
+{
+    return *CHECK_NOTNULL(_static_registries());
+}
 
 CodeobjCallbackRegistry::CodeobjCallbackRegistry(LoadCallback _ld, UnloadCallback _unld)
 : ld_fn(std::move(_ld))
 , unld_fn(std::move(_unld))
 {
-    std::unique_lock<std::mutex> lg(mut);
-    all_registries.insert(this);
+    get_registries().wlock([this](set_type_t& t) { t.insert(this); });
 }
 
 CodeobjCallbackRegistry::~CodeobjCallbackRegistry()
 {
-    std::unique_lock<std::mutex> lg(mut);
-    all_registries.erase(this);
-}
-
-void
-CodeobjCallbackRegistry::Load(rocprofiler_agent_id_t agent,
-                              uint64_t               id,
-                              uint64_t               addr,
-                              uint64_t               size)
-{
-    std::unique_lock<std::mutex> lg(mut);
-    for(auto* reg : all_registries)
-        reg->ld_fn(agent, id, addr, size);
-}
-
-void
-CodeobjCallbackRegistry::Unload(uint64_t id)
-{
-    std::unique_lock<std::mutex> lg(mut);
-    for(auto* reg : all_registries)
-        reg->unld_fn(id);
+    if(auto* reg = _static_registries()) reg->wlock([this](set_type_t& t) { t.erase(this); });
 }
 
 void
 CodeobjCallbackRegistry::IterateLoaded() const
 {
-    std::unique_lock<std::mutex> lg(mut);
-
     rocprofiler::code_object::iterate_loaded_code_objects(
         [&](const rocprofiler::code_object::hsa::code_object& code_object) {
             const auto& data = code_object.rocp_data;
@@ -104,9 +98,12 @@ executable_freeze(hsa_executable_t executable, const char* options)
         [&](const rocprofiler::code_object::hsa::code_object& code_object) {
             if(code_object.hsa_executable != executable) return;
 
-            const auto& data = code_object.rocp_data;
-            CodeobjCallbackRegistry::Load(
-                data.rocp_agent, data.code_object_id, data.load_delta, data.load_size);
+            const auto& co = code_object.rocp_data;
+
+            get_registries().wlock([&](set_type_t& t) {
+                for(auto* reg : t)
+                    reg->ld_fn(co.rocp_agent, co.code_object_id, co.load_delta, co.load_size);
+            });
         });
 
     return HSA_STATUS_SUCCESS;
@@ -117,8 +114,12 @@ executable_destroy(hsa_executable_t executable)
 {
     rocprofiler::code_object::iterate_loaded_code_objects(
         [&](const rocprofiler::code_object::hsa::code_object& code_object) {
-            if(code_object.hsa_executable == executable)
-                CodeobjCallbackRegistry::Unload(code_object.rocp_data.code_object_id);
+            if(code_object.hsa_executable != executable) return;
+
+            get_registries().wlock([&](set_type_t& t) {
+                for(auto* reg : t)
+                    reg->unld_fn(code_object.rocp_data.code_object_id);
+            });
         });
 
     // Call underlying function
@@ -142,6 +143,12 @@ initialize(HsaApiTable* table)
         << "infinite recursion";
 }
 
+void
+finalize()
+{
+    get_registries().wlock([](set_type_t& t) { t.clear(); });
+}
+
 }  // namespace code_object
 }  // namespace thread_trace
 }  // namespace rocprofiler

projects/rocprofiler-sdk/source/lib/rocprofiler-sdk/thread_trace/code_object.hpp

@@ -27,8 +27,6 @@
 #include <hsa/hsa_api_trace.h>
 
 #include <functional>
-#include <mutex>
-#include <set>
 
 namespace rocprofiler
 {
@@ -44,20 +42,18 @@ struct CodeobjCallbackRegistry
     CodeobjCallbackRegistry(LoadCallback ld, UnloadCallback unld);
     virtual ~CodeobjCallbackRegistry();
 
-    void        IterateLoaded() const;
-    static void Load(rocprofiler_agent_id_t agent, uint64_t id, uint64_t addr, uint64_t size);
-    static void Unload(uint64_t id);
+    void IterateLoaded() const;
 
-private:
     LoadCallback   ld_fn;
     UnloadCallback unld_fn;
-
-    static std::mutex                         mut;
-    static std::set<CodeobjCallbackRegistry*> all_registries;
 };
 
 void
 initialize(HsaApiTable* table);
+
+void
+finalize();
+
 }  // namespace code_object
 }  // namespace thread_trace
 }  // namespace rocprofiler

projects/rocprofiler-sdk/tests/thread-trace/multi_dispatch.cpp

@@ -43,6 +43,7 @@
 #include <vector>
 
 constexpr double WAVE_RATIO_TOLERANCE = 0.05;
+constexpr size_t NUM_KERNELS          = 5;
 
 namespace ATTTest
 {
@@ -59,12 +60,12 @@ dispatch_callback(rocprofiler_queue_id_t /* queue_id  */,
                   rocprofiler_user_data_t* dispatch_userdata,
                   void*                    userdata)
 {
-    C_API_BEGIN
+    static std::atomic<size_t> count{0};
+    if(count.fetch_add(1) > NUM_KERNELS) return ROCPROFILER_ATT_CONTROL_NONE;
 
     assert(userdata && "Dispatch callback passed null!");
     dispatch_userdata->ptr = userdata;
 
-    C_API_END
     return ROCPROFILER_ATT_CONTROL_START_AND_STOP;
 }