DEV Community: Lukas Fruntke

Best practices for using AWS StepFunctions

Lukas Fruntke — Sat, 22 Jan 2022 23:30:21 +0000

In this post you will learn some of the best patterns/tricks I have learned during my time creating Step Functions workflows, while working with clients at Accenture.

Table of contents (clickable)

Make use of the Step Functions Workflow Studio
Utilize the service integrations
Use .waitForTaskToken
Make use of the inbuilt retries
Utilize Heartbeats to fail fast
Define a Catch Handler
Further reading

Make use of the Step Functions Workflow Studio

Since it has been introduced in June Step Functions Workflow Studio proved its value to me several times. As its a low-code editor with the most common configurations already baked-in, the effort of writing/designing a workflow with Amazon States Language plummeted. Those minutes and hours handwriting workflows with 100 lines++ are finally over for good, which is something you must not ignore when dealing with Step Functions.

As visualized below, the option to create a workflow from scratch in Workflow Studio directly is already present:

As is the option to edit pre-existing workflows directly in Step Functions:

Utilize the service integrations

While AWS offers some 17 "optimized" service integrations (for the definitive list see here), that include different custom options of integrating with the specific services, AWS has released an option to call the APIs of nearly all AWS services directly, as described in this article. This allows you to scrap some of the utility lambdas one uses to add much-needed functionality-augmentation to a service and go with Step Function instead.

Use `.waitForTaskToken`

By using .waitForTaskToken, you are able to transparently pause the workflow, until a task like a lambda function has finished executing.

Be aware, that you need to specify the Task Token in the payload for the lambda, as Step Function does not inject it automatically for you.

Example

code example
This example shows how to send the task token for success/failure back to Step Functions via AWS JS SDK v3.

import {
  SFNClient,
  SendTaskFailureCommand,
  SendTaskSuccessCommand
} from "@aws-sdk/client-sfn";
const client = new SFNClient();

async function success(taskToken, input) {
  const stepFunctionsCommand = new SendTaskSuccessCommand({
    taskToken,
    output: input
  });
  await client.send(stepFunctionsCommand);
}

async function failure(taskToken, cause, error) {
  const stepFunctionsCommand = new SendTaskFailureCommand({
    taskToken,
    cause,
    error
  });
  await client.send(stepFunctionsCommand);
}

async function main(event) {
  try {
    await success(event.MyTaskToken)
  } catch (error) {
    console.error(error)
    const {
      requestId,
      cfId,
      extendedRequestId
    } = error.$metadata;
    await failure(event.MyTaskToken, {
      requestId,
      cfId,
      extendedRequestId
    })
  }

}

Utilize Heartbeats to fail fast

As StepFunctions can run for up to a year (at least Standard workflows) it is imperative to avoid stuck executions. One way of doing this when integrating with Lambda is the Heartbeat API and specification. This allows developers to specify a max-interval, in which the Heartbeat has to be send back to Step Functions. Failure to meet this deadline leads to termination of the task.

Example

code example
In the following example a task is specified with a max-heartbeat duration of 10 minutes.

{
      "Type": "Task",
      "Resource": "arn:aws:states:::lambda:invoke",
      "HeartbeatSeconds": 600
}

This code-example shows how to send back a heartbeat with AWS JS SDK v3.

import {
  SFNClient,
  SendTaskHeartbeatCommand
} from "@aws-sdk/client-sfn";
const client = new SFNClient();

async function heartbeat(taskToken) {
  const stepFunctionsCommand = new SendTaskHeartbeatCommand({
    taskToken
  });
  await client.send(stepFunctionsCommand);
}

async function main(event) {
    await heartbeat(event.MyTaskToken)

    // some expensive calculation

    await heartbeat(event.MyTaskToken)
}

Define a Catch Handler

Rationale

To quote Werner Vogels, Amazon CTO:

everything fails, all the time

Be prepared for the wildly different and sometimes unexpected errors the AWS APIs can throw, by catching them like you would in a lambda function (if you don't do that we are having a whoole different conversation).

Example

In this pretty simple example, the catch block is used on a lambda task. It works on the other tasks as well. This example uses a catch-all error code, for other error codes see here.

As it is usually a good idea to get notified when an error occurs, this example publishes to a SNS topic, which may have an email subscription. I'd reccomend to use this only for really critical errors, as you may otherwise miss important errors in your then-cluttered inbox.

code example

{
  "Invoke Lambda Task": {
    // [..]
    "Catch": [
      {
        "ErrorEquals": [
          "States.ALL"
        ],
        "Next": "SNS Publish"
      }
    ]
    // [..]
  },
  "SNS Publish": {
    "Type": "Task",
    "Resource": "arn:aws:states:::sns:publish",
    "Parameters": {
      "Message.$": "$"
    },
    "Next": "Fail"
  },
  "Fail": {
    "Type": "Fail"
  }
}

Make use of the inbuilt retries

Instead of catching errors and terminating the flow then, one might as well use retries to try recovery from an error or to await a desired state. One may use this for example to await results from APIs similar to those of a Glue Crawler, which need repeated polling and potentially exponential backoff.

Example

code example

In this example a retry from specific Lambda exceptions is shown. The IntervalSeconds parameter defines an initial offset, which has to pass before the first retry is attempted. The BackoffRate parameter specifies the duration-multiplier which is applied after each unsuccessful attempt. Step Functions will retry after 2,4,8,16,32,64 seconds, limited by the MaxAttempts parameter of 6.

{
  "Invoke Lambda Task": {
    // [..]
    "Retry": [
      {
        "ErrorEquals": [
          "Lambda.ServiceException",
          "Lambda.AWSLambdaException",
          "Lambda.SdkClientException"
        ],
        "IntervalSeconds": 2,
        "MaxAttempts": 6,
        "BackoffRate": 2
      }
    ]
    // [..]
  }
}

Creating an AWS Private Certificate Root Authority with Lambda and Node.js

Lukas Fruntke — Fri, 29 Nov 2019 09:11:31 +0000

According to the Security Perspective of the AWS Cloud Adoption Framework, Data has to be safeguarded during its transit.
While it is a common practice to terminate the HTTPS traffic at the Application Load Balancer and forward it to the Application via HTTP, this does not ensure a continuous encryption of possible sensitive data.

When implementing End-to-End encryption in order to protect the data in transit between the Application Load Balancer and the Application, two alternatives need to be considered:

Passing the HTTPS traffic to the Application, where the private key is available too. This has the downside, that the AWS Certificate Manager does not allow the export of private keys, therefore another solution would be needed to store the private key.
Using the AWS Certificate Authority as a Root Authority to sign own certificates for the communication between Application Load Balancer and the Application. Although there is quite a pricetag attached to this solution - 400$ per month per running Certificate Authority and 0,75$ per certificate for the fist thousand certificates - it is AWS native and does not need any creative way of storing the private certificate.

In this blog post the second alternative was chosen - using the Private Certificate Authority to generate the certificates for the communication between ALB and the Application.

To automate the process of creating the Private Certificate Authority and issuing a certificate, a Lambda function written in JavaScript is utilized here. In the process, a mixture of AWS PCA API calls and OpenSSL calls are used. As AWS Lambda removed the OpenSSL binaries from recent Node.JS Lambda Runtimes, Node.JS 8.10 has to be used, which will reach its End of Life at the 31st December 2019. Therefore, the OpenSSL binary will need to be added as a Lambda Layer if the function is used in 2020. This is pretty straightforward, one just needs to spin up an EC2 instance and zip the executable with the correct permissions and in the folder bin, upload it as a Layer and reconfigure the function to use it.

Creating a PCA with the aws-sdk is pretty straightforward:

const AWS = require('aws-sdk');
const pca = new AWS.ACMPCA();

async function createCA() {
    const caParams = {
      CertificateAuthorityConfiguration: {
        KeyAlgorithm: "RSA_2048",
        SigningAlgorithm: "SHA256WITHRSA",
        Subject: {
          Country: 'DE',
          Organization: 'SPIRIT21',
        }
      },
      CertificateAuthorityType: "ROOT",
      RevocationConfiguration: {
        CrlConfiguration: {
          Enabled: false,
        }
      }
    }
    const {CertificateAuthorityArn} = await pca.createCertificateAuthority(caParams).promise();
}

It takes some time, until the creation is done, so the waitFor() method of the SDK has to be used, in order to await the completion of the PCA's creation like this:

const AWS = require('aws-sdk');
const pca = new AWS.ACMPCA();

async function wait(CertificateAuthorityArn){
    await pca.waitFor('certificateAuthorityCSRCreated', { CertificateAuthorityArn }).promise();
}

When the PCA is created, it should be visible in the console:

As displayed, a CA Certificate needs to be installed, before the PCA is ready to use. To sign the PCA certificate, the Certificate Signing Request of the Certificate Authority is required, which can be retrieved via the AWS SDK:

const AWS = require('aws-sdk');
const pca = new AWS.ACMPCA();

async function getCSR(CertificateAuthorityArn){
    const { CSR } = await pca.getCertificateAuthorityCsr({ CertificateAuthorityArn }).promise();
    return CSR;
}

To sign the Root CA Certificate, a issue request of a root CA Certificate has to be performed against the Certificate Authority. Previous attempts of issuing the certificate with OpenSSL failed, because AWS would not accept the generated certificate as a CA Root Certificate, hence the issuing is done via the API. The Certificate Authority needs the Certificate Signing Request and a few other parameters arranged like this:

const AWS = require('aws-sdk');
const pca = new AWS.ACMPCA();

async function issueRootCertificate(CertificateAuthorityArn, CSR) {
  const CACertParams = {
    CertificateAuthorityArn,
    Csr: Buffer.from(CSR),
    SigningAlgorithm: "SHA256WITHRSA",
    TemplateArn: "arn:aws:acm-pca:::template/RootCACertificate/V1",
    Validity: {
      Type: "YEARS",
      Value: 10
    }
  }

  const {CertificateArn} = await pca.issueCertificate(CACertParams).promise();
}

After the Certificate is issued, it has to be imported in the CA:

const AWS = require('aws-sdk');
const pca = new AWS.ACMPCA();

async function importRootCertificate(CertificateAuthorityArn, CertificateArn) {
  const {Certificate} = await pca.getCertificate({
    CertificateAuthorityArn,
    CertificateArn
  }).promise();

  return await pca.importCertificateAuthorityCertificate({
    CertificateAuthorityArn,
    Certificate: Buffer.from(Certificate)
  }).promise();
}

Now the CA should be ready for usage, which should be visible in the console:

Now, the CA is finally ready to issue certificates, which can be used for encrypting traffic. A certificate is issued like this (with the help of OpenSSL for generating the Certificate Signing Request):

const AWS = require('aws-sdk');
const pca = new AWS.ACMPCA();
const util = require('util');
const exec = util.promisify(require('child_process').exec);
const read = util.promisify(require('fs').readFile);
const write = util.promisify(require('fs').writeFile);

const privateKeyFile = "/tmp/private.key"
const CSRFile = "/tmp/CSR.csr";

// This is important for OpenSSL, otherwise it would exit with an error, because the .rnd File in the old Home dir is not writeable
process.env.HOME = "/tmp";

async function issueCertificate(CertificateAuthorityArn) {

  await exec(`openssl req -nodes -newkey rsa:4096 -days 3600 -keyout ${privateKeyFile} -out ${CSRFile} -subj "/C=DE/O=SPIRIT21/CN=ExampleInternalCA"`)

  const csr = await read(CSRFile);

  const certParams = {
    CertificateAuthorityArn,
    Csr: Buffer.from(csr),
    SigningAlgorithm: "SHA256WITHRSA",
    Validity: {
      Type: "DAYS",
      Value: 3600
    }
  }
  const certData = await pca.issueCertificate(certParams).promise();

  // Sometimes the CA isn't finished with issuing the cert, 
  // which is why we have to wait here, before getting the cert
  await sleep(500);

  const cert = await pca.getCertificate({
    CertificateArn: certData.CertificateArn,
    CertificateAuthorityArn
  }).promise();

  return {
    CertificateArn: certData.CertificateArn,
    Certificate: Buffer.from(cert.Certificate).toString("base64")
  };
}

When everything is tied together and packaged as a handler it might look like this:

const AWS = require('aws-sdk');
const pca = new AWS.ACMPCA();
const util = require('util');
const exec = util.promisify(require('child_process').exec);
const read = util.promisify(require('fs').readFile);
const write = util.promisify(require('fs').writeFile);
const exists = require('fs').existsSync;

const privateKeyFile = "/tmp/private.key"
const CSRFile = "/tmp/CSR.csr";
process.env.HOME = "/tmp";

const caParams = {
  CertificateAuthorityConfiguration: {
    KeyAlgorithm: "RSA_2048",
    SigningAlgorithm: "SHA256WITHRSA",
    Subject: {
      Country: 'DE',
      Organization: 'SPIRIT21',
    }
  },
  CertificateAuthorityType: "ROOT",
  RevocationConfiguration: {
    CrlConfiguration: {
      Enabled: false,
    }
  }
}


async function testPCA(arn) {
  var params = {
    CertificateAuthorityArn: arn
  };
  try {
    await pca.getCertificateAuthorityCsr(params);
    return true;
  } catch (e) {
    return false;
  }

}


const sleep = m => new Promise(r => setTimeout(r, m))

async function handler(event, context) {
  try {
    let CertificateAuthorityArn = "";

    if (event.hasOwnProperty("arn")) {
      CertificateAuthorityArn = event.arn;
    }


    if (!await testPCA(CertificateAuthorityArn)) {
      console.log('Generating PCA', caParams);
      const {
        CertificateAuthorityArn
      } = await pca.createCertificateAuthority(caParams).promise();
      console.log(CertificateAuthorityArn);

      console.log("Waiting for the CSR creation..");
      await pca.waitFor('certificateAuthorityCSRCreated', {
        CertificateAuthorityArn
      }).promise();
      console.log("Getting CA-CSR now...");

      const {
        Csr
      } = await pca.getCertificateAuthorityCsr({
        CertificateAuthorityArn
      }).promise();
      console.log('CA-CSR loaded, generating Root CA Cert');

      const CACertParams = {
        CertificateAuthorityArn,
        Csr: Buffer.from(Csr),
        SigningAlgorithm: "SHA256WITHRSA",
        TemplateArn: "arn:aws:acm-pca:::template/RootCACertificate/V1",
        Validity: {
          Type: "YEARS",
          Value: 10
        }
      }

      const {
        CertificateArn
      } = await pca.issueCertificate(CACertParams).promise();
      console.log("Root CA Cert generated");

      // Sometimes the CA is not done with issuing the cert, which is why we have to wait here, before getting the cert
      await sleep(500);

      const CAcert = await pca.getCertificate({
        CertificateAuthorityArn,
        CertificateArn
      }).promise();
      console.log(CAcert);

      await pca.importCertificateAuthorityCertificate({
        CertificateAuthorityArn,
        Certificate: Buffer.from(CAcert.Certificate)
      }).promise();
      console.log("Root CA Cert imported");
    }

    // END CA GENERATION


    // CERTIFICATE GENERATION

    console.log("Generating CSR for new CA Cert");
    await exec(`openssl req -nodes -newkey rsa:4096 -days 3600 -keyout ${privateKeyFile} -out ${CSRFile} -subj "/C=DE/O=SPIRIT21/CN=ExampleInternalCA-Root"`)

    const csr = await read(CSRFile);

    const certParams = {
      CertificateAuthorityArn,
      Csr: Buffer.from(csr),
      SigningAlgorithm: "SHA256WITHRSA",
      Validity: {
        Type: "DAYS",
        Value: 3600
      }
    }

    console.log("Generating Cert in CA");
    const certData = await pca.issueCertificate(certParams).promise();

    // Again, the CA might not be ready.
    await sleep(500);

    const cert = await pca.getCertificate({
      CertificateArn: certData.CertificateArn,
      CertificateAuthorityArn
    }).promise();
    console.log(cert);

    return {
      CertificateArn: certData.CertificateArn,
      Certificate: Buffer.from(cert.Certificate).toString("base64")
    };

  } catch (e) {
    console.error(e);
  }

}

module.exports = {
  handler
};

Pay attention to setting the Lambda timeout on a greater value than 10 seconds, which was the mean execution time during testing. Also do not forget to set the Runtime to Node.js 8.10 or use a Lambda Layer with OpenSSL. Contrary to what one might expect the certificates issued by the Private Certificate Authority are not visible in the normal AWS Certificate Manager, so it is important to store the ARNs of the Certificates created too.

DEV Community: Lukas Fruntke

Best practices for using AWS StepFunctions

Table of contents (clickable)

Make use of the Step Functions Workflow Studio

Utilize the service integrations

Use .waitForTaskToken

Example

Utilize Heartbeats to fail fast

Example

Define a Catch Handler

Rationale

Example

Make use of the inbuilt retries

Example

Further reading

Creating an AWS Private Certificate Root Authority with Lambda and Node.js

Use `.waitForTaskToken`