Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[ASM] Dont deserialize rcm payloads until they are needed for memory optimization #5296

Merged
merged 5 commits into from
Apr 11, 2024
Merged

[ASM] Dont deserialize rcm payloads until they are needed for memory optimization #5296

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
1a34259
Dont deserialize rcm payloads until they are needed
anna-git Mar 11, 2024
7b63746
have the right version
anna-git Mar 11, 2024
ea12b77
test for change when storing configs
anna-git Mar 26, 2024
562d267
Fix tests
anna-git Mar 26, 2024
ee10989
Document important method
anna-git Apr 10, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 0 additions & 2 deletions tracer/src/Datadog.Trace/AppSec/Rcm/AsmFeaturesProduct.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,9 +6,7 @@
#nullable enable

using System.Collections.Generic;
using System.Linq;
using Datadog.Trace.AppSec.Rcm.Models.AsmFeatures;
using Datadog.Trace.ExtensionMethods;
using Datadog.Trace.RemoteConfigurationManagement;

namespace Datadog.Trace.AppSec.Rcm;
Expand Down
107 changes: 103 additions & 4 deletions tracer/src/Datadog.Trace/AppSec/Rcm/ConfigurationStatus.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,29 +8,38 @@
using System;
using System.Collections.Generic;
using System.Linq;
using Datadog.Trace.AppSec.Coordinator;
using Datadog.Trace.AppSec.Rcm.Models.Asm;
using Datadog.Trace.AppSec.Rcm.Models.AsmData;
using Datadog.Trace.AppSec.Rcm.Models.AsmDd;
using Datadog.Trace.AppSec.Rcm.Models.AsmFeatures;
using Datadog.Trace.AppSec.Waf.Initialization;
using Datadog.Trace.Logging;
using Datadog.Trace.ExtensionMethods;
using Datadog.Trace.RemoteConfigurationManagement;
using Datadog.Trace.Vendors.Newtonsoft.Json.Linq;
using Datadog.Trace.Vendors.Serilog;
using Action = Datadog.Trace.AppSec.Rcm.Models.Asm.Action;

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.
/// </summary>
internal record ConfigurationStatus
{
internal const string WafRulesKey = "rules";
internal const string WafRulesOverridesKey = "rules_override";
internal const string WafExclusionsKey = "exclusions";
internal const string WafRulesDataKey = "rules_data";
internal const string WafCustomRulesKey = "custom_rules";
private readonly IAsmConfigUpdater _asmFeatureProduct = new AsmFeaturesProduct();

private static readonly IDatadogLogger Log = DatadogLogging.GetLoggerFor<ConfigurationStatus>();
private readonly IReadOnlyDictionary<string, IAsmConfigUpdater> _productConfigUpdaters = new Dictionary<string, IAsmConfigUpdater> { { RcmProducts.Asm, new AsmProduct() }, { RcmProducts.AsmDd, new AsmDdProduct() }, { RcmProducts.AsmData, new AsmDataProduct() } };

private readonly string? _embeddedRulesPath;
private Dictionary<string, List<RemoteConfiguration>> _fileUpdates = new();
private Dictionary<string, List<RemoteConfigurationPath>> _fileRemoves = new();

public ConfigurationStatus(string? embeddedRulesPath) => _embeddedRulesPath = embeddedRulesPath;

Expand Down Expand Up @@ -138,6 +147,96 @@ internal Dictionary<string, object> BuildDictionaryForWafAccordingToIncomingUpda
return dictionary;
}

/// <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.
/// </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.
foreach (var updater in _productConfigUpdaters)
{
var fileUpdates = _fileUpdates.TryGetValue(updater.Key, out var value);
if (fileUpdates)
{
updater.Value.ProcessUpdates(this, value!);
}

var fileRemoves = _fileRemoves.TryGetValue(updater.Key, out var valueRemove);
if (fileRemoves)
{
updater.Value.ProcessRemovals(this, valueRemove!);
}
}
}

/// <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.
/// </summary>
/// <param name="configsByProduct">configsByProduct</param>
/// <param name="removedConfigs">removedConfigs</param>
/// <returns>whether or not there is any change, i.e any update/removal</returns>
public bool StoreLastConfigState(Dictionary<string, List<RemoteConfiguration>> configsByProduct, Dictionary<string, List<RemoteConfigurationPath>>? removedConfigs)
{
_fileUpdates.Clear();
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

NIT: Is this thread safe? I'm not sure if a lock wouldn't hurt here and in ApplyStoredFiles()

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So rcm updates are not multi threaded, everything is pulled on one thread, and sequentially every x seconds

_fileRemoves.Clear();
List<RemoteConfiguration> asmFeaturesToUpdate = new();
List<RemoteConfigurationPath> asmFeaturesToRemove = new();
var anyChange = configsByProduct.Count > 0 || removedConfigs?.Count > 0;
if (anyChange)
{
foreach (var configByProduct in configsByProduct)
{
if (configByProduct.Key == RcmProducts.AsmFeatures)
{
asmFeaturesToUpdate.AddRange(configByProduct.Value);
}
else
{
if (_fileUpdates.ContainsKey(configByProduct.Key))
{
_fileUpdates[configByProduct.Key].AddRange(configByProduct.Value);
}
else
{
_fileUpdates[configByProduct.Key] = configByProduct.Value;
}
}
}

if (removedConfigs != null)
{
foreach (var configByProductToRemove in removedConfigs)
{
if (configByProductToRemove.Key == RcmProducts.AsmFeatures)
{
asmFeaturesToRemove.AddRange(configByProductToRemove.Value);
}
else
{
if (_fileRemoves.ContainsKey(configByProductToRemove.Key))
{
_fileRemoves[configByProductToRemove.Key].AddRange(configByProductToRemove.Value);
}
else
{
_fileRemoves[configByProductToRemove.Key] = configByProductToRemove.Value;
}
}
}

// only treat asm_features as it will decide if asm gets toggled on and if we deserialize all the others
_asmFeatureProduct.ProcessUpdates(this, asmFeaturesToUpdate);
_asmFeatureProduct.ProcessRemovals(this, asmFeaturesToRemove);
EnableAsm = !AsmFeaturesByFile.IsEmpty() && AsmFeaturesByFile.All(a => a.Value.Enabled == true);
}
}

return anyChange;
}

public void ResetUpdateMarkers() => IncomingUpdateState.Reset();

internal record IncomingUpdateStatus
Expand Down
2 changes: 1 addition & 1 deletion tracer/src/Datadog.Trace/AppSec/Rcm/Models/AsmFeatures/AsmFeature.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@

#nullable enable

namespace Datadog.Trace.AppSec;
namespace Datadog.Trace.AppSec.Rcm.Models.AsmFeatures;

internal class AsmFeature
{
Expand Down
54 changes: 19 additions & 35 deletions tracer/src/Datadog.Trace/AppSec/Security.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,8 +37,6 @@ internal class Security : IDatadogSecurity, IDisposable
private static bool _globalInstanceInitialized;
private static object _globalInstanceLock = new();
private readonly SecuritySettings _settings;
private readonly IReadOnlyDictionary<string, IAsmConfigUpdater> _productConfigUpdaters = new Dictionary<string, IAsmConfigUpdater> { { RcmProducts.AsmFeatures, new AsmFeaturesProduct() }, { RcmProducts.Asm, new AsmProduct() }, { RcmProducts.AsmDd, new AsmDdProduct() }, { RcmProducts.AsmData, new AsmDataProduct() } };

private readonly ConfigurationStatus _configurationStatus;
private readonly bool _noLocalRules;
private readonly IRcmSubscriptionManager _rcmSubscriptionManager;
Expand Down Expand Up @@ -166,52 +164,39 @@ internal void SubscribeToChanges(params string[] productNames)
}
}

internal ApplyDetails[] UpdateFromRcm(Dictionary<string, List<RemoteConfiguration>> configsByProduct, Dictionary<string, List<RemoteConfigurationPath>>? removedConfigs)
private ApplyDetails[] UpdateFromRcm(Dictionary<string, List<RemoteConfiguration>> configsByProduct, Dictionary<string, List<RemoteConfigurationPath>>? removedConfigs)
{
string? rcmUpdateError = null;
UpdateResult? updateResult = null;

try
{
foreach (var product in _productConfigUpdaters)
{
if (configsByProduct.TryGetValue(product.Key, out var configurations))
{
product.Value.ProcessUpdates(_configurationStatus, configurations);
}

if (removedConfigs?.TryGetValue(product.Key, out var configsForThisProductToRemove) is true)
{
product.Value.ProcessRemovals(_configurationStatus, configsForThisProductToRemove);
}
}

_configurationStatus.EnableAsm = !_configurationStatus.AsmFeaturesByFile.IsEmpty() && _configurationStatus.AsmFeaturesByFile.All(a => a.Value.Enabled == true);
// 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 (_configurationStatus.IncomingUpdateState.SecurityStateChange && _settings.CanBeToggled)
if (securityStateChange && _settings.CanBeToggled)
{
// disable ASM scenario
if (Enabled && _configurationStatus.EnableAsm == false)
{
DisposeWafAndInstrumentations(true);
_configurationStatus.IncomingUpdateState.SecurityStateChange = false;
}
} // enable ASM scenario taking into account rcm changes for other products/data
else if (!Enabled && _configurationStatus.EnableAsm == true)
{
_configurationStatus.ApplyStoredFiles();
InitWafAndInstrumentations(true);
rcmUpdateError = _wafInitResult?.ErrorMessage;
// no point in updating the waf with potentially new rules as it's initialized here with new rules
_configurationStatus.IncomingUpdateState.WafKeysToApply.Remove(ConfigurationStatus.WafRulesKey);
// reapply others
_configurationStatus.IncomingUpdateState.WafKeysToApply.Add(ConfigurationStatus.WafExclusionsKey);
_configurationStatus.IncomingUpdateState.WafKeysToApply.Add(ConfigurationStatus.WafRulesDataKey);
_configurationStatus.IncomingUpdateState.WafKeysToApply.Add(ConfigurationStatus.WafRulesOverridesKey);
_configurationStatus.IncomingUpdateState.SecurityStateChange = false;
if (_wafInitResult?.RuleFileVersion is not null)
{
WafRuleFileVersion = _wafInitResult.RuleFileVersion;
}
}
}

if (Enabled && _configurationStatus.IncomingUpdateState.WafKeysToApply.Any())
} // update asm configuration
else if (Enabled && anyChange)
{
_configurationStatus.ApplyStoredFiles();
updateResult = _waf?.UpdateWafFromConfigurationStatus(_configurationStatus);
if (updateResult?.Success ?? false)
{
Expand Down Expand Up @@ -371,7 +356,7 @@ private void SetRemoteConfigCapabilites()
rcm.SetCapability(RcmCapabilitiesIndices.AsmTrustedIps, _noLocalRules);
}

private void InitWafAndInstrumentations(bool fromRemoteConfig = false)
private void InitWafAndInstrumentations(bool configurationFromRcm = false)
{
// initialization of WafLibraryInvoker
if (_libraryInitializationResult == null)
Expand All @@ -393,8 +378,7 @@ private void InitWafAndInstrumentations(bool fromRemoteConfig = false)
_settings.ObfuscationParameterKeyRegex,
_settings.ObfuscationParameterValueRegex,
_settings.Rules,
_configurationStatus.RulesByFile.Values.FirstOrDefault()?.All,
setupWafSchemaExtraction: _settings.ApiSecurityEnabled,
configurationFromRcm ? _configurationStatus : null,
_settings.UseUnsafeEncoder,
GlobalSettings.Instance.DebugEnabledInternal && _settings.WafDebugEnabled);
if (_wafInitResult.Success)
Expand All @@ -410,13 +394,13 @@ private void InitWafAndInstrumentations(bool fromRemoteConfig = false)
_rateLimiter ??= new(_settings.TraceRateLimit);
Enabled = true;
InitializationError = null;
Log.Information("AppSec is now Enabled, _settings.Enabled is {EnabledValue}, coming from remote config: {EnableFromRemoteConfig}", _settings.Enabled, fromRemoteConfig);
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 (!fromRemoteConfig)
if (!configurationFromRcm)
{
// occurs the first time we initialize the WAF
TelemetryFactory.Metrics.SetWafVersion(_waf!.Version);
Expand Down
27 changes: 18 additions & 9 deletions tracer/src/Datadog.Trace/AppSec/Waf/Waf.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,8 +52,7 @@ internal Waf(IntPtr wafHandle, WafLibraryInvoker wafLibraryInvoker, IEncoder enc
/// <param name="obfuscationParameterValueRegex">the regex that will be used to obfuscate possible sensitive data in values that are highlighted WAF as potentially malicious,
/// empty string means use default embedded in the WAF</param>
/// <param name="embeddedRulesetPath">can be null, means use rules embedded in the manifest </param>
/// <param name="rulesFromRcm">can be null. RemoteConfig rules json. Takes precedence over rulesFile </param>
/// <param name="setupWafSchemaExtraction">should we read the config file for schema extraction</param>
/// <param name="configurationStatus">can be null. RemoteConfig rules json. Takes precedence over rulesFile </param>
/// <param name="useUnsafeEncoder">use legacy encoder</param>
/// <param name="wafDebugEnabled">if debug level logs should be enabled for the WAF</param>
/// <returns>the waf wrapper around waf native</returns>
Expand All @@ -62,8 +61,7 @@ internal static InitResult Create(
string obfuscationParameterKeyRegex,
string obfuscationParameterValueRegex,
string? embeddedRulesetPath = null,
JToken? rulesFromRcm = null,
bool setupWafSchemaExtraction = false,
ConfigurationStatus? configurationStatus = null,
bool useUnsafeEncoder = false,
bool wafDebugEnabled = false)
{
Expand All @@ -72,8 +70,19 @@ internal static InitResult Create(
// set the log level and setup the logger
wafLibraryInvoker.SetupLogging(wafDebugEnabled);

var jtokenRoot = rulesFromRcm ?? WafConfigurator.DeserializeEmbeddedOrStaticRules(embeddedRulesetPath)!;
if (jtokenRoot is null)
object? configurationToEncode = null;
if (configurationStatus is not null)
{
var configFromRcm = configurationStatus.BuildDictionaryForWafAccordingToIncomingUpdate();
if (configFromRcm.Count > 0)
{
configurationToEncode = configFromRcm;
}
}

configurationToEncode ??= WafConfigurator.DeserializeEmbeddedOrStaticRules(embeddedRulesetPath)!;

if (configurationToEncode is null)
{
return InitResult.FromUnusableRuleFile();
}
Expand All @@ -84,15 +93,15 @@ internal static InitResult Create(
configWafStruct.KeyRegex = keyRegex;
configWafStruct.ValueRegex = valueRegex;
IEncoder encoder = useUnsafeEncoder ? new Encoder() : new EncoderLegacy(wafLibraryInvoker);
var result = encoder.Encode(jtokenRoot, applySafetyLimits: false);
var result = encoder.Encode(configurationToEncode, applySafetyLimits: false);
var rulesObj = result.ResultDdwafObject;

var diagnostics = new DdwafObjectStruct { Type = DDWAF_OBJ_TYPE.DDWAF_OBJ_MAP };

try
{
var initResult = wafConfigurator.Configure(ref rulesObj, encoder, configWafStruct, ref diagnostics, rulesFromRcm == null ? embeddedRulesetPath : "RemoteConfig");
initResult.EmbeddedRules = jtokenRoot;
var initResult = wafConfigurator.Configure(ref rulesObj, encoder, configWafStruct, ref diagnostics, configurationStatus == null ? embeddedRulesetPath : "RemoteConfig");
initResult.EmbeddedRules = initResult.EmbeddedRules;
return initResult;
}
finally
Expand Down
2 changes: 1 addition & 1 deletion tracer/test/Datadog.Trace.Security.Unit.Tests/ContextTests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@ public void MultipleContextsRun(bool useUnsafeEncoder)
AppSec.WafEncoding.Encoder.SetPoolSize(0);
}

var initResult = Waf.Create(WafLibraryInvoker, string.Empty, string.Empty, setupWafSchemaExtraction: true, useUnsafeEncoder: useUnsafeEncoder);
var initResult = Waf.Create(WafLibraryInvoker, string.Empty, string.Empty, useUnsafeEncoder: useUnsafeEncoder);
using var waf = initResult.Waf;
waf.Should().NotBeNull();

Expand Down
2 changes: 1 addition & 1 deletion tracer/test/Datadog.Trace.Security.Unit.Tests/WafEphemeralDataTests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -100,7 +100,7 @@ private void Execute(string address, object[] values, bool[] matches, string flo

private void ExecuteInternal(string address, object[] values, bool[] matches, string flow, string rule, bool newEncoder)
{
var initResult = Waf.Create(WafLibraryInvoker, string.Empty, string.Empty, setupWafSchemaExtraction: false, useUnsafeEncoder: newEncoder);
var initResult = Waf.Create(WafLibraryInvoker, string.Empty, string.Empty, useUnsafeEncoder: newEncoder);
using var waf = initResult.Waf;
waf.Should().NotBeNull();
using var context = waf.CreateContext();
Expand Down
2 changes: 1 addition & 1 deletion tracer/test/Datadog.Trace.Security.Unit.Tests/WafTests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -147,7 +147,7 @@ private void ExecuteInternal(string address, object value, string flow, string r
args.Add(AddressesConstants.RequestMethod, "GET");
}

var initResult = Waf.Create(WafLibraryInvoker, string.Empty, string.Empty, setupWafSchemaExtraction: extractSchema, useUnsafeEncoder: newEncoder);
var initResult = Waf.Create(WafLibraryInvoker, string.Empty, string.Empty, useUnsafeEncoder: newEncoder);
using var waf = initResult.Waf;
waf.Should().NotBeNull();
using var context = waf.CreateContext();
Expand Down
Loading