Add rotateAuthKey (aptos-labs#156)

* add rotateAuthKey

* move rotateAuthKey to transactionSubmission

* address comment

src/api/transactionSubmission.ts

@@ -2,7 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 import { AptosConfig } from "./aptosConfig";
-import { Account } from "../core";
+import { Account, PrivateKey } from "../core";
 import { AccountAuthenticator } from "../transactions/authenticator/account";
 import {
   AnyRawTransaction,
@@ -16,10 +16,11 @@ import {
   InputSimulateTransactionData,
   InputGenerateTransactionOptions,
 } from "../transactions/types";
-import { UserTransactionResponse, PendingTransactionResponse, HexInput } from "../types";
+import { UserTransactionResponse, PendingTransactionResponse, HexInput, TransactionResponse } from "../types";
 import {
   generateTransaction,
   publicPackageTransaction,
+  rotateAuthKey,
   signAndSubmitTransaction,
   signTransaction,
   simulateTransaction,
@@ -189,4 +190,18 @@ export class TransactionSubmission {
   }): Promise<InputSingleSignerTransaction> {
     return publicPackageTransaction({ aptosConfig: this.config, ...args });
   }
+
+  /**
+   * Rotate an account's auth key. After rotation, only the new private key can be used to sign txns for
+   * the account.
+   * Note: Only legacy Ed25519 scheme is supported for now.
+   * More info: {@link https://aptos.dev/guides/account-management/key-rotation/}
+   * @param args.fromAccount The account to rotate the auth key for
+   * @param args.toNewPrivateKey The new private key to rotate to
+   *
+   * @returns PendingTransactionResponse
+   */
+  async rotateAuthKey(args: { fromAccount: Account; toNewPrivateKey: PrivateKey }): Promise<TransactionResponse> {
+    return rotateAuthKey({ aptosConfig: this.config, ...args });
+  }
 }

src/internal/transactionSubmission.ts

@@ -6,10 +6,13 @@
  */
 
 import { AptosConfig } from "../api/aptosConfig";
-import { MoveVector } from "../bcs";
+import { MoveVector, U8 } from "../bcs";
 import { postAptosFullNode } from "../client";
 import { Account } from "../core/account";
+import { AccountAddress } from "../core/accountAddress";
+import { PrivateKey } from "../core/crypto";
 import { AccountAuthenticator } from "../transactions/authenticator/account";
+import { RotationProofChallenge } from "../transactions/instances/rotationProofChallenge";
 import {
   buildTransaction,
   generateTransactionPayload,
@@ -25,7 +28,8 @@ import {
   InputSingleSignerTransaction,
   InputGenerateTransactionPayloadDataWithRemoteABI,
 } from "../transactions/types";
-import { UserTransactionResponse, PendingTransactionResponse, MimeType, HexInput } from "../types";
+import { UserTransactionResponse, PendingTransactionResponse, MimeType, HexInput, TransactionResponse } from "../types";
+import { getInfo } from "./account";
 
 /**
  * Generates any transaction by passing in the required arguments
@@ -226,3 +230,55 @@ export async function publicPackageTransaction(args: {
   });
   return transaction as InputSingleSignerTransaction;
 }
+
+/**
+ * TODO: Need to refactor and move this function out of transactionSubmission
+ */
+export async function rotateAuthKey(args: {
+  aptosConfig: AptosConfig;
+  fromAccount: Account;
+  toNewPrivateKey: PrivateKey;
+}): Promise<TransactionResponse> {
+  const { aptosConfig, fromAccount, toNewPrivateKey } = args;
+  const accountInfo = await getInfo({
+    aptosConfig,
+    accountAddress: fromAccount.accountAddress.toString(),
+  });
+
+  const newAccount = Account.fromPrivateKey({ privateKey: toNewPrivateKey, legacy: true });
+
+  const challenge = new RotationProofChallenge({
+    sequenceNumber: BigInt(accountInfo.sequence_number),
+    originator: fromAccount.accountAddress,
+    currentAuthKey: AccountAddress.fromHexInput(accountInfo.authentication_key),
+    newPublicKey: newAccount.publicKey,
+  });
+
+  // Sign the challenge
+  const challengeHex = challenge.bcsToBytes();
+  const proofSignedByCurrentPrivateKey = fromAccount.sign(challengeHex);
+  const proofSignedByNewPrivateKey = newAccount.sign(challengeHex);
+
+  // Generate transaction
+  const rawTxn = await generateTransaction({
+    aptosConfig,
+    sender: fromAccount.accountAddress.toString(),
+    data: {
+      function: "0x1::account::rotate_authentication_key",
+      functionArguments: [
+        new U8(fromAccount.signingScheme.valueOf()), // from scheme
+        MoveVector.U8(fromAccount.publicKey.toUint8Array()),
+        new U8(newAccount.signingScheme.valueOf()), // to scheme
+        MoveVector.U8(newAccount.publicKey.toUint8Array()),
+        MoveVector.U8(proofSignedByCurrentPrivateKey.toUint8Array()),
+        MoveVector.U8(proofSignedByNewPrivateKey.toUint8Array()),
+      ],
+    },
+  });
+  const pendingTxn = await signAndSubmitTransaction({
+    aptosConfig,
+    signer: fromAccount,
+    transaction: rawTxn,
+  });
+  return pendingTxn;
+}

src/transactions/instances/index.ts

@@ -5,6 +5,7 @@ export * from "./chainId";
 export * from "./identifier";
 export * from "./moduleId";
 export * from "./rawTransaction";
+export * from "./rotationProofChallenge";
 export * from "./signedTransaction";
 export * from "./transactionArgument";
 export * from "./transactionPayload";

src/transactions/instances/rotationProofChallenge.ts

@@ -0,0 +1,58 @@
+// Copyright © Aptos Foundation
+// SPDX-License-Identifier: Apache-2.0
+
+import { Serializer, Serializable } from "../../bcs/serializer";
+import { AccountAddress } from "../../core/accountAddress";
+import { AnyNumber } from "../../types";
+import { PublicKey } from "../../core/crypto/asymmetricCrypto";
+import { MoveString, MoveVector, U64, U8 } from "../../bcs";
+
+/**
+ * Representation of the challenge which is needed to sign by owner of the account
+ * to rotate the authentication key.
+ */
+export class RotationProofChallenge extends Serializable {
+  // Resource account address
+  public readonly accountAddress: AccountAddress = AccountAddress.ONE;
+
+  // Module name, i.e: 0x1::account
+  public readonly moduleName: MoveString = new MoveString("account");
+
+  // The rotation proof challenge struct name that live under the module
+  public readonly structName: MoveString = new MoveString("RotationProofChallenge");
+
+  // Signer's address
+  public readonly originator: AccountAddress;
+
+  // Signer's current authentication key
+  public readonly currentAuthKey: AccountAddress;
+
+  // New public key to rotate to
+  public readonly newPublicKey: MoveVector<U8>;
+
+  // Sequence number of the account
+  public readonly sequenceNumber: U64;
+
+  constructor(args: {
+    sequenceNumber: AnyNumber;
+    originator: AccountAddress;
+    currentAuthKey: AccountAddress;
+    newPublicKey: PublicKey;
+  }) {
+    super();
+    this.sequenceNumber = new U64(args.sequenceNumber);
+    this.originator = args.originator;
+    this.currentAuthKey = args.currentAuthKey;
+    this.newPublicKey = MoveVector.U8(args.newPublicKey.toUint8Array());
+  }
+
+  serialize(serializer: Serializer): void {
+    serializer.serialize(this.accountAddress);
+    serializer.serialize(this.moduleName);
+    serializer.serialize(this.structName);
+    serializer.serialize(this.sequenceNumber);
+    serializer.serialize(this.originator);
+    serializer.serialize(this.currentAuthKey);
+    serializer.serialize(this.newPublicKey);
+  }
+}

tests/e2e/api/account.test.ts

@@ -1,7 +1,7 @@
 // Copyright © Aptos Foundation
 // SPDX-License-Identifier: Apache-2.0
 
-import { Account, Aptos, AptosConfig, Network, SigningSchemeInput, U64 } from "../../../src";
+import { Account, Aptos, AptosConfig, Ed25519PrivateKey, Network, SigningSchemeInput, U64 } from "../../../src";
 
 describe("account api", () => {
   const FUND_AMOUNT = 100_000_000;
@@ -302,4 +302,30 @@ describe("account api", () => {
       expect(data.authentication_key).toBe("0x0000000000000000000000000000000000000000000000000000000000000001");
     });
   });
+
+  describe("Key Rotation", () => {
+    test("it should rotate ed25519 to ed25519 auth key correctly", async () => {
+      const config = new AptosConfig({ network: Network.LOCAL });
+      const aptos = new Aptos(config);
+
+      // Current Account
+      const account = Account.generate({ scheme: SigningSchemeInput.Ed25519, legacy: true });
+      await aptos.fundAccount({ accountAddress: account.accountAddress.toString(), amount: 1_000_000_000 });
+
+      // account that holds the new key
+      const rotateToPrivateKey = Ed25519PrivateKey.generate();
+
+      // Rotate the key
+      const pendingTxn = await aptos.rotateAuthKey({ fromAccount: account, toNewPrivateKey: rotateToPrivateKey });
+      await aptos.waitForTransaction({ transactionHash: pendingTxn.hash });
+
+      // lookup original account address
+      const lookupAccountAddress = await aptos.lookupOriginalAccountAddress({
+        authenticationKey: Account.authKey({ publicKey: rotateToPrivateKey.publicKey() }).toString(),
+      });
+
+      // Check if the lookup account address is the same as the original account address
+      expect(lookupAccountAddress.toString()).toBe(account.accountAddress.toString());
+    });
+  });
 });