Handle Cognito as a MultiKey

examples/typescript/package.json

@@ -29,7 +29,9 @@
   "dependencies": {
     "@noble/curves": "^1.4.0",
     "@types/readline-sync": "^1.4.8",
+    "axios": "^1.7.7",
     "dotenv": "^16.3.1",
+    "jwt-decode": "^4.0.0",
     "npm-run-all": "latest",
     "readline-sync": "^1.4.10",
     "superagent": "^8.1.2"

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 { getIssAudAndUidVal, 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 } = getIssAudAndUidVal({ 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,28 @@ 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;
+  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 @@
   // 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,21 @@ 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,
+  KeylessSigner,
+  MultiKeyAccount,
+  ProofFetchCallback,
+} from "../account";
 import { PepperFetchRequest, PepperFetchResponse, ProverRequest, ProverResponse } from "../types/keyless";
 import { lookupOriginalAccountAddress } from "./account";
 import { FederatedKeylessPublicKey } from "../core/crypto/federatedKeyless";
@@ -32,7 +41,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.
@@ -176,7 +185,7 @@ export async function deriveKeylessAccount(args: {
   uidKey?: string;
   pepper?: HexInput;
   proofFetchCallback?: ProofFetchCallback;
-}): Promise<FederatedKeylessAccount>;
+}): Promise<FederatedKeylessAccount | MultiKeyAccount>;
 
 export async function deriveKeylessAccount(args: {
   aptosConfig: AptosConfig;
@@ -186,7 +195,7 @@ export async function deriveKeylessAccount(args: {
   uidKey?: string;
   pepper?: HexInput;
   proofFetchCallback?: ProofFetchCallback;
-}): Promise<KeylessAccount | FederatedKeylessAccount> {
+}): Promise<KeylessSigner> {
   const { aptosConfig, jwt, jwkAddress, uidKey, proofFetchCallback, pepper = await getPepper(args) } = args;
   const { verificationKey, maxExpHorizonSecs } = await getKeylessConfig({ aptosConfig });
 
@@ -200,6 +209,32 @@ export async function deriveKeylessAccount(args: {
 
   // Look up the original address to handle key rotations and then instantiate the account.
   if (jwkAddress !== undefined) {
+    if (isCognito(jwt)) {
+      const multiKey = new MultiKey({
+        publicKeys: [
+          FederatedKeylessPublicKey.fromJwtAndPepperWithoutUnescaping({ jwt, pepper, jwkAddress, uidKey }),
+          FederatedKeylessPublicKey.fromJwtAndPepper({ jwt, pepper, jwkAddress, uidKey }),
+        ],
+        signaturesRequired: 1,
+      });
+      const address = await lookupOriginalAccountAddress({
+        aptosConfig,
+        authenticationKey: multiKey.authKey().derivedAddress(),
+      });
+      const signer = FederatedKeylessAccount.create({
+        ...args,
+        address,
+        proof,
+        pepper,
+        proofFetchCallback,
+        jwkAddress,
+        verificationKey,
+      });
+      return new MultiKeyAccount({
+        multiKey,
+        signers: [signer],
+      });
+    }
     const publicKey = FederatedKeylessPublicKey.fromJwtAndPepper({ jwt, pepper, jwkAddress, uidKey });
     const address = await lookupOriginalAccountAddress({
       aptosConfig,
@@ -225,6 +260,11 @@ export async function deriveKeylessAccount(args: {
   return KeylessAccount.create({ ...args, address, proof, pepper, proofFetchCallback, verificationKey });
 }
 
+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-_]+$/;