AI/rocm-systems
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
* Fix use-after-free for ATT code objects

* Formatting

* Temporary fix for last kernel

---------

Co-authored-by: Giovanni Baraldi <gbaraldi@amd.com>
This commit is contained in:
Baraldi, Giovanni
2024-12-04 19:22:57 +01:00
committed by GitHub
parent 85ad519ca7
commit 21c577ba60
4 changed files with 49 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files
source/lib/rocprofiler-sdk/thread_trace/code_object.cpp
+39 -32
View File
@@ -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