Rework the listeners for AMS_FEATURES (#7158)

Rework the listeners for ASM_FEATURES

dd-java-agent/appsec/build.gradle

@@ -97,6 +97,9 @@ ext {
     'com.datadog.appsec.config.AppSecConfig.AppSecConfigV1',
     'com.datadog.appsec.config.AppSecConfig.AppSecConfigV2',
     'com.datadog.appsec.config.AppSecConfig.NumberJsonAdapter',
+    'com.datadog.appsec.config.AppSecFeatures',
+    'com.datadog.appsec.config.AppSecFeatures.Asm',
+    'com.datadog.appsec.config.AppSecFeatures.ApiSecurity',
     'com.datadog.appsec.event.ReplaceableEventProducerService',
   ]
   excludedClassesBranchCoverage = [

dd-java-agent/appsec/src/main/java/com/datadog/appsec/AppSecSystem.java

@@ -67,18 +67,21 @@ private static void doStart(SubscriptionService gw, SharedCommunicationObjects s
     EventDispatcher eventDispatcher = new EventDispatcher();
     REPLACEABLE_EVENT_PRODUCER.replaceEventProducerService(eventDispatcher);
 
+    ApiSecurityRequestSampler requestSampler = new ApiSecurityRequestSampler(config);
+
     ConfigurationPoller configurationPoller = sco.configurationPoller(config);
     // may throw and abort startup
     APP_SEC_CONFIG_SERVICE =
         new AppSecConfigServiceImpl(
-            config, configurationPoller, () -> reloadSubscriptions(REPLACEABLE_EVENT_PRODUCER));
+            config,
+            configurationPoller,
+            requestSampler,
+            () -> reloadSubscriptions(REPLACEABLE_EVENT_PRODUCER));
     APP_SEC_CONFIG_SERVICE.init();
 
     sco.createRemaining(config);
 
     RateLimiter rateLimiter = getRateLimiter(config, sco.monitoring);
-    ApiSecurityRequestSampler requestSampler =
-        new ApiSecurityRequestSampler(config, configurationPoller);
 
     GatewayBridge gatewayBridge =
         new GatewayBridge(

dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiSecurityRequestSampler.java

@@ -1,56 +1,34 @@
 package com.datadog.appsec.api.security;
 
-import static datadog.remoteconfig.tuf.RemoteConfigRequest.ClientInfo.CAPABILITY_ASM_API_SECURITY_SAMPLE_RATE;
-
-import com.datadog.appsec.config.AppSecFeaturesDeserializer;
-import datadog.remoteconfig.ConfigurationPoller;
-import datadog.remoteconfig.Product;
 import datadog.trace.api.Config;
 import java.util.concurrent.atomic.AtomicLong;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 
 public class ApiSecurityRequestSampler {
 
-  private static final Logger log = LoggerFactory.getLogger(ApiSecurityRequestSampler.class);
-
   private volatile int sampling;
   private final AtomicLong cumulativeCounter = new AtomicLong();
 
   public ApiSecurityRequestSampler(final Config config) {
     sampling = computeSamplingParameter(config.getApiSecurityRequestSampleRate());
   }
 
-  public ApiSecurityRequestSampler(final Config config, ConfigurationPoller configurationPoller) {
-    this(config);
-    if (configurationPoller == null) {
-      return;
+  /**
+   * Sets the new sampling parameter
+   *
+   * @return {@code true} if the value changed
+   */
+  public boolean setSampling(final float newSamplingFloat) {
+    int newSampling = computeSamplingParameter(newSamplingFloat);
+    if (newSampling != sampling) {
+      sampling = newSampling;
+      cumulativeCounter.set(0); // Reset current sampling counter
+      return true;
     }
+    return false;
+  }
 
-    configurationPoller.addListener(
-        Product.ASM_FEATURES,
-        "asm_api_security",
-        AppSecFeaturesDeserializer.INSTANCE,
-        (configKey, newConfig, pollingRateHinter) -> {
-          if (newConfig != null && newConfig.apiSecurity != null) {
-            Float newSamplingFloat = newConfig.apiSecurity.requestSampleRate;
-            if (newSamplingFloat != null) {
-              int newSampling = computeSamplingParameter(newSamplingFloat);
-              if (newSampling != sampling) {
-                sampling = newSampling;
-                cumulativeCounter.set(0); // Reset current sampling counter
-                if (sampling == 0) {
-                  log.info("Api Security is disabled via remote-config");
-                } else {
-                  log.info(
-                      "Api Security changed via remote-config. New sampling rate is {}% of all requests.",
-                      sampling);
-                }
-              }
-            }
-          }
-        });
-    configurationPoller.addCapabilities(CAPABILITY_ASM_API_SECURITY_SAMPLE_RATE);
+  public int getSampling() {
+    return sampling;
   }
 
   public boolean sampleRequest() {
@@ -69,7 +47,7 @@ static int computeSamplingParameter(final float pct) {
       return 100;
     }
     if (pct < 0) {
-      // We don't support disabling Api Security by setting it, so we set it to 100%.
+      // Api security can only be disabled by setting the sampling to zero, so we set it to 100%.
       // TODO: We probably want a warning here.
       return 100;
     }

dd-java-agent/appsec/src/main/java/com/datadog/appsec/config/AppSecConfigServiceImpl.java

@@ -12,8 +12,10 @@
 import static datadog.remoteconfig.tuf.RemoteConfigRequest.ClientInfo.CAPABILITY_ASM_REQUEST_BLOCKING;
 import static datadog.remoteconfig.tuf.RemoteConfigRequest.ClientInfo.CAPABILITY_ASM_TRUSTED_IPS;
 import static datadog.remoteconfig.tuf.RemoteConfigRequest.ClientInfo.CAPABILITY_ASM_USER_BLOCKING;
+import static datadog.trace.api.telemetry.LogCollector.SEND_TELEMETRY;
 
 import com.datadog.appsec.AppSecSystem;
+import com.datadog.appsec.api.security.ApiSecurityRequestSampler;
 import com.datadog.appsec.config.AppSecModuleConfigurer.SubconfigListener;
 import com.datadog.appsec.config.CurrentAppSecConfig.DirtyStatus;
 import com.datadog.appsec.util.AbortStartupException;
@@ -34,7 +36,6 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
-import javax.annotation.Nullable;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -60,27 +61,34 @@ public class AppSecConfigServiceImpl implements AppSecConfigService {
   private final Config tracerConfig;
   private final List<TraceSegmentPostProcessor> traceSegmentPostProcessors = new ArrayList<>();
   private final AppSecModuleConfigurer.Reconfiguration reconfiguration;
+  private final ApiSecurityRequestSampler apiSecurityRequestSampler;
 
   private final ConfigurationEndListener applyWAFChangesAsListener = this::applyWAFChanges;
 
   private boolean hasUserWafConfig;
 
   public AppSecConfigServiceImpl(
       Config tracerConfig,
-      @Nullable ConfigurationPoller configurationPoller,
+      ConfigurationPoller configurationPoller,
+      ApiSecurityRequestSampler apiSecurityRequestSampler,
       AppSecModuleConfigurer.Reconfiguration reconfig) {
     this.tracerConfig = tracerConfig;
     this.configurationPoller = configurationPoller;
     this.reconfiguration = reconfig;
+    this.apiSecurityRequestSampler = apiSecurityRequestSampler;
   }
 
   private void subscribeConfigurationPoller() {
     // see also close() method
     if (tracerConfig.getAppSecActivation() == ProductActivation.ENABLED_INACTIVE) {
       subscribeActivation();
     } else {
-      log.debug("Will not subscribe to ASM_FEATURES (AppSec explicitly enabled)");
+      log.debug(
+          "Will not subscribe to ASM_FEATURES['asm_features_activation'] (AppSec explicitly enabled)");
     }
+
+    subscribeApiSecurity();
+
     if (!hasUserWafConfig) {
       subscribeRulesAndData();
     } else {
@@ -164,8 +172,18 @@ private void subscribeActivation() {
           if (!initialized) {
             throw new IllegalStateException();
           }
-          final boolean newState =
-              newConfig != null && newConfig.asm != null && newConfig.asm.enabled;
+          final boolean newState;
+          if (newConfig == null) {
+            // configuration file was removed, restore the default
+            newState = tracerConfig.getAppSecActivation() == ProductActivation.FULLY_ENABLED;
+          } else if (newConfig.asm == null || newConfig.asm.enabled == null) {
+            // invalid payload from the backend, restore the default
+            log.debug(
+                SEND_TELEMETRY, "Invalid 'asm_features_activation' payload : {}", newConfig.asm);
+            newState = tracerConfig.getAppSecActivation() == ProductActivation.FULLY_ENABLED;
+          } else {
+            newState = newConfig.asm.enabled;
+          }
           if (AppSecSystem.isActive() != newState) {
             log.info("AppSec {} (runtime)", newState ? "enabled" : "disabled");
             AppSecSystem.setActive(newState);
@@ -179,6 +197,43 @@ private void subscribeActivation() {
     this.configurationPoller.addCapabilities(CAPABILITY_ASM_ACTIVATION);
   }
 
+  private void subscribeApiSecurity() {
+    this.configurationPoller.addListener(
+        Product.ASM_FEATURES,
+        "asm_api_security",
+        AppSecFeaturesDeserializer.INSTANCE,
+        (configKey, newConfig, hinter) -> {
+          if (!initialized) {
+            throw new IllegalStateException();
+          }
+          final float newSampling;
+          if (newConfig == null) {
+            // configuration file was removed on the backend, restore the default
+            newSampling = tracerConfig.getApiSecurityRequestSampleRate();
+          } else if (newConfig.apiSecurity == null
+              || newConfig.apiSecurity.requestSampleRate == null) {
+            // invalid payload from the backend, restore the default
+            log.debug(
+                SEND_TELEMETRY, "Invalid 'asm_api_security' payload : {}", newConfig.apiSecurity);
+            newSampling = tracerConfig.getApiSecurityRequestSampleRate();
+          } else {
+            newSampling = newConfig.apiSecurity.requestSampleRate;
+          }
+
+          if (apiSecurityRequestSampler.setSampling(newSampling)) {
+            int pct = apiSecurityRequestSampler.getSampling();
+            if (pct == 0) {
+              log.info("Api Security is disabled via remote-config");
+            } else {
+              log.info(
+                  "Api Security changed via remote-config. New sampling rate is {}% of all requests.",
+                  pct);
+            }
+          }
+        });
+    this.configurationPoller.addCapabilities(CAPABILITY_ASM_API_SECURITY_SAMPLE_RATE);
+  }
+
   private void distributeSubConfigurations(
       Map<String, Object> newConfig, AppSecModuleConfigurer.Reconfiguration reconfiguration) {
     for (Map.Entry<String, SubconfigListener> entry : subconfigListeners.entrySet()) {

dd-java-agent/appsec/src/main/java/com/datadog/appsec/config/AppSecFeatures.java

@@ -7,11 +7,26 @@ public class AppSecFeatures {
   public ApiSecurity apiSecurity;
 
   public static class Asm {
-    public boolean enabled;
+    public Boolean enabled;
+
+    @Override
+    public String toString() {
+      return "Asm{" + "enabled=" + enabled + '}';
+    }
   }
 
   public static class ApiSecurity {
     @com.squareup.moshi.Json(name = "request_sample_rate")
     public Float requestSampleRate;
+
+    @Override
+    public String toString() {
+      return "ApiSecurity{" + "requestSampleRate=" + requestSampleRate + '}';
+    }
+  }
+
+  @Override
+  public String toString() {
+    return "AppSecFeatures{" + "asm=" + asm + ", apiSecurity=" + apiSecurity + '}';
   }
 }

dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/ApiSecurityRequestSamplerTest.groovy

@@ -1,10 +1,5 @@
 package com.datadog.appsec.api.security
 
-import com.datadog.appsec.config.AppSecFeatures
-import com.datadog.appsec.config.AppSecFeaturesDeserializer
-import datadog.remoteconfig.ConfigurationChangesTypedListener
-import datadog.remoteconfig.ConfigurationPoller
-import datadog.remoteconfig.Product
 import datadog.trace.api.Config
 import datadog.trace.test.util.DDSpecification
 import spock.lang.Shared
@@ -46,31 +41,13 @@ class ApiSecurityRequestSamplerTest extends DDSpecification {
     -0.5                     | [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]    // Wrong sample rate - use 100%
   }
 
-  void 'update sample rate via remote-config'() {
+  void 'update sample rate'() {
     given:
-    ConfigurationPoller poller = Mock()
     def config = Spy(Config.get())
-    ConfigurationChangesTypedListener<AppSecFeatures> listener
-    AppSecFeatures newConfig = new AppSecFeatures().tap {
-      asm = new AppSecFeatures.Asm().tap {
-        enabled = true
-      }
-      apiSecurity = new AppSecFeatures.ApiSecurity().tap {
-        requestSampleRate = 0.2
-      }
-    }
-
-    when:
-    def sampler = new ApiSecurityRequestSampler(config, poller)
-
-    then:
-    1 * poller.addListener(Product.ASM_FEATURES, 'asm_api_security', AppSecFeaturesDeserializer.INSTANCE, _) >> {
-      listener = it[3] as ConfigurationChangesTypedListener<AppSecFeatures>
-    }
-    listener != null
+    def sampler = new ApiSecurityRequestSampler(config)
 
     when:
-    listener.accept(null, newConfig, null)
+    sampler.setSampling(0.2)
 
     then:
     sampler.sampling == 20