Skip to content

Commit

Permalink
fix: Improve error handling in AWS account ID parsing logic (#2984)
Browse files Browse the repository at this point in the history
  • Loading branch information
tippmar-nr authored Feb 5, 2025
1 parent 7ee9825 commit 9b66750
Show file tree
Hide file tree
Showing 3 changed files with 85 additions and 32 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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>
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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()
{
Expand Down Expand Up @@ -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; }
Expand Down

0 comments on commit 9b66750

Please sign in to comment.