[ASM] Support removing whole files from RC (#3862)

robertpi · web-flow · commit 9fb7812aeb27 · 2023-03-07T16:58:11.000+01:00
* support removing whole files from RC
diff --git a/tracer/src/Datadog.Trace/AppSec/RcmModels/RemoteConfigurationStatus.cs b/tracer/src/Datadog.Trace/AppSec/RcmModels/RemoteConfigurationStatus.cs
@@ -5,6 +5,7 @@
 
 using System.Collections.Generic;
 using System.Collections.ObjectModel;
+using System.Linq;
 using Datadog.Trace.AppSec.RcmModels.Asm;
 using Datadog.Trace.AppSec.RcmModels.AsmData;
 using Datadog.Trace.Vendors.Newtonsoft.Json.Linq;
@@ -13,11 +14,11 @@ namespace Datadog.Trace.AppSec.RcmModels;
 
 internal class RemoteConfigurationStatus
 {
-    internal List<RuleOverride> RulesOverrides { get; } = new();
+    internal Dictionary<string, RuleOverride[]> RulesOverridesByFile { get; } = new();
 
-    internal List<RuleData> RulesData { get; } = new();
+    internal Dictionary<string, RuleData[]> RulesDataByFile { get; } = new();
 
-    internal List<JToken> Exclusions { get; } = new();
+    internal Dictionary<string, JArray> ExclusionsByFile { get; } = new();
 
     internal IDictionary<string, Action> Actions { get; set; } = new Dictionary<string, Action>();
 
diff --git a/tracer/src/Datadog.Trace/AppSec/Security.cs b/tracer/src/Datadog.Trace/AppSec/Security.cs
@@ -9,6 +9,7 @@
 using System.Linq;
 using System.Threading;
 using Datadog.Trace.AppSec.RcmModels;
+using Datadog.Trace.AppSec.RcmModels.AsmData;
 using Datadog.Trace.AppSec.Waf;
 using Datadog.Trace.AppSec.Waf.Initialization;
 using Datadog.Trace.AppSec.Waf.NativeBindings;
@@ -283,27 +284,28 @@ private void SetRemoteConfigCapabilites()
 
         private void AsmDDProductConfigChanged(object sender, ProductConfigChangedEventArgs e)
         {
-            if (!_enabled) { return; }
-
             var asmDd = e.GetConfigurationAsString().FirstOrDefault();
             if (!string.IsNullOrEmpty(asmDd.TypedFile))
             {
                 _remoteConfigurationStatus.RemoteRulesJson = asmDd.TypedFile;
-                var result = _waf?.UpdateRules(_remoteConfigurationStatus.RemoteRulesJson);
-                WafRuleFileVersion = result?.RuleFileVersion;
-                if (_wafInitResult?.Success ?? false)
+                if (_enabled)
                 {
-                    e.Acknowledge(asmDd.Name);
-                }
-                else
-                {
-                    e.Error(asmDd.Name, "An error happened updating waf rules");
+                    var result = _waf?.UpdateRules(_remoteConfigurationStatus.RemoteRulesJson);
+                    WafRuleFileVersion = result?.RuleFileVersion;
+                    if (_wafInitResult?.Success ?? false)
+                    {
+                        e.Acknowledge(asmDd.Name);
+                    }
+                    else
+                    {
+                        e.Error(asmDd.Name, "An error happened updating waf rules");
+                    }
+
+                    return;
                 }
             }
-            else
-            {
-                e.Acknowledge(asmDd.Name);
-            }
+
+            e.Acknowledge(asmDd.Name);
         }
 
         private void FeaturesProductConfigChanged(object sender, ProductConfigChangedEventArgs e)
@@ -325,30 +327,60 @@ private void FeaturesProductConfigChanged(object sender, ProductConfigChangedEve
             e.Acknowledge(features.Name);
         }
 
-        private void AsmDataProductConfigChanged(object sender, ProductConfigChangedEventArgs e)
+        private void AsmDataProductConfigRemoved(object sender, ProductConfigChangedEventArgs e)
         {
-            if (!_enabled)
+            var asmDataConfigs = e.GetDeserializedConfigurations<RcmModels.AsmData.Payload>();
+            foreach (var asmDataConfig in asmDataConfigs)
             {
-                return;
+                if (_remoteConfigurationStatus.RulesDataByFile.ContainsKey(asmDataConfig.Name))
+                {
+                    _remoteConfigurationStatus.RulesDataByFile.Remove(asmDataConfig.Name);
+                }
+            }
+
+            var updated = true;
+            if (_enabled)
+            {
+                var ruleData = _remoteConfigurationStatus.RulesDataByFile.SelectMany(x => x.Value).ToList();
+                updated = UpdateWafWithRulesData(ruleData);
             }
 
+            foreach (var asmDataConfig in asmDataConfigs)
+            {
+                if (!updated)
+                {
+                    e.Error(asmDataConfig.Name, "Waf could not remove the rules data");
+                }
+                else
+                {
+                    e.Acknowledge(asmDataConfig.Name);
+                }
+            }
+        }
+
+        private void AsmDataProductConfigChanged(object sender, ProductConfigChangedEventArgs e)
+        {
             var asmDataConfigs = e.GetDeserializedConfigurations<RcmModels.AsmData.Payload>();
             foreach (var asmDataConfig in asmDataConfigs)
             {
                 if (asmDataConfig.TypedFile?.RulesData?.Length > 0)
                 {
-                    _remoteConfigurationStatus.RulesData.AddRange(asmDataConfig.TypedFile.RulesData);
+                    _remoteConfigurationStatus.RulesDataByFile[asmDataConfig.Name] = asmDataConfig.TypedFile.RulesData;
                 }
+            }
 
-                e.Acknowledge(asmDataConfig.Name);
+            var updated = true;
+            if (_enabled)
+            {
+                var ruleData = _remoteConfigurationStatus.RulesDataByFile.SelectMany(x => x.Value).ToList();
+                updated = UpdateWafWithRulesData(ruleData);
             }
 
-            var updated = UpdateWafWithRulesData();
             foreach (var asmDataConfig in asmDataConfigs)
             {
                 if (!updated)
                 {
-                    e.Error(asmDataConfig.Name, "Waf could not update the rules");
+                    e.Error(asmDataConfig.Name, "Waf could not update the rules data");
                 }
                 else
                 {
@@ -357,25 +389,64 @@ private void AsmDataProductConfigChanged(object sender, ProductConfigChangedEven
             }
         }
 
-        private void AsmProductConfigChanged(object sender, ProductConfigChangedEventArgs e)
+        private void AsmProductConfigRemoved(object sender, ProductConfigChangedEventArgs e)
         {
-            if (!_enabled) { return; }
-
             var asmConfigs = e.GetDeserializedConfigurations<RcmModels.Asm.Payload>();
 
-            _remoteConfigurationStatus.RulesOverrides.Clear();
-            _remoteConfigurationStatus.Exclusions.Clear();
+            foreach (var asmConfig in asmConfigs)
+            {
+                if (_remoteConfigurationStatus.RulesOverridesByFile.ContainsKey(asmConfig.Name))
+                {
+                    _remoteConfigurationStatus.RulesOverridesByFile.Remove(asmConfig.Name);
+                }
+
+                if (_remoteConfigurationStatus.ExclusionsByFile.ContainsKey(asmConfig.Name))
+                {
+                    _remoteConfigurationStatus.ExclusionsByFile.Remove(asmConfig.Name);
+                }
+            }
+
+            var result = true;
+            if (_enabled)
+            {
+                var overrides = _remoteConfigurationStatus.RulesOverridesByFile.SelectMany(x => x.Value).ToList();
+                var exclusions = _remoteConfigurationStatus.ExclusionsByFile.SelectMany(x => x.Value).ToList();
+
+                result = _waf.UpdateRulesStatus(overrides, exclusions);
+                Log.Debug<bool, int, int>(
+                    "_waf.Update was updated for removal: {Success}, ({RulesOverridesCount} rule status entries), ({ExclusionsCount} exclusion filters)",
+                    result,
+                    overrides.Count,
+                    exclusions.Count);
+            }
+
+            foreach (var asmConfig in asmConfigs)
+            {
+                if (result)
+                {
+                    e.Acknowledge(asmConfig.Name);
+                }
+                else
+                {
+                    e.Error(asmConfig.Name, "waf couldn't be remove with rule asm product");
+                }
+            }
+        }
+
+        private void AsmProductConfigChanged(object sender, ProductConfigChangedEventArgs e)
+        {
+            var asmConfigs = e.GetDeserializedConfigurations<RcmModels.Asm.Payload>();
 
             foreach (var asmConfig in asmConfigs)
             {
                 if (asmConfig.TypedFile.RuleOverrides?.Length > 0)
                 {
-                    _remoteConfigurationStatus.RulesOverrides.AddRange(asmConfig.TypedFile.RuleOverrides);
+                    _remoteConfigurationStatus.RulesOverridesByFile[asmConfig.Name] = asmConfig.TypedFile.RuleOverrides;
                 }
 
                 if (asmConfig.TypedFile.Exclusions?.Count > 0)
                 {
-                    _remoteConfigurationStatus.Exclusions.AddRange(asmConfig.TypedFile.Exclusions);
+                    _remoteConfigurationStatus.ExclusionsByFile[asmConfig.Name] = asmConfig.TypedFile.Exclusions;
                 }
 
                 if (asmConfig.TypedFile.Actions != null)
@@ -395,12 +466,19 @@ private void AsmProductConfigChanged(object sender, ProductConfigChangedEventArg
                 }
             }
 
-            var result = _waf.UpdateRulesStatus(_remoteConfigurationStatus.RulesOverrides, _remoteConfigurationStatus.Exclusions);
-            Log.Debug<bool, int, int>(
-                "_waf.Update was updated: {Success}, ({RulesOverridesCount} rule status entries), ({ExclusionsCount} exclusion filters)",
-                result,
-                _remoteConfigurationStatus.RulesOverrides.Count,
-                _remoteConfigurationStatus.Exclusions.Count);
+            var result = true;
+            if (_enabled)
+            {
+                var overrides = _remoteConfigurationStatus.RulesOverridesByFile.SelectMany(x => x.Value).ToList();
+                var exclusions = _remoteConfigurationStatus.ExclusionsByFile.SelectMany(x => x.Value).ToList();
+
+                result = _waf.UpdateRulesStatus(overrides, exclusions);
+                Log.Debug<bool, int, int>(
+                    "_waf.Update was updated for change: {Success}, ({RulesOverridesCount} rule status entries), ({ExclusionsCount} exclusion filters)",
+                    result,
+                    overrides.Count,
+                    exclusions.Count);
+            }
 
             foreach (var asmConfig in asmConfigs)
             {
@@ -410,12 +488,12 @@ private void AsmProductConfigChanged(object sender, ProductConfigChangedEventArg
                 }
                 else
                 {
-                    e.Error(asmConfig.Name, "waf couldn't be updated with rule overrides");
+                    e.Error(asmConfig.Name, "waf couldn't be updated with asm product");
                 }
             }
         }
 
-        private bool UpdateWafWithRulesData() => _waf?.UpdateRulesData(_remoteConfigurationStatus.RulesData) ?? false;
+        private bool UpdateWafWithRulesData(List<RuleData> ruleData) => _waf?.UpdateRulesData(ruleData) ?? false;
 
         private void InitWafAndInstrumentations(bool fromRemoteConfig = false)
         {
@@ -441,7 +519,8 @@ private void InitWafAndInstrumentations(bool fromRemoteConfig = false)
                 _waf = _wafInitResult.Waf;
                 oldWaf?.Dispose();
                 Log.Debug("Disposed old waf and affected new waf");
-                UpdateWafWithRulesData();
+                var ruleData = _remoteConfigurationStatus.RulesDataByFile.SelectMany(x => x.Value).ToList();
+                UpdateWafWithRulesData(ruleData);
                 AddInstrumentationsAndProducts(fromRemoteConfig);
             }
             else
@@ -464,6 +543,8 @@ private void AddInstrumentationsAndProducts(bool fromRemoteConfig)
             {
                 AsmRemoteConfigurationProducts.AsmDataProduct.ConfigChanged += AsmDataProductConfigChanged;
                 AsmRemoteConfigurationProducts.AsmProduct.ConfigChanged += AsmProductConfigChanged;
+                AsmRemoteConfigurationProducts.AsmDataProduct.ConfigRemoved += AsmDataProductConfigRemoved;
+                AsmRemoteConfigurationProducts.AsmProduct.ConfigRemoved += AsmProductConfigRemoved;
                 AddAppsecSpecificInstrumentations();
 
                 _rateLimiter ??= new AppSecRateLimiter(_settings.TraceRateLimit);
@@ -480,6 +561,8 @@ private void RemoveInstrumentationsAndProducts(bool fromRemoteConfig)
             {
                 AsmRemoteConfigurationProducts.AsmDataProduct.ConfigChanged -= AsmDataProductConfigChanged;
                 AsmRemoteConfigurationProducts.AsmProduct.ConfigChanged -= AsmProductConfigChanged;
+                AsmRemoteConfigurationProducts.AsmDataProduct.ConfigRemoved -= AsmDataProductConfigRemoved;
+                AsmRemoteConfigurationProducts.AsmProduct.ConfigRemoved -= AsmProductConfigRemoved;
                 RemoveAppsecSpecificInstrumentations();
 
                 _enabled = false;
@@ -508,6 +591,8 @@ private void RunShutdown()
         {
             AsmRemoteConfigurationProducts.AsmDataProduct.ConfigChanged -= AsmDataProductConfigChanged;
             AsmRemoteConfigurationProducts.AsmProduct.ConfigChanged -= AsmProductConfigChanged;
+            AsmRemoteConfigurationProducts.AsmDataProduct.ConfigRemoved -= AsmDataProductConfigRemoved;
+            AsmRemoteConfigurationProducts.AsmProduct.ConfigRemoved -= AsmProductConfigRemoved;
             AsmRemoteConfigurationProducts.AsmFeaturesProduct.ConfigChanged -= FeaturesProductConfigChanged;
             AsmRemoteConfigurationProducts.AsmDDProduct.ConfigChanged -= AsmDDProductConfigChanged;
             Dispose();
diff --git a/tracer/src/Datadog.Trace/RemoteConfigurationManagement/RemoteConfigurationManager.cs b/tracer/src/Datadog.Trace/RemoteConfigurationManagement/RemoteConfigurationManager.cs
@@ -322,6 +322,7 @@ private void ProcessResponse(GetRcmResponse response, IDictionary<string, Produc
                 try
                 {
                     var removedConfigurations = GetRemovedConfigurations(product);
+
                     if (removedConfigurations is null)
                     {
                         continue;
diff --git a/tracer/test/Datadog.Trace.Security.IntegrationTests/Rcm/AspNetCore5AsmRulesToggle.cs b/tracer/test/Datadog.Trace.Security.IntegrationTests/Rcm/AspNetCore5AsmRulesToggle.cs
@@ -5,6 +5,7 @@
 
 #if NETCOREAPP3_0_OR_GREATER
 
+using System;
 using System.Collections.Generic;
 using System.Collections.Immutable;
 using System.Threading.Tasks;
@@ -38,28 +39,38 @@ public async Task TestRulesToggling()
             var ruleId = "crs-942-290";
 
             var spans0 = await SendRequestsAsync(agent, url);
-            var acknowledgedId = nameof(TestRulesToggling);
+            var acknowledgedId = nameof(TestRulesToggling) + Guid.NewGuid();
 
             var request1 = await agent.SetupRcmAndWait(Output, new[] { ((object)new Payload { RuleOverrides = new[] { new RuleOverride { Id = ruleId, Enabled = false } } }, acknowledgedId) }, ASMProduct, appliedServiceNames: new[] { acknowledgedId });
-            CheckAckState(request1, ASMProduct, ApplyStates.ACKNOWLEDGED, null, "First RCM call");
+            CheckAckState(request1, ASMProduct,  1, ApplyStates.ACKNOWLEDGED, null, "First RCM call");
 
             var spans1 = await SendRequestsAsync(agent, url);
 
-            var acknowledgedId2 = nameof(TestRulesToggling) + 2;
+            var acknowledgedId2 = nameof(TestRulesToggling) + Guid.NewGuid();
             var request2 = await agent.SetupRcmAndWait(Output, new[] { ((object)new Payload { RuleOverrides = new[] { new RuleOverride { Id = ruleId, Enabled = true } } }, acknowledgedId2) }, ASMProduct, appliedServiceNames: new[] { acknowledgedId2 });
-            CheckAckState(request2, ASMProduct, ApplyStates.ACKNOWLEDGED, null, "Second RCM call");
+            CheckAckState(request2, ASMProduct, 1, ApplyStates.ACKNOWLEDGED, null, "Second RCM call");
             var spans2 = await SendRequestsAsync(agent, url);
 
-            var acknowledgedId3 = nameof(TestRulesToggling) + 3;
-            var request3 = await agent.SetupRcmAndWait(Output, new[] { ((object)new Payload { RuleOverrides = new[] { new RuleOverride { Id = ruleId, Enabled = true, OnMatch = new[] { "block" } } } }, acknowledgedId3) }, ASMProduct, appliedServiceNames: new[] { acknowledgedId3 });
-            CheckAckState(request3, ASMProduct, ApplyStates.ACKNOWLEDGED, null, "Third RCM call");
+            var acknowledgedId3 = nameof(TestRulesToggling) + Guid.NewGuid();
+            var acknowledgedId4 = nameof(TestRulesToggling) + Guid.NewGuid();
+            var payload1 = ((object)new Payload { RuleOverrides = new[] { new RuleOverride { Id = ruleId, Enabled = true } } }, acknowledgedId3);
+            var payload2 = ((object)new Payload { RuleOverrides = new[] { new RuleOverride { Id = ruleId, OnMatch = new[] { "block" } } } }, acknowledgedId4);
+            var request3 = await agent.SetupRcmAndWait(Output, new[] { payload1, payload2 }, ASMProduct, appliedServiceNames: new[] { acknowledgedId3, acknowledgedId4 });
+            CheckAckState(request3, ASMProduct, 2, ApplyStates.ACKNOWLEDGED, null, "Third RCM call");
             var spans3 = await SendRequestsAsync(agent, url);
 
+            var acknowledgedId5 = nameof(TestRulesToggling) + Guid.NewGuid();
+            var payload3 = ((object)new Payload { RuleOverrides = new[] { new RuleOverride { Id = ruleId, Enabled = true } } }, acknowledgedId5);
+            var request4 = await agent.SetupRcmAndWait(Output, new[] { payload3 }, ASMProduct, appliedServiceNames: new[] { acknowledgedId5 });
+            CheckAckState(request4, ASMProduct, 1, ApplyStates.ACKNOWLEDGED, null, "Forth RCM call");
+            var spans4 = await SendRequestsAsync(agent, url);
+
             var spans = new List<MockSpan>();
             spans.AddRange(spans0);
             spans.AddRange(spans1);
             spans.AddRange(spans2);
             spans.AddRange(spans3);
+            spans.AddRange(spans4);
 
             await VerifySpans(spans.ToImmutableList(), settings);
         }
diff --git a/tracer/test/Datadog.Trace.Security.IntegrationTests/Rcm/AspNetCore5AsmToggle.cs b/tracer/test/Datadog.Trace.Security.IntegrationTests/Rcm/AspNetCore5AsmToggle.cs
@@ -69,17 +69,17 @@ public async Task TestSecurityToggling()
 
             var request1 = await agent.SetupRcmAndWait(Output, new[] { ((object)new AsmFeatures { Asm = new AsmFeature { Enabled = false } }, acknowledgedId) }, "ASM_FEATURES", "first", new[] { acknowledgedId });
 
-            RcmBase.CheckAckState(request1, "ASM_FEATURES", expectedState, null, "First RCM call");
+            RcmBase.CheckAckState(request1, "ASM_FEATURES", 1, expectedState, null, "First RCM call");
             request1.Client.State.BackendClientState.Should().Be("first");
 
-            RcmBase.CheckAckState(request1, "ASM_FEATURES", expectedState, null, "First RCM call");
+            RcmBase.CheckAckState(request1, "ASM_FEATURES", 1, expectedState, null, "First RCM call");
 
             var spans2 = await SendRequestsAsync(agent, url);
             var acknowledgedId2 = nameof(TestSecurityToggling) + Guid.NewGuid();
 
             var request2 = await agent.SetupRcmAndWait(Output, new[] { ((object)new AsmFeatures { Asm = new AsmFeature { Enabled = true } }, acknowledgedId2) }, "ASM_FEATURES", "second", new[] { acknowledgedId2 });
 
-            RcmBase.CheckAckState(request2, "ASM_FEATURES", expectedState, null, "Second RCM call");
+            RcmBase.CheckAckState(request2, "ASM_FEATURES", 1, expectedState, null, "Second RCM call");
 
             var request3 = await agent.WaitRcmRequestAndReturnLast(appliedServiceNames: new[] { acknowledgedId2 });
             request3.Client.State.BackendClientState.Should().Be("second");
@@ -116,7 +116,7 @@ public async Task TestRemoteConfigError()
 
             var request = await agent.SetupRcmAndWait(Output, new[] { ((object)"haha, you weren't expect this!", acknowledgedId) }, "ASM_FEATURES", appliedServiceNames: new[] { acknowledgedId });
 
-            RcmBase.CheckAckState(request, "ASM_FEATURES", ApplyStates.ERROR, "Error converting value \"haha, you weren't expect this!\" to type 'Datadog.Trace.AppSec.AsmFeatures'. Path '', line 1, position 32.", "First RCM call");
+            RcmBase.CheckAckState(request, "ASM_FEATURES", 1, ApplyStates.ERROR, "Error converting value \"haha, you weren't expect this!\" to type 'Datadog.Trace.AppSec.AsmFeatures'. Path '', line 1, position 32.", "First RCM call");
 
             await VerifySpans(spans1.ToImmutableList(), settings);
         }
diff --git a/tracer/test/Datadog.Trace.Security.IntegrationTests/Rcm/RcmBase.cs b/tracer/test/Datadog.Trace.Security.IntegrationTests/Rcm/RcmBase.cs
diff --git a/tracer/test/snapshots/Security.AspNetCore5AsmRulesToggle._.verified.txt b/tracer/test/snapshots/Security.AspNetCore5AsmRulesToggle._.verified.txt

Original file line number	Diff line number	Diff line change
`@@ -322,6 +322,7 @@ private void ProcessResponse(GetRcmResponse response, IDictionary<string, Produc`
`322`	`322`	`try`
`323`	`323`	`{`
`324`	`324`	`var removedConfigurations = GetRemovedConfigurations(product);`
	`325`	`+`
`325`	`326`	`if (removedConfigurations is null)`
`326`	`327`	`{`
`327`	`328`	`continue;`