Clean iast state when connection is closed

newrelic · Dec 3, 2024 · 4a6e932 · 4a6e932
1 parent 2d034bb
commit 4a6e932
Showing 1 changed file with 13 additions and 0 deletions.
diff --git a/.../main/java/com/newrelic/agent/security/intcodeagent/apache/httpclient/SecurityClient.java b/.../main/java/com/newrelic/agent/security/intcodeagent/apache/httpclient/SecurityClient.java
@@ -2,9 +2,12 @@
 
 import com.newrelic.agent.security.AgentConfig;
 import com.newrelic.agent.security.AgentInfo;
+import com.newrelic.agent.security.instrumentator.dispatcher.DispatcherPool;
+import com.newrelic.agent.security.instrumentator.httpclient.RestRequestThreadPool;
 import com.newrelic.agent.security.instrumentator.utils.INRSettingsKey;
 import com.newrelic.agent.security.intcodeagent.websocket.EventSendPool;
 import com.newrelic.agent.security.intcodeagent.websocket.EventSender;
+import com.newrelic.api.agent.security.instrumentation.helpers.GrpcClientRequestReplayHelper;
 import com.newrelic.api.agent.security.utils.ConnectionException;
 import com.newrelic.agent.security.intcodeagent.filelogging.FileLoggerThreadPool;
 import com.newrelic.agent.security.util.IUtilConstants;
@@ -147,9 +150,19 @@ public ReadResult send(JSONStreamAware message, String api) throws ApacheHttpExc
 
     @Override
     public void close(String message) {
+        cleanIASTState();
         httpClient.shutdown();
     }
 
+    private static void cleanIASTState() {
+        RestRequestThreadPool.getInstance().resetIASTProcessing();
+        GrpcClientRequestReplayHelper.getInstance().resetIASTProcessing();
+        RestRequestThreadPool.getInstance().getRejectedIds().clear();
+        GrpcClientRequestReplayHelper.getInstance().getRejectedIds().clear();
+        DispatcherPool.getInstance().reset();
+        EventSendPool.getInstance().reset();
+    }
+
     @Override
     public void ping() {
         try {