@aptos-labs/ts-sdk 1.26.0 → 1.27.1-zeta.0

package/src/account/FederatedKeylessAccount.ts

@@ -0,0 +1,328 @@
+// Copyright © Aptos Foundation
+// SPDX-License-Identifier: Apache-2.0
+
+import { JwtPayload, jwtDecode } from "jwt-decode";
+import EventEmitter from "eventemitter3";
+import { EphemeralCertificateVariant, HexInput, SigningScheme } from "../types";
+import { AccountAddress, AccountAddressInput } from "../core/accountAddress";
+import { AnyPublicKey, AnySignature, KeylessSignature, EphemeralCertificate, ZeroKnowledgeSig } from "../core/crypto";
+
+import { Account } from "./Account";
+import { EphemeralKeyPair } from "./EphemeralKeyPair";
+import { Hex } from "../core/hex";
+import { AccountAuthenticatorSingleKey } from "../transactions/authenticator/account";
+import { Deserializer, Serializable, Serializer } from "../bcs";
+import { deriveTransactionType } from "../transactions/transactionBuilder/signingMessage";
+import { AnyRawTransaction } from "../transactions/types";
+import { base64UrlDecode } from "../utils/helpers";
+import { FederatedKeylessPublicKey } from "../core/crypto/federatedKeyless";
+import { KeylessAccount, ProofFetchCallback, ProofFetchEvents, TransactionAndProof } from "./KeylessAccount";
+
+/**
+ * Account implementation for the Keyless authentication scheme.
+ *
+ * Used to represent a Keyless based account and sign transactions with it.
+ *
+ * Use KeylessAccount.fromJWTAndProof to instantiate a KeylessAccount with a JWT, proof and EphemeralKeyPair.
+ *
+ * When the proof expires or the JWT becomes invalid, the KeylessAccount must be instantiated again with a new JWT,
+ * EphemeralKeyPair, and corresponding proof.
+ */
+export class FederatedKeylessAccount extends Serializable implements Account {
+  static readonly PEPPER_LENGTH: number = 31;
+
+  /**
+   * The KeylessPublicKey associated with the account
+   */
+  readonly publicKey: FederatedKeylessPublicKey;
+
+  /**
+   * The EphemeralKeyPair used to generate sign.
+   */
+  readonly ephemeralKeyPair: EphemeralKeyPair;
+
+  /**
+   * The claim on the JWT to identify a user.  This is typically 'sub' or 'email'.
+   */
+  readonly uidKey: string;
+
+  /**
+   * The value of the uidKey claim on the JWT.  This intended to be a stable user identifier.
+   */
+  readonly uidVal: string;
+
+  /**
+   * The value of the 'aud' claim on the JWT, also known as client ID.  This is the identifier for the dApp's
+   * OIDC registration with the identity provider.
+   */
+  readonly aud: string;
+
+  /**
+   * A value contains 31 bytes of entropy that preserves privacy of the account. Typically fetched from a pepper provider.
+   */
+  readonly pepper: Uint8Array;
+
+  /**
+   * Account address associated with the account
+   */
+  readonly accountAddress: AccountAddress;
+
+  /**
+   * The zero knowledge signature (if ready) which contains the proof used to validate the EphemeralKeyPair.
+   */
+  proof: ZeroKnowledgeSig | undefined;
+
+  /**
+   * The proof of the EphemeralKeyPair or a promise that provides the proof.  This is used to allow for awaiting on
+   * fetching the proof.
+   */
+  readonly proofOrPromise: ZeroKnowledgeSig | Promise<ZeroKnowledgeSig>;
+
+  /**
+   * Signing scheme used to sign transactions
+   */
+  readonly signingScheme: SigningScheme;
+
+  /**
+   * The JWT token used to derive the account
+   */
+  readonly jwt: string;
+
+  /**
+   * An event emitter used to assist in handling asycronous proof fetching.
+   */
+  private readonly emitter: EventEmitter<ProofFetchEvents>;
+
+  // Use the static constructor 'create' instead.
+  private constructor(args: {
+    address?: AccountAddress;
+    ephemeralKeyPair: EphemeralKeyPair;
+    iss: string;
+    uidKey: string;
+    uidVal: string;
+    aud: string;
+    pepper: HexInput;
+    jwkAddress: AccountAddress;
+    proof: ZeroKnowledgeSig | Promise<ZeroKnowledgeSig>;
+    proofFetchCallback?: ProofFetchCallback;
+    jwt: string;
+  }) {
+    super();
+    const { address, ephemeralKeyPair, uidKey, uidVal, aud, pepper, proof, proofFetchCallback, jwt } = args;
+    this.ephemeralKeyPair = ephemeralKeyPair;
+    this.publicKey = FederatedKeylessPublicKey.create(args);
+    this.accountAddress = address ? AccountAddress.from(address) : this.publicKey.authKey().derivedAddress();
+    this.uidKey = uidKey;
+    this.uidVal = uidVal;
+    this.aud = aud;
+    this.jwt = jwt;
+    this.emitter = new EventEmitter<ProofFetchEvents>();
+    this.proofOrPromise = proof;
+    if (proof instanceof ZeroKnowledgeSig) {
+      this.proof = proof;
+    } else {
+      if (proofFetchCallback === undefined) {
+        throw new Error("Must provide callback for async proof fetch");
+      }
+      this.emitter.on("proofFetchFinish", async (status) => {
+        await proofFetchCallback(status);
+        this.emitter.removeAllListeners();
+      });
+      this.init(proof);
+    }
+    this.signingScheme = SigningScheme.SingleKey;
+    const pepperBytes = Hex.fromHexInput(pepper).toUint8Array();
+    if (pepperBytes.length !== KeylessAccount.PEPPER_LENGTH) {
+      throw new Error(`Pepper length in bytes should be ${KeylessAccount.PEPPER_LENGTH}`);
+    }
+    this.pepper = pepperBytes;
+  }
+
+  /**
+   * This initializes the asyncronous proof fetch
+   * @return
+   */
+  async init(promise: Promise<ZeroKnowledgeSig>) {
+    try {
+      this.proof = await promise;
+      this.emitter.emit("proofFetchFinish", { status: "Success" });
+    } catch (error) {
+      if (error instanceof Error) {
+        this.emitter.emit("proofFetchFinish", { status: "Failed", error: error.toString() });
+      } else {
+        this.emitter.emit("proofFetchFinish", { status: "Failed", error: "Unknown" });
+      }
+    }
+  }
+
+  serialize(serializer: Serializer): void {
+    serializer.serializeStr(this.jwt);
+    serializer.serializeStr(this.uidKey);
+    serializer.serializeFixedBytes(this.pepper);
+    this.publicKey.jwkAddress.serialize(serializer);
+    this.ephemeralKeyPair.serialize(serializer);
+    if (this.proof === undefined) {
+      throw new Error("Connot serialize - proof undefined");
+    }
+    this.proof.serialize(serializer);
+  }
+
+  static deserialize(deserializer: Deserializer): FederatedKeylessAccount {
+    const jwt = deserializer.deserializeStr();
+    const uidKey = deserializer.deserializeStr();
+    const pepper = deserializer.deserializeFixedBytes(31);
+    const jwkAddress = AccountAddress.deserialize(deserializer);
+    const ephemeralKeyPair = EphemeralKeyPair.deserialize(deserializer);
+    const proof = ZeroKnowledgeSig.deserialize(deserializer);
+    return FederatedKeylessAccount.create({
+      proof,
+      pepper,
+      jwkAddress,
+      uidKey,
+      jwt,
+      ephemeralKeyPair,
+    });
+  }
+
+  /**
+   * Checks if the proof is expired.  If so the account must be rederived with a new EphemeralKeyPair
+   * and JWT token.
+   * @return boolean
+   */
+  isExpired(): boolean {
+    return this.ephemeralKeyPair.isExpired();
+  }
+
+  /**
+   * Sign a message using Keyless.
+   * @param message the message to sign, as binary input
+   * @return the AccountAuthenticator containing the signature, together with the account's public key
+   */
+  signWithAuthenticator(message: HexInput): AccountAuthenticatorSingleKey {
+    const signature = new AnySignature(this.sign(message));
+    const publicKey = new AnyPublicKey(this.publicKey);
+    return new AccountAuthenticatorSingleKey(publicKey, signature);
+  }
+
+  /**
+   * Sign a transaction using Keyless.
+   * @param transaction the raw transaction
+   * @return the AccountAuthenticator containing the signature of the transaction, together with the account's public key
+   */
+  signTransactionWithAuthenticator(transaction: AnyRawTransaction): AccountAuthenticatorSingleKey {
+    const signature = new AnySignature(this.signTransaction(transaction));
+    const publicKey = new AnyPublicKey(this.publicKey);
+    return new AccountAuthenticatorSingleKey(publicKey, signature);
+  }
+
+  /**
+   * Waits for asyncronous proof fetching to finish.
+   * @return
+   */
+  async waitForProofFetch() {
+    if (this.proofOrPromise instanceof Promise) {
+      await this.proofOrPromise;
+    }
+  }
+
+  /**
+   * Sign the given message using Keyless.
+   * @param message in HexInput format
+   * @returns Signature
+   */
+  sign(data: HexInput): KeylessSignature {
+    const { expiryDateSecs } = this.ephemeralKeyPair;
+    if (this.isExpired()) {
+      throw new Error("EphemeralKeyPair is expired");
+    }
+    if (this.proof === undefined) {
+      throw new Error("Proof not defined");
+    }
+    const ephemeralPublicKey = this.ephemeralKeyPair.getPublicKey();
+    const ephemeralSignature = this.ephemeralKeyPair.sign(data);
+
+    return new KeylessSignature({
+      jwtHeader: base64UrlDecode(this.jwt.split(".")[0]),
+      ephemeralCertificate: new EphemeralCertificate(this.proof, EphemeralCertificateVariant.ZkProof),
+      expiryDateSecs,
+      ephemeralPublicKey,
+      ephemeralSignature,
+    });
+  }
+
+  /**
+   * Sign the given transaction with Keyless.
+   * Signs the transaction and proof to guard against proof malleability.
+   * @param transaction the transaction to be signed
+   * @returns KeylessSignature
+   */
+  signTransaction(transaction: AnyRawTransaction): KeylessSignature {
+    if (this.proof === undefined) {
+      throw new Error("Proof not found");
+    }
+    const raw = deriveTransactionType(transaction);
+    const txnAndProof = new TransactionAndProof(raw, this.proof.proof);
+    const signMess = txnAndProof.hash();
+    return this.sign(signMess);
+  }
+
+  /**
+   * Note - This function is currently incomplete and should only be used to verify ownership of the KeylessAccount
+   *
+   * Verifies a signature given the message.
+   *
+   * TODO: Groth16 proof verification
+   *
+   * @param args.message the message that was signed.
+   * @param args.signature the KeylessSignature to verify
+   * @returns boolean
+   */
+  verifySignature(args: { message: HexInput; signature: KeylessSignature }): boolean {
+    const { message, signature } = args;
+    if (this.isExpired()) {
+      return false;
+    }
+    if (!this.ephemeralKeyPair.getPublicKey().verifySignature({ message, signature: signature.ephemeralSignature })) {
+      return false;
+    }
+    return true;
+  }
+
+  static fromBytes(bytes: Uint8Array): KeylessAccount {
+    return KeylessAccount.deserialize(new Deserializer(bytes));
+  }
+
+  static create(args: {
+    address?: AccountAddress;
+    proof: ZeroKnowledgeSig | Promise<ZeroKnowledgeSig>;
+    jwt: string;
+    ephemeralKeyPair: EphemeralKeyPair;
+    pepper: HexInput;
+    jwkAddress: AccountAddressInput;
+    uidKey?: string;
+    proofFetchCallback?: ProofFetchCallback;
+  }): FederatedKeylessAccount {
+    const { address, proof, jwt, ephemeralKeyPair, pepper, jwkAddress, uidKey = "sub", proofFetchCallback } = args;
+
+    const jwtPayload = jwtDecode<JwtPayload & { [key: string]: string }>(jwt);
+    const iss = jwtPayload.iss!;
+    if (typeof jwtPayload.aud !== "string") {
+      throw new Error("aud was not found or an array of values");
+    }
+    const aud = jwtPayload.aud!;
+    const uidVal = jwtPayload[uidKey];
+    return new FederatedKeylessAccount({
+      address,
+      proof,
+      ephemeralKeyPair,
+      iss,
+      uidKey,
+      uidVal,
+      aud,
+      pepper,
+      jwkAddress: AccountAddress.from(jwkAddress),
+      jwt,
+      proofFetchCallback,
+    });
+  }
+}

package/src/account/KeylessAccount.ts

@@ -331,7 +331,7 @@ export class KeylessAccount extends Serializable implements Account {
  * A container class to hold a transaction and a proof.  It implements CryptoHashable which is used to create
  * the signing message for Keyless transactions.  We sign over the proof to ensure non-malleability.
  */
-class TransactionAndProof extends Serializable {
+export class TransactionAndProof extends Serializable {
   /**
    * The transaction to sign.
    */

package/src/api/keyless.ts

@@ -2,11 +2,17 @@
 // SPDX-License-Identifier: Apache-2.0
 
 import { EphemeralKeyPair, KeylessAccount, ProofFetchCallback } from "../account";
-import { ZeroKnowledgeSig } from "../core";
-import { deriveKeylessAccount, getPepper, getProof } from "../internal/keyless";
+import { FederatedKeylessAccount } from "../account/FederatedKeylessAccount";
+import { AccountAddressInput, ZeroKnowledgeSig } from "../core";
+import { deriveFederatedKeylessAccount, deriveKeylessAccount, getPepper, getProof } from "../internal/keyless";
 import { HexInput } from "../types";
 import { AptosConfig } from "./aptosConfig";
 
+
+interface InputGenerateMultiAgentRawTransactionData {
+  jwkAddress: AccountAddressInput;
+}
+
 /**
  * A class to query all `Keyless` related queries on Aptos.
  *
@@ -51,6 +57,27 @@ export class Keyless {
   }): Promise<ZeroKnowledgeSig> {
     return getProof({ aptosConfig: this.config, ...args });
   }
+  
+  async deriveKeylessAccount(
+    args: {
+      jwt: string;
+      ephemeralKeyPair: EphemeralKeyPair;
+      uidKey?: string;
+      pepper?: HexInput;
+      proofFetchCallback?: ProofFetchCallback;
+    },
+  ): Promise<KeylessAccount>;
+
+  async deriveKeylessAccount(
+    args: {
+      jwt: string;
+      ephemeralKeyPair: EphemeralKeyPair;
+      jwkAddress: AccountAddressInput;
+      uidKey?: string;
+      pepper?: HexInput;
+      proofFetchCallback?: ProofFetchCallback;
+    } & InputGenerateMultiAgentRawTransactionData,
+  ): Promise<FederatedKeylessAccount>;
 
   /**
    * Derives the Keyless Account from the JWT token and corresponding EphemeralKeyPair.  It will lookup the pepper from
@@ -69,10 +96,15 @@ export class Keyless {
   async deriveKeylessAccount(args: {
     jwt: string;
     ephemeralKeyPair: EphemeralKeyPair;
+    jwkAddress?: AccountAddressInput;
     uidKey?: string;
     pepper?: HexInput;
     proofFetchCallback?: ProofFetchCallback;
-  }): Promise<KeylessAccount> {
+  }): Promise<KeylessAccount | FederatedKeylessAccount> {
+    const { jwkAddress } = args;
+    if(jwkAddress !== undefined) {
+      return deriveFederatedKeylessAccount({ aptosConfig: this.config, jwkAddress, ...args });
+    }
     return deriveKeylessAccount({ aptosConfig: this.config, ...args });
   }
 }

package/src/bcs/deserializer.ts

@@ -37,6 +37,15 @@ export class Deserializer {
     return bytes;
   }
 
+  /**
+   * Returns the number of bytes remaining in the buffer.
+   *
+   * Useful to tell if there's more data to be read.
+   */
+  remaining(): number {
+    return this.buffer.byteLength - this.offset;
+  }
+
   /**
    * Deserializes a string. UTF8 string is supported. Reads the string's bytes length "l" first,
    * and then reads "l" bytes of content. Decodes the byte array into a string.

package/src/bcs/serializable/moveStructs.ts

@@ -244,7 +244,8 @@ export class MoveString extends Serializable implements TransactionArgument {
 
   serializeForScriptFunction(serializer: Serializer): void {
     // Serialize the string as a fixed byte string, i.e., without the length prefix
-    const fixedStringBytes = this.bcsToBytes().slice(1);
+    const textEncoder = new TextEncoder();
+    const fixedStringBytes = textEncoder.encode(this.value);
     // Put those bytes into a vector<u8> and serialize it as a script function argument
     const vectorU8 = MoveVector.U8(fixedStringBytes);
     vectorU8.serializeForScriptFunction(serializer);

package/src/core/crypto/federatedKeyless.ts

@@ -0,0 +1,134 @@
+// Copyright © Aptos Foundation
+// SPDX-License-Identifier: Apache-2.0
+
+import { AccountPublicKey, PublicKey } from "./publicKey";
+import { Deserializer, Serializer } from "../../bcs";
+import { Hex } from "../hex";
+import { HexInput, AnyPublicKeyVariant, SigningScheme } from "../../types";
+import { AuthenticationKey } from "../authenticationKey";
+import { AccountAddress, AccountAddressInput } from "../accountAddress";
+import { KeylessPublicKey, KeylessSignature } from "./keyless";
+
+/**
+ * Represents the FederatedKeylessPublicKey public key
+ *
+ * These keys use an onchain address as a source of truth for the JWK used to verify signatures.
+ *
+ * FederatedKeylessPublicKey authentication key is represented in the SDK as `AnyPublicKey`.
+ */
+export class FederatedKeylessPublicKey extends AccountPublicKey {
+  /**
+   * The address that contains the JWK set to be used for verification.
+   */
+  readonly jwkAddress: AccountAddress;
+
+  /**
+   * The inner public key which contains the standard Keyless public key.
+   */
+  readonly keylessPublicKey: KeylessPublicKey;
+
+  constructor(jwkAddress: AccountAddressInput, keylessPublicKey: KeylessPublicKey) {
+    super();
+    this.jwkAddress = AccountAddress.from(jwkAddress);
+    this.keylessPublicKey = keylessPublicKey;
+  }
+
+  /**
+   * Get the authentication key for the federated keyless public key
+   *
+   * @returns AuthenticationKey
+   */
+  authKey(): AuthenticationKey {
+    const serializer = new Serializer();
+    serializer.serializeU32AsUleb128(AnyPublicKeyVariant.FederatedKeyless);
+    serializer.serializeFixedBytes(this.bcsToBytes());
+    return AuthenticationKey.fromSchemeAndBytes({
+      scheme: SigningScheme.SingleKey,
+      input: serializer.toUint8Array(),
+    });
+  }
+
+  /**
+   * Get the public key in bytes (Uint8Array).
+   *
+   * @returns Uint8Array representation of the public key
+   */
+  toUint8Array(): Uint8Array {
+    return this.bcsToBytes();
+  }
+
+  /**
+   * Get the public key as a hex string with the 0x prefix.
+   *
+   * @returns string representation of the public key
+   */
+  toString(): string {
+    return Hex.fromHexInput(this.toUint8Array()).toString();
+  }
+
+  /**
+   * Verifies a signed data with a public key
+   *
+   * @param args.message message
+   * @param args.signature The signature
+   * @returns true if the signature is valid
+   */
+  // eslint-disable-next-line @typescript-eslint/no-unused-vars, class-methods-use-this
+  verifySignature(args: { message: HexInput; signature: KeylessSignature }): boolean {
+    throw new Error("Not yet implemented");
+  }
+
+  serialize(serializer: Serializer): void {
+    this.jwkAddress.serialize(serializer);
+    this.keylessPublicKey.serialize(serializer);
+  }
+
+  static deserialize(deserializer: Deserializer): FederatedKeylessPublicKey {
+    const jwkAddress = AccountAddress.deserialize(deserializer);
+    const keylessPublicKey = KeylessPublicKey.deserialize(deserializer);
+    return new FederatedKeylessPublicKey(jwkAddress, keylessPublicKey);
+  }
+
+  static isPublicKey(publicKey: PublicKey): publicKey is FederatedKeylessPublicKey {
+    return publicKey instanceof FederatedKeylessPublicKey;
+  }
+
+  /**
+   * Creates a FederatedKeylessPublicKey from the JWT components plus pepper
+   *
+   * @param args.iss the iss of the identity
+   * @param args.uidKey the key to use to get the uidVal in the JWT token
+   * @param args.uidVal the value of the uidKey in the JWT token
+   * @param args.aud the client ID of the application
+   * @param args.pepper The pepper used to maintain privacy of the account
+   * @returns FederatedKeylessPublicKey
+   */
+  static create(args: {
+    iss: string;
+    uidKey: string;
+    uidVal: string;
+    aud: string;
+    pepper: HexInput;
+    jwkAddress: AccountAddressInput;
+  }): FederatedKeylessPublicKey {
+    return new FederatedKeylessPublicKey(args.jwkAddress, KeylessPublicKey.create(args));
+  }
+
+  static fromJwtAndPepper(args: {
+    jwt: string;
+    pepper: HexInput;
+    jwkAddress: AccountAddressInput;
+    uidKey?: string;
+  }): FederatedKeylessPublicKey {
+    return new FederatedKeylessPublicKey(args.jwkAddress, KeylessPublicKey.fromJwtAndPepper(args));
+  }
+
+  static isInstance(publicKey: PublicKey) {
+    return (
+      "jwkAddress" in publicKey &&
+      publicKey.jwkAddress instanceof AccountAddress &&
+      "keylessPublicKey" in publicKey &&
+      publicKey.keylessPublicKey instanceof KeylessPublicKey
+    );
+  }
+}

package/src/core/crypto/singleKey.ts

@@ -6,6 +6,7 @@ import { AccountPublicKey, PublicKey, VerifySignatureArgs } from "./publicKey";
 import { Secp256k1PublicKey, Secp256k1Signature } from "./secp256k1";
 import { KeylessPublicKey, KeylessSignature } from "./keyless";
 import { Signature } from "./signature";
+import { FederatedKeylessPublicKey } from "./federatedKeyless";
 
 /**
  * Represents any public key supported by Aptos.
@@ -37,6 +38,8 @@ export class AnyPublicKey extends AccountPublicKey {
       this.variant = AnyPublicKeyVariant.Secp256k1;
     } else if (publicKey instanceof KeylessPublicKey) {
       this.variant = AnyPublicKeyVariant.Keyless;
+    } else if (publicKey instanceof FederatedKeylessPublicKey) {
+      this.variant = AnyPublicKeyVariant.FederatedKeyless;
     } else {
       throw new Error("Unsupported public key type");
     }
@@ -91,6 +94,9 @@ export class AnyPublicKey extends AccountPublicKey {
       case AnyPublicKeyVariant.Keyless:
         publicKey = KeylessPublicKey.deserialize(deserializer);
         break;
+      case AnyPublicKeyVariant.FederatedKeyless:
+        publicKey = FederatedKeylessPublicKey.deserialize(deserializer);
+        break;
       default:
         throw new Error(`Unknown variant index for AnyPublicKey: ${variantIndex}`);
     }

package/src/internal/keyless.ts

@@ -10,6 +10,7 @@
 import { AptosConfig } from "../api/aptosConfig";
 import { postAptosPepperService, postAptosProvingService } from "../client";
 import {
+  AccountAddressInput,
   EphemeralSignature,
   Groth16Zkp,
   Hex,
@@ -23,6 +24,8 @@ import { EphemeralKeyPair, KeylessAccount, ProofFetchCallback } from "../account
 import { PepperFetchRequest, PepperFetchResponse, ProverRequest, ProverResponse } from "../types/keyless";
 import { nowInSeconds } from "../utils/helpers";
 import { lookupOriginalAccountAddress } from "./account";
+import { FederatedKeylessPublicKey } from "../core/crypto/federatedKeyless";
+import { FederatedKeylessAccount } from "../account/FederatedKeylessAccount";
 
 export async function getPepper(args: {
   aptosConfig: AptosConfig;
@@ -123,7 +126,33 @@ export async function deriveKeylessAccount(args: {
     authenticationKey: publicKey.authKey().derivedAddress(),
   });
 
-  const keylessAccount = KeylessAccount.create({ ...args, address, proof, pepper, proofFetchCallback });
+  return KeylessAccount.create({ ...args, address, proof, pepper, proofFetchCallback });
+}
+
+export async function deriveFederatedKeylessAccount(args: {
+  aptosConfig: AptosConfig;
+  jwt: string;
+  ephemeralKeyPair: EphemeralKeyPair;
+  jwkAddress: AccountAddressInput;
+  uidKey?: string;
+  pepper?: HexInput;
+  proofFetchCallback?: ProofFetchCallback;
+}): Promise<FederatedKeylessAccount> {
+  const { aptosConfig, jwt, jwkAddress, uidKey, proofFetchCallback, pepper = await getPepper(args) } = args;
+  const proofPromise = getProof({ ...args, pepper });
+  // 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 syncronously.
+  const proof = proofFetchCallback ? proofPromise : await proofPromise;
+
+  // Look up the original address to handle key rotations
+  const publicKey = FederatedKeylessPublicKey.fromJwtAndPepper({ jwt, pepper, jwkAddress, uidKey });
+  const address = await lookupOriginalAccountAddress({
+    aptosConfig,
+    authenticationKey: publicKey.authKey().derivedAddress(),
+  });
 
-  return keylessAccount;
+  return FederatedKeylessAccount.create({ ...args, address, proof, pepper, proofFetchCallback });
 }