Instrumentation and unit test cases for httpclient5 (async)

instrumentation-security-test/src/main/java/com/newrelic/agent/security/introspec/internal/HttpServerRule.java

@@ -30,6 +30,12 @@ protected void after() {
 
     @Override
     public void shutdown() {
+        try {
+            // to prevent socket.io: broken pipe error for async calls
+            Thread.sleep(1000);
+        } catch (InterruptedException e) {
+            throw new RuntimeException(e);
+        }
         server.shutdown();
     }
 

instrumentation-security/httpclient-5.0/src/main/java/com/newrelic/agent/security/instrumentation/httpclient50/BasicRequestProducer_Instrumentation.java

@@ -0,0 +1,24 @@
+/*
+ *
+ *  * Copyright 2023 New Relic Corporation. All rights reserved.
+ *  * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+
+package com.newrelic.agent.security.instrumentation.httpclient50;
+
+import com.newrelic.api.agent.security.NewRelicSecurity;
+import com.newrelic.api.agent.weaver.MatchType;
+import com.newrelic.api.agent.weaver.Weave;
+import org.apache.hc.core5.http.HttpRequest;
+import org.apache.hc.core5.http.nio.AsyncEntityProducer;
+
+import static com.newrelic.agent.security.instrumentation.httpclient50.SecurityHelper.APACHE5_ASYNC_REQUEST_PRODUCER;
+
+@Weave(type=MatchType.BaseClass, originalName = "org.apache.hc.core5.http.nio.support.BasicRequestProducer")
+public class BasicRequestProducer_Instrumentation {
+
+    public BasicRequestProducer_Instrumentation(final HttpRequest request, final AsyncEntityProducer dataProducer) {
+        NewRelicSecurity.getAgent().getSecurityMetaData().addCustomAttribute(APACHE5_ASYNC_REQUEST_PRODUCER+this.hashCode(), request);
+    }
+}

instrumentation-security/httpclient-5.0/src/main/java/com/newrelic/agent/security/instrumentation/httpclient50/HttpAsyncClient_Instrumentation.java

@@ -0,0 +1,76 @@
+/*
+ *
+ *  * Copyright 2023 New Relic Corporation. All rights reserved.
+ *  * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+
+package com.newrelic.agent.security.instrumentation.httpclient50;
+
+import com.newrelic.api.agent.security.NewRelicSecurity;
+import com.newrelic.api.agent.security.instrumentation.helpers.GenericHelper;
+import com.newrelic.api.agent.security.schema.AbstractOperation;
+import com.newrelic.api.agent.weaver.MatchType;
+import com.newrelic.api.agent.weaver.Weave;
+import com.newrelic.api.agent.weaver.Weaver;
+import org.apache.hc.core5.concurrent.FutureCallback;
+import org.apache.hc.core5.http.HttpRequest;
+import org.apache.hc.core5.http.nio.AsyncPushConsumer;
+import org.apache.hc.core5.http.nio.AsyncRequestProducer;
+import org.apache.hc.core5.http.nio.AsyncResponseConsumer;
+import org.apache.hc.core5.http.nio.HandlerFactory;
+import org.apache.hc.core5.http.protocol.HttpContext;
+
+import java.net.URISyntaxException;
+import java.util.concurrent.Future;
+
+import static com.newrelic.agent.security.instrumentation.httpclient50.SecurityHelper.APACHE5_ASYNC_REQUEST_PRODUCER;
+
+@Weave(type = MatchType.Interface, originalName = "org.apache.hc.client5.http.async.HttpAsyncClient")
+public class HttpAsyncClient_Instrumentation {
+
+    public <T> Future<T> execute(
+            AsyncRequestProducer requestProducer,
+            AsyncResponseConsumer<T> responseConsumer,
+            HandlerFactory<AsyncPushConsumer> pushHandlerFactory,
+            HttpContext context,
+            FutureCallback<T> callback) {
+        HttpRequest request = NewRelicSecurity.getAgent().getSecurityMetaData().getCustomAttribute(APACHE5_ASYNC_REQUEST_PRODUCER+requestProducer.hashCode(), HttpRequest.class);
+
+        boolean isLockAcquired = acquireLockIfPossible();
+        AbstractOperation operation = null;
+        // Preprocess Phase
+        if (isLockAcquired) {
+            try {
+                operation = SecurityHelper.preprocessSecurityHook(request, request.getUri().toString(), this.getClass().getName(), SecurityHelper.METHOD_NAME_EXECUTE);
+            } catch (URISyntaxException ignored) {
+            }
+        }
+        Future<T> returnObj = null;
+        // Actual Call
+        try {
+            returnObj = Weaver.callOriginal();
+        } finally {
+            if (isLockAcquired) {
+                releaseLock();
+            }
+        }
+        SecurityHelper.registerExitOperation(isLockAcquired, operation);
+        return returnObj;
+    }
+
+    private void releaseLock() {
+        try {
+            GenericHelper.releaseLock(SecurityHelper.NR_SEC_CUSTOM_ATTRIB_NAME, this.hashCode());
+        } catch (Throwable ignored) {
+        }
+    }
+
+    private boolean acquireLockIfPossible() {
+        try {
+            return GenericHelper.acquireLockIfPossible(SecurityHelper.NR_SEC_CUSTOM_ATTRIB_NAME, this.hashCode());
+        } catch (Throwable ignored) {
+        }
+        return false;
+    }
+}

instrumentation-security/httpclient-5.0/src/main/java/com/newrelic/agent/security/instrumentation/httpclient50/SecurityHelper.java

@@ -15,6 +15,7 @@ public class SecurityHelper {
     public static final String METHOD_NAME_EXECUTE = "execute";
 
     public static final String NR_SEC_CUSTOM_ATTRIB_NAME = "SSRF_OPERATION_LOCK_APACHE5-";
+    public static final String APACHE5_ASYNC_REQUEST_PRODUCER = "APACHE5_ASYNC_REQUEST_PRODUCER_";
 
     public static void registerExitOperation(boolean isProcessingAllowed, AbstractOperation operation) {
         try {

instrumentation-security/httpclient-5.0/src/test/java/com/nr/agent/security/instrumentation/httpclient5/HttpAsyncClientTest.java

@@ -0,0 +1,186 @@
+package com.nr.agent.security.instrumentation.httpclient5;
+
+import com.newrelic.agent.security.introspec.InstrumentationTestConfig;
+import com.newrelic.agent.security.introspec.SecurityInstrumentationTestRunner;
+import com.newrelic.agent.security.introspec.SecurityIntrospector;
+import com.newrelic.agent.security.introspec.internal.HttpServerRule;
+import com.newrelic.api.agent.Trace;
+import com.newrelic.api.agent.security.instrumentation.helpers.GenericHelper;
+import com.newrelic.api.agent.security.instrumentation.helpers.ServletHelper;
+import com.newrelic.api.agent.security.schema.AbstractOperation;
+import com.newrelic.api.agent.security.schema.VulnerabilityCaseType;
+import com.newrelic.api.agent.security.schema.operation.SSRFOperation;
+import org.apache.hc.client5.http.async.methods.SimpleHttpRequest;
+import org.apache.hc.client5.http.async.methods.SimpleHttpResponse;
+import org.apache.hc.client5.http.async.methods.SimpleRequestBuilder;
+import org.apache.hc.client5.http.impl.async.CloseableHttpAsyncClient;
+import org.apache.hc.client5.http.impl.async.HttpAsyncClients;
+import org.apache.hc.core5.http.Message;
+import org.apache.hc.core5.http.nio.entity.BasicAsyncEntityConsumer;
+import org.apache.hc.core5.http.nio.entity.BasicAsyncEntityProducer;
+import org.apache.hc.core5.http.nio.support.BasicRequestProducer;
+import org.apache.hc.core5.http.nio.support.BasicResponseConsumer;
+import org.apache.hc.core5.http.protocol.BasicHttpContext;
+import org.apache.hc.core5.http.protocol.HttpContext;
+import org.junit.Assert;
+import org.junit.ClassRule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.net.URISyntaxException;
+import java.util.List;
+import java.util.Map;
+import java.util.UUID;
+import java.util.concurrent.Future;
+
+@RunWith(SecurityInstrumentationTestRunner.class)
+@InstrumentationTestConfig(includePrefixes = "com.newrelic.agent.security.instrumentation.httpclient50")
+public class HttpAsyncClientTest {
+    @ClassRule
+    public static HttpServerRule server = new HttpServerRule();
+
+    @Test
+    public void testExecute() throws URISyntaxException {
+        String headerValue = String.valueOf(UUID.randomUUID());
+
+        SecurityIntrospector introspector = SecurityInstrumentationTestRunner.getIntrospector();
+        setCSECHeaders(headerValue, introspector);
+
+        callExecute();
+
+        List<AbstractOperation> operations = introspector.getOperations();
+        Assert.assertTrue("No operations detected", operations.size() > 0);
+        SSRFOperation operation = (SSRFOperation) operations.get(0);
+        Map<String, String> headers = server.getHeaders();
+        Assert.assertEquals("Invalid executed parameters.", server.getEndPoint().toString(), operation.getArg());
+        Assert.assertEquals("Invalid event category.", VulnerabilityCaseType.HTTP_REQUEST, operation.getCaseType());
+        Assert.assertEquals("Invalid executed method name.", "execute", operation.getMethodName());
+        verifyHeaders(headerValue, headers);
+    }
+
+    @Test
+    public void testExecute1() throws URISyntaxException {
+
+        String headerValue = String.valueOf(UUID.randomUUID());
+
+        SecurityIntrospector introspector = SecurityInstrumentationTestRunner.getIntrospector();
+        setCSECHeaders(headerValue, introspector);
+        callExecute1();
+
+        List<AbstractOperation> operations = introspector.getOperations();
+        Assert.assertTrue("No operations detected", operations.size() > 0);
+        SSRFOperation operation = (SSRFOperation) operations.get(0);
+        Map<String, String> headers = server.getHeaders();
+        Assert.assertEquals("Invalid executed parameters.", server.getEndPoint().toString(), operation.getArg());
+        Assert.assertEquals("Invalid event category.", VulnerabilityCaseType.HTTP_REQUEST, operation.getCaseType());
+        Assert.assertEquals("Invalid executed method name.", "execute", operation.getMethodName());
+        verifyHeaders(headerValue, headers);
+    }
+
+    @Test
+    public void testExecute2() throws URISyntaxException {
+
+        String headerValue = String.valueOf(UUID.randomUUID());
+
+        SecurityIntrospector introspector = SecurityInstrumentationTestRunner.getIntrospector();
+        setCSECHeaders(headerValue, introspector);
+        callExecute2();
+
+        List<AbstractOperation> operations = introspector.getOperations();
+        Assert.assertTrue("No operations detected", operations.size() > 0);
+        SSRFOperation operation = (SSRFOperation) operations.get(0);
+        Map<String, String> headers = server.getHeaders();
+        Assert.assertEquals("Invalid executed parameters.", server.getEndPoint().toString(), operation.getArg());
+        Assert.assertEquals("Invalid event category.", VulnerabilityCaseType.HTTP_REQUEST, operation.getCaseType());
+        Assert.assertEquals("Invalid executed method name.", "execute", operation.getMethodName());
+        verifyHeaders(headerValue, headers);
+
+    }
+
+    @Test
+    public void testExecute3() throws URISyntaxException {
+
+        String headerValue = String.valueOf(UUID.randomUUID());
+
+        SecurityIntrospector introspector = SecurityInstrumentationTestRunner.getIntrospector();
+        setCSECHeaders(headerValue, introspector);
+        callExecute3();
+
+        List<AbstractOperation> operations = introspector.getOperations();
+        Assert.assertTrue("No operations detected", operations.size() > 0);
+        SSRFOperation operation = (SSRFOperation) operations.get(0);
+        Map<String, String> headers = server.getHeaders();
+        Assert.assertEquals("Invalid executed parameters.", server.getEndPoint().toString(), operation.getArg());
+        Assert.assertEquals("Invalid event category.", VulnerabilityCaseType.HTTP_REQUEST, operation.getCaseType());
+        Assert.assertEquals("Invalid executed method name.", "execute", operation.getMethodName());
+        verifyHeaders(headerValue, headers);
+
+    }
+
+    private void setCSECHeaders(String headerValue, SecurityIntrospector introspector) {
+        introspector.setK2FuzzRequestId(headerValue+"a");
+        introspector.setK2ParentId(headerValue+"b");
+        introspector.setK2TracingData(headerValue);
+    }
+
+    private void verifyHeaders(String headerValue, Map<String, String> headers) {
+        Assert.assertTrue(String.format("Missing K2 header: %s", ServletHelper.CSEC_IAST_FUZZ_REQUEST_ID), headers.containsKey(ServletHelper.CSEC_IAST_FUZZ_REQUEST_ID));
+        Assert.assertEquals(String.format("Invalid K2 header value for:  %s", ServletHelper.CSEC_IAST_FUZZ_REQUEST_ID), headerValue+"a", headers.get(ServletHelper.CSEC_IAST_FUZZ_REQUEST_ID));
+        Assert.assertTrue(String.format("Missing K2 header: %s", GenericHelper.CSEC_PARENT_ID), headers.containsKey(GenericHelper.CSEC_PARENT_ID));
+        Assert.assertEquals(String.format("Invalid K2 header value for:  %s", GenericHelper.CSEC_PARENT_ID), headerValue+"b", headers.get(GenericHelper.CSEC_PARENT_ID));
+        Assert.assertTrue(String.format("Missing K2 header: %s", ServletHelper.CSEC_DISTRIBUTED_TRACING_HEADER), headers.containsKey(ServletHelper.CSEC_DISTRIBUTED_TRACING_HEADER.toLowerCase()));
+        Assert.assertEquals(String.format("Invalid K2 header value for:  %s", ServletHelper.CSEC_DISTRIBUTED_TRACING_HEADER), String.format("%s;DUMMY_UUID/dummy-api-id/dummy-exec-id;",
+                headerValue), headers.get(
+                ServletHelper.CSEC_DISTRIBUTED_TRACING_HEADER.toLowerCase()));
+    }
+
+    @Trace(dispatcher = true)
+    private void callExecute() {
+        try (CloseableHttpAsyncClient httpAsyncClient = HttpAsyncClients.createDefault()) {
+            httpAsyncClient.start();
+            SimpleHttpRequest request = SimpleRequestBuilder.get(server.getEndPoint()).build();
+            Future<SimpleHttpResponse> future = httpAsyncClient.execute(request, null);
+            SimpleHttpResponse response = future.get();
+            System.out.println("response: " + response.getBody());
+        } catch (Exception ignored) {
+        }
+    }
+
+    @Trace(dispatcher = true)
+    private void callExecute1() {
+        try (CloseableHttpAsyncClient httpAsyncClient = HttpAsyncClients.createDefault()) {
+            httpAsyncClient.start();
+            SimpleHttpRequest request = SimpleRequestBuilder.get(server.getEndPoint()).build();
+            HttpContext httpContext = new BasicHttpContext();
+            Future<SimpleHttpResponse> future = httpAsyncClient.execute(request, httpContext, null);
+            SimpleHttpResponse response = future.get();
+            System.out.println("response: " + response.getBody());
+        } catch (Exception ignored) {
+        }
+    }
+
+    @Trace(dispatcher = true)
+    private void callExecute2() {
+        try (CloseableHttpAsyncClient httpAsyncClient = HttpAsyncClients.createDefault()) {
+            httpAsyncClient.start();
+            SimpleHttpRequest request = SimpleRequestBuilder.get(server.getEndPoint()).build();
+            Future<SimpleHttpResponse> future = httpAsyncClient.execute(new BasicRequestProducer(request, new BasicAsyncEntityProducer("test data")), new BasicResponseConsumer(new BasicAsyncEntityConsumer()), null);
+            SimpleHttpResponse response = future.get();
+            System.out.println("response: " + response.getBody());
+        } catch (Exception ignored) {
+        }
+    }
+
+    @Trace(dispatcher = true)
+    public void callExecute3() {
+        try (CloseableHttpAsyncClient httpAsyncClient = HttpAsyncClients.createDefault()) {
+            httpAsyncClient.start();
+            SimpleHttpRequest request = SimpleRequestBuilder.get(server.getEndPoint()).build();
+            HttpContext httpContext = new BasicHttpContext();
+            Future<Message> future = httpAsyncClient.execute(new BasicRequestProducer(request, new BasicAsyncEntityProducer("test data")), new BasicResponseConsumer(new BasicAsyncEntityConsumer()), httpContext, null);
+            Message response = future.get();
+            System.out.println("response: " + response.getBody());
+        } catch (Exception ignored) {
+        }
+    }
+}