[Step Function] 1. Set up logging

serverless/src/index.ts

@@ -1,5 +1,6 @@
 import { getConfigFromCfnMappings, getConfigFromCfnParams, validateParameters, Configuration } from "./lambda/env";
 import { instrumentLambdas } from "./lambda/lambda";
+import { instrumentStateMachines } from "./step_function/step_function";
 import { InputEvent, OutputEvent, SUCCESS, FAILURE } from "./types";
 import log from "loglevel";
 
@@ -37,6 +38,11 @@ export const handler = async (event: InputEvent, _: any): Promise<OutputEvent> =
       return lambdaOutput;
     }
 
+    const stepFunctionOutput = await instrumentStateMachines(event);
+    if (stepFunctionOutput.status === FAILURE) {
+      return stepFunctionOutput;
+    }
+
     return {
       requestId: event.requestId,
       status: SUCCESS,

serverless/src/step_function/env.ts

@@ -0,0 +1,50 @@
+import log from "loglevel";
+
+export interface Configuration {
+  // When set, it will be added to the state machine's log group name.
+  env?: string;
+}
+
+const envEnvVar = "DD_ENV";
+
+// Same interface as Configuration above, except all parameters are optional, since user does
+// not have to provide the values (in which case we will use the default configuration below).
+interface CfnParams extends Partial<Configuration> {}
+
+export const defaultConfiguration: Configuration = {};
+
+/**
+ * Returns the default configuration with any values overwritten by environment variables.
+ */
+export function getConfigFromEnvVars(): Configuration {
+  const config: Configuration = {
+    ...defaultConfiguration,
+  };
+
+  if (envEnvVar in process.env) {
+    config.env = process.env[envEnvVar];
+  }
+
+  return config;
+}
+
+/**
+ * Takes a set of parameters from the CloudFormation template. This could come from either
+ * the Mappings section of the template, or directly from the Parameters under the transform/macro
+ * as the 'params' property under the original InputEvent to the handler in src/index.ts
+ *
+ * Uses these parameters as the Datadog configuration, and for values that are required in the
+ * configuration but not provided in the parameters, uses the default values from
+ * the defaultConfiguration above.
+ */
+export function getConfigFromCfnParams(params: CfnParams) {
+  let datadogConfig = params as Partial<Configuration> | undefined;
+  if (datadogConfig === undefined) {
+    log.debug("No Datadog config found, using the default config");
+    datadogConfig = {};
+  }
+  return {
+    ...getConfigFromEnvVars(),
+    ...datadogConfig,
+  };
+}

serverless/src/step_function/log.ts

@@ -0,0 +1,146 @@
+import { Resources } from "../types";
+import log from "loglevel";
+import { StateMachine } from "step_function/types";
+
+const unsupportedCaseErrorMessage =
+  "Step Function Instrumentation is not supported. \
+Please open a feature request in https://github.com/DataDog/datadog-cdk-constructs.";
+
+const FN_SUB = "Fn::Sub";
+const FN_GET_ATT = "Fn::GetAtt";
+
+/**
+ * Set up logging for the given state machine:
+ * 1. Set log level to ALL
+ * 2. Set includeExecutionData to true
+ * 3. Create a destination log group (if not set already)
+ * 4. Add permissions to the state machine role to log to CloudWatch Logs
+ */
+export function setUpLogging(resources: Resources, stateMachine: StateMachine): void {
+  log.debug(`Setting up logging`);
+  if (!stateMachine.properties.LoggingConfiguration) {
+    stateMachine.properties.LoggingConfiguration = {};
+  }
+
+  const logConfig = stateMachine.properties.LoggingConfiguration;
+
+  logConfig.Level = "ALL";
+  logConfig.IncludeExecutionData = true;
+  if (!logConfig.Destinations) {
+    log.debug(`Log destination not found, creating one`);
+    const logGroupKey = createLogGroup(resources, stateMachine);
+    logConfig.Destinations = [
+      {
+        CloudWatchLogsLogGroup: {
+          LogGroupArn: {
+            "Fn::GetAtt": [logGroupKey, "Arn"],
+          },
+        },
+      },
+    ];
+  } else {
+    log.debug(`Log destination already exists, skipping creating one`);
+  }
+}
+
+function createLogGroup(resources: Resources, stateMachine: StateMachine): string {
+  const logGroupKey = `${stateMachine.resourceKey}LogGroup`;
+  resources[logGroupKey] = {
+    Type: "AWS::Logs::LogGroup",
+    Properties: {
+      LogGroupName: buildLogGroupName(stateMachine, undefined),
+      RetentionInDays: 7,
+    },
+  };
+
+  let role;
+  if (stateMachine.properties.RoleArn) {
+    log.debug(`A role is already defined. Parsing its resource key from the roleArn.`);
+    const roleArn = stateMachine.properties.RoleArn;
+
+    if (typeof roleArn !== "object") {
+      throw new Error(`RoleArn is not an object. ${unsupportedCaseErrorMessage}`);
+    }
+
+    let roleKey;
+    if (roleArn[FN_GET_ATT]) {
+      // e.g.
+      //   Fn::GetAtt: [MyStateMachineRole, "Arn"]
+      roleKey = roleArn[FN_GET_ATT][0];
+    } else if (roleArn[FN_SUB]) {
+      // e.g.
+      //   Fn::Sub: ${StatesExecutionRole.Arn}
+      const arnMatch = roleArn[FN_SUB].match(/^\${(.*)\.Arn}$/);
+      if (arnMatch) {
+        roleKey = arnMatch[1];
+      } else {
+        throw new Error(`Unsupported Fn::Sub format: ${roleArn[FN_SUB]}. ${unsupportedCaseErrorMessage}`);
+      }
+    } else {
+      throw new Error(`Unsupported RoleArn format: ${roleArn}. ${unsupportedCaseErrorMessage}`);
+    }
+    log.debug(`Found State Machine role Key: ${roleKey}`);
+    role = resources[roleKey];
+  } else {
+    log.debug(`No role is defined. Creating one.`);
+    const roleKey = `${stateMachine.resourceKey}Role`;
+    role = {
+      Type: "AWS::IAM::Role",
+      Properties: {
+        AssumeRolePolicyDocument: {
+          Version: "2012-10-17",
+          Statement: [
+            {
+              Effect: "Allow",
+              Principal: {
+                Service: "states.amazonaws.com",
+              },
+              Action: "sts:AssumeRole",
+            },
+          ],
+        },
+        Policies: [],
+      },
+    };
+    resources[roleKey] = role;
+  }
+
+  log.debug(`Add a policy to the role to grant permissions to the log group`);
+  if (!role.Properties.Policies) {
+    role.Properties.Policies = [];
+  }
+  role.Properties.Policies.push({
+    PolicyName: `${stateMachine.resourceKey}LogPolicy`,
+    PolicyDocument: {
+      Version: "2012-10-17",
+      Statement: [
+        {
+          Effect: "Allow",
+          Action: [
+            "logs:CreateLogDelivery",
+            "logs:CreateLogStream",
+            "logs:GetLogDelivery",
+            "logs:UpdateLogDelivery",
+            "logs:DeleteLogDelivery",
+            "logs:ListLogDeliveries",
+            "logs:PutLogEvents",
+            "logs:PutResourcePolicy",
+            "logs:DescribeResourcePolicies",
+            "logs:DescribeLogGroups",
+          ],
+          Resource: "*",
+        },
+      ],
+    },
+  });
+  return logGroupKey;
+}
+
+/**
+ * Builds log group name for a state machine.
+ * @returns log group name like "/aws/vendedlogs/states/MyStateMachine-Logs" (without env)
+ *                           or "/aws/vendedlogs/states/MyStateMachine-Logs-dev" (with env)
+ */
+export const buildLogGroupName = (stateMachine: StateMachine, env: string | undefined): string => {
+  return `/aws/vendedlogs/states/${stateMachine.resourceKey}-Logs${env !== undefined ? "-" + env : ""}`;
+};

serverless/src/step_function/step_function.ts

@@ -0,0 +1,45 @@
+import { InputEvent, OutputEvent, SUCCESS, Resources } from "../types";
+import log from "loglevel";
+import { StateMachine, StateMachineProperties } from "../../src/step_function/types";
+import { setUpLogging } from "../../src/step_function/log";
+
+const STATE_MACHINE_RESOURCE_TYPE = "AWS::StepFunctions::StateMachine";
+
+export async function instrumentStateMachines(event: InputEvent): Promise<OutputEvent> {
+  const fragment = event.fragment;
+  const resources = fragment.Resources;
+
+  const stateMachines = findStateMachines(resources);
+  for (const stateMachine of stateMachines) {
+    instrumentStateMachine(resources, stateMachine);
+  }
+
+  return {
+    requestId: event.requestId,
+    status: SUCCESS,
+    fragment,
+  };
+}
+
+function instrumentStateMachine(resources: Resources, stateMachine: StateMachine): void {
+  log.debug(`Instrumenting State Machine ${stateMachine.resourceKey}`);
+  setUpLogging(resources, stateMachine);
+}
+
+export function findStateMachines(resources: Resources): StateMachine[] {
+  return Object.entries(resources)
+    .map(([key, resource]) => {
+      if (resource.Type !== STATE_MACHINE_RESOURCE_TYPE) {
+        log.debug(`Resource ${key} is not a State Machine, skipping...`);
+        return;
+      }
+
+      const properties: StateMachineProperties = resource.Properties;
+
+      return {
+        properties: properties,
+        resourceKey: key,
+      } as StateMachine;
+    })
+    .filter((resource) => resource !== undefined) as StateMachine[];
+}

serverless/src/step_function/types.ts

@@ -0,0 +1,27 @@
+export interface StateMachine {
+  properties: StateMachineProperties;
+  resourceKey: string;
+}
+
+export interface StateMachineProperties {
+  LoggingConfiguration?: LoggingConfiguration;
+  RoleArn?: string | { [key: string]: any };
+}
+
+export interface LoggingConfiguration {
+  Destinations?: LogDestination[];
+  IncludeExecutionData?: boolean;
+  Level?: string;
+}
+
+export interface LogDestination {
+  CloudWatchLogsLogGroup: CloudWatchLogsLogGroup;
+}
+
+export interface CloudWatchLogsLogGroup {
+  LogGroupArn:
+    | string
+    | {
+        "Fn::GetAtt": string[];
+      };
+}

serverless/test/step_function/helper.ts

@@ -0,0 +1,20 @@
+export function getEmptyStateMachineRole() {
+  return {
+    Type: "AWS::IAM::Role",
+    Properties: {
+      AssumeRolePolicyDocument: {
+        Version: "2012-10-17",
+        Statement: [
+          {
+            Effect: "Allow",
+            Principal: {
+              Service: "states.amazonaws.com",
+            },
+            Action: "sts:AssumeRole",
+          },
+        ],
+      },
+      Policies: [],
+    },
+  };
+}