Handle Cognito as a MultiKey

src/api/keyless.ts

@@ -1,14 +1,15 @@
 // Copyright © Aptos Foundation
 // SPDX-License-Identifier: Apache-2.0
 
-import { Account, EphemeralKeyPair, KeylessAccount, ProofFetchCallback } from "../account";
+import { Account, EphemeralKeyPair, KeylessAccount, ProofFetchCallback, MultiKeyAccount } from "../account";
 import { FederatedKeylessAccount } from "../account/FederatedKeylessAccount";
 import { AccountAddressInput, ZeroKnowledgeSig } from "../core";
 import {
   deriveKeylessAccount,
   getPepper,
   getProof,
   updateFederatedKeylessJwkSetTransaction,
+  deriveCognitoKeylessAccount,
 } from "../internal/keyless";
 import { InputGenerateTransactionOptions, SimpleTransaction } from "../transactions";
 import { HexInput } from "../types";
@@ -203,6 +204,59 @@ export class Keyless {
     return deriveKeylessAccount({ aptosConfig: this.config, ...args });
   }
 
+  /**
+   * Derives a Cognito Keyless Account from the provided JWT token and corresponding EphemeralKeyPair. This does the same as
+   * deriveKeylessAccount but derives the account as a MultiKeyAccount in order to support different JWT formattings.
+   *
+   * @param args - The arguments required to derive the Cognito Keyless Account.
+   * @param args.jwt - The JWT token used for deriving the account.
+   * @param args.ephemeralKeyPair - The EphemeralKeyPair used to generate the nonce in the JWT token.
+   * @param args.jwkAddress - The address the where the JWKs used to verify signatures are found.  Setting the value derives a
+   * FederatedKeylessAccount.
+   * @param args.uidKey - An optional key in the JWT token to set the uidVal in the IdCommitment.
+   * @param args.pepper - An optional pepper value.
+   * @param args.proofFetchCallback - An optional callback function for fetching the proof in the background, allowing for a more
+   * responsive user experience.
+   *
+   * @returns A MultiKeyAccount that can be used to sign transactions.
+   *
+   * @example
+   * ```typescript
+   * import { Aptos, AptosConfig, Network, deriveKeylessAccount } from "@aptos-labs/ts-sdk";
+   *
+   * const config = new AptosConfig({ network: Network.TESTNET });
+   * const aptos = new Aptos(config);
+   *
+   * async function runExample() {
+   *   const jwt = "your_jwt_token"; // replace with a real JWT token
+   *   const ephemeralKeyPair = new EphemeralKeyPair(); // create a new ephemeral key pair
+   *
+   *   // Deriving the Keyless Account
+   *   const keylessAccount = await deriveCognitoKeylessAccount({
+   *     jwt,
+   *     ephemeralKeyPair,
+   *     jwkAddress: "your_jwk_address",
+   *     uidKey: "your_uid_key", // optional
+   *     pepper: "your_pepper", // optional
+   *   });
+   *
+   *   console.log("Keyless Account derived:", keylessAccount);
+   * }
+   * runExample().catch(console.error);
+   * ```
+   * @group Keyless
+   */
+  async deriveCognitoKeylessAccount(args: {
+    jwt: string;
+    ephemeralKeyPair: EphemeralKeyPair;
+    jwkAddress: AccountAddressInput;
+    uidKey?: string;
+    pepper?: HexInput;
+    proofFetchCallback?: ProofFetchCallback;
+  }): Promise<MultiKeyAccount> {
+    return deriveCognitoKeylessAccount({ aptosConfig: this.config, ...args });
+  }
+
   /**
    * This installs a set of FederatedJWKs at an address for a given iss.
    *

src/core/crypto/federatedKeyless.ts

@@ -6,7 +6,7 @@ import { Deserializer, Serializer } from "../../bcs";
 import { HexInput, AnyPublicKeyVariant, SigningScheme } from "../../types";
 import { AuthenticationKey } from "../authenticationKey";
 import { AccountAddress, AccountAddressInput } from "../accountAddress";
-import { KeylessPublicKey, KeylessSignature } from "./keyless";
+import { getIssAudAndUidValWithoutUnescaping, KeylessPublicKey, KeylessSignature } from "./keyless";
 
 /**
  * Represents the FederatedKeylessPublicKey public key
@@ -116,6 +116,20 @@ export class FederatedKeylessPublicKey extends AccountPublicKey {
     return new FederatedKeylessPublicKey(args.jwkAddress, KeylessPublicKey.fromJwtAndPepper(args));
   }
 
+  static fromJwtAndPepperWithoutUnescaping(args: {
+    jwt: string;
+    pepper: HexInput;
+    jwkAddress: AccountAddressInput;
+    uidKey?: string;
+  }): FederatedKeylessPublicKey {
+    const { jwt, pepper, uidKey = "sub" } = args;
+    const { iss, aud, uidVal } = getIssAudAndUidValWithoutUnescaping({ jwt, uidKey });
+    return new FederatedKeylessPublicKey(
+      args.jwkAddress,
+      KeylessPublicKey.create({ iss, uidKey, uidVal, aud, pepper }),
+    );
+  }
+
   static isInstance(publicKey: PublicKey) {
     return (
       "jwkAddress" in publicKey &&

src/core/crypto/keyless.ts

@@ -34,6 +34,7 @@ import { memoizeAsync } from "../../utils/memoize";
 import { AccountAddress, AccountAddressInput } from "../accountAddress";
 import { getErrorMessage } from "../../utils";
 import { KeylessError, KeylessErrorType } from "../../errors";
+import { getClaimWithoutUnescaping } from "./utils";
 
 /**
  * @group Implementation
@@ -264,15 +265,9 @@ export class KeylessPublicKey extends AccountPublicKey {
    */
   static fromJwtAndPepper(args: { jwt: string; pepper: HexInput; uidKey?: string }): KeylessPublicKey {
     const { jwt, pepper, uidKey = "sub" } = args;
-    const jwtPayload = jwtDecode<JwtPayload & { [key: string]: string }>(jwt);
-    if (typeof jwtPayload.iss !== "string") {
-      throw new Error("iss was not found");
-    }
-    if (typeof jwtPayload.aud !== "string") {
-      throw new Error("aud was not found or an array of values");
-    }
-    const uidVal = jwtPayload[uidKey];
-    return KeylessPublicKey.create({ iss: jwtPayload.iss, uidKey, uidVal, aud: jwtPayload.aud, pepper });
+
+    const { iss, aud, uidVal } = getIssAudAndUidVal({ jwt, uidKey });
+    return KeylessPublicKey.create({ iss, uidKey, uidVal, aud, pepper });
   }
 
   /**
@@ -932,6 +927,29 @@ export function getIssAudAndUidVal(args: { jwt: string; uidKey?: string }): {
   return { iss: jwtPayload.iss, aud: jwtPayload.aud, uidVal };
 }
 
+/**
+ * Parses a JWT and returns the 'iss', 'aud', and 'uid' values without unescaping the values.
+ *
+ * @param args - The arguments for parsing the JWT.
+ * @param args.jwt - The JWT to parse.
+ * @param args.uidKey - The key to use for the 'uid' value; defaults to 'sub'.
+ * @returns The 'iss', 'aud', and 'uid' values from the JWT.
+ */
+export function getIssAudAndUidValWithoutUnescaping(args: { jwt: string; uidKey?: string }): {
+  iss: string;
+  aud: string;
+  uidVal: string;
+} {
+  const { jwt, uidKey = "sub" } = args;
+  // This checks that the JWT is valid and that the iss, aud, and uidVal are present
+  getIssAudAndUidVal(args);
+  return {
+    iss: getClaimWithoutUnescaping(jwt, "iss"),
+    aud: getClaimWithoutUnescaping(jwt, "aud"),
+    uidVal: getClaimWithoutUnescaping(jwt, uidKey),
+  };
+}
+
 /**
  * Retrieves the KeylessConfiguration set on chain.
  *

src/core/crypto/utils.ts

@@ -24,3 +24,60 @@ export const convertSigningMessage = (message: HexInput): HexInput => {
   // message is a Uint8Array
   return message;
 };
+
+function b64DecodeUnicode(str: string) {
+  return decodeURIComponent(
+    atob(str).replace(/(.)/g, (m, p) => {
+      let code = (p as string).charCodeAt(0).toString(16).toUpperCase();
+      if (code.length < 2) {
+        code = `0${code}`;
+      }
+      return `%${code}`;
+    }),
+  );
+}
+
+function base64UrlDecode(str: string) {
+  let output = str.replace(/-/g, "+").replace(/_/g, "/");
+  switch (output.length % 4) {
+    case 0:
+      break;
+    case 2:
+      output += "==";
+      break;
+    case 3:
+      output += "=";
+      break;
+    default:
+      throw new Error("base64 string is not of the correct length");
+  }
+
+  try {
+    return b64DecodeUnicode(output);
+  } catch (err) {
+    return atob(output);
+  }
+}
+
+export function getClaimWithoutUnescaping(jwt: string, claim: string): string {
+  const parts = jwt.split(".");
+  const payload = parts[1];
+  const payloadStr = base64UrlDecode(payload);
+  const claimIdx = payloadStr.indexOf(`"${claim}"`) + claim.length + 2;
+  let claimVal = "";
+  let foundStart = false;
+  for (let i = claimIdx; i < payloadStr.length; i += 1) {
+    if (payloadStr[i] === "\"") {
+      if (foundStart) {
+        break;
+      }
+      foundStart = true;
+      // eslint-disable-next-line no-continue
+      continue;
+    }
+    if (foundStart) {
+      claimVal += payloadStr[i];
+    }
+  }
+  return claimVal;
+}

src/internal/keyless.ts

@@ -18,12 +18,20 @@ import {
   Hex,
   KeylessPublicKey,
   MoveJWK,
+  MultiKey,
   ZeroKnowledgeSig,
   ZkProof,
+  getIssAudAndUidVal,
   getKeylessConfig,
 } from "../core";
 import { HexInput, ZkpVariant } from "../types";
-import { Account, EphemeralKeyPair, KeylessAccount, ProofFetchCallback } from "../account";
+import {
+  Account,
+  EphemeralKeyPair,
+  KeylessAccount,
+  MultiKeyAccount,
+  ProofFetchCallback,
+} from "../account";
 import { PepperFetchRequest, PepperFetchResponse, ProverRequest, ProverResponse } from "../types/keyless";
 import { lookupOriginalAccountAddress } from "./account";
 import { FederatedKeylessPublicKey } from "../core/crypto/federatedKeyless";
@@ -32,7 +40,7 @@ import { MoveVector } from "../bcs";
 import { generateTransaction } from "./transactionSubmission";
 import { InputGenerateTransactionOptions, SimpleTransaction } from "../transactions";
 import { KeylessError, KeylessErrorType } from "../errors";
-import { FIREBASE_AUTH_ISS_PATTERN } from "../utils/const";
+import { COGNITO_ISS_PATTERN, FIREBASE_AUTH_ISS_PATTERN } from "../utils/const";
 
 /**
  * Retrieves a pepper value based on the provided configuration and authentication details.
@@ -188,6 +196,9 @@ export async function deriveKeylessAccount(args: {
   proofFetchCallback?: ProofFetchCallback;
 }): Promise<KeylessAccount | FederatedKeylessAccount> {
   const { aptosConfig, jwt, jwkAddress, uidKey, proofFetchCallback, pepper = await getPepper(args) } = args;
+  if (isCognito(jwt)) {
+    throw new Error("This API does not support Cognito JWTs. Please use deriveCognitoKeylessAccount instead.");
+  };
   const { verificationKey, maxExpHorizonSecs } = await getKeylessConfig({ aptosConfig });
 
   const proofPromise = getProof({ ...args, pepper, maxExpHorizonSecs });
@@ -225,6 +236,56 @@ export async function deriveKeylessAccount(args: {
   return KeylessAccount.create({ ...args, address, proof, pepper, proofFetchCallback, verificationKey });
 }
 
+export async function deriveCognitoKeylessAccount(args: {
+  aptosConfig: AptosConfig;
+  jwt: string;
+  ephemeralKeyPair: EphemeralKeyPair;
+  jwkAddress: AccountAddressInput;
+  uidKey?: string;
+  pepper?: HexInput;
+  proofFetchCallback?: ProofFetchCallback;
+}): Promise<MultiKeyAccount> {
+  const { aptosConfig, jwt, jwkAddress, uidKey, proofFetchCallback, pepper = await getPepper(args) } = args;
+  if (!isCognito(jwt)) {
+    throw new Error("JWT is not a Cognito JWT");
+  };
+
+  const { verificationKey, maxExpHorizonSecs } = await getKeylessConfig({ aptosConfig });
+
+  const proofPromise = getProof({ ...args, pepper, maxExpHorizonSecs });
+  // If a callback is provided, pass in the proof as a promise to KeylessAccount.create.  This will make the proof be fetched in the
+  // background and the callback will handle the outcome of the fetch.  This allows the developer to not have to block on the proof fetch
+  // allowing for faster rendering of UX.
+  //
+  // If no callback is provided, the just await the proof fetch and continue synchronously.
+  const proof = proofFetchCallback ? proofPromise : await proofPromise;
+
+  const multiKey = new MultiKey({
+    publicKeys: [
+      FederatedKeylessPublicKey.fromJwtAndPepperWithoutUnescaping({ jwt, pepper, jwkAddress, uidKey }),
+      FederatedKeylessPublicKey.fromJwtAndPepper({ jwt, pepper, jwkAddress, uidKey }),
+    ],
+    signaturesRequired: 1,
+  });
+  const signer = FederatedKeylessAccount.create({
+    ...args,
+    proof,
+    pepper,
+    proofFetchCallback,
+    jwkAddress,
+    verificationKey,
+  });
+  return new MultiKeyAccount({
+    multiKey,
+    signers: [signer],
+  });
+}
+
+function isCognito(jwt: string): boolean {
+  const { iss } = getIssAudAndUidVal({ jwt });
+  return COGNITO_ISS_PATTERN.test(iss);
+}
+
 export interface JWKS {
   keys: MoveJWK[];
 }

src/utils/const.ts

@@ -95,3 +95,5 @@ export enum ProcessorType {
  * where project-id can contain letters, numbers, hyphens, and underscores
  */
 export const FIREBASE_AUTH_ISS_PATTERN = /^https:\/\/securetoken\.google\.com\/[a-zA-Z0-9-_]+$/;
+
+export const COGNITO_ISS_PATTERN = /^https:\/\/cognito-idp\.[a-zA-Z0-9-_]+\.amazonaws\.com\/[a-zA-Z0-9-_]+$/;