Example AWS DynamoDB Adapter

[sahinvardar]

Note

Target version: ^9.0.0
AWS SDK JS Target : ^3.0.0

Warning

No guarantees this is bug-free, if you find a bug please comment on the discussion

Based on the previous work of @SachinShekhar here

I updated the code to use AWS SDK for JavaScript v3 with the use of

• @aws-sdk/client-dynamodb
• @aws-sdk/util-dynamodb

// This work is based on the previous work of Sachin Shekhar and just updated to use AWS SDK v3.
// The original code can be found here: @see https://github.com/panva/node-oidc-provider/blob/v8.x/example/adapters/contributed/dynamodb-gsi.ts
// This adapter is written in Typescript. If you want to get Javascript version without compiling it with tsc, paste this code here: https://babeljs.io/en/repl (choose Typescript preset)

/**
 * Prerequisites:
 *
 * 1. Create a DynamoDB Table with following details:
 *        Partition Key: modelId
 *        TTL Attribute: expiresAt
 *        Three Global Secondary Indexes:
 *            GSI 1:
 *                Index Name: uidIndex
 *                Partition Key: uid
 *            GSI 2:
 *                Index Name: grantIdIndex
 *                Partition Key: grantId
 *            GSI 3:
 *                Index Name: userCodeIndex
 *                Partition Key: userCode
 *
 * 2. Put the Table's name in environment variable TABLE_NAME or simply replace the value of constant TABLE_NAME below.
 *
 * 3. You'll also need to change value of TABLE_REGION constant below if you aren't in AWS compute environment or if DynamoDB Table exists in different region.
 *
 * 4. If you are in AWS' compute environment, nothing more needs to be changed in code.
 *    You just need to give proper IAM permissions of DynamoDB Table.
 *    Required Permissions:
 *        dynamodb:GetItem
 *        dynamodb:ConditionCheckItem
 *        dynamodb:UpdateItem
 *        dynamodb:DeleteItem
 *        dynamodb:Query
 *        dynamodb:BatchWriteItem
 *    If you aren't in AWS' compute environment, you'll also need to configure SDK with proper credentials.
 *    @see https://docs.aws.amazon.com/sdk-for-javascript/v2/developer-guide/configuring-the-jssdk.html
 */

import { Adapter, AdapterPayload } from 'oidc-provider';
import {
  DynamoDBClient,
  DeleteItemCommand,
  GetItemCommand,
  QueryCommand,
  QueryCommandOutput,
  UpdateItemCommand,
  BatchWriteItemCommand,
  AttributeValue
} from '@aws-sdk/client-dynamodb';
import { marshall, unmarshall } from '@aws-sdk/util-dynamodb';

const TABLE_NAME = process.env.TABLE_NAME!;
const TABLE_REGION = process.env.AWS_REGION;

const dynamoClient = new DynamoDBClient({
  region: TABLE_REGION
});

export class DynamoDBAdapter implements Adapter {
  name: string;

  constructor(name: string) {
    this.name = name;
  }

  async upsert(id: string, payload: AdapterPayload, expiresIn?: number): Promise<void> {
    // DynamoDB can recognise TTL values only in seconds
    const expiresAt = expiresIn ? Math.floor(Date.now() / 1000) + expiresIn : null;

    await dynamoClient.send(
      new UpdateItemCommand({
        TableName: TABLE_NAME,
        Key: marshall({ modelId: this.name + '-' + id }),
        UpdateExpression:
          'SET payload = :payload' +
          (expiresAt ? ', expiresAt = :expiresAt' : '') +
          (payload.userCode ? ', userCode = :userCode' : '') +
          (payload.uid ? ', uid = :uid' : '') +
          (payload.grantId ? ', grantId = :grantId' : ''),
        ExpressionAttributeValues: marshall(
          {
            ':payload': payload,
            ...(expiresAt ? { ':expiresAt': expiresAt } : {}),
            ...(payload.userCode ? { ':userCode': payload.userCode } : {}),
            ...(payload.uid ? { ':uid': payload.uid } : {}),
            ...(payload.grantId ? { ':grantId': payload.grantId } : {})
          },
          { removeUndefinedValues: true }
        )
      })
    );
  }

  async find(id: string): Promise<AdapterPayload | undefined> {
    const response = await dynamoClient.send(
      new GetItemCommand({
        TableName: TABLE_NAME,
        Key: marshall({ modelId: this.name + '-' + id }),
        ProjectionExpression: 'payload, expiresAt'
      })
    );

    if (!response.Item) {
      return undefined;
    }

    const result = unmarshall(response.Item);

    // DynamoDB can take upto 48 hours to drop expired items, so a check is required
    if (result.expiresAt && Date.now() > result.expiresAt * 1000) {
      return undefined;
    }

    return result.payload;
  }

  async findByUserCode(userCode: string): Promise<AdapterPayload | undefined> {
    const response = await dynamoClient.send(
      new QueryCommand({
        TableName: TABLE_NAME,
        IndexName: 'userCodeIndex',
        KeyConditionExpression: 'userCode = :userCode',
        ExpressionAttributeValues: marshall(
          {
            ':userCode': userCode
          },
          { removeUndefinedValues: true }
        ),
        Limit: 1,
        ProjectionExpression: 'payload, expiresAt'
      })
    );

    if (!response.Items || !response.Items[0]) {
      return undefined;
    }

    const result = unmarshall(response.Items[0]);

    // DynamoDB can take upto 48 hours to drop expired items, so a check is required
    if (result.expiresAt && Date.now() > result.expiresAt * 1000) {
      return undefined;
    }

    return result.payload;
  }

  async findByUid(uid: string): Promise<AdapterPayload | undefined> {
    const response = await dynamoClient.send(
      new QueryCommand({
        TableName: TABLE_NAME,
        IndexName: 'uidIndex',
        KeyConditionExpression: 'uid = :uid',
        ExpressionAttributeValues: marshall(
          {
            ':uid': uid
          },
          { removeUndefinedValues: true }
        ),
        Limit: 1,
        ProjectionExpression: 'payload, expiresAt'
      })
    );

    if (!response.Items || !response.Items[0]) {
      return undefined;
    }

    const result = unmarshall(response.Items[0]);

    // DynamoDB can take upto 48 hours to drop expired items, so a check is required
    if (result.expiresAt && Date.now() > result.expiresAt * 1000) {
      return undefined;
    }

    return result.payload;
  }

  async consume(id: string): Promise<void> {
    await dynamoClient.send(
      new UpdateItemCommand({
        TableName: TABLE_NAME,
        Key: marshall({ modelId: this.name + '-' + id }),
        UpdateExpression: 'SET #payload.#consumed = :value',
        ExpressionAttributeNames: {
          '#payload': 'payload',
          '#consumed': 'consumed'
        },
        ExpressionAttributeValues: marshall(
          {
            ':value': Math.floor(Date.now() / 1000)
          },
          { removeUndefinedValues: true }
        ),
        ConditionExpression: 'attribute_exists(modelId)'
      })
    );
  }

  async destroy(id: string): Promise<void> {
    await dynamoClient.send(
      new DeleteItemCommand({
        TableName: TABLE_NAME,
        Key: marshall({ modelId: this.name + '-' + id })
      })
    );
  }

  async revokeByGrantId(grantId: string): Promise<void> {
    console.log(`Revoking items with grantId ${grantId} for model ${this.name}`);
    let ExclusiveStartKey: Record<string, AttributeValue> | undefined = undefined;

    do {
      const queryResult: QueryCommandOutput = await dynamoClient.send(
        new QueryCommand({
          TableName: TABLE_NAME,
          IndexName: 'grantIdIndex',
          KeyConditionExpression: 'grantId = :grantId',
          ExpressionAttributeValues: marshall(
            {
              ':grantId': grantId
            },
            { removeUndefinedValues: true }
          ),
          ProjectionExpression: 'modelId',
          Limit: 25,
          ExclusiveStartKey
        })
      );
      ExclusiveStartKey = queryResult.LastEvaluatedKey;

      if (!queryResult.Items || !queryResult.Items.length) {
        return;
      }

      await dynamoClient.send(
        new BatchWriteItemCommand({
          RequestItems: {
            [TABLE_NAME]: queryResult.Items.map((item: Record<string, AttributeValue>) => ({
              DeleteRequest: {
                Key: { modelId: item.modelId }
              }
            }))
          }
        })
      );
    } while (ExclusiveStartKey);
  }
}