First pass at incorporating account ID

newrelic · Oct 28, 2024 · 4340181 · 4340181
1 parent c673718
commit 4340181
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 4 deletions.
diff --git a/src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/AwsSdk/AwsSdkPipelineWrapper.cs b/src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/AwsSdk/AwsSdkPipelineWrapper.cs
@@ -3,6 +3,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Linq;
 using System.Threading.Tasks;
 using NewRelic.Agent.Api;
 using NewRelic.Agent.Extensions.Providers.Wrapper;
@@ -12,6 +13,7 @@ namespace NewRelic.Providers.Wrapper.AwsSdk
     public class AwsSdkPipelineWrapper : IWrapper
     {
         public bool IsTransactionRequired => true;
+        private string _accountId = null;
 
         private const string WrapperName = "AwsSdkPipelineWrapper";
         private static HashSet<string> _unsupportedRequestTypes = new();
@@ -38,6 +40,24 @@ private string GetRegion(IAgent agent, dynamic requestContext)
             return "";
         }
 
+        private string GetAccountId(IAgent agent)
+        {
+            if (_accountId != null)
+            {
+                return _accountId;
+            }
+            _accountId = agent.Configuration.AwsAccountId;
+            if (_accountId == null)
+            {
+                _accountId = "";
+            }
+            else if ((_accountId.Length != 12) || _accountId.Any(c => (c < '0') || (c > '9')))
+            {
+                agent.Logger.Warn("Supplied AWS Account Id appears to be invalid: {0}", _accountId);
+            }
+            return _accountId;
+        }
+
         public AfterWrappedMethodDelegate BeforeWrappedMethod(InstrumentedMethodCall instrumentedMethodCall, IAgent agent, ITransaction transaction)
         {
             // Get the IExecutionContext (the only parameter)
@@ -67,14 +87,15 @@ public AfterWrappedMethodDelegate BeforeWrappedMethod(InstrumentedMethodCall ins
             }
             dynamic request = requestContext.OriginalRequest;
             string requestType = request.GetType().FullName;
+            string accountId = agent.Configuration.AwsAccountId;
 
             if (requestType.StartsWith("Amazon.SQS"))
             {
                 return SQSRequestHandler.HandleSQSRequest(instrumentedMethodCall, agent, transaction, request, isAsync, executionContext);
             }
             else if (requestType == "Amazon.Lambda.Model.InvokeRequest")
             {
-                return LambdaInvokeRequestHandler.HandleInvokeRequest(instrumentedMethodCall, agent, transaction, request, isAsync, GetRegion(agent, requestContext));
+                return LambdaInvokeRequestHandler.HandleInvokeRequest(instrumentedMethodCall, agent, transaction, request, isAsync, GetRegion(agent, requestContext), accountId);
             }
 
             if (_unsupportedRequestTypes.Add(requestType)) // log once per unsupported request type

diff --git a/src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/AwsSdk/LambdaInvokeRequestHandler.cs b/src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/AwsSdk/LambdaInvokeRequestHandler.cs
@@ -47,7 +47,7 @@ private static void SetRequestIdIfAvailable(IAgent agent, ITransaction transacti
             }
         }
 
-        public static AfterWrappedMethodDelegate HandleInvokeRequest(InstrumentedMethodCall instrumentedMethodCall, IAgent agent, ITransaction transaction, dynamic request, bool isAsync, string region)
+        public static AfterWrappedMethodDelegate HandleInvokeRequest(InstrumentedMethodCall instrumentedMethodCall, IAgent agent, ITransaction transaction, dynamic request, bool isAsync, string region, string accountId)
         {
             string functionName = request.FunctionName;
             string qualifier = request.Qualifier;
@@ -65,12 +65,12 @@ public static AfterWrappedMethodDelegate HandleInvokeRequest(InstrumentedMethodC
                 string key = $"{region}:{functionName}";
                 if (!_arnCache.TryGetValue(key, out arn))
                 {
-                    arn = AwsSdkHelpers.ConstructArn(agent, functionName, region, "");
+                    arn = AwsSdkHelpers.ConstructArn(agent, functionName, region, accountId);
                     _arnCache.TryAdd(key, arn);
                 }
             }
             var segment = transaction.StartTransactionSegment(instrumentedMethodCall.MethodCall, "InvokeRequest");
-            segment.GetExperimentalApi().MakeLeaf();
+            //segment.GetExperimentalApi().MakeLeaf();
 
             transaction.AddCloudSdkAttribute("cloud.platform", "aws_lambda");
             transaction.AddCloudSdkAttribute("aws.operation", "InvokeRequest");