Add extension to Lambda handlers using custom runtimes (#219)

* Add extension for custom runtimes

* Add unit tests

* Add integration test for provided_al2 runtime

* Add integration test for provided_al2 ARM handler

v2/integration_tests/stacks/lambda-provided-arm-stack.ts

@@ -0,0 +1,49 @@
+import * as lambda from "aws-cdk-lib/aws-lambda";
+import * as s3 from "aws-cdk-lib/aws-s3";
+import { LambdaRestApi, LogGroupLogDestination } from "aws-cdk-lib/aws-apigateway";
+import { LogGroup } from "aws-cdk-lib/aws-logs";
+import { Stack, StackProps, App } from "aws-cdk-lib";
+import { Datadog } from "../../src/index";
+
+export class ExampleStack extends Stack {
+  constructor(scope: App, id: string, props?: StackProps) {
+    super(scope, id, props);
+
+    const s3Bucket = new s3.Bucket(this, "exampleBucket", {
+      objectOwnership: s3.ObjectOwnership.BUCKET_OWNER_ENFORCED,
+      blockPublicAccess: s3.BlockPublicAccess.BLOCK_ALL,
+    });
+
+    const providedLambda = new lambda.Function(this, "HelloHandler", {
+      runtime: lambda.Runtime.PROVIDED_AL2,
+      code: lambda.Code.fromBucket(s3Bucket, "fake-key-for-test"),
+      handler: "handler.handler",
+      architecture: lambda.Architecture.ARM_64,
+    });
+
+    s3Bucket.grantRead(providedLambda);
+
+    const restLogGroup = new LogGroup(this, "restLogGroup");
+    new LambdaRestApi(this, "rest-test", {
+      handler: providedLambda,
+      deployOptions: {
+        accessLogDestination: new LogGroupLogDestination(restLogGroup),
+      },
+    });
+
+    const datadogCDK = new Datadog(this, "Datadog", {
+      extensionLayerVersion: 49,
+      apiKey: "1234",
+      site: "datadoghq.com",
+      sourceCodeIntegration: false,
+    });
+    datadogCDK.addLambdaFunctions([providedLambda]);
+    datadogCDK.addForwarderToNonLambdaLogGroups([restLogGroup]);
+  }
+}
+
+const app = new App();
+const env = { account: "601427279990", region: "sa-east-1" };
+const stack = new ExampleStack(app, "lambda-provided-arm-stack", { env: env });
+console.log("Stack name: " + stack.stackName);
+app.synth();

v2/integration_tests/stacks/lambda-provided-stack.ts

@@ -0,0 +1,48 @@
+import * as lambda from "aws-cdk-lib/aws-lambda";
+import * as s3 from "aws-cdk-lib/aws-s3";
+import { LambdaRestApi, LogGroupLogDestination } from "aws-cdk-lib/aws-apigateway";
+import { LogGroup } from "aws-cdk-lib/aws-logs";
+import { Stack, StackProps, App } from "aws-cdk-lib";
+import { Datadog } from "../../src/index";
+
+export class ExampleStack extends Stack {
+  constructor(scope: App, id: string, props?: StackProps) {
+    super(scope, id, props);
+
+    const s3Bucket = new s3.Bucket(this, "exampleBucket", {
+      objectOwnership: s3.ObjectOwnership.BUCKET_OWNER_ENFORCED,
+      blockPublicAccess: s3.BlockPublicAccess.BLOCK_ALL,
+    });
+
+    const providedLambda = new lambda.Function(this, "HelloHandler", {
+      runtime: lambda.Runtime.PROVIDED_AL2,
+      code: lambda.Code.fromBucket(s3Bucket, "fake-key-for-test"),
+      handler: "handler.handler",
+    });
+
+    s3Bucket.grantRead(providedLambda);
+
+    const restLogGroup = new LogGroup(this, "restLogGroup");
+    new LambdaRestApi(this, "rest-test", {
+      handler: providedLambda,
+      deployOptions: {
+        accessLogDestination: new LogGroupLogDestination(restLogGroup),
+      },
+    });
+
+    const datadogCDK = new Datadog(this, "Datadog", {
+      extensionLayerVersion: 49,
+      apiKey: "1234",
+      site: "datadoghq.com",
+      sourceCodeIntegration: false,
+    });
+    datadogCDK.addLambdaFunctions([providedLambda]);
+    datadogCDK.addForwarderToNonLambdaLogGroups([restLogGroup]);
+  }
+}
+
+const app = new App();
+const env = { account: "601427279990", region: "sa-east-1" };
+const stack = new ExampleStack(app, "lambda-provided-stack", { env: env });
+console.log("Stack name: " + stack.stackName);
+app.synth();

v2/scripts/run_integration_tests.sh

@@ -11,7 +11,7 @@ set -e
 # To add new tests create a new ts file in the 'integration_tests/stacks' directory, append its file name to the STACK_CONFIGS array.
 # Note: Each ts file will have its respective snapshot built in the snapshots directory, e.g. lambda-function-stack.ts
 #       will generate both snapshots/test-lambda-function-stack-snapshot.json and snapshots/correct-lambda-function-stack-snapshot.json
-STACK_CONFIGS=("lambda-singleton-function-stack" "lambda-function-arm-stack" "lambda-function-stack" "lambda-nodejs-function-stack" "lambda-python-function-stack" "lambda-java-function-stack")
+STACK_CONFIGS=("lambda-provided-stack" "lambda-provided-arm-stack" "lambda-singleton-function-stack" "lambda-function-arm-stack" "lambda-function-stack" "lambda-nodejs-function-stack" "lambda-python-function-stack" "lambda-java-function-stack")
 
 SCRIPT_PATH=${BASH_SOURCE[0]}
 SCRIPTS_DIR=$(dirname $SCRIPT_PATH)

v2/src/constants.ts

@@ -22,6 +22,7 @@ export enum RuntimeType {
   NODE,
   PYTHON,
   JAVA,
+  CUSTOM,
   UNSUPPORTED,
 }
 
@@ -60,6 +61,8 @@ export const runtimeLookup: { [key: string]: RuntimeType } = {
   "java8.al2": RuntimeType.JAVA,
   "java11": RuntimeType.JAVA,
   "java17": RuntimeType.JAVA,
+  "provided": RuntimeType.CUSTOM,
+  "provided.al2": RuntimeType.CUSTOM,
 };
 
 export const runtimeToLayerName: { [key: string]: string } = {

v2/src/redirect.ts

@@ -62,6 +62,8 @@ function getDDHandler(lambdaRuntime: RuntimeType, addLayers: boolean) {
       return addLayers ? JS_HANDLER_WITH_LAYERS : JS_HANDLER;
     case RuntimeType.PYTHON:
       return PYTHON_HANDLER;
+    case RuntimeType.CUSTOM:
+      return null;
     case RuntimeType.JAVA:
       return null;
   }

v2/test/layer.spec.ts

@@ -262,6 +262,63 @@ describe("applyLayers", () => {
     });
   });
 
+  it("adds extension layer to provided runtime", () => {
+    const app = new App();
+    const stack = new Stack(app, "stack", {
+      env: {
+        region: "us-west-2",
+      },
+    });
+    const hello = new lambda.Function(stack, "HelloHandler", {
+      runtime: lambda.Runtime.PROVIDED_AL2,
+      code: lambda.Code.fromAsset("test"),
+      handler: "hello.handler",
+    });
+    const datadogCdk = new Datadog(stack, "Datadog", {
+      extensionLayerVersion: EXTENSION_LAYER_VERSION,
+      apiKmsKey: "1234",
+      addLayers: true,
+      enableDatadogTracing: false,
+      flushMetricsToLogs: true,
+      site: "datadoghq.com",
+    });
+    datadogCdk.addLambdaFunctions([hello]);
+
+    Template.fromStack(stack).hasResourceProperties("AWS::Lambda::Function", {
+      Layers: [`arn:aws:lambda:${stack.region}:${DD_ACCOUNT_ID}:layer:Datadog-Extension:${EXTENSION_LAYER_VERSION}`],
+    });
+  });
+
+  it("adds extension ARM layer to provided runtime using ARM", () => {
+    const app = new App();
+    const stack = new Stack(app, "stack", {
+      env: {
+        region: "us-west-2",
+      },
+    });
+    const hello = new lambda.Function(stack, "HelloHandler", {
+      runtime: lambda.Runtime.PROVIDED_AL2,
+      code: lambda.Code.fromAsset("test"),
+      handler: "hello.handler",
+      architecture: lambda.Architecture.ARM_64,
+    });
+    const datadogCdk = new Datadog(stack, "Datadog", {
+      extensionLayerVersion: EXTENSION_LAYER_VERSION,
+      apiKmsKey: "1234",
+      addLayers: true,
+      enableDatadogTracing: false,
+      flushMetricsToLogs: true,
+      site: "datadoghq.com",
+    });
+    datadogCdk.addLambdaFunctions([hello]);
+
+    Template.fromStack(stack).hasResourceProperties("AWS::Lambda::Function", {
+      Layers: [
+        `arn:aws:lambda:${stack.region}:${DD_ACCOUNT_ID}:layer:Datadog-Extension-ARM:${EXTENSION_LAYER_VERSION}`,
+      ],
+    });
+  });
+
   it("doesn't add layer to container image Lambda with extension and layer versions", () => {
     const app = new App();
     const stack = new Stack(app, "stack", {