Init all loaded module infos at the initialization of dataflow

Load all preloaded modules at dataflow init

Author: daniel-romano-DD

tracer/src/Datadog.Tracer.Native/cor_profiler.cpp

@@ -270,7 +270,8 @@ HRESULT STDMETHODCALLTYPE CorProfiler::Initialize(IUnknown* cor_profiler_info_un
 
         if (isIastEnabled || isRaspEnabled)
         {
-            _dataflow = new iast::Dataflow(info_, rejit_handler, runtime_information_);
+            auto modules = module_ids.Get();
+            _dataflow = new iast::Dataflow(info_, rejit_handler, modules.Ref(), runtime_information_);
         }
         else
         {
@@ -2076,7 +2077,8 @@ int CorProfiler::RegisterIastAspects(WCHAR** aspects, int aspectsLength, UINT32
     if (dataflow == nullptr && IsCallSiteManagedActivationEnabled())
     {
         Logger::Debug("Creating Dataflow.");
-        dataflow = new iast::Dataflow(info_, rejit_handler, runtime_information_);
+        auto modules = module_ids.Get();
+        dataflow = new iast::Dataflow(info_, rejit_handler, modules.Ref(), runtime_information_);
     }
 
     if (dataflow != nullptr)

tracer/src/Datadog.Tracer.Native/iast/dataflow.cpp

@@ -158,7 +158,9 @@ AspectFilter* ModuleAspects::GetFilter(DataflowAspectFilterValue filterValue)
 
 //--------------------
 
-Dataflow::Dataflow(ICorProfilerInfo* profiler, std::shared_ptr<RejitHandler> rejitHandler,
+Dataflow::Dataflow(ICorProfilerInfo* profiler, 
+                   std::shared_ptr<RejitHandler> rejitHandler,
+                   std::vector<ModuleID> moduleIds, 
                    const RuntimeInformation& runtimeInfo) :
     Rejitter(rejitHandler, RejitterPriority::Low, false)
 {
@@ -178,7 +180,8 @@ Dataflow::Dataflow(ICorProfilerInfo* profiler, std::shared_ptr<RejitHandler> rej
         _profiler = nullptr;
         trace::Logger::Error("Dataflow::Dataflow -> Something very wrong happened, as QI on ICorProfilerInfo3 failed. Disabling Dataflow. HRESULT : ", Hex(hr));
     }
-}
+
+    _preLoadedModuleIds = moduleIds;
 
 Dataflow::~Dataflow()
 {
@@ -379,6 +382,19 @@ HRESULT Dataflow::AppDomainShutdown(AppDomainID appDomainId)
 
 HRESULT Dataflow::ModuleLoaded(ModuleID moduleId, ModuleInfo** pModuleInfo)
 {
+    ENTER_FUNC
+
+    // Retrieve all already modules at once to mimic initialization from creation behavior
+    if (_preLoadedModuleIds.size() > 0)
+    {
+        CSGUARD(_cs);
+        for (auto const& id : _preLoadedModuleIds)
+        {
+            GetModuleInfo(id);
+        }
+        _preLoadedModuleIds.clear();
+    }
+
     GetModuleInfo(moduleId);
     return S_OK;
 }
@@ -526,20 +542,21 @@ ModuleInfo* Dataflow::GetModuleInfo(ModuleID id)
     }
 
     // Retrieve module information if not found
+    const int pathLen = 2048;
     LPCBYTE pbBaseLoadAddr;
-    WCHAR wszPath[300];
-    ULONG cchNameIn = 300;
-    ULONG cchNameOut;
+    ULONG pathOut;
+    WCHAR wszName[pathLen];
+    WCHAR wszPath[pathLen];
     AssemblyID assemblyId;
     AppDomainID appDomainId;
     ModuleID modIDDummy;
-    WCHAR wszName[1024];
 
     DWORD dwModuleFlags;
-    HRESULT hr = _profiler->GetModuleInfo2(id, &pbBaseLoadAddr, cchNameIn, &cchNameOut, wszPath, &assemblyId, &dwModuleFlags);
+    pathOut = 0;
+    HRESULT hr = _profiler->GetModuleInfo2(id, &pbBaseLoadAddr, pathLen, &pathOut, wszPath, &assemblyId, &dwModuleFlags);
     if (FAILED(hr))
     {
-        trace::Logger::Error("Dataflow::GetModuleInfo -> GetModuleInfo2 failed for ModuleId ", id);
+        trace::Logger::Error("Dataflow::GetModuleInfo -> GetModuleInfo2 failed for ModuleId ", id, " hr:", Hex(hr));
         _modules[id] = nullptr; 
         return nullptr;
     }
@@ -549,10 +566,12 @@ ModuleInfo* Dataflow::GetModuleInfo(ModuleID id)
         return nullptr;
     } // Ignore any Windows Runtime modules.  We cannot obtain writeable metadata interfaces on them or instrument their IL
 
-    hr = _profiler->GetAssemblyInfo(assemblyId, 1024, nullptr, wszName, &appDomainId, &modIDDummy);
+    pathOut = 0;
+    hr = _profiler->GetAssemblyInfo(assemblyId, pathLen, &pathOut, wszName, &appDomainId, &modIDDummy);
     if (FAILED(hr))
     {
-        trace::Logger::Error("Dataflow::GetModuleInfo -> GetAssemblyInfo failed for ModuleId ", id, " AssemblyId ", assemblyId);
+        trace::Logger::Error("Dataflow::GetModuleInfo -> GetAssemblyInfo failed for ModuleId ", id, " AssemblyId ",
+                             assemblyId, " hr:", Hex(hr));
         _modules[id] = nullptr;
         return nullptr;
     }

tracer/src/Datadog.Tracer.Native/iast/dataflow.h

@@ -52,14 +52,16 @@ namespace iast
         friend class ModuleInfo;
         friend class ModuleAspects;
     public:
-        Dataflow(ICorProfilerInfo* profiler, std::shared_ptr<RejitHandler> rejitHandler, const RuntimeInformation& runtimeInfo);
+        Dataflow(ICorProfilerInfo* profiler, std::shared_ptr<RejitHandler> rejitHandler, std::vector<ModuleID> moduleIds,
+                 const RuntimeInformation& runtimeInfo);
         virtual ~Dataflow();
     private:
         CS _cs;
         ICorProfilerInfo3* _profiler = nullptr;
         COR_PRF_RUNTIME_TYPE m_runtimeType = COR_PRF_DESKTOP_CLR;
         VersionInfo m_runtimeVersion = VersionInfo{4, 0, 0, 0};
 
+        std::vector<ModuleID> _preLoadedModuleIds;
         std::map<ModuleID, ModuleInfo*> _modules;
         std::map<AppDomainID, AppDomainInfo*> _appDomains;