home

Blog


Reducing Cognito machine-to-machine (M2M) token costs


Cognito's machine-to-machine authentication is just the OAuth 2.0 Client Credentials flow. To use it, your service hands Cognito a client ID and a client secret, and it gets back a short-lived access token. It's all pretty handy and every successful token request costs you almost nothing. However, when you have a fleet of services asking for a brand-new token on every single call, you're suddenly paying AWS just to hand you a fresh token a few million times when you already have a valid one. Thank goodness that with a few small changes most of those requests simply disappear.

Forwarding token requests through a Lambda

This first move doesn't save a cent on its own, but everything else stands on top of it: put a single Lambda between your services and Cognito. Instead of each service calling /oauth2/token directly, they call the Lambda (behind a Function URL or API Gateway), and the Lambda is the only thing that ever talks to Cognito's token endpoint.

Why bother with the extra hop? Because it gives you a choke point, a single place to add caching, to watch and throttle. Also, if you keep the function tiny, it will be a rounding error on your Lambda bill. Yes, on its own it issues the exact same number of token requests as before, so hold your horses, the saving comes next.

Caching tokens in DynamoDB

This is where the money stops leaking. Before calling Cognito, the Lambda checks a DynamoDB table for a token it already fetched, keyed by app client and scope. If there's one and it hasn't expired, return it and go home. If there isn't (or it's stale), request a fresh one from Cognito, write it back with an expiry, and return that.

import time, boto3, requests

table = boto3.resource("dynamodb").Table("m2m-tokens")
SAFETY = 60  # seconds of buffer, so we never serve a token about to die

def get_token(client_id, client_secret, scope):
    key = f"{client_id}#{scope}"
    token = table.get_item(Key={"pk": key}).get("Item")
    token_expires_at = token.get("expires_at") if token else None
    token_access_token = token.get("access_token") if token else None

    if token and token_expires_at and token_expires_at > int(time.time()) and token_access_token:
        return token_access_token

    res = requests.post(
        "https://your-domain.auth.us-east-1.amazoncognito.com/oauth2/token",
        data={"grant_type": "client_credentials", "scope": scope},
        auth=(client_id, client_secret),
    ).json()
    expires_at = int(time.time()) + res["expires_in"] - SAFETY
    access_token = res["access_token"]

    table.put_item(Item={
        "pk": key,
        "access_token": access_token,
        "expires_at": expires_at,
        "ttl": expires_at,
    })

    return access_token

Now a token that's valid for an hour means at most one Cognito call per hour per client, no matter how many times a service is hammering the Lambda in between. And since DynamoDB reads are pretty cheap too, with a generous free tier (just like Lambda), you're finally set to stop wasting money on Cognito.

Worth noting

  • The ttl attribute lets DynamoDB's Time To Live (TTL) feature sweep expired items for you, so the table doesn't fill up with dead tokens. But since it's a bit lazy about the exact deletion time, the code still checks for the expires_at anyway;
  • The SAFETY buffer subtracts a minute from expires_in, so you never hand out a token that expires mid-flight;
  • Keying on client_id#scope keeps tokens with different scopes from stepping on each other.

One honest caveat: if a bunch of requests miss the cache at the exact same instant (token just expired), they can all stampede Cognito at once and each buy a token. If that actually shows up in your numbers, it might be mitigated by using strongly consistent reads, but it isn't a silver bullet since DynamoDB has no built-in locking. You may also implement a locking mechanism yourself, if it would be worth the trouble, or look for libraries that do it, so you can still save some more money.

Locking the back door with AWS WAF

Here's the hole in the plan so far: nothing stops a service, or a well-meaning developer in a hurry, from skipping your Lambda and hitting the Cognito token endpoint directly. You'd pay full price and the cache would be dodged entirely, which is exactly the behavior you just spent two sections killing.

Fortunately, you can close that door because AWS WAF can be associated with a Cognito user pool and inspect requests to the hosted UI and the OAuth endpoints (/oauth2/token included). The AWS team even has a ready-to-use sample to solve this exact problem.

# The proxy reads this at runtime and sets it as the x-origin-verify header.
resource "aws_ssm_parameter" "origin_verify" {
  name             = "/cognito-m2m/origin-verify-token"
  type             = "SecureString"
  value_wo         = var.origin_verify_secret
  value_wo_version = 1
}

resource "aws_wafv2_web_acl" "cognito" {
  name  = "cognito-m2m"
  scope = "REGIONAL"

  # Block everything by default; only the allow rule below lets the proxy in.
  default_action {
    block {
      custom_response {
        response_code            = 403
        custom_response_body_key = "cognito-denied"
      }
    }
  }

  custom_response_body {
    key          = "cognito-denied"
    content      = jsonencode({ message = "Direct access to Cognito is blocked" })
    content_type = "APPLICATION_JSON"
  }

  visibility_config {
    cloudwatch_metrics_enabled = true
    metric_name                = "cognito-m2m"
    sampled_requests_enabled   = true
  }

  # The allow rule lives in its own aws_wafv2_web_acl_rule resource below.
  lifecycle {
    ignore_changes = [rule]
  }
}

resource "aws_wafv2_web_acl_rule" "allow_origin_verify" {
  name        = "allow-origin-verify-requests"
  priority    = 0
  web_acl_arn = aws_wafv2_web_acl.cognito.arn

  action {
    allow {}
  }

  statement {
    byte_match_statement {
      search_string         = sensitive(var.origin_verify_secret)
      positional_constraint = "EXACTLY"

      field_to_match {
        single_header {
          name = "x-origin-verify"
        }
      }

      text_transformation {
        priority = 0
        type     = "NONE"
      }
    }
  }

  visibility_config {
    cloudwatch_metrics_enabled = true
    metric_name                = "allow-origin-verify-requests"
    sampled_requests_enabled   = true
  }
}

resource "aws_wafv2_web_acl_association" "cognito" {
  resource_arn = aws_cognito_user_pool.this.arn
  web_acl_arn  = aws_wafv2_web_acl.cognito.arn
}

I've made a few tweaks to the overall solution, but the idea remains the same. Instead of allowing everything and then trying to pick off the bad requests, the web ACL blocks by default (handing back a friendly 403) and adds a single allow rule. That rule only lets through requests carrying a secret x-origin-verify header, which your Lambda reads from Parameter Store and sets on every call it makes to Cognito. Also, you may add a rate-limiting rule to prevent bad actors from skyrocketing your bill in case the secret gets leaked.

Wrap-up

A Lambda in front of token requests, a DynamoDB table storing and pulling valid tokens without touching Cognito and AWS WAF making sure everyone uses the front door. With that, the tokens you already paid for get reused until they actually expire, and Cognito only hears from you when there's genuinely a new token to mint. Same flow, a fraction of the token requests, and a bill that finally stops charging you for asking the same question over and over.

Powered by Simple Blog API