KMS integration with AWS Key Management Service

[panva]

Note

Author: @panva
Target version: ^9.5.0

Warning

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

Caution

ExternalSigningKey support is an experimental feature.

This example enables PS256, ES256, and ML-DSA-44 signatures produced by private keys stored in AWS Key Management Service.

import * as crypto from "node:crypto";

import * as KMS from "@aws-sdk/client-kms";
import { derToJose } from "ecdsa-sig-formatter";
import { Provider, ExternalSigningKey } from "oidc-provider";

const KeySpecs = {
  RSA_2048: 'PS256',
  RSA_3072: 'PS256',
  RSA_4096: 'PS256',
  ECC_NIST_P256: 'ES256',
  ML_DSA_44: 'ML-DSA-44',
};

const SigningAlgorithmSpecs = {
  PS256: KMS.SigningAlgorithmSpec.RSASSA_PSS_SHA_256,
  ES256: KMS.SigningAlgorithmSpec.ECDSA_SHA_256,
  'ML-DSA-44': KMS.SigningAlgorithmSpec.ML_DSA_SHAKE_256,
};

const client = new KMS.KMSClient();

function shake256(data, outputLength) {
  return crypto.hash('shake256', data, {
    outputEncoding: 'buffer',
    outputLength,
  });
}

function sha256(data) {
  return crypto.hash('sha256', data, 'buffer');
}

class AwsKmsKey extends ExternalSigningKey {
  #KeyId;

  #alg;

  #publicKeyObject;

  #SigningAlgorithm;

  constructor(KeyId, publicKey, alg) {
    super();
    this.#KeyId = KeyId;
    this.#alg = alg;
    this.#publicKeyObject = publicKey;
    this.#SigningAlgorithm = SigningAlgorithmSpecs[alg];
  }

  get alg() {
    return this.#alg;
  }

  keyObject() {
    return this.#publicKeyObject;
  }

  #getMessageParams(data) {
    if (this.#SigningAlgorithm === KMS.SigningAlgorithmSpec.ML_DSA_SHAKE_256) {
      const pubKeyBytes = Buffer.from(this.pub, 'base64url');
      const hashedPubKey = shake256(pubKeyBytes, 64);

      const externalMu = shake256(Buffer.concat([
        hashedPubKey,
        new Uint8Array([0x00, 0x00]),
        data,
      ]), 64);

      return {
        MessageType: KMS.MessageType.EXTERNAL_MU,
        Message: externalMu,
      };
    }

    return {
      MessageType: KMS.MessageType.DIGEST,
      Message: sha256(data),
    };
  }

  async sign(data) {
    const signOutput = await client.send(
      new KMS.SignCommand({
        KeyId: this.#KeyId,
        SigningAlgorithm: this.#SigningAlgorithm,
        ...this.#getMessageParams(data),
      }),
    );

    if (!signOutput.Signature) {
      throw new Error('unexpected SignCommand response', { cause: signOutput });
    }

    if (this.#SigningAlgorithm === KMS.SigningAlgorithmSpec.ECDSA_SHA_256) {
      return Buffer.from(derToJose(Buffer.from(signOutput.Signature), this.alg), 'base64url');
    }

    return signOutput.Signature;
  }
}

export async function setupKey(KeyId) {
  const describeKeyOutput = await client.send(
    new KMS.DescribeKeyCommand({
      KeyId,
    }),
  );

  if (!describeKeyOutput.KeyMetadata?.KeyId) {
    throw new Error('unexpected DescribeKeyCommand response', { cause: describeKeyOutput });
  }

  if (!describeKeyOutput.KeyMetadata.Enabled) {
    throw new Error('unexpected DescribeKeyCommandOutput.KeyMetadata.Enabled value', {
      cause: describeKeyOutput,
    });
  }

  const alg = KeySpecs[describeKeyOutput.KeyMetadata.KeySpec];

  if (!alg) {
    throw new Error('unsupported DescribeKeyCommandOutput.KeyMetadata.KeySpec value', {
      cause: describeKeyOutput,
    });
  }

  const expectedSigningAlgorithm = Object
    .fromEntries(Object
      .entries(SigningAlgorithmSpecs)
      .map(([key, value]) => [value, key]))[alg];

  if (!describeKeyOutput.KeyMetadata.SigningAlgorithms?.includes(expectedSigningAlgorithm)) {
    throw new Error(
      'unsupported DescribeKeyCommandOutput.KeyMetadata.SigningAlgorithms values',
      {
        cause: describeKeyOutput,
      },
    );
  }

  switch (describeKeyOutput.KeyMetadata.KeyUsage) {
    case KMS.KeyUsageType.SIGN_VERIFY:
      break;
    default:
      throw new Error('unsupported DescribeKeyCommandOutput.KeyMetadata.KeyUsage value', {
        cause: describeKeyOutput,
      });
  }

  const getPublicKeyOutput = await client.send(
    new KMS.GetPublicKeyCommand({
      KeyId,
    }),
  );

  if (!(getPublicKeyOutput.PublicKey instanceof Uint8Array)) {
    throw new Error('unexpected GetPublicKeyCommandOutput.PublicKey value', {
      cause: getPublicKeyOutput,
    });
  }

  const keyObject = crypto.createPublicKey({
    format: 'der',
    type: 'spki',
    key: getPublicKeyOutput.PublicKey,
  });

  return new AwsKmsKey(KeyId, keyObject, alg);
}

const provider = new Provider("/* your issuer identifier */", {
  jwks: {
    keys: await Promise.all([setupKey("<KeyId 1>"), setupKey("<KeyId 2>")]),
  },
  features: {
    externalSigningSupport: { enabled: true, ack: "experimental-01" },
  },
});