DEV Community: Duarte Nunes

Visual Testing with Chromatic

Duarte Nunes — Mon, 28 Mar 2022 21:56:58 +0000

Have you ever pushed a commit into prod, only to notice later it made your site look straight out of a Hieronymus Bosch panel? If so, we’ve got you covered! This article will explain how introducing visual regression testing with Chromatic into your workflow will help you avoid unintended UI changes.

What is visual testing?

When developing a user interface, there are two important targets for testing: the behavior of the graphical elements, and how they are presented and arranged. The former is usually achieved by unit and E2E tests, while for the latter it’s common to leverage snapshot tests. Snapshot tests work by comparing the output of a test against a version-controlled golden file, failing on a mismatch. Intentional changes include an update to that golden file.

Tools like Jest make it easy to create non-visual snapshots based on the markup of an interface. These tests are useful to alert PR authors that they may be introducing unintended changes, but it makes it hard on reviewers to validate and approve the intended ones: it’s not an easy task to mentally conjure up the exact visual changes from looking at HTML alone. Developers reviewing a changeset need to spin up two versions of the UI and manually track down the changes. If the UI under test can have many states and variations, this can easily turn into a long and laborious task: maybe it was the layout of an error message that changed, or the position of a spinner rendered when the component is loading data. In the context of web development, having a published, trunk-based Storybook is key to this workflow.

To make matters worse, markup snapshots don’t capture externally-defined styles, as is the case with HTML and CSS. This is where visual regression tools like Chromatic really shine, extending the UI snapshots to their complete, rendered state, and layering a review process on top. Figure 1 contains an example of the Chromatic review screen for a visual test.

Fig. 1 - Chromatic review screen for a snapshot

On the left we have the snapshot from a previous build and on the right the snapshot with the changes we’re introducing, highlighted in green. Reviewers can comment on each snapshot, accepting or rejecting the changes. Pretty great, right?

In the following sections we’ll cover how to create these snapshots and integrate Chromatic into a CI pipeline.

Writing snapshots

Chromatic integrates with Storybook, capturing a screenshot of each story on a configurable set of browsers (Chrome by default), for a given set of viewports.

At Umani, we like our Storybook stories to be interactive and expose a bunch of controls. Figure 2 contains an interactive story for an Avatar component.

Fig. 2 - An interactive story with multiple controls

This story is written as:

export default {
    title: "Avatar",
    parameters: {
        chromatic: {
            viewports: [360, breakpoints.desktop],
        },
    },
}

interface AvatarStoryProps extends AvatarProps {
    readonly showContent?: boolean
}

const Template: Story<AvatarStoryProps> = ({ showContent = false, size, ...args }) => {
    return (
        <Avatar size={size} {...args}>
            {showContent ? (
                <Stack space="xs">
                    <Text size="md">Art Vandelay</Text>
                    <Text size="sm" variation="subtle">
                        View profile
                    </Text>
                </Stack>
            ) : null}
        </Avatar>
    )
}

export const Basic = Template.bind({})

Basic.args = {
    showContent: false,
    size: "md",
}

Basic.argTypes = {
    showContent: {
        name: "Show Content",
        description: "Content is shown to the right.",
    },
    size: {
        name: "Size",
        description: "Avatar size.",
    },
}

These stories don’t make up a very good snapshot, which is why we disable Chromatic in their parameters:

Basic.parameters = {
    chromatic: {
        disabled: true,
    },
}

The stories we are interested in capturing with Chromatic visual tests are fine-grained and uninteractive. We’ll usually include (a sensible version of) the cartesian product of all variations of a given component within a snapshot. For example, the snapshot story for our Avatar component is defined as:

export const Snapshot: Story = () => {
    const stories: Story[] = []
    for (const showContent of [true, false]) {
        for (const size of ["sm", "md"] as const) {
            const props = { showContent, size }
            const story: Story = () => <Template {...props} fallback="" />
            story.storyName = `Avatar with photo, with${!showContent ? `out` : ``} content and size ${size}`
            stories.push(story)
        }
    }
    return <StoryGroup stories={stories} />
}

Figure 3 contains the rendered snapshot story.

Fig. 3 - A rendered snapshot story containing multiple variations

The reason we bundle different variations into the same story is so we don’t blow up our snapshot budget. Similarly, we’ll strive to minimize duplicate snapshots: if the variations of a component like Avatar have already been tested in isolation, we may not need to include them when using that component in a composite story. Minimizing stories is helpful to remain within limits and also curb the time it takes to review changes.

Notice that we configure Chromatic to produce two snapshots at two different viewports with

chromatic: {
    viewports: [360, breakpoints.desktop],
}

This is useful for responsive components and pages.

Snapshotting CSS states like hover and focus often require using Storybook play functions or the ability to trigger those states from component props.

Setting up Chromatic with Github Actions

At Umani we use Github Actions for our CI pipeline. Integrating Chromatic is very easy, but not without its subtleties. This is our workflow job which builds and publishes the Storybook into Chromatic:

    storybook:
        name: Storybook

        runs-on: ubuntu-latest

        steps:
            - name: Checkout
              uses: actions/checkout@v2
              with:
                  fetch-depth: 0
                  ref: ${{ github.event.pull_request.head.sha }}

            - uses: ./.github/actions/load-node-modules

            - name: Create snapshots
              run: yarn chromatic --only-changed --skip 'dependabot/**'
              env:
                  CHROMATIC_PROJECT_TOKEN: ${{ secrets.CHROMATIC_PROJECT_TOKEN }}
                  CHROMATIC_SHA: ${{ github.event.pull_request.head.sha }}
                  CHROMATIC_BRANCH: ${{ github.event.pull_request.head.ref }}

There’s a few things to unpack here, but the important bits are straightforward: we check out the PR’s code (with full history, which is required by Chromatic), use a composite action to load the node modules from cache, and invoke Chromatic. (There’s an official Github Action, but we are not yet leveraging it.)

This job generates a unique build in Chromatic. A branch/PR can have many builds and, unless otherwise specified, snapshots are checked for differences against their counterparts from a previous build either on the same branch or belonging to an ancestor commit. Chromatic’s documentation goes into detail about how baselines are calculated. For us, that baseline is either a build within the same branch or a build for the main branch. Since we’re not using Chromatic’s UI Review tool and we squash our PRs, there’s no association between the merge commit and the commits on the merged PR. This means Chromatic can’t establish the builds of a merged PR as the baseline for new PRs. To explicitly associate a build with a merge commit, we run a separate action on push:

name: Publish Storybook

on:
    push:
        branches:
            - main

jobs:
    storybook:
        name: Storybook

        runs-on: ubuntu-latest

        steps:
            - name: Checkout
              uses: actions/checkout@v2
              with:
                  fetch-depth: 0

            - uses: ./.github/actions/load-node-modules

            - name: Create snapshots
              run: yarn chromatic --only-changed --auto-accept-changes
              env:
                  CHROMATIC_PROJECT_TOKEN: ${{ secrets.CHROMATIC_PROJECT_TOKEN }}
                  CHROMATIC_SHA: ${{ github.event.after }}
                  CHROMATIC_BRANCH: main

This time we’re specifying the --auto-accept-changes flag to automatically accept the changes, as they have already been reviewed in the context of the PR.

We’re enabling Chromatic’s TurboSnap with the --only-changed flag. TurboSnap uses Webpack’s dependency graph to determine which stories have changed, thus minimizing the amount of snapshots needed per PR. That is especially desirable in the context of a monorepo like ours, since many PRs don’t touch the UI and don’t need to trigger any snapshots. TurboSnap errors on the side of caution though, and if there are changes to package.json, all stories will be considered changed. Since our dependency updates are automated, we use Chromatic’s skip option to mark the visual tests as passed without actually creating any snapshots. It’s possible that updating a dependency will cause UI changes that will go undetected by Chromatic, but right now we’re preferring to conserve the snapshot budget. Note that because we use vanilla-extract for styling, the dependency graph can trace CSS changes to specific stories.

Limitations and pitfalls

As with all tools, there are some non-obvious usages that leave us scratching our heads. These are the ones we repeatedly encounter:

Snapshotted stories need to be written deterministically to avoid false positives. This means ensuring the absence of randomness and the stability of things like element order and dates: a story that uses Date.now() or shuffles the images in a carousel will always require approval (if snapshotted). Stories are easy enough to fix, but sometimes the non-determinism comes from deep within a component. To help with those, we can tell whether we’re running under Chromatic by using the isChromatic() function.
Chromatic doesn’t capture animations. Instead, videos and CSS/SVG animations are paused automatically and reset to their initial state. JavaScript animations must be disabled explicitly (isChromatic() is useful here as well). Alternatively, Chromatic can be configured with a delay to allow animations to complete before a snapshot is taken. This doesn’t always solve the problem though. If you’re creating a looping animation (so adding a delay isn’t useful) with a library like framer-motion, which doesn’t expose a way to globally disable animations, then you may need to instruct Chromatic to ignore a DOM element.
Finally, if using TurboSnap, it’s important to be aware of its limitations. We already mentioned that changes to package.json trigger full snapshots. Another situation that can lead to more snapshots being taken than expected is when stories (or intermediary files) import components through an index file. If any (transitive) import in that index file was changed, then all importers of the index file will be considered changed as well.

Conclusion

Visual regression testing is essential to confidently make changes to a user interface. Front-end development is sufficiently complex that most changes can only be noticed by comparing the rendered interface in a specific viewport and browser. Chromatic makes this very easy by integrating with Storybook, a nearly ubiquitous tool in the JavaScript ecosystem, and layering on top a great review workflow that lets developers comment on and approve or reject changes to an application’s UI.

Indexing S3 files in DynamoDB, Serverlessly

Duarte Nunes — Sat, 26 Mar 2022 21:55:01 +0000

This article will explore different architectures for uploading files, like images, to a serverless backend. The files are stored in S3 and indexed by DynamoDB, where they are associated with an arbitrary domain entity. When a client queries an entity, they obtain the IDs of the files from which URLs to S3 can be derived.

In the absence of a distributed transaction between these systems, the problem we're interested in solving is how to maintain consistency between the file, in S3, and the index, in DynamoDB.

Our file management system contemplates both uploads and deletions.

A file uploading API

Let's consider the following GraphQL API for uploading files. The upload mutation will be used to obtain an S3 signed URL where we can PUT the file.

input UploadFileInput {
    entity: ID!
    md5: String!
}

type Header {
    key: String!
    value: String!
}

type UploadFilePayload {
    uploadURL: AWSURL!
    uploadHeaders: [Header!]!
}

type Mutation {
    upload(input: UploadFileInput!): UploadFilePayload!
}

The mutation accepts the file's MD5 (so the signed URL can only be used to upload files with a matching digest) and the entity ID where we'll index it. In the reply we get the upload URL and the headers to include in the PUT request. The entity ID needs to be carried as metadata throughout the system, and is contained in those headers.

Obtaining the signed URL and issuing a PUT request containing the file is the synchronous part of the architecture depicted in Figure 1.

Fig. 1 - File upload architecture

In the background, when the file is uploaded, an S3 notification¹ triggers a Lambda function which will index the file. We know what entity to associate the file with because its ID is stored in the file metadata in S3 (it got there through the headers used in the PUT request).

Deletions and race conditions

So far so good! Let's now consider the mutation to delete a file.

input DeleteImageInput {
    url: AWSURL!
}

type DeleteImagePayload {
    _: Boolean
}

type Mutation {
    deleteImage(input: DeleteImageInput!): DeleteImagePayload!
}

A pretty simple API, but what about the backend implementation?

Ideally, uploading the file and indexing it would be an atomic, distributed transaction. But since it is done in two discrete steps, there are race conditions galore when trying to delete the file, especially if the deletion happens right after the upload.

One approach could be to delete the file in S3 and use the subsequent notification to delete the index entry in DynamoDB. The issue here is that S3 doesn't guarantee the order of notifications. There is, however, a sequencer contained in the notification payload that can be used to order the notifications for a given object. In our case, we can simply synchronize the notifications using the index, as we'll shortly describe.

S3 guarantees at-least-once delivery of notifications, so we also have to account for retries: processing the notifications must be idempotent. We will guarantee idempotence by storing a tombstone when the file is deleted. The tombstone will live for a grace period (a couple of days should be sufficient) and must be cleaned up asynchronously (either through another process or through a TTL on the file item). This will handle the following scenario:

The upload notification is processed and the file is indexed, but the Lambda fails for some reason;
A deletion happens right after and the deletion notification is processed, replacing the key by a tombstone;
The upload notification is retried; it now encounters a tombstone and does nothing².

Note that if we clean up the tombstone in the upload notification, we open ourselves up to the possibility of the upload being retried (either from S3 or from the user) and the file being incorrectly indexed.

If we ensure each file has a unique key in S3, an upload never wins over a deletion. Were we to allow uploads to the same key after a deletion, we would need to store the sequencer in the index (which would act as the timestamp in last-write-wins systems like Cassandra and ScyllaDB).

Alternatively, we can handle the deletion by writing the tombstone directly to DynamoDB without going through the S3 notification, since it is the index that is being used for synchronization. With this approach the deletion effect is immediately visible to an observer, instead of eventually. Asynchronously, we must clean up the tombstone and delete the file from S3 (we can periodically run a compaction process or set a TTL on the item and use DynamoDB streams to later delete the file). This is the architecture depicted in Figure 2.

Fig. 2 - File deletion architecture

Pretty complex, huh? This complexity is mostly arising from us wanting to handle a deletion in close succession to the upload, such that the notifications can overlap. Maybe if we gate the deletion on the presence of the index we can simplify our architecture. That would entail returning an error to the user if they attempt to delete the file after uploading it, while it has not been indexed yet. This isn't great UX though. We can instead store the deletion intent in a queue and keep trying to process it until the file is indexed. Here we're replacing the tombstone cleanup process with a queue, but it is perhaps simpler. There's also a period where the image can appear to resurrect, after it is indexed but before the deletion completes. This is the architecture in Figure 3.

Fig. 3 - File deletion architecture, but with a queue

Would a REST API be simpler?

Since there is no built-in mechanism in GraphQL to upload files, common wisdom suggests that a REST endpoint should be used for uploading instead. Will that really make a difference though? Consider the architecture in Figure 4, where an API Gateway handler receives the file to upload, writes it to S3, and indexes it in DynamoDB.

Fig. 4 - Uploading files with REST

Seems pretty straightforward, right? But there's a wrinkle: the handler may fail between uploading the file to S3 and indexing it in DynamoDB. This means two things:

The handler must be idempotent to handle retries; and
A retry may not happen, so the handler should clean up after itself.

We can't simply put the request in a queue due to message size limitations, so we need a more elaborate solution. We're right back at the approaches of the previous section.

Conclusion

Indexing S3 files in DynamoDB seems straightforward, but bringing deletions into the picture sure complicates things due to potential race conditions. We looked at multiple strategies to deal with those races, balancing UX and complexity. As with most problems in distributed systems, there isn't a perfect solution.

It can take over a minute for the file to be indexed and available for queries. ↩
We must use DynamoDB conditional statements to ensure consistency. ↩

Building an AWS Lambda extension with Rust

Duarte Nunes — Wed, 04 Nov 2020 22:41:56 +0000

AWS Lambda extensions were recently announced. An extension is a long running process, executed alongside your Lambda function, helping with use cases like:

capturing telemetry;
doing work outside the invocation path, like refreshing configuration settings, secrets, or feature flags;
providing a language-agnostic way to implement some common behavior;
etc.

As an example, imagine a framework that provides utilities for writing integration tests as Lambda functions in arbitrary languages. Those utilities can help with sending captured production traffic to a particular Lambda version, or with comparing invocation responses against historical data.

We'll be building something much simpler though.

Before we delve into the code, a brief overview of how Lambda extensions behave.

Anatomy of an extension

An extension is a program distributed as a Lambda layer that extracts to a special directory named "extensions". Programs in that directory are executed by Lambda and have three distinct phases:

The init phase, covering initialization logic and extension registration;
The invoke phase, in which the extension polls for incoming invocation events;
The shutdown phase, for cleanup logic.

An extension interacts with the Lambda service through the Extensions API, an HTTP API similar to the Runtime API, which allows custom runtimes to receive invocation events from Lambda. (An example of such a runtime is the one we'll be using, the aws-lambda-rust-runtime.)

The documentation for Lambda extensions contains the following helpful diagram:

You'll note that the diagram includes two extensions of different types. There can be up to 10 extensions for a function, and an extension can be internal or external. Internal extensions run as separate threads within the runtime process, which starts and stops them. As such, they don't handle the shutdown event. External extensions - which is what we'll be building - run as independent processes in the execution environment. This is what allows them to be written in a different language than the function.

The documentation suggests external extensions be implemented as a compiled language so they are compatible with all runtimes. Let's start! 🦀

A Rusty extension

Our extension needs to make HTTP requests, parse JSON, and deal with errors. These are the crates we're using for those tasks:

[dependencies]
anyhow = "1.0"
serde = "1.0"
serde_json = "1.0"
reqwest = { version = "0.10.8", default-features = false, features = ["blocking", "json"] }

Our entry point will look like this:

fn main() -> Result<()> {
    let client = Client::builder().timeout(None).build()?;
    let r = register(&client)?;
    loop {
        std::thread::sleep(time::Duration::from_secs(1));
        println!("Waiting for event...");
        match next_event(&client, &r.extension_id) {
            // process event, that is, print some stuff
        }
    }
}

The entry point consists of initialization logic and a main event loop (this is similar to how we structure a Lambda function, with initialization logic living outside of the event handler). In this example, we just initialize the HTTP client and register the extension. If you're wondering whether it's a good idea to use an infinite timeout for our HTTP calls, the answer is yes: the extension can be suspended for an arbitrary period of time until there is an event to return, which is also why we're using blocking I/O for our interaction with the Extension API.

The init phase completes when we enter the event loop and request an event for the first time; the Lambda service now knows we're ready to process events. Events are only allowed to come in when the Lambda function is done initializing and all extensions reach this point. This means that the initialization we do in our extensions directly impacts cold start times. (On the topic of cold starts, I recommend this video by Marc Brooker on virtualization technology underpinning Lambda!)

Notice that we're adding a 1 second sleep, as it'll make it easier to point out some behavior when we look at logs.

Registering the extension is a simple HTTP POST:

#[derive(Debug)]
struct RegisterResponse {
    pub extension_id: String,
}

fn register(client: &reqwest::blocking::Client) -> Result<RegisterResponse> {
    let mut map = HashMap::new();
    map.insert("events", vec!["INVOKE", "SHUTDOWN"]);
    let url = format!("{}/register", base_url()?);
    let res = client
        .post(&url)
        .header(EXTENSION_NAME_HEADER, EXTENSION_NAME)
        .json(&map)
        .send()?;

    ensure!(
        res.status() == StatusCode::OK,
        "Unable to register extension",
    );

    let ext_id = res.headers().get(EXTENSION_ID_HEADER).unwrap().to_str()?;

    Ok(RegisterResponse {
        extension_id: ext_id.into(),
    })
}

The most important thing about this operation is that the extension name must match the filename of the binary, or else we'll get back a 403 response.

The base_url() returns the endpoint for the Extension API, at http://$AWS_LAMBDA_RUNTIME_API/2020-01-01/extension.

As an external extension, we register for the shutdown event. We get back the extension ID, which we flow back on subsequent API calls.

Requesting the next event is also a simple HTTP operation:

fn next_event(client: &reqwest::blocking::Client, ext_id: &str) -> Result<NextEventResponse> {
    let url = format!("{}/event/next", base_url()?);
    Ok(client
        .get(&url)
        .header(EXTENSION_ID_HEADER, ext_id)
        .send()?
        .json()?)
}

Where NextEventResponse is define as:

#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
struct Tracing {
    pub r#type: String,
    pub value: String,
}

#[derive(Debug, Deserialize)]
#[serde(rename_all = "UPPERCASE", tag = "eventType")]
enum NextEventResponse {
    #[serde(rename_all = "camelCase")]
    Invoke {
        deadline_ms: u64,
        request_id: String,
        invoked_function_arn: String,
        tracing: Tracing,
    },
    #[serde(rename_all = "camelCase")]
    Shutdown {
        shutdown_reason: String,
        deadline_ms: u64,
    },
}

As I'm sure you've noticed, the event's payload is conspicuously missing from these definitions. That's because the invocation event sent to each extension contains only metadata. The only way for the extension to get the payload is by communicating with the function, which we'll cover in the next section. But first, let's look at CloudWatch logs for an invocation of a simple test Lambda that's deployed using this extension:

timestamp	message
17:26:52.325-03:00	START RequestId: d11df8ca-1955-4797-a19d-5de248d0dc86 Version: $LATEST
17:26:52.325-03:00	Waiting for event...
17:26:52.325-03:00	EXTENSION Name: sidecar State: Ready Events: [INVOKE,SHUTDOWN]
17:26:52.327-03:00	Invoke event d11df8ca-1955-4797-a19d-5de248d0dc86; deadline: 1464625
17:26:52.344-03:00	Hello, world!
17:26:53.327-03:00	Waiting for event... (1 second after)
17:26:53.327-03:00	END RequestId: d11df8ca-1955-4797-a19d-5de248d0dc86
17:26:53.327-03:00	REPORT RequestId: d11df8ca-1955-4797-a19d-5de248d0dc86 Duration: 1001.13 ms Billed Duration: 2100 ms Memory Size: 128 MB Max Memory Used: 34 MB Init Duration: 1039.65 ms
17:37:51.351-03:00	Exiting: spindown

The first line is logged by the Lambda service right before the extension logs it's waiting for an event. We're not seeing the initial 1 second delay in the logs, but it is there: notice the last log line, where the init duration is 1039.65 ms. We can see that indeed, Lambda waits for extensions to become initialized before handling incoming requests.

The function logs "Hello, world!", which is all that it does. After 1 second, the extension signals it's ready for more events. Notice that it's only after this delay that the Lambda service prints the "END" log entry. Lambda waits until all extensions have finished processing an event before sending out the Lambda function's response. Extensions signal that they have finished processing an event by requesting the next one.

An important thing to keep in mind is that an extension can add cold start latency and add latency to each request.

Roughly 10 minutes after that singleton event is processed, the Lambda service sends the extension a shutdown event with the "spindown" reason. During shutdown, the extension has two seconds to execute any cleanup logic.

What we have now isn't terribly useful. Let's add a little bit more code.

Extension <-> Function IPC

For some use-cases it's useful for the extension to communicate with the function, but there's no built-in or established way to do that. Fortunately, the extension and the function share the memory and /tmp disk storage, so there's plenty of ways to fashion an IPC mechanism for this purpose.

We could create a FIFO, which is a simple approach for unidirectional communication (and a pair of FIFOs could be used for full-duplex communication). We could use a domain socket, or, similarly, the extension could start an HTTP server. This would fit well as a simple but flexible way for the integration test Lambda to communicate with the helpful, but imaginary extension we gave as an example in the beginning of the post.

Yet another alternative is to use the filesystem, which is better suited for unidirectional communication. The extension can expect that the lambda writes a file named with the request ID after each invocation, containing some result. This is exactly what we're going to implement.

The extension must know when the file is created and when the function is finished writing to it. It must also tolerate failures and deal with the absence of the file. We could use inotify or some sophisticated approach to deal with all potential corner cases, but we'll instead assume the system tolerates a 10 minute delay between the function ending and the results becoming available to the extension, and just process the previous event's result when the next one arrives.

Our extension entry point becomes:

fn main() -> Result<()> {
    let client = Client::builder().timeout(None).build()?;
    let r = register(&client)?;
    let mut prev_request: Option<String> = Option::None;
    loop {
        std::thread::sleep(time::Duration::from_secs(1));
        println!("Waiting for event...");
        let evt = next_event(&client, &r.extension_id);
        prev_request.map(process_result);
        match evt {
            Ok(evt) => match evt {
                NextEventResponse::Invoke {
                    request_id,
                    deadline_ms,
                    ..
                } => {
                    println!("Invoke event {}; deadline: {}", request_id, deadline_ms);
                    prev_request = Some(request_id);
                }
                NextEventResponse::Shutdown {
                    shutdown_reason, ..
                } => {
                    println!("Exiting: {}", shutdown_reason);
                    return Ok(());
                }
            },
            Err(err) => {
                eprintln!("Error: {:?}", err);
                println!("Exiting");
                return Err(err);
            }
        }
    }
}

The relevant addition is the prev_request variable, which hold the ID of the previous event. process_result() is defined as

#[derive(Deserialize)]
struct InvocationResult {
    payload: Value,
}

fn read_result(req_id: String) -> Result<InvocationResult> {
    let filename = format!("/tmp/{}", req_id);
    let f = fs::File::open(filename)?;
    let reader = BufReader::new(f);
    let res = serde_json::from_reader(reader)?;
    Ok(res)
}

fn process_result(req_id: String) {
    match read_result(req_id) {
        Ok(InvocationResult { payload }) => println!("Payload: {}", payload),
        Err(e) => eprintln!("Error processing invocation result: {:?}", e),
    }
}

Basically, it reads a file at /tmp/{id} containing a JSON object with the "payload" key mapping to the invocation's payload. The Lambda function writes the files:

fn handler(v: Value, ctx: Context) -> Result<(), HandlerError> {
    println!("Hello, world!");
    let wrapped = json!({ "payload": v });
    match File::create(format!("/tmp/{}", ctx.aws_request_id)) {
        Ok(mut file) => match file.write_all(wrapped.to_string().as_bytes()) {
            Ok(_) => Ok(()),
            Err(e) => Err(e.to_string().as_str().into()),
        },
        Err(e) => Err(e.to_string().as_str().into()),
    }
}

When we invoke this function with

aws lambda invoke --function-name rust-test-function --cli-binary-format raw-in-base64-out --payload '{"hello": "world"}'

We can see in the logs, 10 minutes after:

timestamp	message
18:27:32.101-03:00	Payload: {"hello":"world"}
18:27:32.101-03:00	Exiting: spindown

That's it: the Lambda function wrote a file with the payload, which the extension read and logged to CloudWatch.

Building and deploying

Since we're writing native code, we have to worry about which platform to target when compiling, and which libraries to link against. Fortunately, there's the lambci/lambda:build-provided.al2 Docker image, which we use to ensure we link against the exact library versions that exist in the AWS Lambda environment. As for compiling, we're using rustup to fetch and install the appropriate Rust toolchain for the platform. Pretty simple.

We're using the AWS Cloud Development Kit to deploy the stack:

const app = new App()
const s = new Stack(app, "Test", {
    env: {
        account: process.env.CDK_DEFAULT_ACCOUNT,
        region: process.env.CDK_DEFAULT_REGION,
    },
})

const l = new Lambda(this, "Function", {
    functionName: "rust-test-function",
    runtime: Runtime.PROVIDED_AL2,
    handler: "doesnt.matter",
    code: Code.fromAsset(app.node.tryGetContext("lambda")),
    currentVersionOptions: {
        removalPolicy: RemovalPolicy.DESTROY,
        retryAttempts: 2,
    },
    timeout: Duration.minutes(1),
})

l.addLayers([
    new LayerVersion(s, "Sidecar", {
        code: Code.fromAsset(app.node.tryGetContext("sidecar")),
    }),
])

The Test stack is deployed with:

cdk deploy --context lambda=./lambda.zip --context sidecar=./sidecar.zip Test

Where ./lambda.zip is an archive containing the function's binary, a single file named "bootstrap", which is a requirement when using a custom runtime. The extension is packaged in ./sidecar.zip, containing a binary named "sidecar" - the same name we give the Lambda service when registering the extension -, inside the "extensions" folder.

Final thoughts

We've shown how easy it is to build a Lambda extension: it's a single binary that can be written in any (compiled) language, which uses a simple HTTP-based protocol to communicate with the Lambda service. Surely the Extension API will become richer and more flexible, which will unlock more use cases for extensions - for example, to implement a good and efficient IPC mechanism it would be really helpful to know when a function has finished processing an event.

As with all things serverless, latency and cold starts are a concern, especially if you plan to write extensions for a wider audience.

Tangentially, we showed how easy and cool it is to use Rust to write Lambda functions, and how effortless it is to deploy using the CDK.

Passwordless Authentication with Cognito

Duarte Nunes — Fri, 22 May 2020 16:53:50 +0000

Passwordless authentication is a broad term for any authentication method that doesn't rely on passwords. Implementations typically perform proof of identity based on something that is uniquely associated with a user, such as an e-mail address, a phone, a software one-time password (OTP) generator, or a hardware authentication device like a YubiKey: the user inputs the secret that the system shares with one of those methods, proving its ownership. A passwordless flow essentially skips the "what you know" factor of two-factor authentication. We're epistemological skeptics now.

Cognito is an infamous AWS service for identity management. It can do much for us: authentication, access control, and integration with external identity providers. A big advantage over competing solutions is that it integrates well with other AWS services, like AppSync.

In this post we'll look at how we can implement passwordless authentication using Cognito to satisfy our requirements, which aren't quite as described in this AWS blog post. We assume a passing familiarity with Cognito and AWS Lambda. The code examples are Lambda functions written in Go.

"As a user, I want..."

The app we're going to work on authenticates users according to the following UX:

A user inputs their e-mail or phone number.
We send an OTP over e-mail or SMS.
The user inputs the code.
If this is their first sign-in:
1. We collect additional data, including a secondary authentication method.
2. We verify it using a second OTP.

This is very convenient: users can sign-in using information they know by heart, without needing to use password managers or, worse, their memory (they don't even need to remember whether they've already signed-up).

These are our requirements:

[ ] The user must be able to sign-in with their email or phone number
[ ] There is no separate sign-up process, as it should follow directly from the sign-in
[ ] We mustn't leak user existence errors through the public API, including Cognito's
[ ] OTPs are valid only for a brief period of time
[ ] Users have more than one attempt to input the correct OTP

What features does Cognito offer that help us meet these requirements? Let's find out.

Cognito Custom Auth

Cognito doesn't have a built-in passwordless feature, but it supports a custom authentication flow implemented as a state machine based on three Lambda function triggers that we need to implement.

The define auth challenge Lambda is the state machine coordinator. It returns instructions to Cognito on how the flow should progress. The options are: a challenge is required; the authentication failed; or the authentication succeeded and tokens can be emitted. It is invoked by Cognito when a client calls the InititateAuth API.
If a challenge is required, the create auth challenge Lambda creates it. It's invoked right after the previous Lambda¹, and is responsible for delivering the challenge to the user. This can be done through public challenge parameters returned by InitiateAuth, or through some other means². Private challenge parameters are shared with the verify auth challenge response Lambda, typically including the challenge answer.
The verify auth challenge response Lambda is invoked by Cognito following the RespondToAuthChallenge API call, containing the challenge answer provided by the user. It returns whether the answer is correct, and Cognito invokes the define auth challenge Lambda again to decide whether the authentication can terminate or should continue with more challenges (in which case the create auth challenge Lambda gets called again).

The custom auth state machine, from the documentation.

A custom authentication flow can be composed of many challenges. State is kept in an encrypted session object, which is initially returned by the InititateAuth API call, and then threaded back to Cognito when calling RespondToAuthChallenge. It expires after 3 minutes, which already allows us to tick off a requirement!

[ ] The user must be able to sign-in with their email or phone number
[ ] There is not separate sign-up process, as it should follow directly from the sign-in
[ ] We mustn't leak user existence errors through the public API, including Cognito's
[x] OTPs are valid only for a brief period of time
[ ] Users have more than one attempt to input the correct OTP

Configuring the Cognito user pool

Our journey first tasks us with configuring a Cognito user pool in such a way that e-mails and phone numbers can be used as aliases. Thanks to the CDK, we don't have to confront the maze of darkness that is CloudFormation YAML:



this.pool = new UserPool(stack, "UserPool", {
  userPoolName: props.userPoolName,
  signInAliases: { username: true, email: true, phone: true },
  lambdaTriggers: {
    defineAuthChallenge: props.triggers.defineAuthChallenge,
    createAuthChallenge: props.triggers.createAuthChallenge,
    verifyAuthChallengeResponse: props.triggers.verifyAuthChallengeResponse,
  },
  selfSignUpEnabled: false,
})

This translates to a user pool that's part of a CloudFormation stack, which:

Allows users to sign-in using their e-mail address and phone-number;
Is configured with the ARN of the Lambdas that will carry out our passwordless authentication flow;
Doesn't allow users to sign themselves up, because the SignUp API can be used as an oracle for whether a user exists in a pool. This means we must use the AdminCreateUser API.

We also configure the user pool client to only authenticate users through the custom authentication flow, thus disallowing password-based authentication. Clients are allowed to refresh user tokens (tokens have a validity and need to be refreshed for the session to be extended without users having to go through the sign-in process again).



this.pool.addClient(`UserPoolClient${clientName}`, {
  userPoolClientName: clientName,
  authFlows: {
    custom: true,
    refreshToken: true,
  },
})

Verifying an email or phone number

Implementing the flow using the extension Lambda functions seems straightforward:

The define auth challenge Lambda checks for a successfully completed challenge; if there isn't any, it either fails the authentication if there are over N attempts, or keeps the state machine running.
The create auth challenge Lambda sends an e-mail or SMS, depending on what was used to sign-in.
The verify auth challenge response receives the answer from the client and marks the challenge as correct or not.

Alas, it's not so simple. We need to pass a bit of information from the client to the create auth challenge Lambda, so it can decide between sending an e-mail or an SMS. Cognito, however, doesn't tell the Lambda what alias was used to start the flow, nor does it thread the ClientMetadata parameter of InititateAuth to the define or create auth challenge Lambdas.

How can we work around this? We could pass that information through a data store like DynamoDB, but a better solution is to leverage that ClientMetadata parameter. InititateAuth doesn't pass it to the Lambda, but RespondToAuthChallenge certainly does.

An RPC by any other name

These authentication challenges actually comprise a complete request-response protocol over which the client can communicate with the state machine. Clients can answer the challenges while passing along arbitrary data to the Lambdas, which can keep the protocol going as long as needed. With this in mind, we can solve the problem by using an extra challenge to specify what the actual method we're verifying is:

The first challenge delivers no code, and expects email or phone to be passed in the ClientMetadata. The second challenge is the actual OTP.

The metadata returned by the create auth challenge Lambda is kept in the encrypted session object, and is used to communicate with the define auth challenge Lambda as well as with future invocations of the create auth challenge Lambda. The private parameters are shared only with the verify auth challenge response Lambda.

This custom authentication flow now behaves as a primitive that lets us validate a user attribute, namely an e-mail or phone number.

Implementation

The verify auth challenge response Lambda has the easiest task:



func handler(ctx context.Context, event *events.CognitoEventUserPoolsVerifyAuthChallenge,
) (*events.CognitoEventUserPoolsVerifyAuthChallenge, error) {
    code := event.Request.PrivateChallengeParameters["code"]
    event.Response.AnswerCorrect = code == "" || code == event.Request.ChallengeAnswer
    return event, nil
}

If there is no code in the private challenge parameters it's because this is the first step of the protocol, in which case the answer is always accepted. Otherwise, we match the code against the user-provided answer.

The create auth challenge Lambda calculates and delivers the challenge. We encode a challenge as:



type AttributeType string

const (
    Email          AttributeType = "email"
    Phone          AttributeType = "phone_number"
    NoVerification AttributeType = "none"
)

type Challenge struct {
    Attr AttributeType
    Code string
}

This challenge is persisted in the metadata shared between all Lambdas in the state machine and stored in the session object. The private challenge parameters are shared with the verify auth challenge response Lambda. If our challenge had public parameters, we would specify them here and they would be sent to the client (in plaintext).



func handler(ctx context.Context, event *events.CognitoEventUserPoolsCreateAuthChallenge,
) (*events.CognitoEventUserPoolsCreateAuthChallenge, error) {
    challenge := calculateAndDeliverChallenge(&event.Request, userTest, &metadata)
    event.Response = triggers.CognitoEventUserPoolsCreateAuthChallengeResponse{
        PrivateChallengeParameters: map[string]string{
            "code": challenge.Code,
        },
        ChallengeMetadata: auth.Serialize(challenge),
    }
    return event, nil
}

We generate an empty code for the first challenge and an OTP for the second one. In the second challenge we expect the client metadata to contain the attribute type to verify (the type, not the value), which we use to select the the attribute value from the current user attributes, delivering the code by e-mail or SMS.³ The client essentially sends us a pointer to the actual attribute to verify.



func calculateAndDeliverChallenge(req *events.CognitoEventUserPoolsCreateAuthChallengeRequest,
) auth.Challenge {
    sessionSize := len(request.Session)
    if sessionSize == 0 {
        return auth.Challenge{auth.NoVerification, ""}
    }
    lastSession := request.Session[sessionSize-1]
    lastChallenge := auth.Deserialize(lastSession.ChallengeMetadata)
    if !lastSession.ChallengeResult {
        return lastChallenge
    }
    attr := auth.AttributeToVerify(request.ClientMetadata, request.UserAttributes)
    code := auth.GenerateOTP()
    deliverCode(attr, code)
    return auth.Challenge{attr.Type, code}
}

If the previous challenge failed, that is, if the user introduced the wrong code, we reuse it.

The define auth challenge Lambda coordinates these two:



type ChallengeStatus struct {
  Challenge      auth.Challenge
    AttemptCount int
    Passed       bool
}

func handler(ctx context.Context, event *events.CognitoEventUserPoolsDefineAuthChallenge,
) (*events.CognitoEventUserPoolsDefineAuthChallenge, error) {
    status := calculateChallengeStatus(event.Request.Session)
    challengeName := ""
    done := status.Passed && status.Challenge.Attr != auth.NoVerification
    if done {
      // ...
    } else {
        challengeName = "CUSTOM_CHALLENGE"
    }
    event.Response = events.CognitoEventUserPoolsDefineAuthChallengeResponse{
        IssueTokens:        done,
        FailAuthentication: status.AttemptCount >= 4,
        ChallengeName:      challengeName,
    }
    return event, nil
}

The state machine keeps turning until we reach the maximum allowed failed attempts or the user provides us with the correct OTP. We tick off two requirements 🎉

[x] The user must be able to sign-in with their email or phone number
[ ] There is not separate sign-up process, as it should follow directly from the sign-in
[ ] We mustn't leak user existence errors through the public API, including Cognito's
[x] OTPs are valid only for a brief period of time
[x] Users have more than one attempt to input the correct OTP

Unfortunately, this isn't sufficient to support our authentication flow. Recall that we have no separate sign-up process, so we might not have a registered user for a run of this protocol. However, we need the attribute to verify to be present in the user attributes, which are only populated for existing users (also, Cognito doesn't issue tokens for users absent from the pool). We must ensure a user exists, even if a partial one.

Sign-In

Users aren't allowed to sign themselves up (to prevent using the SignUp API as a user-existence oracle), so we need something that does that on their behalf, say, a Lambda resolver for a GraphQL API on AppSync.

We can call the AdminCreateUser API to ensure the user exists, return to the client and let it handle the rest of the flow, but we'd be optimizing for the less common scenario of a user not existing in the pool.

Instead, we try to InititateAuth and if that fails, we issue the AdminCreateUser request. Does this mean we need to mediate all requests between the client and Cognito? No! We can take the session object returned by Cognito and return it to the client to be used in the RespondToAuthChallenge call.

Since we're already here though, we can even go ahead and execute the first step of the protocol, telling Cognito which attribute we want to verify:



func handler(ctx context.Context, input SignInInput) (*SignInPayload, error) {
  session, op := auth.InitiateAuth(cognito, input.EmailOrPhone, input.Method)
    if op.Error == nil {
        return &SignInPayload{string(session)}, op
    }
    op = auth.CreateUser(cognito, input.EmailOrPhone, input.Method)
    if op.Error != nil {
        return nil, op.Error
    }
  session, op = auth.InitiateAuth(cognito, input.EmailOrPhone, input.Method)
    return &SignInPayload{string(session)}, op.Error
}

The CreateUser function calls the AdminCreateUser API, generating a random username and password and marking the email or phone number attribute as verified so it can be used as an alias. ⁴ The InitiateAuth function calls Cognito's own InitiateAuth and then the first RespondToAuthChallenge. The rest is up to the client.

We encapsulate the AdminCreateUser API and behave the same regardless of the user existing before the request or not. Yep, another one down:

[x] The user must be able to sign-in with their email or phone number
[ ] There is not separate sign-up process, as it should follow directly from the sign-in
[x] We mustn't leak user existence errors through the public API, including Cognito's
[x] OTPs are valid only for a brief period of time
[x] Users have more than one attempt to input the correct OTP

Sign-Up

Uff. The user is signed-in. However, there's a chance the JWT token our client got back from Cognito contains only a single claim for either the user's email or phone number, and we need more! In that case, our UI proceeds to collect some user data: their name and the second authentication method.

We have another attribute to verify, so we need to run another instance of our protocol. Lets add a GraphQL resolver that writes the user data into Cognito by updating the user attributes.⁵



func handler(ctx context.Context, input SignUpInput) (SignUpPayload, error) {
    auth.SetAttributes(cognito, input.Username, map[string]string{
        input.Data.Method.String(): input.Data.EmailOrPhone,
        fmt.Sprintf("%s_verified", input.Data.Method.String()): "false",
        "given_name":  input.Data.GivenName,
        "family_name": input.Data.FamilyName,
    }, auth.UserPoolID()))
  session, op := auth.InitiateAuth(cognito, input.Username, input.Data.Method)
    return SignUpPayload{string(session)}, op.Error
}

Note that we mark that second method as not verified, because Cognito only allows one user to have a given email or phone number marked as verified and thus use it as an alias ⁶. We want to avoid the following scenario:

Alice uses the verified email E;
Bob signs-in using their phone number, and then signs-up with E;
Bob doesn't verify E, but we already marked the attribute as verified;
Alice signs-in using E, sees Bob's account.

Our system should consider a user as valid only when it has both attributes verified, but who marks the second method as such? That's a job for our define auth challenge Lambda:



if done {
  attr := fmt.Sprintf("%s_verified", string(status.Attr))
    if event.Request.UserAttributes[attr] == "false" {
            auth.SetAttributes(cognito, event.UserName, map[string]string{
            attr: "true",
        })
    }
} // ...

Finally, note that this very same logic can be used to implement an API for changing a user's e-mail or phone number.

Refreshing the token

We're almost done, but there's one missing detail: Cognito isn't prepared for us modifying a user's attributes during the authentication process, so the token the client ends up with after the sign-up doesn't have the second attribute marked as verified. There's no clever solution here: we must call the InitiateAuth API using the REFRESH_TOKEN_AUTH flow instead of CUSTOM_AUTH.

We got them all.

[x] The user must be able to sign-in with their email or phone number
[x] There is not separate sign-up process, as it should follow directly from the sign-in
[x] We mustn't leak user existence errors through the public API, including Cognito's
[x] OTPs are valid only for a brief period of time
[x] Users have more than one attempt to input the correct OTP

Conclusion

Cognito isn't easy to work with and I find the client libraries and documentation to be somewhat lacking, but at the end of the say it's solving a huge problem on our behalf. It has enough flexibility to be subject to some unorthodox usage, and I hope it keeps on improving.

For this particular use case I missed some additional toggles, which could be added without loss of generality: it could flow the client metadata from InitiateAuth to the authentication Lambdas, and it could allow the define auth challenge Lambda to override some token claims, much like the pre-token generation Lambda can. Another easy improvement is allowing the custom auth Lambdas to return errors, which should be returned from the API calls (even if as some generic error). By contacting AWS support I was able to turn the first of these in a feature request, but of course, without any estimate when it will be available.

Thanks to @_jwnx for reviewing this post.

Yes, an InitiateAuth call can hit 2 Lambda cold-boots. ↩
For example, by sending an e-mail, SMS, or push notification using some AWS service. ↩
Cognito doesn't support returning custom errors from the extension Lambdas. Errors must be propagated in the state machine in such a way that authentication will end up failing. ↩
It also calls the AdminSetUserPassword API to move the user from the FORCE_CHANGE_PASSWORD state to CONFIRMED. ↩
This resolver is using Cognito authorization, so only signed-in users can use it. The Username points to the user making the request. ↩
Cognito understands the email_verified and phone_number_verified attributes. ↩