Ensure cached subscriptions are cleared on reconfiguration via RC

Author: manuel-alvarez-alvarez

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

@@ -39,6 +39,7 @@ public class AppSecSystem {
   private static final Map<AppSecModule, String> STARTED_MODULES_INFO = new HashMap<>();
   private static AppSecConfigServiceImpl APP_SEC_CONFIG_SERVICE;
   private static ReplaceableEventProducerService REPLACEABLE_EVENT_PRODUCER; // testing
+  private static Runnable STOP_SUBSCRIPTION_SERVICE;
   private static Runnable RESET_SUBSCRIPTION_SERVICE;
 
   public static void start(SubscriptionService gw, SharedCommunicationObjects sco) {
@@ -90,7 +91,8 @@ private static void doStart(SubscriptionService gw, SharedCommunicationObjects s
     loadModules(eventDispatcher, sco.monitoring);
 
     gatewayBridge.init();
-    RESET_SUBSCRIPTION_SERVICE = gatewayBridge::stop;
+    STOP_SUBSCRIPTION_SERVICE = gatewayBridge::stop;
+    RESET_SUBSCRIPTION_SERVICE = gatewayBridge::reset;
 
     setActive(appSecEnabledConfig == ProductActivation.FULLY_ENABLED);
 
@@ -127,7 +129,8 @@ public static void stop() {
       return;
     }
     REPLACEABLE_EVENT_PRODUCER = null;
-    RESET_SUBSCRIPTION_SERVICE.run();
+    STOP_SUBSCRIPTION_SERVICE.run();
+    STOP_SUBSCRIPTION_SERVICE = null;
     RESET_SUBSCRIPTION_SERVICE = null;
     Blocking.setBlockingService(BlockingService.NOOP);
 
@@ -176,6 +179,10 @@ private static void reloadSubscriptions(
     newEd.subscribeDataAvailable(dataSubscriptionSet);
 
     replaceableEventProducerService.replaceEventProducerService(newEd);
+
+    if (RESET_SUBSCRIPTION_SERVICE != null) {
+      RESET_SUBSCRIPTION_SERVICE.run();
+    }
   }
 
   public static boolean isStarted() {

dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java

@@ -173,6 +173,30 @@ public void init() {
     }
   }
 
+  /**
+   * This method clears all the cached subscriptions, should be used everytime the configuration
+   * changes and new addresses might appear or disappear from the config.
+   */
+  public void reset() {
+    initialReqDataSubInfo = null;
+    rawRequestBodySubInfo = null;
+    requestBodySubInfo = null;
+    pathParamsSubInfo = null;
+    respDataSubInfo = null;
+    grpcServerMethodSubInfo = null;
+    grpcServerRequestMsgSubInfo = null;
+    graphqlServerRequestMsgSubInfo = null;
+    requestEndSubInfo = null;
+    dbSqlQuerySubInfo = null;
+    ioNetUrlSubInfo = null;
+    ioFileSubInfo = null;
+    sessionIdSubInfo = null;
+    userIdSubInfo = null;
+    loginEventSubInfo.clear();
+    execCmdSubInfo = null;
+    shellCmdSubInfo = null;
+  }
+
   private Flow<Void> onUser(
       final RequestContext ctx_, final UserIdCollectionMode mode, final String originalUser) {
     if (mode == DISABLED) {

dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy

@@ -73,6 +73,10 @@ class GatewayBridgeSpecification extends DDSpecification {
     i
   }()
 
+  EventProducerService.DataSubscriberInfo emptyDsInfo = Stub() {
+    isEmpty() >> true
+  }
+
   TraceSegmentPostProcessor pp = Mock()
   GatewayBridge bridge = new GatewayBridge(ig, eventDispatcher, null, [pp])
 
@@ -1280,4 +1284,28 @@ class GatewayBridgeSpecification extends DDSpecification {
 
     0 * eventDispatcher.publishDataEvent
   }
+
+  void 'test configuration updates should reset cached subscriptions'() {
+    when:
+    requestSessionCB.apply(ctx, UUID.randomUUID().toString())
+
+    then:
+    1 * eventDispatcher.getDataSubscribers(KnownAddresses.SESSION_ID) >> emptyDsInfo
+    0 * eventDispatcher.publishDataEvent
+
+    when:
+    requestSessionCB.apply(ctx, UUID.randomUUID().toString())
+
+    then:
+    0 * eventDispatcher.getDataSubscribers
+    0 * eventDispatcher.publishDataEvent
+
+    when:
+    bridge.reset()
+    requestSessionCB.apply(ctx, UUID.randomUUID().toString())
+
+    then:
+    1 * eventDispatcher.getDataSubscribers(KnownAddresses.SESSION_ID) >> nonEmptyDsInfo
+    1 * eventDispatcher.publishDataEvent(_, _, _, _)
+  }
 }