From df77b7b707fc92efd3d609c4dcf9f7f4c9a18822 Mon Sep 17 00:00:00 2001 From: lovesh-ap Date: Mon, 18 Nov 2024 12:07:21 +0530 Subject: [PATCH] Experimental changes for skipping complete trace if transitive API call is present in skip list --- .../instrumentator/dispatcher/Dispatcher.java | 5 ++ .../dispatcher/DispatcherPool.java | 4 ++ .../httpclient/RestRequestProcessor.java | 9 ++- .../javaagent/IntCodeControlCommand.java | 9 +++ .../schedulers/SchedulerHelper.java | 20 ++++++ .../utils/IastExclusionUtils.java | 66 +++++++++++++++++++ .../utils/RestrictionUtility.java | 21 ++++-- .../security/intcodeagent/utils/TTLMap.java | 54 +++++++++++++++ .../newrelic/api/agent/security/Agent.java | 1 + 9 files changed, 183 insertions(+), 6 deletions(-) create mode 100644 newrelic-security-agent/src/main/java/com/newrelic/agent/security/intcodeagent/utils/IastExclusionUtils.java create mode 100644 newrelic-security-agent/src/main/java/com/newrelic/agent/security/intcodeagent/utils/TTLMap.java diff --git a/newrelic-security-agent/src/main/java/com/newrelic/agent/security/instrumentator/dispatcher/Dispatcher.java b/newrelic-security-agent/src/main/java/com/newrelic/agent/security/instrumentator/dispatcher/Dispatcher.java index 0acc7370e..35835467e 100644 --- a/newrelic-security-agent/src/main/java/com/newrelic/agent/security/instrumentator/dispatcher/Dispatcher.java +++ b/newrelic-security-agent/src/main/java/com/newrelic/agent/security/instrumentator/dispatcher/Dispatcher.java @@ -8,6 +8,8 @@ import com.newrelic.agent.security.instrumentator.utils.CallbackUtils; import com.newrelic.agent.security.instrumentator.utils.INRSettingsKey; import com.newrelic.agent.security.intcodeagent.filelogging.FileLoggerThreadPool; +import com.newrelic.agent.security.intcodeagent.utils.IastExclusionUtils; +import com.newrelic.agent.security.intcodeagent.utils.RestrictionUtility; import com.newrelic.api.agent.security.Agent; import com.newrelic.api.agent.security.utils.logging.LogLevel; import com.newrelic.agent.security.intcodeagent.logging.DeployedApplication; @@ -180,6 +182,9 @@ public Object call() throws Exception { break; case HTTP_REQUEST: SSRFOperation ssrfOperationalBean = (SSRFOperation) operation; + if(RestrictionUtility.skippedApiDetected(AgentConfig.getInstance().getAgentMode().getSkipScan(), ((SSRFOperation) operation).getArg())) { + IastExclusionUtils.getInstance().registerSkippedTrace(NewRelic.getAgent().getTraceMetadata().getTraceId()); + } eventBean = prepareSSRFEvent(eventBean, ssrfOperationalBean); break; case XPATH: diff --git a/newrelic-security-agent/src/main/java/com/newrelic/agent/security/instrumentator/dispatcher/DispatcherPool.java b/newrelic-security-agent/src/main/java/com/newrelic/agent/security/instrumentator/dispatcher/DispatcherPool.java index 0f9db44a7..7297009b1 100644 --- a/newrelic-security-agent/src/main/java/com/newrelic/agent/security/instrumentator/dispatcher/DispatcherPool.java +++ b/newrelic-security-agent/src/main/java/com/newrelic/agent/security/instrumentator/dispatcher/DispatcherPool.java @@ -5,6 +5,7 @@ import com.newrelic.agent.security.intcodeagent.executor.CustomFutureTask; import com.newrelic.agent.security.intcodeagent.executor.CustomThreadPoolExecutor; import com.newrelic.agent.security.intcodeagent.filelogging.FileLoggerThreadPool; +import com.newrelic.agent.security.intcodeagent.utils.IastExclusionUtils; import com.newrelic.api.agent.security.utils.logging.LogLevel; import com.newrelic.agent.security.intcodeagent.logging.IAgentConstants; import com.newrelic.agent.security.intcodeagent.models.javaagent.EventStats; @@ -217,6 +218,9 @@ public void dispatchEvent(AbstractOperation operation, SecurityMetaData security TraceMetadata traceMetadata = NewRelic.getAgent().getTraceMetadata(); securityMetaData.addCustomAttribute(NR_APM_TRACE_ID, traceMetadata.getTraceId()); securityMetaData.addCustomAttribute(NR_APM_SPAN_ID, traceMetadata.getSpanId()); + + IastExclusionUtils.getInstance().addEncounteredTrace(traceMetadata.getTraceId(), operation.getApiID()); + this.executor.submit(new Dispatcher(operation, new SecurityMetaData(securityMetaData))); AgentInfo.getInstance().getJaHealthCheck().getEventStats().getDispatcher().incrementSubmitted(); diff --git a/newrelic-security-agent/src/main/java/com/newrelic/agent/security/instrumentator/httpclient/RestRequestProcessor.java b/newrelic-security-agent/src/main/java/com/newrelic/agent/security/instrumentator/httpclient/RestRequestProcessor.java index 618d33a0f..94bd988eb 100644 --- a/newrelic-security-agent/src/main/java/com/newrelic/agent/security/instrumentator/httpclient/RestRequestProcessor.java +++ b/newrelic-security-agent/src/main/java/com/newrelic/agent/security/instrumentator/httpclient/RestRequestProcessor.java @@ -7,6 +7,7 @@ import com.newrelic.agent.security.instrumentator.utils.CallbackUtils; import com.newrelic.agent.security.intcodeagent.filelogging.FileLoggerThreadPool; import com.newrelic.agent.security.intcodeagent.logging.IAgentConstants; +import com.newrelic.agent.security.intcodeagent.utils.IastExclusionUtils; import com.newrelic.api.agent.security.NewRelicSecurity; import com.newrelic.api.agent.security.schema.ServerConnectionConfiguration; import com.newrelic.api.agent.security.utils.logging.LogLevel; @@ -41,7 +42,7 @@ public class RestRequestProcessor implements Callable { public static final String JSON_PARSING_ERROR_WHILE_PROCESSING_FUZZING_REQUEST_S = "JSON parsing error while processing fuzzing request : %s"; private static final int MAX_REPETITION = 3; public static final String ENDPOINT_LOCALHOST_S = "%s://localhost:%s"; - private static final String IAST_REQUEST_HAS_NO_ARGUMENTS = "IAST request has no arguments : %s"; + private static final String IAST_REQUEST_HAS_NO_ARGUMENTS = "IAST request has incomplete arguments : %s"; public static final String AGENT_IS_NOT_ACTIVE = "Agent is not active"; public static final String WS_RECONNECTING = "Websocket reconnecting failing for control command id: %s"; private IntCodeControlCommand controlCommand; @@ -72,6 +73,12 @@ public Boolean call() throws InterruptedException { return false; } + //Skip if API ID is registered under skip list + if(IastExclusionUtils.getInstance().skipTraceApi(controlCommand.getApiId())) { + logger.log(LogLevel.FINER, String.format("Skip the control command %s, the record already exist under Skip list.", controlCommand.getId()), RestRequestProcessor.class.getSimpleName()); + return true; + } + FuzzRequestBean httpRequest = null; try { if (WSUtils.getInstance().isReconnecting()) { diff --git a/newrelic-security-agent/src/main/java/com/newrelic/agent/security/intcodeagent/models/javaagent/IntCodeControlCommand.java b/newrelic-security-agent/src/main/java/com/newrelic/agent/security/intcodeagent/models/javaagent/IntCodeControlCommand.java index 38be9b467..deec11279 100644 --- a/newrelic-security-agent/src/main/java/com/newrelic/agent/security/intcodeagent/models/javaagent/IntCodeControlCommand.java +++ b/newrelic-security-agent/src/main/java/com/newrelic/agent/security/intcodeagent/models/javaagent/IntCodeControlCommand.java @@ -32,6 +32,7 @@ public class IntCodeControlCommand { private Object data; private List arguments; private Map reflectedMetaData; + private String apiId; public IntCodeControlCommand() { } @@ -98,6 +99,14 @@ public Map getReflectedMetaData() { return reflectedMetaData; } + public String getApiId() { + return apiId; + } + + public void setApiId(String apiId) { + this.apiId = apiId; + } + public void setReflectedMetaData(Map reflectedMetaData) { this.reflectedMetaData = reflectedMetaData; } diff --git a/newrelic-security-agent/src/main/java/com/newrelic/agent/security/intcodeagent/schedulers/SchedulerHelper.java b/newrelic-security-agent/src/main/java/com/newrelic/agent/security/intcodeagent/schedulers/SchedulerHelper.java index 5e1ea25b9..4599c957e 100644 --- a/newrelic-security-agent/src/main/java/com/newrelic/agent/security/intcodeagent/schedulers/SchedulerHelper.java +++ b/newrelic-security-agent/src/main/java/com/newrelic/agent/security/intcodeagent/schedulers/SchedulerHelper.java @@ -84,6 +84,26 @@ public ScheduledFuture scheduleApplicationRuntimeErrorPosting(Runnable comman return future; } + public ScheduledFuture scheduleTTLMapCleanup(Runnable command, + long initialDelay, + long period, + TimeUnit unit, + String mapId){ + if(scheduledFutureMap.containsKey("ttl-map-cleanup-"+mapId)){ + throw new IllegalArgumentException("TTL map cleanup already scheduled for mapId: "+mapId); + } + ScheduledFuture future = commonExecutor.scheduleWithFixedDelay(command, initialDelay, period, unit); + scheduledFutureMap.put("ttl-map-cleanup-"+mapId, future); + return future; + } + + public void cancelTTLMapCleanup(String mapId){ + if(scheduledFutureMap.containsKey("ttl-map-cleanup-"+mapId)){ + ScheduledFuture currentFuture = scheduledFutureMap.get("ttl-map-cleanup-"+mapId); + currentFuture.cancel(false); + } + } + public ScheduledFuture scheduleDailyLogRollover(Runnable command) { if(LogFileHelper.isDailyRollover()) { diff --git a/newrelic-security-agent/src/main/java/com/newrelic/agent/security/intcodeagent/utils/IastExclusionUtils.java b/newrelic-security-agent/src/main/java/com/newrelic/agent/security/intcodeagent/utils/IastExclusionUtils.java new file mode 100644 index 000000000..6d20dfff8 --- /dev/null +++ b/newrelic-security-agent/src/main/java/com/newrelic/agent/security/intcodeagent/utils/IastExclusionUtils.java @@ -0,0 +1,66 @@ +package com.newrelic.agent.security.intcodeagent.utils; + +import com.newrelic.agent.security.intcodeagent.filelogging.FileLoggerThreadPool; + +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; + +public class IastExclusionUtils { + + private static final FileLoggerThreadPool logger = FileLoggerThreadPool.getInstance(); + + private final TTLMap> encounteredTraces = new TTLMap<>("encounteredTraces"); + + private final TTLMap skippedTraces = new TTLMap<>("skippedTraces"); + + private final Set skippedTraceApis = ConcurrentHashMap.newKeySet(); + + private IastExclusionUtils() { + } + + public boolean skippedTrace(String traceId) { + return skippedTraces.containsKey(traceId); + } + + public boolean skipTraceApi(String id) { + return skippedTraceApis.contains(id); + } + + private static final class InstanceHolder { + static final IastExclusionUtils instance = new IastExclusionUtils(); + } + + public static IastExclusionUtils getInstance() { + return InstanceHolder.instance; + } + + public void addEncounteredTrace(String traceId, String operationApiId) { + Set operationApiIds = encounteredTraces.get(traceId); + if (operationApiIds == null) { + operationApiIds = ConcurrentHashMap.newKeySet(); + encounteredTraces.put(traceId, operationApiIds); + } + operationApiIds.add(operationApiId); + updateSkippedTraceApis(traceId, operationApiId); + } + + public void registerSkippedTrace(String traceId) { + skippedTraces.put(traceId, true); + updateSkippedTraceApis(traceId); + } + + private void updateSkippedTraceApis(String traceId) { + Set operationApiIds = encounteredTraces.get(traceId); + if (operationApiIds != null) { + skippedTraceApis.addAll(operationApiIds); + } + } + + private void updateSkippedTraceApis(String traceId, String operationApiId) { + if (skippedTraces.containsKey(traceId)) { + skippedTraceApis.add(operationApiId); + } + } + + +} diff --git a/newrelic-security-agent/src/main/java/com/newrelic/agent/security/intcodeagent/utils/RestrictionUtility.java b/newrelic-security-agent/src/main/java/com/newrelic/agent/security/intcodeagent/utils/RestrictionUtility.java index f0a0c4518..44516dc62 100644 --- a/newrelic-security-agent/src/main/java/com/newrelic/agent/security/intcodeagent/utils/RestrictionUtility.java +++ b/newrelic-security-agent/src/main/java/com/newrelic/agent/security/intcodeagent/utils/RestrictionUtility.java @@ -6,6 +6,7 @@ import com.fasterxml.jackson.databind.node.ArrayNode; import com.newrelic.agent.security.intcodeagent.exceptions.RestrictionModeException; import com.newrelic.agent.security.intcodeagent.filelogging.FileLoggerThreadPool; +import com.newrelic.api.agent.NewRelic; import com.newrelic.api.agent.security.instrumentation.helpers.ServletHelper; import com.newrelic.api.agent.security.schema.HttpRequest; import com.newrelic.api.agent.security.schema.policy.RestrictionCriteria; @@ -42,19 +43,29 @@ public class RestrictionUtility { private static final FileLoggerThreadPool logger = FileLoggerThreadPool.getInstance(); public static boolean skippedApiDetected(SkipScan skipScan, HttpRequest httpRequest) { - if (skipScan == null) { - return false; - } + if (httpRequest == null) { return false; } - if(skipScan.getApiRoutes().isEmpty()) { + return skippedApiDetected(skipScan, httpRequest.getUrl()); + } + + public static boolean skippedApiDetected(SkipScan skipScan, String url) { + + String uri = StringUtils.substringBefore(url, SEPARATOR_CHARS_QUESTION_MARK); + + if (skipScan == null || skipScan.getApiRoutes().isEmpty()) { return false; } + if(IastExclusionUtils.getInstance().skippedTrace(NewRelic.getAgent().getTraceMetadata().getTraceId())) { + logger.log(LogLevel.FINER, String.format("Skipping the scan for url %s due to traceId %s", uri, NewRelic.getAgent().getTraceMetadata().getTraceId()), RestrictionUtility.class.getName()); + return true; + } + for (Pattern pattern : skipScan.getApiRoutes()) { - if (pattern.matcher(httpRequest.getUrl()).matches()) { + if (pattern.matcher(uri).matches()) { return true; } } diff --git a/newrelic-security-agent/src/main/java/com/newrelic/agent/security/intcodeagent/utils/TTLMap.java b/newrelic-security-agent/src/main/java/com/newrelic/agent/security/intcodeagent/utils/TTLMap.java new file mode 100644 index 000000000..3cd3bb70f --- /dev/null +++ b/newrelic-security-agent/src/main/java/com/newrelic/agent/security/intcodeagent/utils/TTLMap.java @@ -0,0 +1,54 @@ +package com.newrelic.agent.security.intcodeagent.utils; + +import com.newrelic.agent.security.intcodeagent.schedulers.SchedulerHelper; + +import java.util.concurrent.ConcurrentHashMap; + +public class TTLMap { + + private final ConcurrentHashMap map = new ConcurrentHashMap<>(); + private final ConcurrentHashMap timestamps = new ConcurrentHashMap<>(); + private final long TTL; + private final String id; + + public TTLMap(String id) { + this(id, 300_000);// 5 minutes in milliseconds + } + + public TTLMap(String id, long ttl) { + this.id = id; + TTL = ttl; + } + + public void put(K key, V value) { + map.put(key, value); + timestamps.put(key, System.currentTimeMillis()); + } + + public V get(K key) { + return map.get(key); + } + + public void remove(K key) { + map.remove(key); + timestamps.remove(key); + } + + private void removeExpiredEntries() { + long now = System.currentTimeMillis(); + for (K key : timestamps.keySet()) { + if (now - timestamps.get(key) >= TTL) { + map.remove(key); + timestamps.remove(key); + } + } + } + + public void shutdown() { + SchedulerHelper.getInstance().cancelTTLMapCleanup(this.id); + } + + public boolean containsKey(K traceId) { + return map.containsKey(traceId); + } +} diff --git a/newrelic-security-agent/src/main/java/com/newrelic/api/agent/security/Agent.java b/newrelic-security-agent/src/main/java/com/newrelic/api/agent/security/Agent.java index 3a99caef4..302fe4f11 100644 --- a/newrelic-security-agent/src/main/java/com/newrelic/api/agent/security/Agent.java +++ b/newrelic-security-agent/src/main/java/com/newrelic/api/agent/security/Agent.java @@ -376,6 +376,7 @@ public void registerOperation(AbstractOperation operation) { SecurityMetaData securityMetaData = NewRelicSecurity.getAgent().getSecurityMetaData(); if(RestrictionUtility.skippedApiDetected(AgentConfig.getInstance().getAgentMode().getSkipScan(), securityMetaData.getRequest())){ + IastExclusionUtils.getInstance().registerSkippedTrace(NewRelic.getAgent().getTraceMetadata().getTraceId()); logger.log(LogLevel.FINER, String.format(SKIPPING_THE_API_S_AS_IT_IS_PART_OF_THE_SKIP_SCAN_LIST, securityMetaData.getRequest().getUrl()), Agent.class.getName()); return; }