From 7376508a0a4dc4443a5a8b538e74272989c95c87 Mon Sep 17 00:00:00 2001 From: Daniel Romano <108014683+daniel-romano-DD@users.noreply.github.com> Date: Tue, 1 Oct 2024 15:33:16 +0200 Subject: [PATCH] [IAST] Fix for VS Edit and Continue (#6097) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary of changes When VS enables Edit and Continue, it sets this env var `COMPLUS_ForceEnc=true` and some CallSite functions are not instrumented. ## Reason for change When `COMPLUS_ForceEnc=true` the profiler can skip the call to some Rejit events, leaving some functions (specially in callsite) uninstrumented ## Implementation details If this env var is detected, dataflow will set CallSite edited functionsIL in the JIT event as well as in the REJIT event, if this occurs. ## Test coverage Added integration tests ## Other details --- .../Datadog.Tracer.Native/cor_profiler.cpp | 2 +- .../src/Datadog.Tracer.Native/cor_profiler.h | 3 +- .../environment_variables.h | 6 ++ .../environment_variables_util.cpp | 13 +++ .../environment_variables_util.h | 1 + .../Datadog.Tracer.Native/iast/dataflow.cpp | 39 ++++--- .../src/Datadog.Tracer.Native/iast/dataflow.h | 3 +- .../AspNetBase.cs | 2 +- .../IAST/AspNetCore5IastTests.cs | 102 ++++++++++-------- 9 files changed, 111 insertions(+), 60 deletions(-) diff --git a/tracer/src/Datadog.Tracer.Native/cor_profiler.cpp b/tracer/src/Datadog.Tracer.Native/cor_profiler.cpp index 738a7324a620..9441009da13a 100644 --- a/tracer/src/Datadog.Tracer.Native/cor_profiler.cpp +++ b/tracer/src/Datadog.Tracer.Native/cor_profiler.cpp @@ -345,7 +345,7 @@ HRESULT STDMETHODCALLTYPE CorProfiler::Initialize(IUnknown* cor_profiler_info_un if (isIastEnabled || isRaspEnabled) { - _dataflow = new iast::Dataflow(info_, rejit_handler); + _dataflow = new iast::Dataflow(info_, rejit_handler, runtime_information_); if (FAILED(_dataflow->Init())) { Logger::Error("Callsite Dataflow failed to initialize"); diff --git a/tracer/src/Datadog.Tracer.Native/cor_profiler.h b/tracer/src/Datadog.Tracer.Native/cor_profiler.h index a59de0e629ce..99aad8f9fb51 100644 --- a/tracer/src/Datadog.Tracer.Native/cor_profiler.h +++ b/tracer/src/Datadog.Tracer.Native/cor_profiler.h @@ -170,10 +170,10 @@ class CorProfiler : public CorProfilerBase HRESULT STDMETHODCALLTYPE ProfilerDetachSucceeded() override; HRESULT STDMETHODCALLTYPE JITInlining(FunctionID callerId, FunctionID calleeId, BOOL* pfShouldInline) override; + // // ReJIT Methods // - HRESULT STDMETHODCALLTYPE GetReJITParameters(ModuleID moduleId, mdMethodDef methodId, ICorProfilerFunctionControl* pFunctionControl) override; @@ -251,7 +251,6 @@ class CorProfiler : public CorProfilerBase // // Getters for exception filter // - bool IsCallTargetBubbleUpExceptionTypeAvailable() const; bool IsCallTargetBubbleUpFunctionAvailable() const; }; diff --git a/tracer/src/Datadog.Tracer.Native/environment_variables.h b/tracer/src/Datadog.Tracer.Native/environment_variables.h index a89a8a7359b5..b977ca839aeb 100644 --- a/tracer/src/Datadog.Tracer.Native/environment_variables.h +++ b/tracer/src/Datadog.Tracer.Native/environment_variables.h @@ -127,6 +127,12 @@ namespace environment // Enables the workaround for dotnet issue 77973 (https://github.com/dotnet/runtime/issues/77973) const shared::WSTRING internal_workaround_77973_enabled = WStr("DD_INTERNAL_WORKAROUND_77973_ENABLED"); + // IDE Edit and Continue. If enabled, profiler behavior is modified slightly + const shared::WSTRING ide_edit_and_continue_core = WStr("COMPLUS_ForceEnc"); + + // IDE Edit and Continue. If enabled, profiler behavior is modified slightly + const shared::WSTRING ide_edit_and_continue_netfx = WStr("DOTNET_ForceEnc"); + } // namespace environment } // namespace trace diff --git a/tracer/src/Datadog.Tracer.Native/environment_variables_util.cpp b/tracer/src/Datadog.Tracer.Native/environment_variables_util.cpp index 3079949f6e24..cca86cdd2848 100644 --- a/tracer/src/Datadog.Tracer.Native/environment_variables_util.cpp +++ b/tracer/src/Datadog.Tracer.Native/environment_variables_util.cpp @@ -68,4 +68,17 @@ bool IsRaspEnabled() return IsRaspSettingEnabled() && IsAsmSettingEnabled(); } +bool IsEditAndContinueEnabledCore() +{ + ToBooleanWithDefault(shared::GetEnvironmentValue(environment::ide_edit_and_continue_core), false); +} +bool IsEditAndContinueEnabledNetFx() +{ + ToBooleanWithDefault(shared::GetEnvironmentValue(environment::ide_edit_and_continue_netfx), false); +} +bool IsEditAndContinueEnabled() +{ + return IsEditAndContinueEnabledCore() || IsEditAndContinueEnabledNetFx(); +} + } // namespace trace \ No newline at end of file diff --git a/tracer/src/Datadog.Tracer.Native/environment_variables_util.h b/tracer/src/Datadog.Tracer.Native/environment_variables_util.h index 10d7d567c663..e459932927d8 100644 --- a/tracer/src/Datadog.Tracer.Native/environment_variables_util.h +++ b/tracer/src/Datadog.Tracer.Native/environment_variables_util.h @@ -60,6 +60,7 @@ bool IsAzureFunctionsEnabled(); bool IsVersionCompatibilityEnabled(); bool IsIastEnabled(); bool IsRaspEnabled(); +bool IsEditAndContinueEnabled(); } // namespace trace diff --git a/tracer/src/Datadog.Tracer.Native/iast/dataflow.cpp b/tracer/src/Datadog.Tracer.Native/iast/dataflow.cpp index fb7ddbf92d87..c145459fb65f 100644 --- a/tracer/src/Datadog.Tracer.Native/iast/dataflow.cpp +++ b/tracer/src/Datadog.Tracer.Native/iast/dataflow.cpp @@ -10,6 +10,7 @@ #include #include #include "../../../../shared/src/native-src/com_ptr.h" +#include "../environment_variables_util.h" using namespace std::chrono; @@ -153,23 +154,26 @@ AspectFilter* ModuleAspects::GetFilter(DataflowAspectFilterValue filterValue) //-------------------- -Dataflow::Dataflow(ICorProfilerInfo* profiler, std::shared_ptr rejitHandler) : +Dataflow::Dataflow(ICorProfilerInfo* profiler, std::shared_ptr rejitHandler, + const RuntimeInformation& runtimeInfo) : Rejitter(rejitHandler, RejitterPriority::Low) { - HRESULT hr = profiler->QueryInterface(__uuidof(ICorProfilerInfo3), (void**) &_profiler); - if (_profiler != nullptr) + m_runtimeType = runtimeInfo.runtime_type; + m_runtimeVersion = VersionInfo{runtimeInfo.major_version, runtimeInfo.minor_version, runtimeInfo.build_version, 0}; + trace::Logger::Info("Dataflow::Dataflow -> Detected runtime version : ", m_runtimeVersion.ToString()); + + this->_setILOnJit = trace::IsEditAndContinueEnabled(); + if (this->_setILOnJit) { - USHORT major; - USHORT minor; - USHORT build; + trace::Logger::Info("Dataflow detected Edit and Continue feature (COMPLUS_ForceEnc != 0) : Enabling SetILCode in JIT event."); + } - if (SUCCEEDED(_profiler->GetRuntimeInformation(nullptr, &m_runtimeType, &major, &minor, &build, nullptr, 0, - nullptr, nullptr))) - { - m_runtimeVersion = VersionInfo{major, minor, build, 0}; - } + HRESULT hr = profiler->QueryInterface(__uuidof(ICorProfilerInfo3), (void**) &_profiler); + if (FAILED(hr)) + { + _profiler = nullptr; + trace::Logger::Error("Dataflow::Dataflow -> Something very wrong happened, as QI on ICorProfilerInfo3 failed. Disabling Dataflow. HRESULT : ", Hex(hr)); } - trace::Logger::Info("Dataflow::Dataflow -> Detected runtime version : ", m_runtimeVersion.ToString()); } Dataflow::~Dataflow() { @@ -182,6 +186,10 @@ HRESULT Dataflow::Init() { return S_FALSE; } + if (_profiler == nullptr) + { + return E_FAIL; + } HRESULT hr = S_OK; try { @@ -654,8 +662,9 @@ HRESULT Dataflow::RewriteMethod(MethodInfo* method, trace::FunctionControlWrappe } } } - if (!pFunctionControl && written) + if (!pFunctionControl && written && !_setILOnJit) { + // We are in JIT event. If _setILOnJit is false, we abort the commit and request a rejit context.aborted = true; } method->SetInstrumented(written); @@ -668,6 +677,10 @@ HRESULT Dataflow::RewriteMethod(MethodInfo* method, trace::FunctionControlWrappe { hr = method->ApplyFinalInstrumentation((ICorProfilerFunctionControl*) pFunctionControl); } + else if (_setILOnJit) + { + hr = method->ApplyFinalInstrumentation(); + } else { std::vector modulesVector = {module->_id}; diff --git a/tracer/src/Datadog.Tracer.Native/iast/dataflow.h b/tracer/src/Datadog.Tracer.Native/iast/dataflow.h index 804f00e40ee1..6ffe1c70bac2 100644 --- a/tracer/src/Datadog.Tracer.Native/iast/dataflow.h +++ b/tracer/src/Datadog.Tracer.Native/iast/dataflow.h @@ -52,7 +52,7 @@ namespace iast friend class ModuleInfo; friend class ModuleAspects; public: - Dataflow(ICorProfilerInfo* profiler, std::shared_ptr rejitHandler); + Dataflow(ICorProfilerInfo* profiler, std::shared_ptr rejitHandler, const RuntimeInformation& runtimeInfo); virtual ~Dataflow(); private: CS _cs; @@ -75,6 +75,7 @@ namespace iast protected: bool _initialized = false; bool _loaded = false; + bool _setILOnJit = false; std::vector _aspectClasses; std::vector _aspects; diff --git a/tracer/test/Datadog.Trace.Security.IntegrationTests/AspNetBase.cs b/tracer/test/Datadog.Trace.Security.IntegrationTests/AspNetBase.cs index 9f61a0912a99..7237367bece8 100644 --- a/tracer/test/Datadog.Trace.Security.IntegrationTests/AspNetBase.cs +++ b/tracer/test/Datadog.Trace.Security.IntegrationTests/AspNetBase.cs @@ -409,7 +409,7 @@ protected IImmutableList WaitForSpans(MockTracerAgent agent, int expec agent.SpanFilters.Add(s => s.Tags.ContainsKey("http.url") && s.Tags["http.url"].IndexOf(url, StringComparison.InvariantCultureIgnoreCase) > -1); } - var spans = agent.WaitForSpans(expectedSpans, minDateTime: minDateTime); + var spans = agent.WaitForSpans(expectedSpans, minDateTime: minDateTime, assertExpectedCount: false); if (spans.Count != expectedSpans) { Output?.WriteLine($"spans.Count: {spans.Count} != expectedSpans: {expectedSpans}, this is phase: {phase}"); diff --git a/tracer/test/Datadog.Trace.Security.IntegrationTests/IAST/AspNetCore5IastTests.cs b/tracer/test/Datadog.Trace.Security.IntegrationTests/IAST/AspNetCore5IastTests.cs index 65fa1d9eb3a6..e185ea999d85 100644 --- a/tracer/test/Datadog.Trace.Security.IntegrationTests/IAST/AspNetCore5IastTests.cs +++ b/tracer/test/Datadog.Trace.Security.IntegrationTests/IAST/AspNetCore5IastTests.cs @@ -350,7 +350,7 @@ await VerifyHelper.VerifySpans(spansFiltered, settings) } } -// Class to test particular features +// Classes to test particular features public class AspNetCore5IastTestsStackTraces : AspNetCore5IastTests { public AspNetCore5IastTestsStackTraces(AspNetCoreTestFixture fixture, ITestOutputHelper outputHelper) @@ -406,6 +406,65 @@ await VerifyHelper.VerifySpans(spans, settings) } } +public class AspNetCore5IastTestsCompatEditAndContinue : AspNetCore5IastTests +{ + public AspNetCore5IastTestsCompatEditAndContinue(AspNetCoreTestFixture fixture, ITestOutputHelper outputHelper) + : base(fixture, outputHelper, enableIast: true, testName: "AspNetCore5IastTestsCompat", samplingRate: 100, isIastDeduplicationEnabled: false, vulnerabilitiesPerRequest: 200, redactionEnabled: true) + { + SetEnvironmentVariable("COMPLUS_ForceEnc", "1"); + } + + [SkippableFact] + [Trait("RunOnWindows", "True")] + public async Task TestIastCustomSpanRequestManual() + { + var filename = "Iast.CustomManual.AspNetCore5." + (IastEnabled ? "IastEnabled" : "IastDisabled"); + if (RedactionEnabled is true) { filename += ".RedactionEnabled"; } + var url = "/Iast/CustomManual?userName=Vicent"; + IncludeAllHttpSpans = true; + await TryStartApp(); + var agent = Fixture.Agent; + var spans = await SendRequestsAsync(agent, 3, new string[] { url }); + var spansFiltered = spans.Where(s => !s.Resource.StartsWith("CREATE TABLE") && !s.Resource.StartsWith("INSERT INTO")).ToList(); + + var settings = VerifyHelper.GetSpanVerifierSettings(); + settings.AddIastScrubbing(); + await VerifyHelper.VerifySpans(spansFiltered, settings) + .UseFileName(filename) + .DisableRequireUniquePrefix(); + } +} + +public class AspNetCore5IastTestsCompat : AspNetCore5IastTestsCompatEditAndContinue +{ + public AspNetCore5IastTestsCompat(AspNetCoreTestFixture fixture, ITestOutputHelper outputHelper) + : base(fixture, outputHelper) + { + SetEnvironmentVariable("COMPLUS_ForceEnc", "0"); + } + + [SkippableFact] + [Trait("RunOnWindows", "True")] + public async Task TestIastCustomSpanRequestAttribute() + { + var filename = "Iast.CustomAttribute.AspNetCore5." + (IastEnabled ? "IastEnabled" : "IastDisabled"); + if (RedactionEnabled is true) { filename += ".RedactionEnabled"; } + var url = "/Iast/CustomAttribute?userName=Vicent"; + IncludeAllHttpSpans = true; + await TryStartApp(); + var agent = Fixture.Agent; + var spans = await SendRequestsAsync(agent, 2, new string[] { url }); + var spansFiltered = spans.Where(s => !s.Resource.StartsWith("CREATE TABLE") && !s.Resource.StartsWith("INSERT INTO")).ToList(); + + var settings = VerifyHelper.GetSpanVerifierSettings(); + settings.AddIastScrubbing(); + settings.AddRegexScrubber(new Regex(@"_dd.agent_psr: .{1,3},"), string.Empty); + await VerifyHelper.VerifySpans(spansFiltered, settings) + .UseFileName(filename) + .DisableRequireUniquePrefix(); + } +} + public class AspNetCore5IastTestsSpanTelemetryIastEnabled : AspNetCore5IastTests { public AspNetCore5IastTestsSpanTelemetryIastEnabled(AspNetCoreTestFixture fixture, ITestOutputHelper outputHelper) @@ -1169,47 +1228,6 @@ await VerifyHelper.VerifySpans(spansFiltered, settings) .DisableRequireUniquePrefix(); } - [SkippableFact] - [Trait("RunOnWindows", "True")] - public async Task TestIastCustomSpanRequestAttribute() - { - var filename = "Iast.CustomAttribute.AspNetCore5." + (IastEnabled ? "IastEnabled" : "IastDisabled"); - if (RedactionEnabled is true) { filename += ".RedactionEnabled"; } - var url = "/Iast/CustomAttribute?userName=Vicent"; - IncludeAllHttpSpans = true; - await TryStartApp(); - var agent = Fixture.Agent; - var spans = await SendRequestsAsync(agent, 3, new string[] { url }); - var spansFiltered = spans.Where(s => !s.Resource.StartsWith("CREATE TABLE") && !s.Resource.StartsWith("INSERT INTO")).ToList(); - - var settings = VerifyHelper.GetSpanVerifierSettings(); - settings.AddIastScrubbing(); - settings.AddRegexScrubber(new Regex(@"_dd.agent_psr: .{1,3},"), string.Empty); - await VerifyHelper.VerifySpans(spansFiltered, settings) - .UseFileName(filename) - .DisableRequireUniquePrefix(); - } - - [SkippableFact] - [Trait("RunOnWindows", "True")] - public async Task TestIastCustomSpanRequestManual() - { - var filename = "Iast.CustomManual.AspNetCore5." + (IastEnabled ? "IastEnabled" : "IastDisabled"); - if (RedactionEnabled is true) { filename += ".RedactionEnabled"; } - var url = "/Iast/CustomManual?userName=Vicent"; - IncludeAllHttpSpans = true; - await TryStartApp(); - var agent = Fixture.Agent; - var spans = await SendRequestsAsync(agent, 3, new string[] { url }); - var spansFiltered = spans.Where(s => !s.Resource.StartsWith("CREATE TABLE") && !s.Resource.StartsWith("INSERT INTO")).ToList(); - - var settings = VerifyHelper.GetSpanVerifierSettings(); - settings.AddIastScrubbing(); - await VerifyHelper.VerifySpans(spansFiltered, settings) - .UseFileName(filename) - .DisableRequireUniquePrefix(); - } - [SkippableFact] [Trait("RunOnWindows", "True")] public async Task TestNHibernateSqlInjection()