Linking AWS MQ infrastructure entities

agent-bridge/src/main/java/com/newrelic/agent/bridge/CollectionFactory.java

@@ -8,6 +8,7 @@
 package com.newrelic.agent.bridge;
 
 import java.util.Map;
+import java.util.function.Function;
 
 /**
  * Allows instrumentation and bridge API implementations to use collections from third partly libraries without
@@ -28,6 +29,16 @@ public interface CollectionFactory {
      */
     <K, V> Map<K, V> createConcurrentWeakKeyedMap();
 
+    /**
+     * Wraps the provided function into one that will cache the results for future calls.
+     * @param loader the function that calculates the value.
+     * @param maxSize the max number of items to be cached.
+     * @return the cached item, or the result of the loader call.
+     * @param <K> the type of key
+     * @param <V> the type of value stored/returned
+     */
+    <K, V> Function<K, V> memoize(Function<K, V> loader, int maxSize);
+
     /**
      * Create a time based eviction cache in which an entry's age is determined on a last-write basis.
      *

agent-bridge/src/main/java/com/newrelic/agent/bridge/DefaultCollectionFactory.java

@@ -11,14 +11,33 @@
 import java.util.HashMap;
 import java.util.Map;
 import java.util.WeakHashMap;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Function;
 
+/**
+ * This implementation of {@link CollectionFactory} will only be used if the agent-bridge
+ * is being used by an application and the agent is NOT being loaded. Thus, it is unlikely
+ * that the objects created by this implementation are going to receive much use.
+ * So methods in this implementation do not need to implement all functional requirements
+ * of the methods in the interface, but they should not break under low use.
+ */
 public class DefaultCollectionFactory implements CollectionFactory {
 
     @Override
     public <K, V> Map<K, V> createConcurrentWeakKeyedMap() {
         return Collections.synchronizedMap(new WeakHashMap<K, V>());
     }
 
+    @Override
+    public <K, V> Function<K, V> memoize(Function<K, V> loader, int maxSize) {
+        Map<K, V> map = new ConcurrentHashMap<>();
+        return k -> map.computeIfAbsent(k, k1 -> {
+            if (map.size() >= maxSize) {
+                map.remove(map.keySet().iterator().next());
+            }
+            return loader.apply(k1);
+        });
+    }
     /**
      * Note: In this implementation, this method will return a simple concurrent map since an eviction
      * cache can't be easily created with just vanilla JDK Map SDKs.

agent-bridge/src/main/java/com/newrelic/agent/bridge/NoOpPrivateApi.java

@@ -7,6 +7,8 @@
 
 package com.newrelic.agent.bridge;
 
+import com.newrelic.api.agent.DestinationType;
+
 import javax.management.MBeanServer;
 import java.io.Closeable;
 import java.util.Map;
@@ -71,6 +73,11 @@ public void setInstanceName(String instanceName) {
     public void addTracerParameter(String key, String value) {
     }
 
+    @Override
+    public void addTracerParameter(String key, String value, boolean addToSpan) {
+
+    }
+
     @Override
     public void addTracerParameter(String key, Map<String, String> values) {
     }

agent-bridge/src/main/java/com/newrelic/agent/bridge/PrivateApi.java

@@ -7,6 +7,8 @@
 
 package com.newrelic.agent.bridge;
 
+import com.newrelic.api.agent.DestinationType;
+
 import javax.management.MBeanServer;
 import java.io.Closeable;
 import java.util.Map;
@@ -57,6 +59,8 @@ public interface PrivateApi {
 
     void addTracerParameter(String key, String value);
 
+    void addTracerParameter(String key, String value, boolean addToSpan);
+
     void addTracerParameter(String key, Map<String, String> values);
 
     /**
@@ -122,4 +126,5 @@ public interface PrivateApi {
      */
     @Deprecated
     void setInstanceName(String instanceName);
+
 }

agent-bridge/src/main/java/com/newrelic/agent/bridge/external/MessageConsumeParameters.java

@@ -26,7 +26,7 @@ public class MessageConsumeParameters extends com.newrelic.api.agent.MessageCons
     @Deprecated
     protected MessageConsumeParameters(String library, DestinationType destinationType, String destinationName,
             InboundHeaders inboundHeaders) {
-        super(library, destinationType.toApiDestinationType(), destinationName, inboundHeaders);
+        super(library, destinationType.toApiDestinationType(), destinationName, inboundHeaders, null, null, null);
     }
 
 }

agent-bridge/src/main/java/com/newrelic/agent/bridge/external/MessageProduceParameters.java

@@ -16,7 +16,7 @@
 public class MessageProduceParameters extends com.newrelic.api.agent.MessageProduceParameters implements ExternalParameters {
 
     /**
-     * @Deprecated Do not use. Use {@link com.newrelic.api.agent.MessageProduceParameters#MessageProduceParameters} instead.
+     * @Deprecated Do not use. Use the fluent builder {@link com.newrelic.api.agent.MessageProduceParameters#library(String)} instead.
      *
      * @param library
      * @param destinationType
@@ -26,7 +26,7 @@ public class MessageProduceParameters extends com.newrelic.api.agent.MessageProd
     @Deprecated
     protected MessageProduceParameters(String library, DestinationType destinationType, String destinationName,
             OutboundHeaders outboundHeaders) {
-        super(library, destinationType.toApiDestinationType(), destinationName, outboundHeaders);
+        super(library, destinationType.toApiDestinationType(), destinationName, outboundHeaders, null, null, null);
     }
 
 }

agent-bridge/src/main/java/com/newrelic/agent/bridge/messaging/BrokerInstance.java

@@ -0,0 +1,29 @@
+/*
+ *
+ *  * Copyright 2024 New Relic Corporation. All rights reserved.
+ *  * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+package com.newrelic.agent.bridge.messaging;
+
+public class BrokerInstance {
+    private final String hostName;
+    private final Integer port;
+
+    public static BrokerInstance empty() {
+        return new BrokerInstance(null, null);
+    }
+
+    public BrokerInstance(String host, Integer port) {
+        this.hostName = host;
+        this.port = port;
+    }
+
+    public String getHostName() {
+        return hostName;
+    }
+
+    public Integer getPort() {
+        return port;
+    }
+}

agent-bridge/src/main/java/com/newrelic/agent/bridge/messaging/JmsProperties.java

@@ -0,0 +1,11 @@
+/*
+ *
+ *  * Copyright 2024 New Relic Corporation. All rights reserved.
+ *  * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+package com.newrelic.agent.bridge.messaging;
+
+public class JmsProperties {
+    public static final String NR_JMS_BROKER_INSTANCE_PROPERTY = "NR_JMS_BROKER_INSTANCE_PROPERTY";
+}

agent-bridge/src/main/java/com/newrelic/agent/util/AwsAccountUtil.java

@@ -0,0 +1,68 @@
+/*
+ *
+ *  * Copyright 2024 New Relic Corporation. All rights reserved.
+ *  * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+
+package com.newrelic.agent.util;
+
+import com.newrelic.agent.bridge.AgentBridge;
+
+import java.util.function.Function;
+
+
+public class AwsAccountUtil {
+    private static AwsAccountUtil INSTANCE = AwsAccountUtil.create();
+
+    private final Function<String, Long> CACHE = AgentBridge.collectionFactory.memoize(this::doDecodeAccount, 32);
+
+    public static AwsAccountUtil get() {
+        return INSTANCE;
+    }
+
+    public Long decodeAccount(String accessKey) {
+        return CACHE.apply(accessKey);
+    }
+
+    private Long doDecodeAccount(String awsAccessKeyId) {
+        String accessKeyWithoutPrefix = awsAccessKeyId.substring(4).toLowerCase();
+        long encodedAccount = base32Decode(accessKeyWithoutPrefix);
+        // magic number
+        long mask = 140737488355200L;
+        // magic incantation to find out the account
+        return (encodedAccount & mask) >> 7;
+    }
+
+    /**
+     * Character range is A-Z, 2-7. 'A' being 0 and '7', 31.
+     * Characters outside of this range will be considered 0.
+     * @param src the string to be decoded. Must be at least 10 characters.
+     * @return a long containing first 6 bytes of the base 32 decoded data.
+     * @throws ArrayIndexOutOfBoundsException if src has less than 10 characters
+     */
+    private long base32Decode(String src) {
+        long base = 0;
+        char[] chars = src.toCharArray();
+        // each char is 5 bits, we need 48 bits
+        for (int i = 0; i < 10; i++) {
+            char c = chars[i];
+            base <<= 5;
+            if (c >= 'a' && c <= 'z') {
+                base += c - 'a';
+            } else if (c >= '2' && c <= '7') {
+                base += c - '2' + 26;
+            }
+        }
+        // 50 bits were read, dropping the lowest 2
+        return base >> 2;
+    }
+
+    private AwsAccountUtil() {
+        // prevent instantiation of utility class
+    }
+
+    private static AwsAccountUtil create() {
+        return new AwsAccountUtil();
+    }
+}

agent-bridge/src/test/java/com/newrelic/agent/util/AwsAccountUtilTest.java

@@ -0,0 +1,25 @@
+/*
+ *
+ *  * Copyright 2024 New Relic Corporation. All rights reserved.
+ *  * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+
+package com.newrelic.agent.util;
+
+import org.junit.Test;
+
+import static org.junit.Assert.*;
+
+public class AwsAccountUtilTest {
+
+    @Test
+    public void decodeAccount() {
+        Long accountId = AwsAccountUtil.get().decodeAccount("FKKY6RVFFB77ZZZZZZZZ");
+        assertEquals(999999999999L, accountId.longValue());
+
+        accountId = AwsAccountUtil.get().decodeAccount("FKKYQAAAAAAAZZZZZZZZ");
+        assertEquals(1L, accountId.longValue());
+    }
+
+}

functional_test/src/test/java/test/newrelic/test/agent/api/ApiTest.java

@@ -15,9 +15,9 @@
 import com.newrelic.agent.TransactionListener;
 import com.newrelic.agent.bridge.AgentBridge;
 import com.newrelic.agent.bridge.TransactionNamePriority;
-import com.newrelic.agent.browser.BrowserConfigTest;
 import com.newrelic.agent.config.AgentConfigImpl;
 import com.newrelic.agent.config.ConfigConstant;
+import com.newrelic.agent.config.Hostname;
 import com.newrelic.agent.dispatchers.WebRequestDispatcher;
 import com.newrelic.agent.environment.AgentIdentity;
 import com.newrelic.agent.errors.ErrorService;
@@ -38,7 +38,6 @@
 import com.newrelic.agent.tracers.servlet.MockHttpResponse;
 import com.newrelic.agent.transaction.PriorityTransactionName;
 import com.newrelic.agent.transaction.TransactionThrowable;
-import com.newrelic.agent.util.Obfuscator;
 import com.newrelic.api.agent.DatastoreParameters;
 import com.newrelic.api.agent.DestinationType;
 import com.newrelic.api.agent.ExtendedRequest;
@@ -77,13 +76,11 @@
 import java.net.URISyntaxException;
 import java.net.URL;
 import java.nio.charset.StandardCharsets;
-import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
 /* (non-javadoc)
@@ -95,6 +92,7 @@ public class ApiTest implements TransactionListener {
     ApiTestHelper apiTestHelper = new ApiTestHelper();
     private static final String CAT_CONFIG_FILE = "configs/cross_app_tracing_test.yml";
     private static final String HIGH_SECURITY_CONFIG_FILE = "configs/high_security_config.yml";
+    public static String HOSTNAME = Hostname.getHostname(ServiceFactory.getConfigService().getDefaultAgentConfig());
     private static final ClassLoader CLASS_LOADER = ApiTest.class.getClassLoader();
 
     @Before
@@ -2007,6 +2005,27 @@ public void testMessagingAPI() throws Exception {
         }
     }
 
+    @Test
+    public void testMessagingAPIWithHostAndPort() throws Exception {
+        // override default agent config to disabled distributed tracing and use CAT instead
+        EnvironmentHolder holder = setupEnvironmentHolder(CAT_CONFIG_FILE, "cat_enabled_dt_disabled_test");
+        MessagingTestServer server = new MessagingTestServer(8088);
+
+        try {
+            server.start();
+            runTestMessagingAPIWithHostAndPort();
+            String messageBrokerMetric = "MessageBroker/JMS/Queue/Consume/Temp";
+            Assert.assertTrue("The following metric should exist: " + messageBrokerMetric, apiTestHelper.tranStats.getScopedStats().getStatsMap().containsKey(messageBrokerMetric));
+        } catch (IOException e) {
+            e.printStackTrace();
+            Assert.fail();
+        } finally {
+            Transaction.clearTransaction();
+            server.closeAllConnections();
+            holder.close();
+        }
+    }
+
     @Trace(dispatcher = true)
     private void runTestMessagingAPI() {
         URL myURL = null;
@@ -2048,6 +2067,48 @@ private void runTestMessagingAPI() {
         }
     }
 
+    @Trace(dispatcher = true)
+    private void runTestMessagingAPIWithHostAndPort() {
+        URL myURL = null;
+        try {
+            Thread.sleep(600);
+            myURL = new URL("http://localhost:8088");
+            HttpUriRequest request = RequestBuilder.get().setUri(myURL.toURI()).build();
+
+            ApiTestHelper.OutboundWrapper outboundRequestWrapper = new ApiTestHelper.OutboundWrapper(request, HeaderType.MESSAGE);
+
+            // MessageProducer
+            ExternalParameters messageProduceParameters = MessageProduceParameters
+                    .library("JMS")
+                    .destinationType(DestinationType.NAMED_QUEUE)
+                    .destinationName("MessageDestination")
+                    .outboundHeaders(outboundRequestWrapper)
+                    .instance(myURL.getHost(), myURL.getPort())
+                    .build();
+            NewRelic.getAgent().getTracedMethod().reportAsExternal(messageProduceParameters);
+
+            Assert.assertTrue(request.getHeaders("NewRelicID").length != 0);
+            Assert.assertTrue(request.getHeaders("NewRelicTransaction").length != 0);
+
+            CloseableHttpClient connection = HttpClientBuilder.create().build();
+            CloseableHttpResponse response = connection.execute(request);
+
+            // MessageConsumer
+            ExternalParameters messageResponseParameters = MessageConsumeParameters
+                    .library("JMS")
+                    .destinationType(DestinationType.TEMP_QUEUE)
+                    .destinationName("MessageDestination")
+                    .inboundHeaders(new ApiTestHelper.InboundWrapper(response, HeaderType.MESSAGE))
+                    .build();
+            NewRelic.getAgent().getTracedMethod().reportAsExternal(messageResponseParameters);
+
+            Assert.assertTrue(response.getHeaders("NewRelicAppData").length != 0);
+        } catch (Exception e) {
+            e.printStackTrace();
+            Assert.fail();
+        }
+    }
+
     @Test
     public void testNotNull() {
         Assert.assertNotNull(NewRelic.getAgent().getTransaction().getTracedMethod());