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/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, 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, 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_HOST_AND_PORT_PROPERTY = "com.nr.agent.instrumentation.jms.BrokerInstance";
+}

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());
+    }
+
+}