[IAST] Fix for VS Edit and Continue (#6097)

## 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
<!-- Fixes #{issue} -->

<!-- ⚠️ Note: where possible, please obtain 2 approvals prior to
merging. Unless CODEOWNERS specifies otherwise, for external teams it is
typically best to have one review from a team member, and one review
from apm-dotnet. Trivial changes do not require 2 reviews. -->

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");

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;
 };

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
 

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

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
 

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

@@ -10,6 +10,7 @@
 #include <fstream>
 #include <chrono>
 #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> rejitHandler) :
+Dataflow::Dataflow(ICorProfilerInfo* profiler, std::shared_ptr<RejitHandler> 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<ModuleID> modulesVector = {module->_id};

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> rejitHandler);
+        Dataflow(ICorProfilerInfo* profiler, std::shared_ptr<RejitHandler> 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<DataflowAspectClass*> _aspectClasses;
         std::vector<DataflowAspect*> _aspects;

tracer/test/Datadog.Trace.Security.IntegrationTests/AspNetBase.cs

@@ -409,7 +409,7 @@ protected IImmutableList<MockSpan> 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}");

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()