Refines keyless error handling, handles vkey rotations, updates serialization format of KeylessAccounts

CHANGELOG.md

@@ -7,7 +7,8 @@ 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.
+- Adds explicit error handling Keyless accounts using `KeylessError`. Handles JWK rotations and Verifying Key rotations.
+- Includes the address in the `AbstractKeylessAccount` serialization to prevent information loss for key rotated accounts.
 
 # 1.31.0 (2024-10-24)
 

src/account/AbstractKeylessAccount.ts

@@ -3,7 +3,7 @@
 
 import EventEmitter from "eventemitter3";
 import { jwtDecode } from "jwt-decode";
-import { EphemeralCertificateVariant, HexInput, KeylessError, KeylessErrorType, SigningScheme } from "../types";
+import { EphemeralCertificateVariant, HexInput, SigningScheme } from "../types";
 import { AccountAddress } from "../core/accountAddress";
 import {
   AnyPublicKey,
@@ -13,20 +13,22 @@ import {
   EphemeralCertificate,
   ZeroKnowledgeSig,
   ZkProof,
-  getPatchedJWKs,
+  getKeylessJWKs,
   MoveJWK,
+  getKeylessConfig,
 } from "../core/crypto";
 
 import { EphemeralKeyPair } from "./EphemeralKeyPair";
 import { Hex } from "../core/hex";
 import { AccountAuthenticatorSingleKey } from "../transactions/authenticator/account";
-import { Serializable, Serializer } from "../bcs";
+import { Deserializer, Serializable, Serializer } from "../bcs";
 import { deriveTransactionType, generateSigningMessage } from "../transactions/transactionBuilder/signingMessage";
 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";
+import { KeylessError, KeylessErrorType } from "../errors";
 
 /**
  * An interface which defines if an Account utilizes Keyless signing.
@@ -103,6 +105,12 @@ export abstract class AbstractKeylessAccount extends Serializable implements Key
    */
   readonly jwt: string;
 
+  /**
+   * The hash of the verification key used to verify the proof. This is optional and can be used to check verifying key
+   * rotations which may invalidate the proof.
+   */
+  readonly verificationKeyHash?: Uint8Array;
+
   /**
    * An event emitter used to assist in handling asynchronous proof fetching.
    */
@@ -125,6 +133,7 @@ export abstract class AbstractKeylessAccount extends Serializable implements Key
    * @param args.proof - A Zero Knowledge Signature or a promise that resolves to one.
    * @param args.proofFetchCallback - Optional callback function for fetching proof.
    * @param args.jwt - A JSON Web Token used for authentication.
+   * @param args.verificationKeyHash Optional 32-byte verification key hash as hex input used to check proof validity.
    */
   protected constructor(args: {
     address?: AccountAddress;
@@ -138,9 +147,22 @@ export abstract class AbstractKeylessAccount extends Serializable implements Key
     proof: ZeroKnowledgeSig | Promise<ZeroKnowledgeSig>;
     proofFetchCallback?: ProofFetchCallback;
     jwt: string;
+    verificationKeyHash?: HexInput;
   }) {
     super();
-    const { address, ephemeralKeyPair, publicKey, uidKey, uidVal, aud, pepper, proof, proofFetchCallback, jwt } = args;
+    const {
+      address,
+      ephemeralKeyPair,
+      publicKey,
+      uidKey,
+      uidVal,
+      aud,
+      pepper,
+      proof,
+      proofFetchCallback,
+      jwt,
+      verificationKeyHash,
+    } = args;
     this.ephemeralKeyPair = ephemeralKeyPair;
     this.publicKey = publicKey;
     this.accountAddress = address ? AccountAddress.from(address) : this.publicKey.authKey().derivedAddress();
@@ -169,6 +191,12 @@ export abstract class AbstractKeylessAccount extends Serializable implements Key
       throw new Error(`Pepper length in bytes should be ${AbstractKeylessAccount.PEPPER_LENGTH}`);
     }
     this.pepper = pepperBytes;
+    if (verificationKeyHash !== undefined) {
+      if (Hex.hexInputToUint8Array(verificationKeyHash).length !== 32) {
+        throw new Error("verificationKeyHash must be 32 bytes");
+      }
+      this.verificationKeyHash = Hex.hexInputToUint8Array(verificationKeyHash);
+    }
   }
 
   /**
@@ -195,6 +223,7 @@ export abstract class AbstractKeylessAccount extends Serializable implements Key
    * @param serializer - The serializer instance used to convert the jwt data into bytes.
    */
   serialize(serializer: Serializer): void {
+    this.accountAddress.serialize(serializer);
     serializer.serializeStr(this.jwt);
     serializer.serializeStr(this.uidKey);
     serializer.serializeFixedBytes(this.pepper);
@@ -203,6 +232,27 @@ export abstract class AbstractKeylessAccount extends Serializable implements Key
       throw new Error("Cannot serialize - proof undefined");
     }
     this.proof.serialize(serializer);
+    serializer.serializeOption(this.verificationKeyHash, 32);
+  }
+
+  static partialDeserialize(deserializer: Deserializer): {
+    address: AccountAddress;
+    jwt: string;
+    uidKey: string;
+    pepper: Uint8Array;
+    ephemeralKeyPair: EphemeralKeyPair;
+    proof: ZeroKnowledgeSig;
+    verificationKeyHash?: Uint8Array;
+  } {
+    const address = AccountAddress.deserialize(deserializer);
+    const jwt = deserializer.deserializeStr();
+    const uidKey = deserializer.deserializeStr();
+    const pepper = deserializer.deserializeFixedBytes(31);
+    const ephemeralKeyPair = EphemeralKeyPair.deserialize(deserializer);
+    const proof = ZeroKnowledgeSig.deserialize(deserializer);
+    const verificationKeyHash = deserializer.deserializeOption("fixedBytes", 32);
+
+    return { address, jwt, uidKey, pepper, ephemeralKeyPair, proof, verificationKeyHash };
   }
 
   /**
@@ -269,6 +319,19 @@ export abstract class AbstractKeylessAccount extends Serializable implements Key
         details: "checkKeylessAccountValidity failed. JWT is missing 'kid' in header. This should never happen.",
       });
     }
+    if (this.verificationKeyHash !== undefined) {
+      const { verificationKey } = await getKeylessConfig({ aptosConfig });
+      if (Hex.hexInputToString(verificationKey.hash()) !== Hex.hexInputToString(this.verificationKeyHash)) {
+        throw KeylessError.fromErrorType({
+          type: KeylessErrorType.INVALID_PROOF_VERIFICATION_KEY_NOT_FOUND,
+        });
+      }
+    } else {
+      // eslint-disable-next-line no-console
+      console.warn(
+        "[Aptos SDK] The verification key hash was not set. Proof may be invalid if the verification key has rotated.",
+      );
+    }
     await AbstractKeylessAccount.fetchJWK({ aptosConfig, publicKey: this.publicKey, kid: header.kid });
   }
 
@@ -360,18 +423,20 @@ export abstract class AbstractKeylessAccount extends Serializable implements Key
     const keylessPubKey = publicKey instanceof KeylessPublicKey ? publicKey : publicKey.keylessPublicKey;
     const { iss } = keylessPubKey;
 
-    let patchedJWKs: Map<string, MoveJWK[]>;
+    let allJWKs: Map<string, MoveJWK[]>;
+    const jwkAddr = publicKey instanceof FederatedKeylessPublicKey ? publicKey.jwkAddress : undefined;
     try {
-      patchedJWKs = await getPatchedJWKs({ aptosConfig });
+      allJWKs = await getKeylessJWKs({ aptosConfig, jwkAddr });
     } catch (error) {
       throw KeylessError.fromErrorType({
-        type: KeylessErrorType.JWK_FETCH_FAILED,
-        details: `Failed to fetch patched JWKs: ${error}`,
+        type: KeylessErrorType.FULL_NODE_JWKS_LOOKUP_ERROR,
+        error,
+        details: `Failed to fetch ${jwkAddr ? "Federated" : "Patched"}JWKs ${jwkAddr ? `for address ${jwkAddr}` : "0x1"}`,
       });
     }
 
     // Find the corresponding JWK set by `iss`
-    const jwksForIssuer = patchedJWKs.get(iss);
+    const jwksForIssuer = allJWKs.get(iss);
 
     if (jwksForIssuer === undefined) {
       throw KeylessError.fromErrorType({
@@ -386,7 +451,7 @@ export abstract class AbstractKeylessAccount extends Serializable implements Key
     if (jwk === undefined) {
       throw KeylessError.fromErrorType({
         type: KeylessErrorType.INVALID_JWT_JWK_NOT_FOUND,
-        details: `JWK with kid ${kid} for issuer ${iss} not found.`,
+        details: `JWK with kid '${kid}' for issuer '${iss}' not found.`,
       });
     }
 

src/account/FederatedKeylessAccount.ts

@@ -1,15 +1,15 @@
 // Copyright © Aptos Foundation
 // SPDX-License-Identifier: Apache-2.0
 
-import { JwtPayload, jwtDecode } from "jwt-decode";
 import { HexInput } from "../types";
 import { AccountAddress, AccountAddressInput } from "../core/accountAddress";
-import { ZeroKnowledgeSig } from "../core/crypto";
+import { getIssAudAndUidVal, Groth16VerificationKey, ZeroKnowledgeSig } from "../core/crypto";
 
 import { EphemeralKeyPair } from "./EphemeralKeyPair";
 import { Deserializer, Serializer } from "../bcs";
 import { FederatedKeylessPublicKey } from "../core/crypto/federatedKeyless";
 import { AbstractKeylessAccount, ProofFetchCallback } from "./AbstractKeylessAccount";
+import { Hex } from "../core";
 
 /**
  * Account implementation for the FederatedKeyless authentication scheme.
@@ -54,6 +54,7 @@ export class FederatedKeylessAccount extends AbstractKeylessAccount {
     proof: ZeroKnowledgeSig | Promise<ZeroKnowledgeSig>;
     proofFetchCallback?: ProofFetchCallback;
     jwt: string;
+    verificationKeyHash?: HexInput;
   }) {
     const publicKey = FederatedKeylessPublicKey.create(args);
     super({ publicKey, ...args });
@@ -67,15 +68,8 @@ export class FederatedKeylessAccount extends AbstractKeylessAccount {
    * @param serializer - The serializer instance used to convert the transaction data into bytes.
    */
   serialize(serializer: Serializer): void {
-    if (this.proof === undefined) {
-      throw new Error("Cannot serialize - proof undefined");
-    }
-    serializer.serializeStr(this.jwt);
-    serializer.serializeStr(this.uidKey);
-    serializer.serializeFixedBytes(this.pepper);
+    super.serialize(serializer);
     this.publicKey.jwkAddress.serialize(serializer);
-    this.ephemeralKeyPair.serialize(serializer);
-    this.proof.serialize(serializer);
   }
 
   /**
@@ -86,19 +80,22 @@ export class FederatedKeylessAccount extends AbstractKeylessAccount {
    * @returns A KeylessAccount instance created from the deserialized data.
    */
   static deserialize(deserializer: Deserializer): FederatedKeylessAccount {
-    const jwt = deserializer.deserializeStr();
-    const uidKey = deserializer.deserializeStr();
-    const pepper = deserializer.deserializeFixedBytes(31);
+    const { address, proof, ephemeralKeyPair, jwt, uidKey, pepper, verificationKeyHash } =
+      AbstractKeylessAccount.partialDeserialize(deserializer);
     const jwkAddress = AccountAddress.deserialize(deserializer);
-    const ephemeralKeyPair = EphemeralKeyPair.deserialize(deserializer);
-    const proof = ZeroKnowledgeSig.deserialize(deserializer);
-    return FederatedKeylessAccount.create({
+    const { iss, aud, uidVal } = getIssAudAndUidVal({ jwt, uidKey });
+    return new FederatedKeylessAccount({
+      address,
       proof,
-      pepper,
-      jwkAddress,
+      ephemeralKeyPair,
+      iss,
       uidKey,
+      uidVal,
+      aud,
+      pepper,
       jwt,
-      ephemeralKeyPair,
+      verificationKeyHash,
+      jwkAddress,
     });
   }
 
@@ -108,8 +105,8 @@ export class FederatedKeylessAccount extends AbstractKeylessAccount {
    * @param bytes The bytes being interpreted.
    * @returns
    */
-  static fromBytes(bytes: Uint8Array): FederatedKeylessAccount {
-    return FederatedKeylessAccount.deserialize(new Deserializer(bytes));
+  static fromBytes(bytes: HexInput): FederatedKeylessAccount {
+    return FederatedKeylessAccount.deserialize(new Deserializer(Hex.hexInputToUint8Array(bytes)));
   }
 
   /**
@@ -135,29 +132,34 @@ export class FederatedKeylessAccount extends AbstractKeylessAccount {
     jwkAddress: AccountAddressInput;
     uidKey?: string;
     proofFetchCallback?: ProofFetchCallback;
+    verificationKey?: Groth16VerificationKey;
   }): FederatedKeylessAccount {
-    const { address, proof, jwt, ephemeralKeyPair, pepper, jwkAddress, uidKey = "sub", proofFetchCallback } = args;
+    const {
+      address,
+      proof,
+      jwt,
+      ephemeralKeyPair,
+      pepper,
+      jwkAddress,
+      uidKey = "sub",
+      proofFetchCallback,
+      verificationKey,
+    } = 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];
+    const { iss, aud, uidVal } = getIssAudAndUidVal({ jwt, uidKey });
     return new FederatedKeylessAccount({
       address,
       proof,
       ephemeralKeyPair,
-      iss: jwtPayload.iss,
+      iss,
       uidKey,
       uidVal,
-      aud: jwtPayload.aud,
+      aud,
       pepper,
       jwkAddress: AccountAddress.from(jwkAddress),
       jwt,
       proofFetchCallback,
+      verificationKeyHash: verificationKey ? verificationKey.hash() : undefined,
     });
   }
 }