feat: implement the jwks infrastructure using pulumi

index.ts

@@ -0,0 +1 @@
+export * as Stack from "./src";

src/composition-root.ts

@@ -0,0 +1,137 @@
+import * as aws from "@pulumi/aws";
+import { JWKS } from ".";
+import * as Infrastructure from "./infrastructure";
+import * as InfrastructureAdapters from "./infrastructure-adapters";
+import * as pulumi from "@pulumi/pulumi";
+
+export interface Configuration<AccessTokenPayload, RefreshTokenPayload> {
+  customDomain?: aws.apigateway.DomainName;
+  jwksFileKey?: string;
+  kidVersionStagePrefix?: string;
+  kidVersionStagePrefixSeparator?: string;
+  replicas: Infrastructure.Region[];
+  secretsRecoveryWindowInDays: number;
+  region: string;
+  resourcesPrefix: string;
+  s3Bucket?: aws.s3.BucketV2;
+}
+
+export interface Output {
+  accessTokenSecretARN: pulumi.Output<string>;
+  accessTokenSecretID: pulumi.Output<string>;
+  accessTokenSecretName: pulumi.Output<string>;
+  jwksUri: pulumi.Output<string>;
+  keySetRotationProcessorARN: pulumi.Output<string>;
+  keySetRotationProcessorID: pulumi.Output<string>;
+  keySetRotationProcessorName: pulumi.Output<string>;
+  kidVersionStagePrefix: string;
+  kidVersionStagePrefixSeparator: string;
+  pulumiProjectName: string;
+  pulumiStackName: string;
+  refreshTokenSecretARN: pulumi.Output<string>;
+  refreshTokenSecretID: pulumi.Output<string>;
+  refreshTokenSecretName: pulumi.Output<string>;
+  s3BucketARN: pulumi.Output<string>;
+  s3BucketID: pulumi.Output<string>;
+  s3BucketName: pulumi.Output<string>;
+  tableARN: pulumi.Output<string>;
+  tableName: pulumi.Output<string>;
+  tableID: pulumi.Output<string>;
+}
+
+export const compose = <AccessTokenPayload, RefreshTokenPayload>(
+  configuration: Configuration<AccessTokenPayload, RefreshTokenPayload>,
+): Output => {
+  const { resourcesPrefix, region, replicas, secretsRecoveryWindowInDays } =
+    configuration;
+
+  const kidVersionStagePrefix = configuration.kidVersionStagePrefix ?? "kid";
+  const kidVersionStagePrefixSeparator =
+    configuration.kidVersionStagePrefixSeparator ?? "#";
+
+  const jwksFileKey = configuration.jwksFileKey ?? ".well-known/jwks.json";
+
+  const jwksTable = Infrastructure.DynamoDB.jwksTable({
+    prefix: resourcesPrefix,
+    replicas: replicas,
+  });
+
+  const s3Bucket = Infrastructure.S3.configureJwksBucket({
+    jwksFileKey,
+    prefix: configuration.resourcesPrefix,
+    providedBucket: configuration.s3Bucket,
+  });
+
+  const jwksModuleProps = {
+    jwksFileKey,
+  };
+
+  const jwksModule: JWKS.API.JWKSModule = JWKS.Bindings.create({
+    jwksTable: jwksTable.table,
+    region,
+    s3Bucket,
+  })(jwksModuleProps);
+
+  const keySetRotationProcessorProps = {
+    jwksFileKey,
+    prefix: configuration.resourcesPrefix,
+  };
+
+  const keySetRotationProcessor =
+    InfrastructureAdapters.KeySetRotationProcessor.create({
+      jwksModule,
+      table: jwksTable,
+      s3Bucket,
+    })(keySetRotationProcessorProps);
+
+  const privateKeyStore = Infrastructure.SecretsManager.create({
+    prefix: resourcesPrefix,
+    recoveryWindowInDays: secretsRecoveryWindowInDays,
+    replicas: replicas,
+  });
+  const accessTokenSecret = privateKeyStore.provisionSecret({
+    name: "AccessTokenPrivateKey",
+    description: "Secret for the access token private key",
+  });
+  const refreshTokenSecret = privateKeyStore.provisionSecret({
+    name: "RefreshTokenPrivateKey",
+    description: "Secret for the refresh token private key",
+  });
+
+  const jwksUri = (() => {
+    if (configuration.customDomain) {
+      return configuration.customDomain.domainName.apply((domainName) => {
+        return pulumi.output(`https://${domainName}/${jwksFileKey}`);
+      });
+    } else {
+      return s3Bucket.bucket.apply((bucket) => {
+        return pulumi.output(
+          `https://${bucket}.s3.amazonaws.com/${jwksFileKey}`,
+        );
+      });
+    }
+  })();
+
+  return {
+    accessTokenSecretARN: accessTokenSecret.arn,
+    accessTokenSecretID: accessTokenSecret.id,
+    accessTokenSecretName: accessTokenSecret.name,
+    jwksUri,
+    keySetRotationProcessorARN: keySetRotationProcessor.arn,
+    keySetRotationProcessorID: keySetRotationProcessor.id,
+    keySetRotationProcessorName: keySetRotationProcessor.name,
+    kidVersionStagePrefix,
+    kidVersionStagePrefixSeparator,
+    pulumiProjectName: pulumi.getProject(),
+    pulumiStackName: pulumi.getStack(),
+    refreshTokenSecretARN: refreshTokenSecret.arn,
+    refreshTokenSecretID: refreshTokenSecret.id,
+    refreshTokenSecretName: refreshTokenSecret.name,
+    s3BucketARN: s3Bucket.arn,
+    s3BucketID: s3Bucket.id,
+    s3BucketName: s3Bucket.bucket,
+    tableARN: jwksTable.table.arn,
+    tableName: jwksTable.table.name,
+    tableID: jwksTable.table.id,
+  };
+};

src/index.ts

@@ -0,0 +1,2 @@
+export * as CompositionRoot from "./composition-root";
+export * as JWKS from "./module";

src/infrastructure-adapters/index.ts

@@ -0,0 +1,2 @@
+export * as KeySetRotationProcessor from "./key-set-rotation-processor";
+export { PrivateKeyStore } from "./private-key-store";

src/infrastructure-adapters/key-set-rotation-processor.ts

@@ -0,0 +1,176 @@
+import * as aws from "@pulumi/aws";
+import { JWKS } from "..";
+import * as pulumi from "@pulumi/pulumi";
+
+export interface Dependencies {
+  jwksModule: JWKS.API.JWKSModule;
+  s3Bucket: aws.s3.BucketV2;
+  table: {
+    table: aws.dynamodb.Table;
+    policy: aws.iam.Policy;
+  };
+}
+
+export interface Props {
+  jwksFileKey: string;
+  prefix: string;
+}
+
+export const create = (deps: Dependencies) => (props: Props) => {
+  const { jwksModule, s3Bucket, table } = deps;
+  const { jwksFileKey, prefix } = props;
+
+  // Step 2: Create the Lambda Function
+  const lambdaRole = new aws.iam.Role(
+    `${prefix}-jwksKeySetRotationProcessorLambda`,
+    {
+      assumeRolePolicy: aws.iam.assumeRolePolicyForPrincipal({
+        Service: "lambda.amazonaws.com",
+      }),
+    },
+  );
+
+  const streamPolicy = new aws.iam.RolePolicy(
+    `${prefix}-jwksKeySetRotationProcessorLambdaStreamPolicy`,
+    {
+      role: lambdaRole.id,
+      policy: pulumi.output({
+        Version: "2012-10-17",
+        Statement: [
+          {
+            Effect: "Allow",
+            Action: [
+              "dynamodb:ListStreams",
+              "dynamodb:DescribeStream",
+              "dynamodb:GetRecords",
+              "dynamodb:GetShardIterator",
+            ],
+            Resource: table.table.streamArn,
+          },
+          {
+            Effect: "Allow",
+            Action: [
+              "dynamodb:ListStreams",
+              "dynamodb:DescribeStream",
+              "dynamodb:GetRecords",
+              "dynamodb:GetShardIterator",
+            ],
+            Resource: table.table.streamArn,
+          },
+          {
+            Action: "logs:*",
+            Effect: "Allow",
+            Resource: "arn:aws:logs:*:*:*",
+          },
+        ],
+      }),
+    },
+    {
+      dependsOn: [table.table, table.policy],
+    },
+  );
+
+  const dynamoDBPolicy = new aws.iam.RolePolicy(
+    `${prefix}-jwksKeySetRotationProcessorLambdaDynamoDBPolicy`,
+    {
+      role: lambdaRole.id,
+      policy: pulumi.output({
+        Version: "2012-10-17",
+        Statement: [
+          {
+            Effect: "Allow",
+            Action: [
+              "dynamodb:GetItem",
+              "dynamodb:BatchGetItem",
+              "dynamodb:Scan",
+              "dynamodb:Query",
+              "dynamodb:ConditionCheckItem",
+            ],
+            Resource: table.table.arn,
+          },
+        ],
+      }),
+    },
+    {
+      dependsOn: [table.table, table.policy],
+    },
+  );
+
+  const s3Policy = new aws.iam.RolePolicy(
+    `${prefix}-jwksKeySetRotationProcessorLambdaS3Policy`,
+    {
+      role: lambdaRole.id,
+      policy: s3Bucket.bucket.apply((bucketName) =>
+        JSON.stringify({
+          Version: "2012-10-17",
+          Statement: [
+            {
+              Effect: "Allow",
+              Action: [
+                "s3:ListBucket",
+                "s3:PutObject",
+                "s3:PutObjectAcl",
+                "s3:PutObjectVersionAcl",
+              ],
+              Resource: `arn:aws:s3:::${bucketName}/${jwksFileKey}`,
+            },
+          ],
+        }),
+      ),
+    },
+    {
+      dependsOn: [table.table, table.policy],
+    },
+  );
+
+  const dynamoDBPolicyAttachment = new aws.iam.RolePolicyAttachment(
+    `${prefix}-lambdaPolicyAttachment`,
+    {
+      role: lambdaRole.name, // Use the Lambda function's role
+      policyArn: table.policy.arn, // The ARN of the policy to attach
+    },
+    {
+      dependsOn: [table.table, table.policy, lambdaRole, streamPolicy],
+    },
+  );
+
+  const lambda = new aws.lambda.CallbackFunction(
+    `${prefix}-jwks-key-set-rotation-processor`,
+    {
+      architectures: ["arm64"],
+      callback: async () => {
+        await jwksModule.keySetRotationProcessor.execute();
+      },
+      role: lambdaRole.arn,
+    },
+  );
+
+  // Step 3: Create an Event Source Mapping
+  const eventSourceMapping = new aws.lambda.EventSourceMapping(
+    `${prefix}-eventSourceMapping`,
+    {
+      eventSourceArn: table.table.streamArn,
+      functionName: lambda.arn,
+      startingPosition: "TRIM_HORIZON", // To avoid missing any events
+    },
+    {
+      dependsOn: [table.table, table.policy, dynamoDBPolicyAttachment],
+    },
+  );
+
+  // Step 4: Manage Permissions
+  const lambdaInvokePermission = new aws.lambda.Permission(
+    `${prefix}-lambdaInvokePermission`,
+    {
+      action: "lambda:InvokeFunction",
+      function: lambda.name,
+      principal: "dynamodb.amazonaws.com",
+      sourceArn: table.table.streamArn,
+    },
+    {
+      dependsOn: [table.table, table.policy],
+    },
+  );
+
+  return lambda;
+};

src/infrastructure-adapters/private-key-store.ts

@@ -0,0 +1,15 @@
+import * as aws from "@pulumi/aws";
+
+export interface Props {
+  name: string;
+  description: string;
+  secret: string;
+  versionStages: string[];
+}
+
+export interface PrivateKeyStore {
+  provisionSecret: (props: {
+    name: string;
+    description: string;
+  }) => aws.secretsmanager.Secret;
+}

src/infrastructure/dynamo-db.ts

@@ -0,0 +1,55 @@
+import * as aws from "@pulumi/aws";
+import * as pulumi from "@pulumi/pulumi";
+import { Region } from ".";
+
+export interface TableProps {
+  prefix: string;
+  replicas: Region[];
+}
+
+export const jwksTable = (
+  props: TableProps,
+): { table: aws.dynamodb.Table; policy: aws.iam.Policy } => {
+  const { prefix, replicas } = props;
+  const table = new aws.dynamodb.Table(`${prefix}-jwks`, {
+    attributes: [{ name: "pk", type: "S" }],
+    hashKey: "pk",
+    billingMode: "PAY_PER_REQUEST",
+
+    // Enabled point in time recovery
+    pointInTimeRecovery: { enabled: true },
+
+    ttl: { attributeName: "delete_at", enabled: true },
+
+    streamEnabled: true,
+    streamViewType: "KEYS_ONLY",
+
+    // Specify the replicas for the global table
+    replicas: replicas.map((region) => ({ regionName: region })),
+  });
+
+  const policy = new aws.iam.Policy(`${prefix}-dynamoPolicy`, {
+    description: "IAM policy for DynamoDB read/write access",
+    policy: pulumi.output({
+      Version: "2012-10-17",
+      Statement: [
+        {
+          Effect: "Allow",
+          Action: [
+            "dynamodb:GetItem",
+            "dynamodb:PutItem",
+            "dynamodb:UpdateItem",
+            "dynamodb:DeleteItem",
+            "dynamodb:Scan",
+            "dynamodb:Query",
+            "dynamodb:BatchGetItem",
+            "dynamodb:BatchWriteItem",
+          ],
+          Resource: table.arn,
+        },
+      ],
+    }),
+  });
+
+  return { table, policy };
+};

src/infrastructure/index.ts

@@ -0,0 +1,4 @@
+export * as DynamoDB from "./dynamo-db";
+export { Region } from "./region";
+export * as S3 from "./s3";
+export * as SecretsManager from "./secrets-manager";