Skip to content

Commit

Permalink
Experimental changes for skipping complete trace if transitive API ca…
Browse files Browse the repository at this point in the history
…ll is present in skip list
  • Loading branch information
lovesh-ap committed Nov 18, 2024
1 parent 4b5b84c commit df77b7b
Show file tree
Hide file tree
Showing 9 changed files with 183 additions and 6 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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();

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -41,7 +42,7 @@ public class RestRequestProcessor implements Callable<Boolean> {
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;
Expand Down Expand Up @@ -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()) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ public class IntCodeControlCommand {
private Object data;
private List<String> arguments;
private Map<String, String> reflectedMetaData;
private String apiId;

public IntCodeControlCommand() {
}
Expand Down Expand Up @@ -98,6 +99,14 @@ public Map<String, String> getReflectedMetaData() {
return reflectedMetaData;
}

public String getApiId() {
return apiId;
}

public void setApiId(String apiId) {
this.apiId = apiId;
}

public void setReflectedMetaData(Map<String, String> reflectedMetaData) {
this.reflectedMetaData = reflectedMetaData;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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()) {
Expand Down
Original file line number Diff line number Diff line change
@@ -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<String, Set<String>> encounteredTraces = new TTLMap<>("encounteredTraces");

private final TTLMap<String, Boolean> skippedTraces = new TTLMap<>("skippedTraces");

private final Set<String> 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<String> 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<String> operationApiIds = encounteredTraces.get(traceId);
if (operationApiIds != null) {
skippedTraceApis.addAll(operationApiIds);
}
}

private void updateSkippedTraceApis(String traceId, String operationApiId) {
if (skippedTraces.containsKey(traceId)) {
skippedTraceApis.add(operationApiId);
}
}


}
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
}
}
Expand Down
Original file line number Diff line number Diff line change
@@ -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<K, V> {

private final ConcurrentHashMap<K, V> map = new ConcurrentHashMap<>();
private final ConcurrentHashMap<K, Long> 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);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand Down

0 comments on commit df77b7b

Please sign in to comment.