Skip to content

Commit

Permalink
aws-lambda: changing ARN assembly to accept different types of functi…
Browse files Browse the repository at this point in the history
…on references
  • Loading branch information
meiao committed Oct 30, 2024
1 parent fc8cf8f commit 62e9dbd
Show file tree
Hide file tree
Showing 9 changed files with 360 additions and 426 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -10,14 +10,26 @@
import com.newrelic.agent.bridge.AgentBridge;
import com.newrelic.api.agent.CloudAccountInfo;
import com.newrelic.api.agent.CloudParameters;
import com.newrelic.api.agent.NewRelic;

import java.util.function.Function;
import java.util.logging.Level;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

public class LambdaUtil {

private static final String PLATFORM = "aws_lambda";
private static final String NULL_ARN = "";
private static final FunctionProcessedData NULL_DATA = new FunctionProcessedData(NULL_ARN, NULL_ARN);
private static final String PREFIX = "arn:aws:lambda:";
private static final Pattern FUNC_REF_PATTERN = Pattern.compile(
"(arn:(aws[a-zA-Z-]*)?:lambda:)?" + // arn prefix
"((?<region>[a-z]{2}((-gov)|(-iso([a-z]?)))?-[a-z]+-\\d{1}):)?" + // region
"((?<accountId>\\d{12}):)?" + // account id
"(function:)?" + // constant
"(?<functionName>[a-zA-Z0-9-\\.]+)" + // function name (only required part)
"(:(?<qualifier>\\$LATEST|[a-zA-Z0-9-]+))?"); // qualifier: version or alias
private static final Function<FunctionRawData, FunctionProcessedData> CACHE =
AgentBridge.collectionFactory.createAccessTimeBasedCache(3600, 8, LambdaUtil::processData);

Expand Down Expand Up @@ -53,66 +65,47 @@ public static CloudParameters getCloudParameters(FunctionRawData functionRawData
*/
// Visible for testing
static FunctionProcessedData processData(FunctionRawData data) {
String functionRef = data.getFunctionRef();

String[] parts = functionRef.split(":");

String functionName = NULL_ARN;
String arn = NULL_ARN;

if (parts.length == 1) {
// function name: {function-name}
String accountId = getAccountId(data.getSdkClient());
if (accountId != null) {
String qualifier = data.getQualifier();
if (qualifier == null) {
arn = PREFIX + data.getRegion() + ":" + accountId + ":function:" + functionRef;
} else {
arn = PREFIX + data.getRegion() + ":" + accountId + ":function:" + functionRef + ":" + qualifier;
}
}
functionName = functionRef;
String functionRef = data.getFunctionRef();
Matcher matcher = FUNC_REF_PATTERN.matcher(functionRef);
if (!matcher.matches()) {
return NULL_DATA;
}
String region = matcher.group("region");
String accountId = matcher.group("accountId");
String qualifier = matcher.group("qualifier");
String functionName = matcher.group("functionName");

if (functionName == null) {
// will not be able to add any data
NewRelic.getAgent().getLogger().log(Level.INFO, "aws-lambda: Unable to assemble ARN: " + functionRef);
return NULL_DATA;
}

} else if (parts.length == 2) {
// function name + qualifier: {function-name}:{qualifier}
String accountId = getAccountId(data.getSdkClient());
if (accountId != null) {
arn = PREFIX + data.getRegion() + ":" + accountId + ":function:" + functionRef;
}
functionName = parts[0];
if (region == null) {
// if region is not provided, we will try to get it from the SDK client
region = data.getRegion();
}

if (accountId == null) {
// if account id is not provided, we will try to get it from the config
accountId = getAccountId(data.getSdkClient());
}

} else if (parts.length == 3) {
// partial ARN: {account-id}:function:{function-name}
functionName = parts[2];
String qualifier = data.getQualifier();
if (region != null && accountId != null) {
if (qualifier == null) {
arn = PREFIX + data.getRegion() + ":" + functionRef;
} else {
arn = PREFIX + data.getRegion() + ":" + functionRef + ":" + qualifier;
qualifier = data.getQualifier();
}

} else if (parts.length == 4) {
// partial ARN with qualifier: {account-id}:function:{function-name}:{qualifier}
functionName = parts[2];
arn = PREFIX + data.getRegion() + ":" + functionRef;

} else if (parts.length == 7) {
// full ARN: arn:aws:lambda:{region}:{account-id}:function:{function-name}
functionName = parts[6];
String qualifier = data.getQualifier();
if (qualifier == null) {
arn = functionRef;
if (qualifier == null || qualifier.isEmpty() || "$LATEST".equals(qualifier)) {
arn = PREFIX + region + ":" + accountId + ":function:" + functionName;
} else {
arn = functionRef + ":" + qualifier;
arn = PREFIX + region + ":" + accountId + ":function:" + functionName + ":" + qualifier;
}

} else if (parts.length == 8) {
// full ARN with qualifier: arn:aws:lambda:{region}:{account-id}:function:{function-name}:{qualifier}
functionName = parts[6];
arn = functionRef;
} else {
NewRelic.getAgent().getLogger().log(Level.INFO, "aws-lambda: Missing information to assemble ARN.");
}
// reference should be invalid if the number of parts do not match any of the expected cases

return new FunctionProcessedData(functionName, arn);
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ public abstract class AWSLambdaAsyncClient_Instrumentation {
protected abstract String getSigningRegion();

public Future<InvokeResult> invokeAsync(final InvokeRequest request, AsyncHandler<InvokeRequest, InvokeResult> asyncHandler) {
FunctionRawData functionRawData = new FunctionRawData(request.getFunctionName(), request.getQualifier(), getSigningRegion(), this);
FunctionRawData functionRawData = new FunctionRawData(request.getFunctionName(), request.getQualifier(), this.getSigningRegion(), this);
CloudParameters cloudParameters = LambdaUtil.getCloudParameters(functionRawData);
String functionName = LambdaUtil.getSimpleFunctionName(functionRawData);
Segment segment = NewRelic.getAgent().getTransaction().startSegment("Lambda", "invoke/" + functionName);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ public abstract class AWSLambdaClient_Instrumentation {

@Trace(leaf = true)
public InvokeResult invoke(InvokeRequest invokeRequest) {
FunctionRawData functionRawData = new FunctionRawData(invokeRequest.getFunctionName(), invokeRequest.getQualifier(), getSigningRegion(), this);
FunctionRawData functionRawData = new FunctionRawData(invokeRequest.getFunctionName(), invokeRequest.getQualifier(), this.getSigningRegion(), this);
CloudParameters cloudParameters = LambdaUtil.getCloudParameters(functionRawData);
TracedMethod tracedMethod = NewRelic.getAgent().getTracedMethod();
tracedMethod.reportAsExternal(cloudParameters);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -64,163 +64,6 @@ public void testGetCloudParamArnQualifier() {
assertEquals("arn:aws:lambda:us-east-1:123456789012:function:my-function:alias", cloudParameters.getResourceId());
}

@Test
public void testGetArnFunctionName() {
FunctionProcessedData data = LambdaUtil.processData(data("my-function", null));
assertEquals("", data.getArn());
assertEquals("my-function", data.getFunctionName());
}

@Test
public void testGetArnFunctionNameClient() {
mockCloudApiClient();
FunctionProcessedData data = LambdaUtil.processData(data("my-function", null));
assertEquals("arn:aws:lambda:us-east-1:123456789012:function:my-function", data.getArn());
assertEquals("my-function", data.getFunctionName());
}

@Test
public void testGetArnFunctionNameWithAlias() {
FunctionProcessedData data = LambdaUtil.processData(data("my-function:alias", null));
assertEquals("", data.getArn());
assertEquals("my-function", data.getFunctionName());
}

@Test
public void testGetArnFunctionNameWithAliasClient() {
mockCloudApiClient();
FunctionProcessedData data = LambdaUtil.processData(data("my-function:alias", null));
assertEquals("arn:aws:lambda:us-east-1:123456789012:function:my-function:alias", data.getArn());
assertEquals("my-function", data.getFunctionName());
}

@Test
public void testGetArnFunctionNameWithVersion() {
FunctionProcessedData data = LambdaUtil.processData(data("my-function:123", null));
assertEquals("", data.getArn());
assertEquals("my-function", data.getFunctionName());
}

@Test
public void testGetArnFunctionNameWithVersionClient() {
mockCloudApiClient();
FunctionProcessedData data = LambdaUtil.processData(data("my-function:123", null));
assertEquals("arn:aws:lambda:us-east-1:123456789012:function:my-function:123", data.getArn());
assertEquals("my-function", data.getFunctionName());
}

@Test
public void testGetArnFunctionNameAndAliasQualifier() {
FunctionProcessedData data = LambdaUtil.processData(data("my-function", "alias"));
assertEquals("", data.getArn());
assertEquals("my-function", data.getFunctionName());
}

@Test
public void testGetArnFunctionNameAndAliasQualifierClient() {
mockCloudApiClient();
FunctionProcessedData data = LambdaUtil.processData(data("my-function", "alias"));
assertEquals("arn:aws:lambda:us-east-1:123456789012:function:my-function:alias", data.getArn());
assertEquals("my-function", data.getFunctionName());
}

@Test
public void testGetArnFunctionNameAndVersionQualifier() {
FunctionProcessedData data = LambdaUtil.processData(data("my-function", "123"));
assertEquals("", data.getArn());
assertEquals("my-function", data.getFunctionName());
}

@Test
public void testGetArnFunctionNameAndVersionQualifierClient() {
mockCloudApiClient();
FunctionProcessedData data = LambdaUtil.processData(data("my-function", "123"));
assertEquals("arn:aws:lambda:us-east-1:123456789012:function:my-function:123", data.getArn());
assertEquals("my-function", data.getFunctionName());
}

@Test
public void testGetArnPartialArn() {
FunctionProcessedData data = LambdaUtil.processData(data("123456789012:function:my-function", null));
assertEquals("arn:aws:lambda:us-east-1:123456789012:function:my-function", data.getArn());
assertEquals("my-function", data.getFunctionName());
}

@Test
public void testGetArnPartialArnWithAlias() {
FunctionProcessedData data = LambdaUtil.processData(data("123456789012:function:my-function:alias", null));
assertEquals("arn:aws:lambda:us-east-1:123456789012:function:my-function:alias", data.getArn());
assertEquals("my-function", data.getFunctionName());
}

@Test
public void testGetArnPartialArnWithVersion() {
FunctionProcessedData data = LambdaUtil.processData(data("123456789012:function:my-function:123", null));
assertEquals("arn:aws:lambda:us-east-1:123456789012:function:my-function:123", data.getArn());
assertEquals("my-function", data.getFunctionName());
}

@Test
public void testGetArnPartialArnAndAliasQualifier() {
FunctionProcessedData data = LambdaUtil.processData(data("123456789012:function:my-function", "alias"));
assertEquals("arn:aws:lambda:us-east-1:123456789012:function:my-function:alias", data.getArn());
assertEquals("my-function", data.getFunctionName());
}

@Test
public void testGetArnPartialArnAndVersionQualifier() {
FunctionProcessedData data = LambdaUtil.processData(data("123456789012:function:my-function", "123"));
assertEquals("arn:aws:lambda:us-east-1:123456789012:function:my-function:123", data.getArn());
assertEquals("my-function", data.getFunctionName());
}

@Test
public void testGetArnFullArn() {
FunctionProcessedData data = LambdaUtil.processData(data("arn:aws:lambda:us-east-1:123456789012:function:my-function", null));
assertEquals("arn:aws:lambda:us-east-1:123456789012:function:my-function", data.getArn());
assertEquals("my-function", data.getFunctionName());
}

@Test
public void testGetArnFullArnWithAlias() {
FunctionProcessedData data = LambdaUtil.processData(data("arn:aws:lambda:us-east-1:123456789012:function:my-function:alias", null));
assertEquals("arn:aws:lambda:us-east-1:123456789012:function:my-function:alias", data.getArn());
assertEquals("my-function", data.getFunctionName());
}

@Test
public void testGetArnFullArnWithVersion() {
FunctionProcessedData data = LambdaUtil.processData(data("arn:aws:lambda:us-east-1:123456789012:function:my-function:123", null));
assertEquals("arn:aws:lambda:us-east-1:123456789012:function:my-function:123", data.getArn());
assertEquals("my-function", data.getFunctionName());
}

@Test
public void testGetArnFullArnAndAliasQualifier() {
FunctionProcessedData data = LambdaUtil.processData(data("arn:aws:lambda:us-east-1:123456789012:function:my-function", "alias"));
assertEquals("arn:aws:lambda:us-east-1:123456789012:function:my-function:alias", data.getArn());
assertEquals("my-function", data.getFunctionName());
}

@Test
public void testGetArnFullArnAndVersionQualifier() {
FunctionProcessedData data = LambdaUtil.processData(data("arn:aws:lambda:us-east-1:123456789012:function:my-function", "123"));
assertEquals("arn:aws:lambda:us-east-1:123456789012:function:my-function:123", data.getArn());
assertEquals("my-function", data.getFunctionName());
}

@Test
public void testGetArnDifferentRegion() {
FunctionProcessedData data = LambdaUtil.processData(data("arn:aws:lambda:us-west-2:123456789012:function:my-function", null));
assertEquals("arn:aws:lambda:us-west-2:123456789012:function:my-function", data.getArn());
assertEquals("my-function", data.getFunctionName());
}

private static void mockCloudApiClient() {
when(AgentBridge.cloud.getAccountInfo(any(), eq(CloudAccountInfo.AWS_ACCOUNT_ID)))
.thenReturn("123456789012");
}

private FunctionRawData data(String functionRef, String qualifier) {
return new FunctionRawData(functionRef, qualifier, Regions.US_EAST_1.getName(), this);
}
Expand Down
Loading

0 comments on commit 62e9dbd

Please sign in to comment.