[ASM] Fix one flakiness in rcm asm data integration tests and simplify some asm rcm code

tracer/missing-nullability-files.csv

@@ -270,7 +270,6 @@ src/Datadog.Trace/Agent/Transports/MimeTypes.cs
 src/Datadog.Trace/Agent/Transports/SocketHandlerRequestFactory.cs
 src/Datadog.Trace/AppSec/Concurrency/ReaderWriterLock.Core.cs
 src/Datadog.Trace/AppSec/Concurrency/ReaderWriterLock.Framework.cs
-src/Datadog.Trace/AppSec/Waf/IWaf.cs
 src/Datadog.Trace/AppSec/Waf/WafConstants.cs
 src/Datadog.Trace/AppSec/Waf/WafReturnCode.cs
 src/Datadog.Trace/Ci/Agent/ApmAgentWriter.cs

tracer/src/Datadog.Trace/AppSec/Rcm/AsmDdProduct.cs

@@ -40,18 +40,9 @@ public void ProcessUpdates(ConfigurationStatus configurationStatus, List<RemoteC
 
     public void ProcessRemovals(ConfigurationStatus configurationStatus, List<RemoteConfigurationPath> removedConfigsForThisProduct)
     {
-        var oneRemoved = false;
         foreach (var removedConfig in removedConfigsForThisProduct)
         {
-            oneRemoved |= configurationStatus.RulesByFile.Remove(removedConfig.Path);
-        }
-
-        if (configurationStatus.RulesByFile.Count == 0)
-        {
-            configurationStatus.IncomingUpdateState.FallbackToEmbeddedRuleset();
-        }
-        else if (oneRemoved)
-        {
+            configurationStatus.RulesByFile.Remove(removedConfig.Path);
             configurationStatus.IncomingUpdateState.WafKeysToApply.Add(ConfigurationStatus.WafRulesKey);
         }
     }

tracer/src/Datadog.Trace/AppSec/Rcm/ConfigurationStatus.cs

@@ -13,7 +13,6 @@
 using Datadog.Trace.AppSec.Rcm.Models.AsmDd;
 using Datadog.Trace.AppSec.Rcm.Models.AsmFeatures;
 using Datadog.Trace.AppSec.Waf.Initialization;
-using Datadog.Trace.ExtensionMethods;
 using Datadog.Trace.Logging;
 using Datadog.Trace.RemoteConfigurationManagement;
 using Datadog.Trace.Vendors.Newtonsoft.Json.Linq;
@@ -48,8 +47,6 @@ internal record ConfigurationStatus
 
     public ConfigurationStatus(string? embeddedRulesPath) => _embeddedRulesPath = embeddedRulesPath;
 
-    internal RuleSet? FallbackEmbeddedRuleSet { get; set; }
-
     internal bool? EnableAsm { get; set; } = null;
 
     internal string? AutoUserInstrumMode { get; set; } = null;
@@ -104,67 +101,73 @@ internal static List<RuleData> MergeRuleData(IEnumerable<RuleData> res)
         return finalRuleData;
     }
 
-    internal Dictionary<string, object> BuildDictionaryForWafAccordingToIncomingUpdate()
+    internal object? BuildDictionaryForWafAccordingToIncomingUpdate(string? embeddedRulesetPath)
     {
-        var dictionary = new Dictionary<string, object>();
+        var configuration = new Dictionary<string, object>();
 
         if (IncomingUpdateState.WafKeysToApply.Contains(WafExclusionsKey))
         {
             var exclusions = ExclusionsByFile.SelectMany(x => x.Value).ToList();
-            dictionary.Add(WafExclusionsKey, new JArray(exclusions));
+            configuration.Add(WafExclusionsKey, new JArray(exclusions));
         }
 
         if (IncomingUpdateState.WafKeysToApply.Contains(WafRulesOverridesKey))
         {
             var overrides = RulesOverridesByFile.SelectMany(x => x.Value).ToList();
-            dictionary.Add(WafRulesOverridesKey, overrides.Select(r => r.ToKeyValuePair()).ToArray());
+            configuration.Add(WafRulesOverridesKey, overrides.Select(r => r.ToKeyValuePair()).ToArray());
         }
 
         if (IncomingUpdateState.WafKeysToApply.Contains(WafRulesDataKey))
         {
             var rulesData = MergeRuleData(RulesDataByFile.SelectMany(x => x.Value));
-            dictionary.Add(WafRulesDataKey, rulesData.Select(r => r.ToKeyValuePair()).ToArray());
+            configuration.Add(WafRulesDataKey, rulesData.Select(r => r.ToKeyValuePair()).ToArray());
         }
 
         if (IncomingUpdateState.WafKeysToApply.Contains(WafExclusionsDataKey))
         {
             var rulesData = MergeRuleData(ExclusionsDataByFile.SelectMany(x => x.Value));
-            dictionary.Add(WafExclusionsDataKey, rulesData.Select(r => r.ToKeyValuePair()).ToArray());
+            configuration.Add(WafExclusionsDataKey, rulesData.Select(r => r.ToKeyValuePair()).ToArray());
         }
 
         if (IncomingUpdateState.WafKeysToApply.Contains(WafActionsKey))
         {
             var actions = ActionsByFile.SelectMany(x => x.Value).ToList();
-            dictionary.Add(WafActionsKey, actions.Select(r => r.ToKeyValuePair()).ToArray());
+            configuration.Add(WafActionsKey, actions.Select(r => r.ToKeyValuePair()).ToArray());
         }
 
         if (IncomingUpdateState.WafKeysToApply.Contains(WafCustomRulesKey))
         {
             var customRules = CustomRulesByFile.SelectMany(x => x.Value).ToList();
             var mergedCustomRules = new JArray(customRules);
-            dictionary.Add(WafCustomRulesKey, mergedCustomRules);
+            configuration.Add(WafCustomRulesKey, mergedCustomRules);
         }
 
-        if (IncomingUpdateState.FallbackToEmbeddedRulesetAtNextUpdate)
+        // if there's incoming rules or empty rules, or if asm is to be activated, we also want the rules key in waf arguments
+        if (IncomingUpdateState.WafKeysToApply.Contains(WafRulesKey) || (IncomingUpdateState.SecurityStateChange && (EnableAsm ?? false)))
         {
-            if (FallbackEmbeddedRuleSet == null)
+            var rulesetFromRcm = RulesByFile.Values.FirstOrDefault();
+            // should deserialize from LocalRuleFile
+            if (rulesetFromRcm is null)
             {
-                var result = WafConfigurator.DeserializeEmbeddedOrStaticRules(_embeddedRulesPath);
-                if (result != null)
+                var deserializedFromLocalRules = WafConfigurator.DeserializeEmbeddedOrStaticRules(embeddedRulesetPath);
+                if (deserializedFromLocalRules is not null)
                 {
-                    FallbackEmbeddedRuleSet = RuleSet.From(result);
+                    if (configuration.Count == 0)
+                    {
+                        return deserializedFromLocalRules;
+                    }
+
+                    var ruleSet = RuleSet.From(deserializedFromLocalRules);
+                    ruleSet.AddToDictionaryAtRoot(configuration);
                 }
             }
-
-            FallbackEmbeddedRuleSet?.AddToDictionaryAtRoot(dictionary);
-        }
-        else if (IncomingUpdateState.WafKeysToApply.Contains(WafRulesKey))
-        {
-            var rulesetFromRcm = RulesByFile.Values.FirstOrDefault();
-            rulesetFromRcm?.AddToDictionaryAtRoot(dictionary);
+            else
+            {
+                rulesetFromRcm?.AddToDictionaryAtRoot(configuration);
+            }
         }
 
-        return dictionary;
+        return configuration.Count > 0 ? configuration : null;
     }
 
     /// <summary>
@@ -247,7 +250,7 @@ public bool StoreLastConfigState(Dictionary<string, List<RemoteConfiguration>> c
                     }
                 }
 
-                // only treat asm_features as it will decide if asm gets toggled on and if we deserialize all the others
+                // only deserialize and apply asm_features as it will decide if asm gets toggled on and if we deserialize all the others
                 // (the enable of auto user instrumentation as added to asm_features)
                 _asmFeatureProduct.ProcessUpdates(this, asmFeaturesToUpdate);
                 _asmFeatureProduct.ProcessRemovals(this, asmFeaturesToRemove);
@@ -282,19 +285,14 @@ internal record IncomingUpdateStatus
     {
         internal HashSet<string> WafKeysToApply { get; } = new();
 
-        internal bool FallbackToEmbeddedRulesetAtNextUpdate { get; private set; }
-
         internal bool SecurityStateChange { get; set; }
 
         public void Reset()
         {
-            FallbackToEmbeddedRulesetAtNextUpdate = false;
             WafKeysToApply.Clear();
             SecurityStateChange = false;
         }
 
-        public void FallbackToEmbeddedRuleset() => FallbackToEmbeddedRulesetAtNextUpdate = true;
-
         public void SignalSecurityStateChange() => SecurityStateChange = true;
     }
 }

tracer/src/Datadog.Trace/AppSec/Rcm/Models/AsmDd/RuleSet.cs

@@ -1,4 +1,4 @@
-// <copyright file="RuleSet.cs" company="Datadog">
+// <copyright file="RuleSet.cs" company="Datadog">
 // Unless explicitly stated otherwise all files in this repository are licensed under the Apache 2 License.
 // This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2017 Datadog, Inc.
 // </copyright>
@@ -22,21 +22,32 @@ internal class RuleSet
     [JsonProperty("processors")]
     internal JToken? Processors { get; set; }
 
+    [JsonProperty("actions")]
+    internal JToken? Actions { get; set; }
+
     [JsonProperty("scanners")]
     internal JToken? Scanners { get; set; }
 
-    public JToken? All { get; set; }
+    [JsonProperty("exclusions")]
+    internal JToken? Exclusions { get; set; }
+
+    [JsonProperty("custom_rules")]
+    internal JToken? CustomRules { get; set; }
 
     public static RuleSet From(JToken result)
     {
+        // can rules from rc contains exclusions and custom rules?
+
         var ruleset = new RuleSet
         {
             Version = result["version"]?.ToString(),
             Metadata = result["metadata"],
             Rules = result["rules"],
             Processors = result["processors"],
             Scanners = result["scanners"],
-            All = result
+            Actions = result["actions"],
+            Exclusions = result["exclusions"],
+            CustomRules = result["custom_rules"]
         };
         return ruleset;
     }
@@ -49,27 +60,37 @@ public void AddToDictionaryAtRoot(Dictionary<string, object> dictionary)
     {
         if (Rules != null)
         {
-            dictionary.Add("rules", Rules);
+            dictionary["rules"] = Rules;
         }
 
         if (Metadata != null)
         {
-            dictionary.Add("metadata", Metadata);
+            dictionary["metadata"] = Metadata;
         }
 
         if (Version != null)
         {
-            dictionary.Add("version", Version);
+            dictionary["version"] = Version;
         }
 
         if (Processors != null)
         {
-            dictionary.Add("processors", Processors);
+            dictionary["processors"] = Processors;
         }
 
         if (Scanners != null)
         {
-            dictionary.Add("scanners", Scanners);
+            dictionary["scanners"] = Scanners;
+        }
+
+        if (Exclusions is not null)
+        {
+            dictionary["exclusions"] = Exclusions;
+        }
+
+        if (CustomRules is not null)
+        {
+            dictionary["custom_rules"] = CustomRules;
         }
     }
 }

tracer/src/Datadog.Trace/AppSec/Security.cs

@@ -208,7 +208,6 @@ private ApplyDetails[] UpdateFromRcm(Dictionary<string, List<RemoteConfiguration
                 // store the last config state, clearing any previous state, without deserializing any payloads yet.
                 var anyChange = _configurationStatus.StoreLastConfigState(configsByProduct, removedConfigs);
                 var securityStateChange = Enabled != _configurationStatus.EnableAsm;
-
                 // normally CanBeToggled should not need a check as asm_features capacity is only sent if AppSec env var is null, but still guards it in case
                 if (securityStateChange && _settings.CanBeToggled)
                 {
@@ -232,7 +231,7 @@ private ApplyDetails[] UpdateFromRcm(Dictionary<string, List<RemoteConfiguration
                 else if (Enabled && anyChange)
                 {
                     _configurationStatus.ApplyStoredFiles();
-                    updateResult = _waf?.UpdateWafFromConfigurationStatus(_configurationStatus);
+                    updateResult = _waf?.UpdateWafFromConfigurationStatus(_configurationStatus, _settings.Rules);
                     if (updateResult?.Success ?? false)
                     {
                         if (!string.IsNullOrEmpty(updateResult.RuleFileVersion))
@@ -539,7 +538,7 @@ private void InitWafAndInstrumentations(bool configurationFromRcm = false)
                 _wafLibraryInvoker!,
                 _settings.ObfuscationParameterKeyRegex,
                 _settings.ObfuscationParameterValueRegex,
-                _settings.Rules,
+                embeddedRulesetPath: _settings.Rules,
                 configurationFromRcm ? _configurationStatus : null,
                 _settings.UseUnsafeEncoder,
                 GlobalSettings.Instance.DebugEnabledInternal && _settings.WafDebugEnabled);
@@ -557,10 +556,6 @@ private void InitWafAndInstrumentations(bool configurationFromRcm = false)
                 Enabled = true;
                 InitializationError = null;
                 Log.Information("AppSec is now Enabled, _settings.Enabled is {EnabledValue}, coming from remote config: {EnableFromRemoteConfig}", _settings.Enabled, configurationFromRcm);
-                if (_wafInitResult.EmbeddedRules != null)
-                {
-                    _configurationStatus.FallbackEmbeddedRuleSet ??= RuleSet.From(_wafInitResult.EmbeddedRules);
-                }
 
                 if (!configurationFromRcm)
                 {
@@ -590,7 +585,7 @@ private void RemoveInstrumentationsAndProducts(bool fromRemoteConfig)
             {
                 if (_rcmSubscription != null)
                 {
-                    var newKeys = _rcmSubscription.ProductKeys.Except(new[] { RcmProducts.AsmData, RcmProducts.Asm }).ToArray();
+                    var newKeys = _rcmSubscription.ProductKeys.Except([RcmProducts.AsmData, RcmProducts.Asm]).ToArray();
                     if (newKeys.Length > 0)
                     {
                         var newSubscription = new Subscription(UpdateFromRcm, newKeys);

tracer/src/Datadog.Trace/AppSec/Waf/IWaf.cs

@@ -2,7 +2,7 @@
 // Unless explicitly stated otherwise all files in this repository are licensed under the Apache 2 License.
 // This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2017 Datadog, Inc.
 // </copyright>
-
+#nullable enable
 using System;
 using Datadog.Trace.AppSec.Rcm;
 using Datadog.Trace.AppSec.Waf.NativeBindings;
@@ -14,11 +14,11 @@ internal interface IWaf : IDisposable
     {
         public string Version { get; }
 
-        public IContext CreateContext();
+        public IContext? CreateContext();
 
         internal unsafe WafReturnCode Run(IntPtr contextHandle, DdwafObjectStruct* rawPersistentData, DdwafObjectStruct* rawEphemeralData, ref DdwafResultStruct retNative, ulong timeoutMicroSeconds);
 
-        UpdateResult UpdateWafFromConfigurationStatus(ConfigurationStatus configurationStatus);
+        UpdateResult UpdateWafFromConfigurationStatus(ConfigurationStatus configurationStatus, string? staticRulesFilePath = null);
 
         public string[] GetKnownAddresses();
 

tracer/src/Datadog.Trace/AppSec/Waf/Initialization/WafConfigurator.cs

@@ -82,6 +82,13 @@ private static void LogRuleDetailsIfDebugEnabled(JToken root)
             return File.OpenRead(rulesFile);
         }
 
+        /// <summary>
+        /// Deserialize rules for the waf as Jtoken
+        /// If null is passed, will deserialize embedded rule file in the app
+        /// If a path is given but file isn't found, it won't fallback on the embedded rule file
+        /// </summary>
+        /// <param name="rulesFilePath">if null, will fallback on embedded rules file</param>
+        /// <returns>the rules, might be null if file not found</returns>
         internal static JToken? DeserializeEmbeddedOrStaticRules(string? rulesFilePath)
         {
             JToken root;

tracer/src/Datadog.Trace/AppSec/Waf/ReturnTypes.Managed/InitResult.cs

@@ -10,19 +10,17 @@
 using Datadog.Trace.AppSec.WafEncoding;
 using Datadog.Trace.Logging;
 using Datadog.Trace.Vendors.Newtonsoft.Json;
-using Datadog.Trace.Vendors.Newtonsoft.Json.Linq;
 
 namespace Datadog.Trace.AppSec.Waf.ReturnTypes.Managed
 {
     internal class InitResult
     {
         private static readonly IDatadogLogger Log = DatadogLogging.GetLoggerFor(typeof(InitResult));
 
-        private InitResult(ushort failedToLoadRules, ushort loadedRules, string ruleFileVersion, IReadOnlyDictionary<string, object> errors, JToken? embeddedRules = null, bool unusableRuleFile = false, IntPtr? wafHandle = null, WafLibraryInvoker? wafLibraryInvoker = null, IEncoder? encoder = null, bool shouldEnableWaf = true, bool incompatibleWaf = false)
+        private InitResult(ushort failedToLoadRules, ushort loadedRules, string ruleFileVersion, IReadOnlyDictionary<string, object> errors, bool unusableRuleFile = false, IntPtr? wafHandle = null, WafLibraryInvoker? wafLibraryInvoker = null, IEncoder? encoder = null, bool shouldEnableWaf = true, bool incompatibleWaf = false)
         {
             HasErrors = errors.Count > 0;
             Errors = errors;
-            EmbeddedRules = embeddedRules;
             FailedToLoadRules = failedToLoadRules;
             LoadedRules = loadedRules;
             RuleFileVersion = ruleFileVersion;
@@ -55,8 +53,6 @@ private InitResult(ushort failedToLoadRules, ushort loadedRules, string ruleFile
 
         internal IReadOnlyDictionary<string, object> Errors { get; }
 
-        public JToken? EmbeddedRules { get; set; }
-
         internal string ErrorMessage { get; }
 
         internal bool HasErrors { get; }