Refactor remote config application for ASM

Author: andrewlock

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

@@ -7,6 +7,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Diagnostics.CodeAnalysis;
 using System.Linq;
 using Datadog.Trace.AppSec.Rcm.Models.Asm;
 using Datadog.Trace.AppSec.Rcm.Models.AsmData;
@@ -26,8 +27,10 @@ namespace Datadog.Trace.AppSec.Rcm;
 /// <summary>
 /// This class represents the state of RCM for ASM.
 /// It has 2 possible status:
-/// - ASM is not activated, and _fileUpdates/_fileRemoves contain some pending non-deserialized changes to apply when ASM_FEATURES activate ASM. Every time an RC payload is received here, pending changes are reset to the last ones
-/// - ASM is activated, stored configs in _fileUpdates/_fileRemoves are applied every time.
+/// - ASM is not activated, and _pendingByProduct accumulates pending non-deserialized changes (updates and removes)
+///   across RC polls until ASM_FEATURES activates ASM, at which point they are applied. RCM only forwards deltas
+///   (new/changed/removed configs), so the pending state must persist between polls.
+/// - ASM is activated, the pending state in _pendingByProduct is applied every time.
 /// </summary>
 internal sealed record ConfigurationState
 {
@@ -41,8 +44,13 @@ internal sealed record ConfigurationState
 
     private readonly string? _rulesPath;
     private readonly bool _canBeToggled;
-    private readonly Dictionary<string, List<RemoteConfiguration>> _fileUpdates = new();
-    private readonly Dictionary<string, List<RemoteConfigurationPath>> _fileRemoves = new();
+
+    // Pending RCM operations, grouped by product. Inner key is the full config path; the value is
+    // either an update payload or a remove marker. Using a dictionary keyed by path means each path
+    // can only hold one pending operation at a time — a later update or remove for the same path
+    // overwrites whatever was there, which is exactly the RCM "latest delta wins" semantics.
+    private readonly Dictionary<string, Dictionary<string, PendingOperation>> _pendingByProduct = new();
+
     private bool _defaultRulesetApplied;
 
     public ConfigurationState(SecuritySettings settings, IConfigurationTelemetry telemetry, bool wafIsNull)
@@ -136,9 +144,19 @@ internal RemoteConfigWafFiles GetWafConfigurations(bool updating = false)
         var configurations = new Dictionary<string, object>();
         List<string>? removes = null;
 
-        if (updating && _fileRemoves is { Count: > 0 })
+        if (updating)
         {
-            removes = _fileRemoves.SelectMany(p => p.Value).Where(p => p.Product != RcmProducts.AsmFeatures).Select(v => v.Path).ToList();
+            foreach (var entry in _pendingByProduct)
+            {
+                if (entry.Key == RcmProducts.AsmFeatures) { continue; }
+                foreach (var op in entry.Value.Values)
+                {
+                    if (op.IsRemove)
+                    {
+                        (removes ??= new()).Add(op.Path.Path);
+                    }
+                }
+            }
         }
 
         if (AsmConfigs is { Count: > 0 })
@@ -192,11 +210,12 @@ internal RemoteConfigWafFiles GetWafConfigurations(bool updating = false)
 
         return new(configurations.Count > 0 ? configurations : null, removes);
 
+        // True if `path` has a pending update (not a remove) in any product slot.
         bool IsNewUpdate(string path)
         {
-            foreach (var productUpdates in _fileUpdates)
+            foreach (var pending in _pendingByProduct.Values)
             {
-                if (productUpdates.Value.Any(u => u.Path.Path == path))
+                if (pending.TryGetValue(path, out var op) && !op.IsRemove)
                 {
                     return true;
                 }
@@ -207,99 +226,111 @@ bool IsNewUpdate(string path)
     }
 
     /// <summary>
-    /// Calls each product updater to deserialize all remote config payloads and store them properly in dictionaries which might involve various logical merges
-    /// This method deserializes everything stored in _fileUpdates. ConfigurationStatus will have a *bigger* memory footprint.
+    /// Calls each product updater to deserialize all remote config payloads and store them properly in dictionaries which might involve various logical merges.
     /// </summary>
     public void ApplyStoredFiles()
     {
-        // no need to clear _fileUpdates / _fileRemoves after they've been applied, as when we receive a new config, `StoreLastConfigState` method will clear anything remaining anyway.
+        // The pending state is NOT cleared here — GetWafConfigurations still reads it to identify which
+        // configs are new for the WAF update path. Re-application is idempotent: each product updater
+        // writes into its destination collection keyed by config path, so replaying the same pending
+        // entry on a later turn produces the same end state.
         foreach (var updater in _productConfigUpdaters)
         {
-            var fileRemoves = _fileRemoves.TryGetValue(updater.Key, out var removes);
-            var fileUpdates = _fileUpdates.TryGetValue(updater.Key, out var updates);
-            if (fileRemoves || fileUpdates)
+            if (!_pendingByProduct.TryGetValue(updater.Key, out var pending) || pending.Count == 0)
             {
-                updater.Value.ProcessUpdates(this, removes, updates);
+                continue;
             }
+
+            var (removes, updates) = SplitPending(pending);
+            updater.Value.ProcessUpdates(this, removes, updates);
         }
     }
 
     /// <summary>
-    /// This method just stores the config state without deserializing anything, this state will be ready to use and deserialized if ASM is enabled later on.
-    /// This method considers that RC sends us everything again, the whole state together. That's why it's clearing all unapplied updates / removals before processing the last ones received.
-    /// In case ASM remained disabled, we discard previous updates and removals stored here that were never applied.
+    /// Merges an incoming RCM delta into the per-product pending state without deserializing payloads.
+    /// If ASM is enabled (or this delta turns it on), the merged state is applied immediately; otherwise it stays pending until ASM activates.
     /// </summary>
-    /// <param name="configsByProduct">configsByProduct</param>
-    /// <param name="removedConfigs">removedConfigs</param>
+    /// <remarks>
+    /// Background: RCM forwards only the DELTA versus what we last acknowledged — not the full backend state.
+    /// A config received in an earlier poll is NOT re-sent on the next poll if it hasn't changed.
+    /// So while ASM is disabled, we must accumulate deltas across polls, otherwise configs received
+    /// before activation would be silently dropped (they're already in RCM's _appliedConfigurations
+    /// and won't be re-forwarded). Each path holds exactly one pending op in _pendingByProduct;
+    /// a later update or remove for the same path overwrites whatever was there.
+    /// </remarks>
     public void ReceivedNewConfig(Dictionary<string, List<RemoteConfiguration>> configsByProduct, Dictionary<string, List<RemoteConfigurationPath>>? removedConfigs)
     {
-        _fileUpdates.Clear();
-        _fileRemoves.Clear();
-        // if we just have asm features, it might only be to toggle appsec. Other products bring actual configurations to the waf.
-        var hasUpdateConfigurations = false;
-        var anyChange = configsByProduct.Count > 0 || removedConfigs?.Count > 0;
-        if (anyChange)
+        if (configsByProduct.Count == 0 && (removedConfigs is null || removedConfigs.Count == 0))
         {
-            if (removedConfigs != null)
-            {
-                foreach (var configByProductToRemove in removedConfigs)
-                {
-                    if (configByProductToRemove.Key != RcmProducts.AsmFeatures)
-                    {
-                        hasUpdateConfigurations = true;
-                    }
+            return;
+        }
 
-                    if (_fileRemoves.TryGetValue(configByProductToRemove.Key, out var remove))
-                    {
-                        remove.AddRange(configByProductToRemove.Value);
-                    }
-                    else
-                    {
-                        _fileRemoves[configByProductToRemove.Key] = configByProductToRemove.Value;
-                    }
-                }
-            }
+        // Tracks whether this delta touches any product OTHER than ASM_FEATURES. A pure ASM_FEATURES
+        // delta only toggles AppSec on/off — it doesn't carry WAF rules — so when that's all we got,
+        // there's no reason to push a config update through to the WAF below.
+        var hasUpdateConfigurations = false;
 
-            foreach (var configByProduct in configsByProduct)
+        if (removedConfigs is not null)
+        {
+            foreach (var entry in removedConfigs)
             {
-                if (configByProduct.Key != RcmProducts.AsmFeatures)
+                var productKey = entry.Key;
+                hasUpdateConfigurations |= productKey != RcmProducts.AsmFeatures;
+                if (!_pendingByProduct.TryGetValue(productKey, out var pending))
                 {
-                    hasUpdateConfigurations = true;
+                    pending = new Dictionary<string, PendingOperation>();
+                    _pendingByProduct[productKey] = pending;
                 }
 
-                if (_fileUpdates.TryGetValue(configByProduct.Key, out var update))
-                {
-                    update.AddRange(configByProduct.Value);
-                }
-                else
+                foreach (var removed in entry.Value)
                 {
-                    _fileUpdates[configByProduct.Key] = configByProduct.Value;
+                    // Overwrites any pending update OR pending remove for this path. Later op wins.
+                    pending[removed.Path] = new PendingOperation(removed, update: null);
                 }
             }
+        }
 
-            ApplyAsmFeatures(AppsecEnabled);
-            IncomingUpdateState.ShouldUpdateAppsec = !IncomingUpdateState.ShouldInitAppsec && AppsecEnabled && hasUpdateConfigurations;
+        foreach (var entry in configsByProduct)
+        {
+            var productKey = entry.Key;
+            hasUpdateConfigurations |= productKey != RcmProducts.AsmFeatures;
+            if (!_pendingByProduct.TryGetValue(productKey, out var pending))
+            {
+                pending = new Dictionary<string, PendingOperation>();
+                _pendingByProduct[productKey] = pending;
+            }
 
-            if (IncomingUpdateState.ShouldUpdateAppsec || IncomingUpdateState.ShouldInitAppsec)
+            foreach (var update in entry.Value)
             {
-                ApplyStoredFiles();
+                // Same overwrite semantics — replaces any prior op for this path.
+                pending[update.Path.Path] = new PendingOperation(update.Path, update);
             }
         }
+
+        // ASM_FEATURES is processed eagerly because it can flip AppsecEnabled and decide whether
+        // we initialize the WAF, update an already-running WAF, or do nothing.
+        ApplyAsmFeatures(AppsecEnabled);
+        IncomingUpdateState.ShouldUpdateAppsec = !IncomingUpdateState.ShouldInitAppsec && AppsecEnabled && hasUpdateConfigurations;
+
+        // If ASM just turned on, or it was already on and we got real WAF config, drain pending state
+        // into the per-product output dictionaries (RulesetConfigs, AsmConfigs, AsmDataConfigs).
+        if (IncomingUpdateState.ShouldUpdateAppsec || IncomingUpdateState.ShouldInitAppsec)
+        {
+            ApplyStoredFiles();
+        }
     }
 
     private void ApplyAsmFeatures(bool appsecCurrentlyEnabled)
     {
         // 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)
-        var change = _fileRemoves.TryGetValue(RcmProducts.AsmFeatures, out var removals);
-        change |= _fileUpdates.TryGetValue(RcmProducts.AsmFeatures, out var updates);
-
-        if (change)
+        if (!_pendingByProduct.TryGetValue(RcmProducts.AsmFeatures, out var pending) || pending.Count == 0)
         {
-            _asmFeatureProduct.ProcessUpdates(this, removals, updates);
+            return;
         }
 
-        if (!change) { return; }
+        var (removals, updates) = SplitPending(pending);
+        _asmFeatureProduct.ProcessUpdates(this, removals, updates);
 
         // 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 (_canBeToggled)
@@ -341,6 +372,41 @@ private void ApplyAsmFeatures(bool appsecCurrentlyEnabled)
         }
     }
 
+    // Splits the per-product pending map into the (removes, updates) shape that IAsmConfigUpdater.ProcessUpdates expects.
+    // Either list may be null when that side is empty — matching the interface contract.
+    private static (List<RemoteConfigurationPath>? Removes, List<RemoteConfiguration>? Updates) SplitPending(Dictionary<string, PendingOperation> pending)
+    {
+        List<RemoteConfigurationPath>? removes = null;
+        List<RemoteConfiguration>? updates = null;
+        foreach (var op in pending.Values)
+        {
+            if (op.IsRemove)
+            {
+                removes ??= [];
+                removes.Add(op.Path);
+            }
+            else
+            {
+                updates ??= [];
+                updates.Add(op.Update);
+            }
+        }
+
+        return (removes, updates);
+    }
+
+    // A pending RCM operation for one config path: either an update carrying a fresh payload,
+    // or a removal marker. The Update field is null iff this entry represents a removal.
+    private readonly struct PendingOperation(RemoteConfigurationPath path, RemoteConfiguration? update)
+    {
+        public RemoteConfigurationPath Path { get; } = path;
+
+        public RemoteConfiguration? Update { get; } = update;
+
+        [MemberNotNullWhen(false, nameof(Update))]
+        public bool IsRemove => Update is null;
+    }
+
     internal sealed record IncomingUpdateStatus : IDisposable
     {
         internal bool ShouldInitAppsec { get; set; }