feat(lambda): add support for log retention

packages/@aws-cdk/aws-cloudtrail/lib/index.ts

@@ -58,9 +58,9 @@ export interface CloudTrailProps {
 
   /**
    * How long to retain logs in CloudWatchLogs. Ignored if sendToCloudWatchLogs is false
-   *  @default LogRetention.OneYear
+   *  @default logs.RetentionDays.OneYear
    */
-  readonly cloudWatchLogsRetentionTimeDays?: LogRetention;
+  readonly cloudWatchLogsRetentionTimeDays?: logs.RetentionDays;
 
   /** The AWS Key Management Service (AWS KMS) key ID that you want to use to encrypt CloudTrail logs.
    * @default none
@@ -90,26 +90,6 @@ export enum ReadWriteType {
   All = "All"
 }
 
-// TODO: This belongs in a CWL L2
-export enum LogRetention {
-  OneDay = 1,
-  ThreeDays = 3,
-  FiveDays = 5,
-  OneWeek = 7,
-  TwoWeeks =  14,
-  OneMonth = 30,
-  TwoMonths = 60,
-  ThreeMonths = 90,
-  FourMonths = 120,
-  FiveMonths = 150,
-  HalfYear = 180,
-  OneYear = 365,
-  FourHundredDays = 400,
-  EighteenMonths = 545,
-  TwoYears = 731,
-  FiveYears = 1827,
-  TenYears = 3653
-}
 /**
  * Cloud trail allows you to log events that happen in your AWS account
  * For example:
@@ -145,7 +125,7 @@ export class CloudTrail extends cdk.Construct {
     let logsRole: iam.IRole | undefined;
     if (props.sendToCloudWatchLogs) {
       logGroup = new logs.CfnLogGroup(this, "LogGroup", {
-        retentionInDays: props.cloudWatchLogsRetentionTimeDays || LogRetention.OneYear
+        retentionInDays: props.cloudWatchLogsRetentionTimeDays || logs.RetentionDays.OneYear
       });
 
       logsRole = new iam.Role(this, 'LogsRole', { assumedBy: new iam.ServicePrincipal(cloudTrailPrincipal) });

packages/@aws-cdk/aws-cloudtrail/package.json

@@ -72,9 +72,10 @@
   "homepage": "https://github.com/awslabs/aws-cdk",
   "peerDependencies": {
     "@aws-cdk/aws-kms": "^0.26.0",
+    "@aws-cdk/aws-logs": "^0.26.0",
     "@aws-cdk/cdk": "^0.26.0"
   },
   "engines": {
     "node": ">= 8.10.0"
   }
-}
+}

packages/@aws-cdk/aws-cloudtrail/test/test.cloudtrail.ts

@@ -1,7 +1,8 @@
 import { expect, haveResource, not, SynthUtils } from '@aws-cdk/assert';
+import { RetentionDays } from '@aws-cdk/aws-logs';
 import { Stack } from '@aws-cdk/cdk';
 import { Test } from 'nodeunit';
-import { CloudTrail, LogRetention, ReadWriteType } from '../lib';
+import { CloudTrail, ReadWriteType } from '../lib';
 
 const ExpectedBucketPolicyProperties = {
   PolicyDocument: {
@@ -105,7 +106,7 @@ export = {
         const stack = getTestStack();
         new CloudTrail(stack, 'MyAmazingCloudTrail', {
           sendToCloudWatchLogs: true,
-          cloudWatchLogsRetentionTimeDays: LogRetention.OneWeek
+          cloudWatchLogsRetentionTimeDays: RetentionDays.OneWeek
         });
 
         expect(stack).to(haveResource("AWS::CloudTrail::Trail"));

packages/@aws-cdk/aws-lambda/lib/function.ts

@@ -1,6 +1,7 @@
 import cloudwatch = require('@aws-cdk/aws-cloudwatch');
 import ec2 = require('@aws-cdk/aws-ec2');
 import iam = require('@aws-cdk/aws-iam');
+import logs = require('@aws-cdk/aws-logs');
 import sqs = require('@aws-cdk/aws-sqs');
 import cdk = require('@aws-cdk/cdk');
 import { Code } from './code';
@@ -9,6 +10,7 @@ import { FunctionBase, FunctionImportProps, IFunction } from './function-base';
 import { Version } from './lambda-version';
 import { CfnFunction } from './lambda.generated';
 import { ILayerVersion } from './layers';
+import { LogRetention } from './log-retention';
 import { Runtime } from './runtime';
 
 /**
@@ -198,6 +200,15 @@ export interface FunctionProps {
    * You can also add event sources using `addEventSource`.
    */
   readonly events?: IEventSource[];
+
+  /**
+   * The number of days log events are kept in CloudWatch Logs. When updating
+   * this property, unsetting it doesn't remove the log retention policy. To
+   * remove the retention policy, set the value to `Infinity`.
+   *
+   * @default logs never expire
+   */
+  readonly logRetentionDays?: logs.RetentionDays;
 }
 
 /**
@@ -395,6 +406,14 @@ export class Function extends FunctionBase {
     for (const event of props.events || []) {
       this.addEventSource(event);
     }
+
+    // Log retention
+    if (props.logRetentionDays) {
+      new LogRetention(this, 'LogRetention', {
+        logGroupName: `/aws/lambda/${this.functionName}`,
+        retentionDays: props.logRetentionDays
+      });
+    }
   }
 
   /**

packages/@aws-cdk/aws-lambda/lib/index.ts

@@ -10,6 +10,7 @@ export * from './lambda-version';
 export * from './singleton-lambda';
 export * from './event-source';
 export * from './event-source-mapping';
+export * from './log-retention';
 
 // AWS::Lambda CloudFormation Resources:
 export * from './lambda.generated';

packages/@aws-cdk/aws-lambda/lib/log-retention-provider/index.ts

@@ -0,0 +1,97 @@
+// tslint:disable:no-console
+import AWS = require('aws-sdk');
+
+/**
+ * Creates a log group and doesn't throw if it exists.
+ *
+ * @param logGroupName the name of the log group to create
+ */
+async function createLogGroupSafe(logGroupName: string) {
+  try { // Try to create the log group
+    const cloudwatchlogs = new AWS.CloudWatchLogs({ apiVersion: '2014-03-28' });
+    await cloudwatchlogs.createLogGroup({ logGroupName }).promise();
+  } catch (e) {
+    if (e.code !== 'ResourceAlreadyExistsException') {
+      throw e;
+    }
+  }
+}
+
+/**
+ * Puts or deletes a retention policy on a log group.
+ *
+ * @param logGroupName the name of the log group to create
+ * @param retentionInDays the number of days to retain the log events in the specified log group.
+ */
+async function setRetentionPolicy(logGroupName: string, retentionInDays?: number) {
+  const cloudwatchlogs = new AWS.CloudWatchLogs({ apiVersion: '2014-03-28' });
+  if (!retentionInDays) {
+    await cloudwatchlogs.deleteRetentionPolicy({ logGroupName }).promise();
+  } else {
+    await cloudwatchlogs.putRetentionPolicy({ logGroupName, retentionInDays }).promise();
+  }
+}
+
+export async function handler(event: AWSLambda.CloudFormationCustomResourceEvent, context: AWSLambda.Context) {
+  try {
+    console.log(JSON.stringify(event));
+
+    // The target log group
+    const logGroupName = event.ResourceProperties.LogGroupName;
+
+    if (event.RequestType === 'Create' || event.RequestType === 'Update') {
+      // Act on the target log group
+      await createLogGroupSafe(logGroupName);
+      await setRetentionPolicy(logGroupName, parseInt(event.ResourceProperties.RetentionInDays, 10));
+
+      if (event.RequestType === 'Create') {
+        // Set a retention policy of 1 day on the logs of this function. The log
+        // group for this function should already exist at this stage because we
+        // already logged the event but due to the async nature of Lambda logging
+        // there could be a race condition. So we also try to create the log group
+        // of this function first.
+        await createLogGroupSafe(`/aws/lambda/${context.functionName}`);
+        await setRetentionPolicy(`/aws/lambda/${context.functionName}`, 1);
+      }
+    }
+
+    await respond('SUCCESS', 'OK', logGroupName);
+  } catch (e) {
+    console.log(e);
+
+    await respond('FAILED', e.message, event.ResourceProperties.LogGroupName);
+  }
+
+  function respond(responseStatus: string, reason: string, physicalResourceId: string) {
+    const responseBody = JSON.stringify({
+      Status: responseStatus,
+      Reason: reason,
+      PhysicalResourceId: physicalResourceId,
+      StackId: event.StackId,
+      RequestId: event.RequestId,
+      LogicalResourceId: event.LogicalResourceId,
+      Data: {}
+    });
+
+    console.log('Responding', responseBody);
+
+    const parsedUrl = require('url').parse(event.ResponseURL);
+    const requestOptions = {
+      hostname: parsedUrl.hostname,
+      path: parsedUrl.path,
+      method: 'PUT',
+      headers: { 'content-type': '', 'content-length': responseBody.length }
+    };
+
+    return new Promise((resolve, reject) => {
+      try {
+        const request = require('https').request(requestOptions, resolve);
+        request.on('error', reject);
+        request.write(responseBody);
+        request.end();
+      } catch (e) {
+        reject(e);
+      }
+    });
+  }
+}

packages/@aws-cdk/aws-lambda/lib/log-retention.ts

@@ -0,0 +1,64 @@
+import iam = require('@aws-cdk/aws-iam');
+import logs = require('@aws-cdk/aws-logs');
+import cdk = require('@aws-cdk/cdk');
+import path = require('path');
+import { Code } from './code';
+import { Runtime } from './runtime';
+import { SingletonFunction } from './singleton-lambda';
+
+/**
+ * Construction properties for a LogRetention.
+ */
+export interface LogRetentionProps {
+  /**
+   * The log group name.
+   */
+  readonly logGroupName: string;
+
+  /**
+   * The number of days log events are kept in CloudWatch Logs.
+   */
+  readonly retentionDays: logs.RetentionDays;
+}
+
+/**
+ * Creates a custom resource to control the retention policy of a CloudWatch Logs
+ * log group. The log group is created if it doesn't already exist. The policy
+ * is removed when `retentionDays` is `undefined` or equal to `Infinity`.
+ */
+export class LogRetention extends cdk.Construct {
+  constructor(scope: cdk.Construct, id: string, props: LogRetentionProps) {
+    super(scope, id);
+
+    // Custom resource provider
+    const provider = new SingletonFunction(this, 'Provider', {
+      code: Code.asset(path.join(__dirname, 'log-retention-provider')),
+      runtime: Runtime.NodeJS810,
+      handler: 'index.handler',
+      uuid: 'aae0aa3c-5b4d-4f87-b02d-85b201efdd8a',
+      lambdaPurpose: 'LogRetention',
+    });
+
+    if (provider.role && !provider.role.node.tryFindChild('DefaultPolicy')) { // Avoid duplicate statements
+      provider.role.addToPolicy(
+        new iam.PolicyStatement()
+          .addActions('logs:PutRetentionPolicy', 'logs:DeleteRetentionPolicy')
+          // We need '*' here because we will also put a retention policy on
+          // the log group of the provider function. Referencing it's name
+          // creates a CF circular dependency.
+          .addAllResources()
+      );
+    }
+
+    // Need to use a CfnResource here to prevent lerna dependency cycles
+    // @aws-cdk/aws-cloudformation -> @aws-cdk/aws-lambda -> @aws-cdk/aws-cloudformation
+    new cdk.CfnResource(this, 'Resource', {
+      type: 'Custom::LogRetention',
+      properties: {
+        ServiceToken: provider.functionArn,
+        LogGroupName: props.logGroupName,
+        RetentionInDays: props.retentionDays === Infinity ? undefined : props.retentionDays
+      }
+    });
+  }
+}