Skip to content

[IAST] Fix racing conditions in Dataflow #7838

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 5 commits into
base: master
Choose a base branch
from
Open

[IAST] Fix racing conditions in Dataflow #7838

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
41f968b
Fix racing conditions
NachoEchevarria Nov 18, 2025
d196232
make _initialized atomic
NachoEchevarria Nov 18, 2025
e742073
Fix toctou issues
NachoEchevarria Nov 18, 2025
4828694
Add include to avoid transitive includes
NachoEchevarria Nov 19, 2025
64ff1b8
Add additional Initialized check
NachoEchevarria Nov 19, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
64 changes: 48 additions & 16 deletions tracer/src/Datadog.Tracer.Native/iast/dataflow.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -182,7 +182,9 @@ Dataflow::Dataflow(ICorProfilerInfo* profiler, std::shared_ptr<RejitHandler> rej

Dataflow::~Dataflow()
{
_initialized = false;
CSGUARD(_cs);
Copy link
Contributor

@daniel-romano-DD daniel-romano-DD Nov 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Dataflow is destroyed on profiler unload and this is the only place where _initialized is set to false. This CS is unnecessary (although does not hurt).

SetInitialized(false);

REL(_profiler);
DEL_MAP_VALUES(_modules);
DEL_MAP_VALUES(_appDomains);
Expand All @@ -191,9 +193,9 @@ Dataflow::~Dataflow()

void Dataflow::LoadAspects(WCHAR** aspects, int aspectsLength, UINT32 enabledCategories, UINT32 platform)
{
if (!_initialized)
if (!IsInitialized())
{
_initialized = true;
SetInitialized(true);
Copy link
Member

@andrewlock andrewlock Nov 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Silly question, but could this also be part of the problem? 🤔

We set initialized = true here before we've finished the initialization. That means that all the other places where we do if (!IsInitialized()) and bail out will stop bailing 🤔 Also, we don't have a CSGUARD in here, should we?

I feel like we probably need something more like this?

CSGUARD(_cs);
if (!IsInitialized())
{
    // Do all the other work in the if block

    SetInitialized(true);
}

Copy link
Member

@tonyredondo tonyredondo Nov 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yes I do think you need to guard in this one as well.

Copy link
Contributor

@daniel-romano-DD daniel-romano-DD Nov 19, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The main problem is here, in LoadAspects, and this function is the one that has to be well covered. _isInitialized should be set to true as the last thing done here, not in the begining of the function.


// Init aspects
DBG("Dataflow::LoadAspects -> Processing aspects... ", aspectsLength, " Enabled categories: ", enabledCategories, " Platform: ", platform);
Expand Down Expand Up @@ -237,9 +239,15 @@ void Dataflow::LoadAspects(WCHAR** aspects, int aspectsLength, UINT32 enabledCat

LoadSecurityControls();

auto moduleAspects = _moduleAspects;
_moduleAspects.clear();
DEL_MAP_VALUES(moduleAspects);
// Move the map out under the lock
std::map<ModuleID, ModuleAspects*> oldModuleAspects;
{
CSGUARD(_cs);
oldModuleAspects.swap(_moduleAspects);
}

// Destroy outside the lock
DEL_MAP_VALUES(oldModuleAspects)

trace::Logger::Info("Dataflow::LoadAspects -> read ", _aspects.size(), " aspects");
m_rejitHandler->RegisterRejitter(this);
Expand Down Expand Up @@ -369,12 +377,12 @@ void Dataflow::LoadSecurityControls()

HRESULT Dataflow::AppDomainShutdown(AppDomainID appDomainId)
{
if (!_initialized)
CSGUARD(_cs);
Copy link
Member

@andrewlock andrewlock Nov 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we want to use double checking here to avoid contention? I'm not sure it's a big problem for AppDomainShutdown really, and how often would we not be initialized? 🤔

// Initial check outside the lock
if (!IsInitialized())
{
    return S_OK;
}

CSGUARD(_cs);
// Double check inside the lock
if (!IsInitialized())
{
    return S_OK;
}

Copy link
Member

@tonyredondo tonyredondo Nov 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If we add the guard for each IsInitialized() call we should consider to just use a normal bool instead of the atomic bool.

Copy link
Contributor

@daniel-romano-DD daniel-romano-DD Nov 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We are not going to find the situation where a threar reads true while it has been set to false. The problem is the oposite, in initialization. IMO the whole thing could be solved with deferring the set to true in _isInitialized to the end of the LoadAspects function.

if (!IsInitialized())
{
return S_OK;
}

CSGUARD(_cs);
auto it = _appDomains.find(appDomainId);
if (it != _appDomains.end())
{
Expand All @@ -388,7 +396,7 @@ HRESULT Dataflow::AppDomainShutdown(AppDomainID appDomainId)

HRESULT Dataflow::ModuleLoaded(ModuleID moduleId, ModuleInfo** pModuleInfo)
{
if (!_initialized)
if (!IsInitialized())
Copy link
Member

@andrewlock andrewlock Nov 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This still has the CSGUARD after the IsInitialized() check (because it takes the lock in GetModuleInfo... should we move this check inside GetModuleInfo() instead?

Copy link
Collaborator Author

@NachoEchevarria NachoEchevarria Nov 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

GetModuleInfo() already checks for that, this is just an early check.

{
return S_OK;
}
Expand All @@ -399,12 +407,12 @@ HRESULT Dataflow::ModuleLoaded(ModuleID moduleId, ModuleInfo** pModuleInfo)

HRESULT Dataflow::ModuleUnloaded(ModuleID moduleId)
{
if (!_initialized)
CSGUARD(_cs);
if (!IsInitialized())
{
return S_OK;
}

CSGUARD(_cs);
{
auto it = _moduleAspects.find(moduleId);
if (it != _moduleAspects.end())
Expand Down Expand Up @@ -510,6 +518,11 @@ ICorProfilerInfo* Dataflow::GetCorProfilerInfo()
AppDomainInfo* Dataflow::GetAppDomain(AppDomainID id)
{
CSGUARD(_cs);
if (!IsInitialized())
{
return nullptr;
}

auto found = _appDomains.find(id);
if (found != _appDomains.end())
{
Expand Down Expand Up @@ -538,6 +551,11 @@ AppDomainInfo* Dataflow::GetAppDomain(AppDomainID id)
ModuleInfo* Dataflow::GetModuleInfo(ModuleID id)
{
CSGUARD(_cs);
if (!IsInitialized())
{
return nullptr;
}

auto found = _modules.find(id);
if (found != _modules.end())
{
Expand Down Expand Up @@ -595,7 +613,11 @@ ModuleInfo* Dataflow::GetModuleInfo(ModuleID id)

ModuleInfo* Dataflow::GetAspectsModule(AppDomainID id)
{
CSGUARD(_cs);
if (!IsInitialized())
{
return nullptr;
}

ModuleID moduleId = trace::profiler->GetProfilerAssemblyModuleId(id);

if (moduleId > 0)
Expand All @@ -618,7 +640,7 @@ MethodInfo* Dataflow::GetMethodInfo(ModuleID moduleId, mdMethodDef methodId)

bool Dataflow::IsInlineEnabled(ModuleID calleeModuleId, mdToken calleeMethodId)
{
if (!_initialized)
if (!IsInitialized())
Copy link
Member

@andrewlock andrewlock Nov 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Seems like there's a bunch of places here without the guards... is that safe?

Copy link
Collaborator Author

@NachoEchevarria NachoEchevarria Nov 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It is safe as long as we dont perform sensitive operations. In this case, IsInitialized() is already safe because we are using an atomic bool and we will recheck that later in a locked envirinment.
Methods like IsInlineEnabled(), JITProcessMethod(), and JITCompilationStarted() check IsInitialized() outside locks because the atomic read itself is safe, they perform no "sensitive operations" (directly accessing shared state _modules, _appDomains, _moduleAspects) and the real protection is at the leaf level. When these methods call GetModuleInfo() or GetAppDomain(), those methods acquire the lock first, check IsInitialized() inside the lock (TOCTOU-safe) and only then access shared maps.

{
return true;
}
Expand All @@ -632,7 +654,7 @@ bool Dataflow::IsInlineEnabled(ModuleID calleeModuleId, mdToken calleeMethodId)
}
bool Dataflow::JITCompilationStarted(ModuleID moduleId, mdToken methodId)
{
if (!_initialized)
if (!IsInitialized())
{
return false;
}
Expand All @@ -642,7 +664,7 @@ bool Dataflow::JITCompilationStarted(ModuleID moduleId, mdToken methodId)
}
MethodInfo* Dataflow::JITProcessMethod(ModuleID moduleId, mdToken methodId, trace::FunctionControlWrapper* pFunctionControl)
{
if (!_initialized)
if (!IsInitialized())
{
return nullptr;
}
Expand Down Expand Up @@ -678,6 +700,10 @@ HRESULT Dataflow::RewriteMethod(MethodInfo* method, trace::FunctionControlWrappe
HRESULT hr = S_OK;

CSGUARD(_cs);
if (!IsInitialized())
{
return S_FALSE;
}

if (!pFunctionControl)
{
Expand Down Expand Up @@ -741,6 +767,12 @@ HRESULT Dataflow::RewriteMethod(MethodInfo* method, trace::FunctionControlWrappe

std::vector<DataflowAspectReference*> Dataflow::GetAspects(ModuleInfo* module)
{
CSGUARD(_cs);
if (!IsInitialized())
{
return std::vector<DataflowAspectReference*>();
}

auto value = _moduleAspects.find(module->_id);
if (value != _moduleAspects.end())
{
Expand Down Expand Up @@ -783,7 +815,7 @@ void Dataflow::AddNGenInlinerModule(ModuleID moduleId)

HRESULT Dataflow::RejitMethod(trace::FunctionControlWrapper& functionControl)
{
if (!_initialized)
if (!IsInitialized())
{
return S_FALSE;
}
Expand Down
13 changes: 12 additions & 1 deletion tracer/src/Datadog.Tracer.Native/iast/dataflow.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
#pragma once
#include <atomic>
#include "../../../../shared/src/native-src/pal.h"
#include "iast_util.h"
#include "aspect.h"
Expand Down Expand Up @@ -65,7 +66,7 @@ namespace iast

void LoadSecurityControls();
protected:
bool _initialized = false;
std::atomic<bool> _initialized{false};
bool _setILOnJit = false;

std::vector<DataflowAspectClass*> _aspectClasses;
Expand All @@ -77,6 +78,16 @@ namespace iast

std::vector<DataflowAspectReference*> GetAspects(ModuleInfo* module);
static bool InstrumentInstruction(DataflowContext& context, std::vector<DataflowAspectReference*>& aspects);

inline bool IsInitialized() const noexcept
{
return _initialized.load(std::memory_order_acquire);
}

inline void SetInitialized(bool value) noexcept
{
_initialized.store(value, std::memory_order_release);
}

public:
HRESULT AppDomainShutdown(AppDomainID appDomainId);
Expand Down
Loading