fix: Improve error handling in AWS account ID parsing logic (#2984)

Author: tippmar-nr

src/Agent/NewRelic/Agent/Extensions/NewRelic.Agent.Extensions/Caching/WeakReferenceKey.cs

@@ -11,11 +11,13 @@ namespace NewRelic.Agent.Extensions.Caching
     /// <typeparam name="T"></typeparam>
     public class WeakReferenceKey<T> where T : class
     {
+        private readonly int _hashCode;
         private WeakReference<T> WeakReference { get; }
 
         public WeakReferenceKey(T cacheKey)
         {
             WeakReference = new WeakReference<T>(cacheKey);
+            _hashCode = cacheKey.GetHashCode(); // store the hashcode since we use it in the Equals method and the object could have been GC'd by the time we need to look for it
         }
 
         public override bool Equals(object obj)
@@ -27,19 +29,17 @@ public override bool Equals(object obj)
                 {
                     return ReferenceEquals(thisTarget, otherTarget);
                 }
+
+                // if one of the targets has been garbage collected, we can still compare the hashcodes
+                return otherKey.GetHashCode() == _hashCode;
             }
 
             return false;
         }
 
         public override int GetHashCode()
         {
-            if (WeakReference.TryGetTarget(out var target))
-            {
-                return target.GetHashCode();
-            }
-
-            return 0;
+            return WeakReference.TryGetTarget(out var target) ? target.GetHashCode() : _hashCode;
         }
 
         /// <summary>

src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/AwsSdk/AmazonServiceClientWrapper.cs

@@ -11,6 +11,8 @@ namespace NewRelic.Providers.Wrapper.AwsSdk;
 
 public class AmazonServiceClientWrapper : IWrapper
 {
+    private bool _disableServiceClientWrapper;
+
     private const int LRUCapacity = 100;
     // cache the account id per instance of AmazonServiceClient.Config
     public static LRUCache<WeakReferenceKey<object>, string> AwsAccountIdByClientConfigCache = new(LRUCapacity);
@@ -27,32 +29,81 @@ public CanWrapResponse CanWrap(InstrumentedMethodInfo instrumentedMethodInfo)
 
     public AfterWrappedMethodDelegate BeforeWrappedMethod(InstrumentedMethodCall instrumentedMethodCall, IAgent agent, ITransaction transaction)
     {
+        if (_disableServiceClientWrapper) // something bad happened on a previous call, so we're disabling this wrapper
+            return Delegates.NoOp;
+
         object client = instrumentedMethodCall.MethodCall.InvocationTarget;
 
         var weakReferenceKey = new WeakReferenceKey<object>(client);
-        if (AmazonServiceClientInstanceCache.Contains(weakReferenceKey)) // don't do anything if we've already seen this client instance
+
+        // don't do anything if we've already seen this client instance -- the account ID is cached already
+        if (AmazonServiceClientInstanceCache.Contains(weakReferenceKey)) 
             return Delegates.NoOp;
 
+        // add this client instance to cache so we don't process it again
         AmazonServiceClientInstanceCache.Add(weakReferenceKey);
 
-        string awsAccountId;
+        // retrieve and cache the account id
+        string awsAccountId = null;
         try
         {
             // get the AWSCredentials parameter
-            dynamic awsCredentials = instrumentedMethodCall.MethodCall.MethodArguments[0];
+            if (instrumentedMethodCall.MethodCall.MethodArguments.Length > 0)
+            {
+                dynamic awsCredentials = instrumentedMethodCall.MethodCall.MethodArguments[0];
+                if (awsCredentials != null)
+                {
+                    dynamic immutableCredentials = awsCredentials.GetCredentials();
+                    if (immutableCredentials != null)
+                    {
+                        string accessKey = immutableCredentials.AccessKey;
 
-            dynamic immutableCredentials = awsCredentials.GetCredentials();
-            string accessKey = immutableCredentials.AccessKey;
+                        if (!string.IsNullOrEmpty(accessKey))
+                        {
+                            try
+                            {
+                                // convert the access key to an account id
+                                awsAccountId = AwsAccountIdDecoder.GetAccountId(accessKey);
+                            }
+                            catch (Exception e)
+                            {
+                                agent.Logger.Debug(e, "Unexpected exception parsing AWS Account ID from AccessKey.");
+                            }
+                        }
+                        else
+                            agent.Logger.Debug("Unable to parse AWS Account ID from AWSCredentials because AccessKey was null.");
+                    }
+                    else
+                        agent.Logger.Debug("Unable to parse AWS Account ID from AWSCredentials because GetCredentials() returned null.");
+                }
+                else
+                    agent.Logger.Debug("Unable to parse AWS Account ID from AWSCredentials because AWSCredentials was null.");
+            }
+            else
+                agent.Logger.Debug("Unable to parse AWS Account ID from AWSCredentials because there were no arguments in the method call.");
 
-            // convert the access key to an account id
-            awsAccountId = AwsAccountIdDecoder.GetAccountId(accessKey);
+            // fall back to configuration if we didn't get an account id from the credentials
+            if (string.IsNullOrEmpty(awsAccountId))
+            {
+                agent.Logger.Debug("Using AccountId from configuration.");
+                awsAccountId = agent.Configuration.AwsAccountId;
+            }
         }
         catch (Exception e)
         {
-            agent.Logger.Info($"Unable to parse AWS Account ID from AccessKey. Using AccountId from configuration instead. Exception: {e.Message}");
+            agent.Logger.Debug(e, "Unexpected exception in AmazonServiceClientWrapper.BeforeWrappedMethod(). Using AccountId from configuration.");
             awsAccountId = agent.Configuration.AwsAccountId;
         }
 
+        // disable the wrapper if we get this far and there's no account id
+        if (string.IsNullOrEmpty(awsAccountId))
+        {
+            agent.Logger.Warn("Unable to parse AWS Account ID from AWSCredentials or configuration. Further AWS Account ID parsing will be disabled.");
+            _disableServiceClientWrapper = true;
+
+            return Delegates.NoOp;
+        }
+
         return Delegates.GetDelegateFor(onComplete: () =>
         {
             // get the _config field from the client

tests/Agent/UnitTests/NewRelic.Agent.Extensions.Tests/Cache/WeakReferenceKeyTests.cs

@@ -93,24 +93,6 @@ public void GetHashCode_ShouldReturnDifferentHashCodeForDifferentObjects()
             Assert.That(hashCode1, Is.Not.EqualTo(hashCode2));
         }
 
-        [Test]
-        public async Task GetHashCode_ShouldReturnZeroIfTargetIsGarbageCollected()
-        {
-            // Arrange
-            var weakRefKey = GetWeakReferenceKey();
-
-            // Act
-            Assert.That(weakRefKey.Value, Is.Not.Null);
-            // force garbage collection
-            GC.Collect();
-            GC.WaitForPendingFinalizers();
-            await Task.Delay(500);
-            GC.Collect(); // Force another collection
-
-            // Assert
-            Assert.That(weakRefKey.GetHashCode(), Is.EqualTo(0));
-        }
-
         [Test]
         public async Task Value_ShouldReturnNullIfTargetIsGarbageCollected()
         {
@@ -168,6 +150,26 @@ public async Task Equals_ShouldReturnFalseIfTargetIsGarbageCollected()
             Assert.That(weakRefKey1.Equals(weakRefKey2), Is.False);
         }
 
+        [Test]
+        public async Task GetHashCode_ShouldReturnOriginalHashcodeIfTargetIsGarbageCollected()
+        {
+            // Arrange
+            var weakRefKey = GetWeakReferenceKey();
+            var originalHashCode = weakRefKey.GetHashCode();
+
+            // Act
+            Assert.That(weakRefKey.Value, Is.Not.Null);
+
+            // force garbage collection
+            GC.Collect();
+            GC.WaitForPendingFinalizers();
+            await Task.Delay(500);
+            GC.Collect(); // Force another collection
+
+            // Assert
+            Assert.That(weakRefKey.Value, Is.Null); // ensure GC really happened
+            Assert.That(weakRefKey.GetHashCode(), Is.EqualTo(originalHashCode));
+        }
         private class Foo
         {
             public string Bar { get; set; }