Skip to content

Commit

Permalink
Fix body interception in mule
Browse files Browse the repository at this point in the history
  • Loading branch information
@IshikaDawda
IshikaDawda committed Oct 7, 2024
1 parent 4bb0270 commit 5c60b32
Show file tree
Hide file tree
Showing 12 changed files with 52 additions and 46 deletions.
23 changes: 11 additions & 12 deletions ...mule-3.6/src/main/java/com/newrelic/agent/security/instrumentation/mule36/MuleHelper.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,8 +27,8 @@ public class MuleHelper {
public static final String MULE_36 = "MULE-3.6";
private static final String MULE_LOCK_CUSTOM_ATTRIB_NAME = "MULE_LOCK-";
public static final String MULE_SERVER_PORT_ATTRIB_NAME = "MULE_SERVER_PORT";
private static final String RESPONSE_OUTPUTSTREAM_HASH = "RESPONSE_OUTPUTSTREAM_HASH";
private static final String REQUEST_INPUTSTREAM_HASH = "REQUEST_INPUTSTREAM_HASH";
public static final String RESPONSE_OUTPUTSTREAM_HASH = "RESPONSE_OUTPUTSTREAM_HASH";
public static final String REQUEST_INPUTSTREAM_HASH = "REQUEST_INPUTSTREAM_HASH";
public static final String TRANSFORM_METHOD = "transform";
public static final String HANDLE_REQUEST_METHOD = "handleRequest";
private static final String EMPTY = "";
Expand Down Expand Up @@ -135,26 +135,25 @@ public static Map<Integer, String> getHandlerMap() {
return handlerMap;
}

public static void registerOutputStreamHashIfNeeded(int outputStreamHash){
public static void registerStreamHashIfNeeded(int streamHash, String key){
try {
Set<Integer> hashSet = NewRelicSecurity.getAgent().getSecurityMetaData().getCustomAttribute(RESPONSE_OUTPUTSTREAM_HASH, Set.class);
Set<Integer> hashSet = NewRelicSecurity.getAgent().getSecurityMetaData().getCustomAttribute(key, Set.class);
if (hashSet == null) {
hashSet = new HashSet<>();
NewRelicSecurity.getAgent().getSecurityMetaData().addCustomAttribute(RESPONSE_OUTPUTSTREAM_HASH, hashSet);
NewRelicSecurity.getAgent().getSecurityMetaData().addCustomAttribute(key, hashSet);
}
hashSet.add(outputStreamHash);
hashSet.add(streamHash);
} catch (Throwable ignored) {}
}

public static void registerInputStreamHashIfNeeded(int inputStreamHash){
public static boolean preprocessStream(int streamHash, String key){
try {
Set<Integer> hashSet = NewRelicSecurity.getAgent().getSecurityMetaData().getCustomAttribute(REQUEST_INPUTSTREAM_HASH, Set.class);
if(hashSet == null){
hashSet = new HashSet<>();
NewRelicSecurity.getAgent().getSecurityMetaData().addCustomAttribute(REQUEST_INPUTSTREAM_HASH, hashSet);
Set<Integer> hashSet = NewRelicSecurity.getAgent().getSecurityMetaData().getCustomAttribute(key, Set.class);
if(hashSet != null && hashSet.contains(streamHash)){
return true;
}
hashSet.add(inputStreamHash);
} catch (Throwable ignored) {}
return false;
}

public static void processHttpResponseHeaders(com.newrelic.api.agent.security.schema.HttpResponse securityResponse, HttpResponse response){
Expand Down
11 changes: 7 additions & 4 deletions ...c/main/java/org/mule/module/http/internal/domain/ByteArrayHttpEntity_Instrumentation.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import com.newrelic.agent.security.instrumentation.mule36.MuleHelper;
import com.newrelic.api.agent.security.NewRelicSecurity;
import com.newrelic.api.agent.security.instrumentation.helpers.GenericHelper;
import com.newrelic.api.agent.security.schema.SecurityMetaData;
import com.newrelic.api.agent.security.utils.logging.LogLevel;
import com.newrelic.api.agent.weaver.Weave;
import com.newrelic.api.agent.weaver.Weaver;
Expand Down Expand Up @@ -33,10 +34,12 @@ private void extractResponseBody(byte[] content) throws IOException {
encoding = Charsets.UTF_8.name();
}
String body = IOUtils.toString(content, encoding);
if (Objects.equals(NewRelicSecurity.getAgent().getSecurityMetaData().getCustomAttribute(MuleHelper.getNrSecCustomAttribName(MuleHelper.RESPONSE_ENTITY_STREAM), Integer.class), this.hashCode())) {
NewRelicSecurity.getAgent().getSecurityMetaData().getResponse().getResponseBody().append(body);
} else if (Objects.equals(NewRelicSecurity.getAgent().getSecurityMetaData().getCustomAttribute(MuleHelper.getNrSecCustomAttribName(MuleHelper.REQUEST_ENTITY_STREAM), Integer.class), this.hashCode())) {
NewRelicSecurity.getAgent().getSecurityMetaData().getRequest().getBody().append(body);

SecurityMetaData securityMetaData = NewRelicSecurity.getAgent().getSecurityMetaData();
if (MuleHelper.preprocessStream(this.hashCode(), MuleHelper.RESPONSE_ENTITY_STREAM)) {
securityMetaData.getResponse().getResponseBody().append(body);
} else if (MuleHelper.preprocessStream(this.hashCode(), MuleHelper.REQUEST_ENTITY_STREAM)) {
securityMetaData.getRequest().getBody().append(body);
}
}
}
Expand Down
12 changes: 8 additions & 4 deletions ...main/java/org/mule/module/http/internal/domain/InputStreamHttpEntity_Instrumentation.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import com.newrelic.agent.security.instrumentation.mule36.MuleHelper;
import com.newrelic.api.agent.security.NewRelicSecurity;
import com.newrelic.api.agent.security.instrumentation.helpers.GenericHelper;
import com.newrelic.api.agent.security.schema.SecurityMetaData;
import com.newrelic.api.agent.security.utils.logging.LogLevel;
import com.newrelic.api.agent.weaver.Weave;
import com.newrelic.api.agent.weaver.Weaver;
Expand All @@ -27,10 +28,13 @@ public InputStream getInputStream() {
private void extractResponseBody(InputStream stream) {
if (NewRelicSecurity.isHookProcessingActive() && stream != null) {
// check if it is an input or output stream
if (Objects.equals(NewRelicSecurity.getAgent().getSecurityMetaData().getCustomAttribute(MuleHelper.getNrSecCustomAttribName(MuleHelper.RESPONSE_ENTITY_STREAM), Integer.class), this.hashCode())) {
MuleHelper.registerOutputStreamHashIfNeeded(stream.hashCode());
} else if (Objects.equals(NewRelicSecurity.getAgent().getSecurityMetaData().getCustomAttribute(MuleHelper.getNrSecCustomAttribName(MuleHelper.REQUEST_ENTITY_STREAM), Integer.class), this.hashCode())) {
MuleHelper.registerInputStreamHashIfNeeded(stream.hashCode());
// outputBody stream
if (MuleHelper.preprocessStream(this.hashCode(), MuleHelper.RESPONSE_ENTITY_STREAM)) {
MuleHelper.registerStreamHashIfNeeded(stream.hashCode(), MuleHelper.RESPONSE_OUTPUTSTREAM_HASH);
}
// inputBody stream
else if (MuleHelper.preprocessStream(this.hashCode(), MuleHelper.REQUEST_ENTITY_STREAM)) {
MuleHelper.registerStreamHashIfNeeded(stream.hashCode(), MuleHelper.REQUEST_INPUTSTREAM_HASH);
}
}
}
Expand Down
2 changes: 1 addition & 1 deletion ...va/org/mule/module/http/internal/domain/response/HttpResponseBuilder_Instrumentation.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ private void postProcessSecurityHook(HttpResponse response) {
return;
}
if (body != null) {
NewRelicSecurity.getAgent().getSecurityMetaData().addCustomAttribute(MuleHelper.getNrSecCustomAttribName(MuleHelper.getNrSecCustomAttribName(MuleHelper.RESPONSE_ENTITY_STREAM)), body.hashCode());
MuleHelper.registerStreamHashIfNeeded(body.hashCode(), MuleHelper.RESPONSE_ENTITY_STREAM);
}
com.newrelic.api.agent.security.schema.HttpResponse securityResponse = NewRelicSecurity.getAgent().getSecurityMetaData().getResponse();

Expand Down
2 changes: 1 addition & 1 deletion ...n/java/org/mule/module/http/internal/listener/HttpRequestToMuleEvent_Instrumentation.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,7 @@ private static void preprocessSecurityHook(HttpRequestContext requestContext) {

HttpRequest httpRequest = requestContext.getRequest();
if (httpRequest.getEntity() != null) {
securityMetaData.addCustomAttribute(MuleHelper.getNrSecCustomAttribName(MuleHelper.REQUEST_ENTITY_STREAM), httpRequest.getEntity().hashCode());
MuleHelper.registerStreamHashIfNeeded(httpRequest.getEntity().hashCode(), MuleHelper.REQUEST_ENTITY_STREAM);
}
securityRequest.setMethod(httpRequest.getMethod());
securityRequest.setClientIP(requestContext.getRemoteHostAddress().toString());
Expand Down
2 changes: 1 addition & 1 deletion ...ain/java/org/mule/module/http/internal/listener/async/RequestHandler_Instrumentation.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,7 @@ private void preprocessSecurityHook(HttpRequestContext requestContext) {
HttpRequest httpRequest = requestContext.getRequest();

if (httpRequest.getEntity() != null) {
securityMetaData.addCustomAttribute(MuleHelper.getNrSecCustomAttribName(MuleHelper.REQUEST_ENTITY_STREAM), httpRequest.getEntity().hashCode());
MuleHelper.registerStreamHashIfNeeded(httpRequest.getEntity().hashCode(), MuleHelper.REQUEST_ENTITY_STREAM);
}
securityRequest.setMethod(httpRequest.getMethod());
securityRequest.setClientIP(requestContext.getRemoteHostAddress().toString());
Expand Down
23 changes: 11 additions & 12 deletions ...mule-3.7/src/main/java/com/newrelic/agent/security/instrumentation/mule37/MuleHelper.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,13 +27,13 @@ public class MuleHelper {
public static final String MULE_37 = "MULE-3.7";
private static final String MULE_LOCK_CUSTOM_ATTRIB_NAME = "MULE_LOCK-";
public static final String MULE_SERVER_PORT_ATTRIB_NAME = "MULE_SERVER_PORT";
private static final String RESPONSE_OUTPUTSTREAM_HASH = "RESPONSE_OUTPUTSTREAM_HASH";
public static final String RESPONSE_OUTPUTSTREAM_HASH = "RESPONSE_OUTPUTSTREAM_HASH";
public static final String TRANSFORM_METHOD = "transform";
public static final String HANDLE_REQUEST_METHOD = "handleRequest";
private static final String EMPTY = "";
public static final String LIBRARY_NAME = "MULE-SERVER";
private static final Map<Integer, String> handlerMap = new HashMap<>();
private static final String REQUEST_INPUTSTREAM_HASH = "REQUEST_INPUTSTREAM_HASH";
public static final String REQUEST_INPUTSTREAM_HASH = "REQUEST_INPUTSTREAM_HASH";
public static final String RESPONSE_ENTITY_STREAM = "RESPONSE_ENTITY_STREAM";
public static final String REQUEST_ENTITY_STREAM = "REQUEST_ENTITY_STREAM";
public static final String MULE_ENCODING = "MULE_ENCODING";
Expand Down Expand Up @@ -134,26 +134,25 @@ public static Map<Integer, String> getHandlerMap() {
return handlerMap;
}

public static void registerOutputStreamHashIfNeeded(int outputStreamHash){
public static void registerStreamHashIfNeeded(int streamHash, String key){
try {
Set<Integer> hashSet = NewRelicSecurity.getAgent().getSecurityMetaData().getCustomAttribute(RESPONSE_OUTPUTSTREAM_HASH, Set.class);
Set<Integer> hashSet = NewRelicSecurity.getAgent().getSecurityMetaData().getCustomAttribute(key, Set.class);
if (hashSet == null) {
hashSet = new HashSet<>();
NewRelicSecurity.getAgent().getSecurityMetaData().addCustomAttribute(RESPONSE_OUTPUTSTREAM_HASH, hashSet);
NewRelicSecurity.getAgent().getSecurityMetaData().addCustomAttribute(key, hashSet);
}
hashSet.add(outputStreamHash);
hashSet.add(streamHash);
} catch (Throwable ignored) {}
}

public static void registerInputStreamHashIfNeeded(int inputStreamHash){
public static boolean preprocessStream(int streamHash, String key){
try {
Set<Integer> hashSet = NewRelicSecurity.getAgent().getSecurityMetaData().getCustomAttribute(REQUEST_INPUTSTREAM_HASH, Set.class);
if(hashSet == null){
hashSet = new HashSet<>();
NewRelicSecurity.getAgent().getSecurityMetaData().addCustomAttribute(REQUEST_INPUTSTREAM_HASH, hashSet);
Set<Integer> hashSet = NewRelicSecurity.getAgent().getSecurityMetaData().getCustomAttribute(key, Set.class);
if(hashSet != null && hashSet.contains(streamHash)){
return true;
}
hashSet.add(inputStreamHash);
} catch (Throwable ignored) {}
return false;
}

public static void processHttpResponseHeaders(com.newrelic.api.agent.security.schema.HttpResponse securityResponse, HttpResponse response){
Expand Down
5 changes: 2 additions & 3 deletions ...c/main/java/org/mule/module/http/internal/domain/ByteArrayHttpEntity_Instrumentation.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,6 @@
import org.mule.util.IOUtils;

import java.io.IOException;
import java.util.Objects;

@Weave(originalName = "org.mule.module.http.internal.domain.ByteArrayHttpEntity")
public class ByteArrayHttpEntity_Instrumentation {
Expand All @@ -36,9 +35,9 @@ private void extractResponseBody(byte[] content) throws IOException {
String body = IOUtils.toString(content, encoding);

SecurityMetaData securityMetaData = NewRelicSecurity.getAgent().getSecurityMetaData();
if (Objects.equals(securityMetaData.getCustomAttribute(MuleHelper.getNrSecCustomAttribName(MuleHelper.RESPONSE_ENTITY_STREAM), Integer.class), this.hashCode())) {
if (MuleHelper.preprocessStream(this.hashCode(), MuleHelper.RESPONSE_ENTITY_STREAM)) {
securityMetaData.getResponse().getResponseBody().append(body);
} else if (Objects.equals(securityMetaData.getCustomAttribute(MuleHelper.getNrSecCustomAttribName(MuleHelper.REQUEST_ENTITY_STREAM), Integer.class), this.hashCode())) {
} else if (MuleHelper.preprocessStream(this.hashCode(), MuleHelper.REQUEST_ENTITY_STREAM)) {
securityMetaData.getRequest().getBody().append(body);
}
}
Expand Down
12 changes: 7 additions & 5 deletions ...main/java/org/mule/module/http/internal/domain/InputStreamHttpEntity_Instrumentation.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,6 @@
import com.newrelic.api.agent.weaver.Weaver;

import java.io.InputStream;
import java.util.Objects;

@Weave(originalName = "org.mule.module.http.internal.domain.InputStreamHttpEntity")
public class InputStreamHttpEntity_Instrumentation {
Expand All @@ -29,10 +28,13 @@ private void extractResponseBody(InputStream stream) {
if (NewRelicSecurity.isHookProcessingActive() && stream != null) {
// check if it is an input or output stream
SecurityMetaData securityMetaData = NewRelicSecurity.getAgent().getSecurityMetaData();
if (Objects.equals(securityMetaData.getCustomAttribute(MuleHelper.getNrSecCustomAttribName(MuleHelper.RESPONSE_ENTITY_STREAM), Integer.class), this.hashCode())) {
MuleHelper.registerOutputStreamHashIfNeeded(stream.hashCode());
} else if (Objects.equals(securityMetaData.getCustomAttribute(MuleHelper.getNrSecCustomAttribName(MuleHelper.REQUEST_ENTITY_STREAM), Integer.class), this.hashCode())) {
MuleHelper.registerInputStreamHashIfNeeded(stream.hashCode());
// outputBody stream
if (MuleHelper.preprocessStream(this.hashCode(), MuleHelper.RESPONSE_ENTITY_STREAM)) {
MuleHelper.registerStreamHashIfNeeded(stream.hashCode(), MuleHelper.RESPONSE_OUTPUTSTREAM_HASH);
}
// inputBody stream
else if (MuleHelper.preprocessStream(this.hashCode(), MuleHelper.REQUEST_ENTITY_STREAM)) {
MuleHelper.registerStreamHashIfNeeded(stream.hashCode(), MuleHelper.REQUEST_INPUTSTREAM_HASH);
}
}
}
Expand Down
2 changes: 1 addition & 1 deletion ...va/org/mule/module/http/internal/domain/response/HttpResponseBuilder_Instrumentation.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ private void postProcessSecurityHook(HttpResponse response) {
return;
}
if (body != null) {
NewRelicSecurity.getAgent().getSecurityMetaData().addCustomAttribute(MuleHelper.getNrSecCustomAttribName(MuleHelper.getNrSecCustomAttribName(MuleHelper.RESPONSE_ENTITY_STREAM)), body.hashCode());
MuleHelper.registerStreamHashIfNeeded(body.hashCode(), MuleHelper.RESPONSE_ENTITY_STREAM);
}
com.newrelic.api.agent.security.schema.HttpResponse securityResponse = NewRelicSecurity.getAgent().getSecurityMetaData().getResponse();
MuleHelper.processHttpResponseHeaders(securityResponse, response);
Expand Down
2 changes: 1 addition & 1 deletion ...n/java/org/mule/module/http/internal/listener/HttpRequestToMuleEvent_Instrumentation.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,7 @@ private static void preprocessSecurityHook(HttpRequestContext requestContext) {

HttpRequest httpRequest = requestContext.getRequest();
if (httpRequest.getEntity() != null) {
securityMetaData.addCustomAttribute(MuleHelper.getNrSecCustomAttribName(MuleHelper.REQUEST_ENTITY_STREAM), httpRequest.getEntity().hashCode());
MuleHelper.registerStreamHashIfNeeded(httpRequest.getEntity().hashCode(), MuleHelper.REQUEST_ENTITY_STREAM);
}
securityRequest.setMethod(httpRequest.getMethod());
securityRequest.setClientIP(requestContext.getClientConnection().getRemoteHostAddress().toString());
Expand Down
2 changes: 1 addition & 1 deletion ...ain/java/org/mule/module/http/internal/listener/async/RequestHandler_Instrumentation.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,7 @@ private void preprocessSecurityHook(HttpRequestContext requestContext) {

HttpRequest httpRequest = requestContext.getRequest();
if (httpRequest.getEntity() != null) {
securityMetaData.addCustomAttribute(MuleHelper.getNrSecCustomAttribName(MuleHelper.REQUEST_ENTITY_STREAM), httpRequest.getEntity().hashCode());
MuleHelper.registerStreamHashIfNeeded(httpRequest.getEntity().hashCode(), MuleHelper.REQUEST_ENTITY_STREAM);
}
securityRequest.setMethod(httpRequest.getMethod());
securityRequest.setClientIP(requestContext.getClientConnection().getRemoteHostAddress().toString());
Expand Down

0 comments on commit 5c60b32

Please sign in to comment.