Throw error with descriptive message when a Keyless Signature is inva…

…lid due to JWK rotation (#541)

* add this

* update

* add check

* update

* update

* update

* update

* fix

* update

* update

* update iport

* update changelog

CHANGELOG.md

@@ -7,6 +7,7 @@ All notable changes to the Aptos TypeScript SDK will be captured in this file. T
 - [`Breaking`] Updated `AccountAddress.fromString` and `AccountAddress.from` to only accept SHORT strings that are 60-64 characters long by default (with the exception of special addresses). This can be adjusted using `maxMissingChars` which is set to `4` by default. If you would like to keep the previous behavior, set `maxMissingChars` to `63` for relaxed parsing.
 - Add support for AIP-80 compliant private key imports and exports through `toAIP80String`
 - Add `PrivateKey` helpers for AIP-80: `PrivateKey.parseHexInput`, `PrivateKey.formatPrivateKey`, and `PrivateKey.AIP80_PREFIXES`.
+- Adds explicit error handling Keyless accounts using `KeylessError`.  Handles JWK rotations.
 
 # 1.31.0 (2024-10-24)
 

src/account/AbstractKeylessAccount.ts

@@ -2,7 +2,8 @@
 // SPDX-License-Identifier: Apache-2.0
 
 import EventEmitter from "eventemitter3";
-import { EphemeralCertificateVariant, HexInput, SigningScheme } from "../types";
+import { jwtDecode } from "jwt-decode";
+import { EphemeralCertificateVariant, HexInput, KeylessError, KeylessErrorType, SigningScheme } from "../types";
 import { AccountAddress } from "../core/accountAddress";
 import {
   AnyPublicKey,
@@ -12,9 +13,10 @@ import {
   EphemeralCertificate,
   ZeroKnowledgeSig,
   ZkProof,
+  getPatchedJWKs,
+  MoveJWK,
 } from "../core/crypto";
 
-import { Account } from "./Account";
 import { EphemeralKeyPair } from "./EphemeralKeyPair";
 import { Hex } from "../core/hex";
 import { AccountAuthenticatorSingleKey } from "../transactions/authenticator/account";
@@ -23,12 +25,25 @@ import { deriveTransactionType, generateSigningMessage } from "../transactions/t
 import { AnyRawTransaction, AnyRawTransactionInstance } from "../transactions/types";
 import { base64UrlDecode } from "../utils/helpers";
 import { FederatedKeylessPublicKey } from "../core/crypto/federatedKeyless";
+import { Account } from "./Account";
+import { AptosConfig } from "../api/aptosConfig";
+
+/**
+ * An interface which defines if an Account utilizes Keyless signing.
+ */
+export interface KeylessSigner extends Account {
+  checkKeylessAccountValidity(aptosConfig: AptosConfig): Promise<void>;
+}
+
+export function isKeylessSigner(obj: any): obj is KeylessSigner {
+  return obj !== null && obj !== undefined && typeof obj.checkKeylessAccountValidity === "function";
+}
 
 /**
  * Account implementation for the Keyless authentication scheme.  This abstract class is used for standard Keyless Accounts
  * and Federated Keyless Accounts.
  */
-export abstract class AbstractKeylessAccount extends Serializable implements Account {
+export abstract class AbstractKeylessAccount extends Serializable implements KeylessSigner {
   static readonly PEPPER_LENGTH: number = 31;
 
   /**
@@ -231,6 +246,32 @@ export abstract class AbstractKeylessAccount extends Serializable implements Acc
     }
   }
 
+  /**
+   * Validates that the Keyless Account can be used to sign transactions.
+   * @return
+   */
+  async checkKeylessAccountValidity(aptosConfig: AptosConfig): Promise<void> {
+    if (this.isExpired()) {
+      throw KeylessError.fromErrorType({
+        type: KeylessErrorType.EPHEMERAL_KEY_PAIR_EXPIRED,
+      });
+    }
+    await this.waitForProofFetch();
+    if (this.proof === undefined) {
+      throw KeylessError.fromErrorType({
+        type: KeylessErrorType.ASYNC_PROOF_FETCH_FAILED,
+      });
+    }
+    const header = jwtDecode(this.jwt, { header: true });
+    if (header.kid === undefined) {
+      throw KeylessError.fromErrorType({
+        type: KeylessErrorType.JWT_PARSING_ERROR,
+        details: "checkKeylessAccountValidity failed. JWT is missing 'kid' in header. This should never happen.",
+      });
+    }
+    await AbstractKeylessAccount.fetchJWK({ aptosConfig, publicKey: this.publicKey, kid: header.kid });
+  }
+
   /**
    * Sign the given message using Keyless.
    * @param message in HexInput format
@@ -239,10 +280,15 @@ export abstract class AbstractKeylessAccount extends Serializable implements Acc
   sign(message: HexInput): KeylessSignature {
     const { expiryDateSecs } = this.ephemeralKeyPair;
     if (this.isExpired()) {
-      throw new Error("EphemeralKeyPair is expired");
+      throw KeylessError.fromErrorType({
+        type: KeylessErrorType.EPHEMERAL_KEY_PAIR_EXPIRED,
+      });
     }
     if (this.proof === undefined) {
-      throw new Error("Proof not found - make sure to call `await account.waitForProofFetch()` before signing.");
+      throw KeylessError.fromErrorType({
+        type: KeylessErrorType.PROOF_NOT_FOUND,
+        details: "Proof not found - make sure to call `await account.checkKeylessAccountValidity()` before signing.",
+      });
     }
     const ephemeralPublicKey = this.ephemeralKeyPair.getPublicKey();
     const ephemeralSignature = this.ephemeralKeyPair.sign(message);
@@ -264,7 +310,10 @@ export abstract class AbstractKeylessAccount extends Serializable implements Acc
    */
   signTransaction(transaction: AnyRawTransaction): KeylessSignature {
     if (this.proof === undefined) {
-      throw new Error("Proof not found - make sure to call `await account.waitForProofFetch()` before signing.");
+      throw KeylessError.fromErrorType({
+        type: KeylessErrorType.PROOF_NOT_FOUND,
+        details: "Proof not found - make sure to call `await account.checkKeylessAccountValidity()` before signing.",
+      });
     }
     const raw = deriveTransactionType(transaction);
     const txnAndProof = new TransactionAndProof(raw, this.proof.proof);
@@ -293,6 +342,56 @@ export abstract class AbstractKeylessAccount extends Serializable implements Acc
     }
     return true;
   }
+
+  /**
+   * Fetches the JWK from the issuer's well-known JWKS endpoint.
+   *
+   * @param args.publicKey The keyless public key to query
+   * @param args.kid The kid of the JWK to fetch
+   * @returns A JWK matching the `kid` in the JWT header.
+   * @throws {KeylessError} If the JWK cannot be fetched
+   */
+  static async fetchJWK(args: {
+    aptosConfig: AptosConfig;
+    publicKey: KeylessPublicKey | FederatedKeylessPublicKey;
+    kid: string;
+  }): Promise<MoveJWK> {
+    const { aptosConfig, publicKey, kid } = args;
+    const keylessPubKey = publicKey instanceof KeylessPublicKey ? publicKey : publicKey.keylessPublicKey;
+    const { iss } = keylessPubKey;
+
+    let patchedJWKs: Map<string, MoveJWK[]>;
+    try {
+      patchedJWKs = await getPatchedJWKs({ aptosConfig });
+    } catch (error) {
+      throw KeylessError.fromErrorType({
+        type: KeylessErrorType.JWK_FETCH_FAILED,
+        details: `Failed to fetch patched JWKs: ${error}`,
+      });
+    }
+
+    // Find the corresponding JWK set by `iss`
+    const jwksForIssuer = patchedJWKs.get(iss);
+
+    if (jwksForIssuer === undefined) {
+      throw KeylessError.fromErrorType({
+        type: KeylessErrorType.INVALID_JWT_ISS_NOT_RECOGNIZED,
+        details: `JWKs for issuer ${iss} not found.`,
+      });
+    }
+
+    // Find the corresponding JWK by `kid`
+    const jwk = jwksForIssuer.find((key) => key.kid === kid);
+
+    if (jwk === undefined) {
+      throw KeylessError.fromErrorType({
+        type: KeylessErrorType.INVALID_JWT_JWK_NOT_FOUND,
+        details: `JWK with kid ${kid} for issuer ${iss} not found.`,
+      });
+    }
+
+    return jwk;
+  }
 }
 
 /**

src/account/MultiKeyAccount.ts

@@ -7,7 +7,8 @@ import { AccountAddress, AccountAddressInput } from "../core/accountAddress";
 import { HexInput, SigningScheme } from "../types";
 import { AccountAuthenticatorMultiKey } from "../transactions/authenticator/account";
 import { AnyRawTransaction } from "../transactions/types";
-import { AbstractKeylessAccount } from "./AbstractKeylessAccount";
+import { AbstractKeylessAccount, KeylessSigner } from "./AbstractKeylessAccount";
+import { AptosConfig } from "../api/aptosConfig";
 
 /**
  * Arguments required to verify a multi-key signature against a given message.
@@ -28,7 +29,7 @@ export interface VerifyMultiKeySignatureArgs {
  *
  * Note: Generating a signer instance does not create the account on-chain.
  */
-export class MultiKeyAccount implements Account {
+export class MultiKeyAccount implements Account, KeylessSigner {
   /**
    * Public key associated with the account
    */
@@ -154,9 +155,21 @@ export class MultiKeyAccount implements Account {
   }
 
   /**
-   * Sign the given data using the MultiKeyAccount's signers.
-   * @param data - The data to be signed in HexInput format.
-   * @returns MultiKeySignature - The resulting multi-key signature.
+   * Validates that the Keyless Account can be used to sign transactions.
+   * @return
+   */
+  async checkKeylessAccountValidity(aptosConfig: AptosConfig): Promise<void> {
+    const keylessSigners = this.signers.filter(
+      (signer) => signer instanceof AbstractKeylessAccount,
+    ) as AbstractKeylessAccount[];
+    const promises = keylessSigners.map((signer) => signer.checkKeylessAccountValidity(aptosConfig));
+    await Promise.all(promises);
+  }
+
+  /**
+   * Sign the given message using the MultiKeyAccount's signers
+   * @param message in HexInput format
+   * @returns MultiKeySignature
    */
   sign(data: HexInput): MultiKeySignature {
     const signatures = [];

src/client/core.ts

@@ -2,9 +2,17 @@
 // SPDX-License-Identifier: Apache-2.0
 
 import { AptosConfig } from "../api/aptosConfig";
-import { AptosApiError, AptosResponse } from "./types";
 import { VERSION } from "../version";
-import { AnyNumber, AptosRequest, Client, ClientRequest, ClientResponse, MimeType } from "../types";
+import {
+  AnyNumber,
+  AptosApiError,
+  AptosRequest,
+  AptosResponse,
+  Client,
+  ClientRequest,
+  ClientResponse,
+  MimeType,
+} from "../types";
 import { AptosApiType } from "../utils";
 
 /**

src/client/get.ts

@@ -1,7 +1,6 @@
 import { AptosConfig } from "../api/aptosConfig";
 import { aptosRequest } from "./core";
-import { AptosResponse } from "./types";
-import { AnyNumber, ClientConfig, MimeType } from "../types";
+import { AptosResponse, AnyNumber, ClientConfig, MimeType } from "../types";
 import { AptosApiType } from "../utils/const";
 
 /**

src/client/index.ts

@@ -4,4 +4,3 @@
 export * from "./core";
 export * from "./get";
 export * from "./post";
-export * from "./types";

src/client/post.ts

@@ -3,8 +3,7 @@
 
 import { AptosConfig } from "../api/aptosConfig";
 import { aptosRequest } from "./core";
-import { AptosResponse } from "./types";
-import { AnyNumber, ClientConfig, MimeType } from "../types";
+import { AptosResponse, AnyNumber, ClientConfig, MimeType } from "../types";
 import { AptosApiType } from "../utils/const";
 
 /**