Skip to content

Commit 7ddcd51

Browse files
tonyredondotonyredondo
authored
Shared Loader - Add support for more rewriting points (DataDog#2487)
* Add support for more rewriting points * Fixes
1 parent ac518f5 commit 7ddcd51

File tree

1 file changed

+109
-50
lines changed

1 file changed

+109
-50
lines changed

‎shared/src/native-src/loader.cpp‎

Lines changed: 109 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -39,9 +39,14 @@ namespace shared
3939
WStr("Datadog.AutoInstrumentation.Profiler.Managed"),
4040
};
4141

42-
const WSTRING SpecificTypeToInjectName = WStr("System.AppDomain");
43-
const WSTRING SpecificMethodToInjectName = WStr("IsCompatibilitySwitchSet");
44-
42+
const std::vector<std::tuple<WSTRING, WSTRING, WSTRING>> _specificMethodsToRewrite = {
43+
std::tuple<WSTRING, WSTRING, WSTRING>(WStr("mscorlib"), WStr("System.AppDomain"), WStr("IsCompatibilitySwitchSet")),
44+
std::tuple<WSTRING, WSTRING, WSTRING>(WStr("System.Private.CoreLib"), WStr("System.AppDomain"), WStr("IsCompatibilitySwitchSet")),
45+
std::tuple<WSTRING, WSTRING, WSTRING>(WStr("System.Web"), WStr("System.Web.Hosting.ProcessHost"), WStr("GetProcessHost")),
46+
std::tuple<WSTRING, WSTRING, WSTRING>(WStr("System.Web"), WStr("System.Web.Hosting.ProcessHost"), WStr("CreateAppHost")),
47+
std::tuple<WSTRING, WSTRING, WSTRING>(WStr("System.Web"), WStr("System.Web.Hosting.HostingEnvironment"), WStr(".ctor")),
48+
std::tuple<WSTRING, WSTRING, WSTRING>(WStr("System.Web"), WStr("System.Web.HttpRuntime"), WStr("FirstRequestInit")),
49+
};
4550

4651
static Enumerator<mdMethodDef> EnumMethodsWithName(
4752
const ComPtr<IMetaDataImport2>& metadata_import,
@@ -276,7 +281,8 @@ namespace shared
276281
//
277282
// Check and store assembly metadata if the corlib is found.
278283
//
279-
if (assemblyNameString == WStr("mscorlib") || assemblyNameString == WStr("System.Private.CoreLib"))
284+
const auto& isCorLibAssembly = assemblyNameString == WStr("mscorlib") || assemblyNameString == WStr("System.Private.CoreLib");
285+
if (isCorLibAssembly)
280286
{
281287
if (_loaderOptions.LogDebugIsEnabled)
282288
{
@@ -322,52 +328,97 @@ namespace shared
322328
ToString(_corlibMetadata.Metadata.usBuildNumber) +
323329
"]");
324330
}
331+
}
332+
333+
if (_loaderOptions.RewriteMSCorLibMethods && _loaderOptions.IsNet46OrGreater && _corlibMetadata.Token != mdAssemblyNil)
334+
{
335+
mdAssemblyRef corlibAssemblyRef = mdAssemblyRefNil;
336+
mdTypeRef systemObjectTypeRef = mdTypeRefNil;
337+
mdTypeDef loaderTypeDef = mdTypeDefNil;
338+
mdMethodDef loaderMethodDef = mdMethodDefNil;
339+
mdMemberRef securitySafeCriticalCtorMemberRef = mdMemberRefNil;
325340

326-
if (_loaderOptions.RewriteMSCorLibMethods && _loaderOptions.IsNet46OrGreater)
341+
for (const auto& specMethodPair : _specificMethodsToRewrite)
327342
{
343+
// Check assembly name
344+
const auto& specificAssemblyInjectName = std::get<0>(specMethodPair);
345+
if (specificAssemblyInjectName != assemblyNameString)
346+
{
347+
continue;
348+
}
349+
350+
const auto& specificTypeToInjectName = std::get<1>(specMethodPair);
351+
const auto& specificMethodToInjectName = std::get<2>(specMethodPair);
352+
328353
mdTypeDef appDomainTypeDef;
329-
hr = metadataImport->FindTypeDefByName(SpecificTypeToInjectName.c_str(), mdTokenNil, &appDomainTypeDef);
354+
hr = metadataImport->FindTypeDefByName(specificTypeToInjectName.c_str(), mdTokenNil, &appDomainTypeDef);
330355
if (FAILED(hr))
331356
{
332-
Debug("Loader::InjectLoaderToModuleInitializer: " + ToString(SpecificTypeToInjectName) + " not found.");
333-
return S_FALSE;
357+
Debug("Loader::InjectLoaderToModuleInitializer: " + ToString(specificTypeToInjectName) +
358+
" not found.");
359+
continue;
334360
}
335361

336-
auto enumMethods = EnumMethodsWithName(metadataImport, appDomainTypeDef, SpecificMethodToInjectName.c_str());
362+
auto enumMethods =
363+
EnumMethodsWithName(metadataImport, appDomainTypeDef, specificMethodToInjectName.c_str());
337364
auto enumIterator = enumMethods.begin();
338-
if (enumIterator != enumMethods.end()) {
365+
if (enumIterator != enumMethods.end())
366+
{
339367
auto methodDef = *enumIterator;
340368

341-
//
342-
// get a TypeRef for System.Object
343-
//
344-
mdTypeRef systemObjectTypeRef;
345-
hr = metadataEmit->DefineTypeRefByName(corLibAssembly, WStr("System.Object"), &systemObjectTypeRef);
346-
if (FAILED(hr))
369+
if (loaderMethodDef == mdMethodDefNil)
347370
{
348-
Error("Loader::InjectLoaderToModuleInitializer: failed to define typeref: System.Object");
349-
return hr;
350-
}
351371

352-
//
353-
// Define a new TypeDef DD_LoaderMethodsType that extends System.Object
354-
//
355-
mdTypeDef newTypeDef;
356-
hr = metadataEmit->DefineTypeDef(WStr("DD_LoaderMethodsType"), tdAbstract | tdSealed, systemObjectTypeRef, NULL, &newTypeDef);
357-
if (FAILED(hr)) {
358-
Error("Loader::InjectLoaderToModuleInitializer: failed to define typedef: DD_LoaderMethodsType");
359-
return hr;
360-
}
372+
if (corlibAssemblyRef == mdAssemblyRefNil)
373+
{
374+
hr = assmeblyEmit->DefineAssemblyRef(
375+
_corlibMetadata.pPublicKey, _corlibMetadata.PublicKeyLength,
376+
_corlibMetadata.Name.c_str(), &_corlibMetadata.Metadata, NULL, 0, _corlibMetadata.Flags,
377+
&corlibAssemblyRef);
378+
if (FAILED(hr))
379+
{
380+
Error(
381+
"Loader::InjectLoaderToModuleInitializer: Error creating assembly reference to "
382+
"mscorlib.");
383+
return hr;
384+
}
385+
}
361386

362-
//
363-
// Emit the DD_LoaderMethodsType.DD_LoadInitializationAssemblies() mdMethodDef
364-
//
365-
mdMethodDef loaderMethodDef;
366-
mdMemberRef securitySafeCriticalCtorMemberRef;
367-
hr = EmitDDLoadInitializationAssembliesMethod(moduleId, newTypeDef, assemblyNameString, &loaderMethodDef, &securitySafeCriticalCtorMemberRef);
368-
if (FAILED(hr))
369-
{
370-
return hr;
387+
if (systemObjectTypeRef == mdTypeRefNil)
388+
{
389+
//
390+
// get a TypeRef for System.Object
391+
//
392+
hr = metadataEmit->DefineTypeRefByName(corlibAssemblyRef, WStr("System.Object"), &systemObjectTypeRef);
393+
if (FAILED(hr))
394+
{
395+
Error(
396+
"Loader::InjectLoaderToModuleInitializer: failed to define typeref: System.Object");
397+
return hr;
398+
}
399+
}
400+
401+
if (loaderTypeDef == mdTypeDefNil)
402+
{
403+
//
404+
// Define a new TypeDef DD_LoaderMethodsType that extends System.Object
405+
//
406+
hr = metadataEmit->DefineTypeDef(WStr("DD_LoaderMethodsType"), tdAbstract | tdSealed, systemObjectTypeRef, NULL, &loaderTypeDef);
407+
if (FAILED(hr))
408+
{
409+
Error("Loader::InjectLoaderToModuleInitializer: failed to define typedef: DD_LoaderMethodsType");
410+
return hr;
411+
}
412+
}
413+
414+
//
415+
// Emit the DD_LoaderMethodsType.DD_LoadInitializationAssemblies() mdMethodDef
416+
//
417+
hr = EmitDDLoadInitializationAssembliesMethod(moduleId, loaderTypeDef, assemblyNameString, &loaderMethodDef, &securitySafeCriticalCtorMemberRef);
418+
if (FAILED(hr))
419+
{
420+
return hr;
421+
}
371422
}
372423

373424
//
@@ -377,21 +428,16 @@ namespace shared
377428
if (SUCCEEDED(hr))
378429
{
379430
Info("Loader::InjectLoaderToModuleInitializer: Loader injected successfully (in " +
380-
ToString(SpecificTypeToInjectName) + "." + ToString(SpecificMethodToInjectName) + "). [ModuleID=" + moduleIdHex +
381-
", AssemblyID=" + assemblyIdHex +
382-
", AssemblyName=" + ToString(assemblyNameString) +
383-
", AppDomainID=" + appDomainIdHex +
384-
", methodDef=" + HexStr(methodDef) +
385-
", loaderMethodDef=" + HexStr(loaderMethodDef) +
386-
"]");
431+
ToString(specificTypeToInjectName) + "." + ToString(specificMethodToInjectName) +
432+
"). [ModuleID=" + moduleIdHex + ", AssemblyID=" + assemblyIdHex +
433+
", AssemblyName=" + ToString(assemblyNameString) + ", AppDomainID=" + appDomainIdHex +
434+
", methodDef=" + HexStr(methodDef) + ", loaderMethodDef=" + HexStr(loaderMethodDef) + "]");
387435
}
388436
}
389437
}
390-
391-
return hr;
392438
}
393439

394-
if (_loaderOptions.RewriteModulesEntrypoint)
440+
if (!isCorLibAssembly && _loaderOptions.RewriteModulesEntrypoint)
395441
{
396442
//
397443
// Rewrite Module EntryPoint
@@ -514,7 +560,7 @@ namespace shared
514560
}
515561
}
516562

517-
if (_loaderOptions.RewriteModulesInitializers)
563+
if (!isCorLibAssembly && _loaderOptions.RewriteModulesInitializers)
518564
{
519565
// **************************************************************************************************************
520566
//
@@ -659,7 +705,13 @@ namespace shared
659705
return S_FALSE;
660706
}
661707

662-
if (SpecificMethodToInjectName != functionName)
708+
auto methodFound = false;
709+
for (const auto& specMethodPair : _specificMethodsToRewrite)
710+
{
711+
methodFound |= std::get<2>(specMethodPair) == functionName;
712+
}
713+
714+
if (!methodFound)
663715
{
664716
if (ExtraVerboseLogging && _loaderOptions.LogDebugIsEnabled)
665717
{
@@ -682,7 +734,14 @@ namespace shared
682734
return S_FALSE;
683735
}
684736

685-
if (SpecificTypeToInjectName != typeName)
737+
738+
auto typeFound = false;
739+
for (const auto& specMethodPair : _specificMethodsToRewrite)
740+
{
741+
typeFound |= std::get<1>(specMethodPair) == typeName;
742+
}
743+
744+
if (!typeFound)
686745
{
687746
if (ExtraVerboseLogging && _loaderOptions.LogDebugIsEnabled)
688747
{

0 commit comments

Comments
 (0)