diff --git a/yarn-project/aztec.js/src/account/manager/deploy_account_sent_tx.ts b/yarn-project/aztec.js/src/account/manager/deploy_account_sent_tx.ts index c34051e894a..321d8632ae5 100644 --- a/yarn-project/aztec.js/src/account/manager/deploy_account_sent_tx.ts +++ b/yarn-project/aztec.js/src/account/manager/deploy_account_sent_tx.ts @@ -1,7 +1,9 @@ import { FieldsOf } from '@aztec/circuits.js'; import { TxHash, TxReceipt } from '@aztec/types'; -import { SentTx, WaitOpts, Wallet } from '../../index.js'; +import { DefaultWaitOpts, SentTx, WaitOpts } from '../../contract/sent_tx.js'; +import { Wallet } from '../../wallet/index.js'; +import { waitForAccountSynch } from './util.js'; /** Extends a transaction receipt with a wallet instance for the newly deployed contract. */ export type DeployAccountTxReceipt = FieldsOf & { @@ -32,8 +34,9 @@ export class DeployAccountSentTx extends SentTx { * @param opts - Options for configuring the waiting for the tx to be mined. * @returns The transaction receipt with the wallet for the deployed account contract. */ - public async wait(opts?: WaitOpts): Promise { + public async wait(opts: WaitOpts = DefaultWaitOpts): Promise { const receipt = await super.wait(opts); + await waitForAccountSynch(this.pxe, this.wallet.getCompleteAddress(), opts); return { ...receipt, wallet: this.wallet }; } } diff --git a/yarn-project/aztec.js/src/account/manager/index.ts b/yarn-project/aztec.js/src/account/manager/index.ts index 69f73ebea8d..ace65bf7999 100644 --- a/yarn-project/aztec.js/src/account/manager/index.ts +++ b/yarn-project/aztec.js/src/account/manager/index.ts @@ -2,6 +2,7 @@ import { PublicKey, getContractDeploymentInfo } from '@aztec/circuits.js'; import { Fr } from '@aztec/foundation/fields'; import { CompleteAddress, GrumpkinPrivateKey, PXE } from '@aztec/types'; +import { DefaultWaitOpts } from '../../contract/sent_tx.js'; import { AccountWalletWithPrivateKey, ContractDeployer, @@ -12,6 +13,7 @@ import { import { AccountContract, Salt } from '../index.js'; import { AccountInterface } from '../interface.js'; import { DeployAccountSentTx } from './deploy_account_sent_tx.js'; +import { waitForAccountSynch } from './util.js'; /** * Manages a user account. Provides methods for calculating the account's address, deploying the account contract, @@ -88,11 +90,12 @@ export class AccountManager { * Registers this account in the PXE Service and returns the associated wallet. Registering * the account on the PXE Service is required for managing private state associated with it. * Use the returned wallet to create Contract instances to be interacted with from this account. + * @param opts - Options to wait for the account to be synched. * @returns A Wallet instance. */ - public async register(): Promise { - const completeAddress = await this.getCompleteAddress(); - await this.pxe.registerAccount(this.encryptionPrivateKey, completeAddress.partialAddress); + public async register(opts: WaitOpts = DefaultWaitOpts): Promise { + const address = await this.#register(); + await waitForAccountSynch(this.pxe, address, opts); return this.getWallet(); } @@ -105,7 +108,7 @@ export class AccountManager { public async getDeployMethod() { if (!this.deployMethod) { if (!this.salt) throw new Error(`Cannot deploy account contract without known salt.`); - await this.register(); + await this.#register(); const encryptionPublicKey = await this.getEncryptionPublicKey(); const deployer = new ContractDeployer(this.accountContract.getContractArtifact(), this.pxe, encryptionPublicKey); const args = await this.accountContract.getDeploymentArgs(); @@ -138,8 +141,14 @@ export class AccountManager { * @param opts - Options to wait for the tx to be mined. * @returns A Wallet instance. */ - public async waitDeploy(opts: WaitOpts = {}): Promise { + public async waitDeploy(opts: WaitOpts = DefaultWaitOpts): Promise { await this.deploy().then(tx => tx.wait(opts)); return this.getWallet(); } + + async #register(): Promise { + const completeAddress = await this.getCompleteAddress(); + await this.pxe.registerAccount(this.encryptionPrivateKey, completeAddress.partialAddress); + return completeAddress; + } } diff --git a/yarn-project/aztec.js/src/account/manager/util.ts b/yarn-project/aztec.js/src/account/manager/util.ts new file mode 100644 index 00000000000..2c111eca4fe --- /dev/null +++ b/yarn-project/aztec.js/src/account/manager/util.ts @@ -0,0 +1,29 @@ +import { CompleteAddress, PXE, WaitOpts, retryUntil } from '../../index.js'; + +/** + * Waits for the account to finish synchronizing with the PXE Service. + * @param pxe - PXE instance + * @param address - Address to wait for synch + * @param opts - Wait options + */ +export async function waitForAccountSynch( + pxe: PXE, + address: CompleteAddress, + { interval, timeout }: WaitOpts, +): Promise { + const publicKey = address.publicKey.toString(); + await retryUntil( + async () => { + const status = await pxe.getSyncStatus(); + const accountSynchedToBlock = status.notes[publicKey]; + if (typeof accountSynchedToBlock === 'undefined') { + return false; + } else { + return accountSynchedToBlock >= status.blocks; + } + }, + 'waitForAccountSynch', + timeout, + interval, + ); +} diff --git a/yarn-project/aztec.js/src/contract/index.ts b/yarn-project/aztec.js/src/contract/index.ts index a2a3bf8db63..0a2d9ea8994 100644 --- a/yarn-project/aztec.js/src/contract/index.ts +++ b/yarn-project/aztec.js/src/contract/index.ts @@ -37,6 +37,6 @@ */ export * from './contract.js'; export * from './contract_function_interaction.js'; -export * from './sent_tx.js'; +export { SentTx, WaitOpts } from './sent_tx.js'; export * from './contract_base.js'; export * from './batch_call.js'; diff --git a/yarn-project/aztec.js/src/contract/sent_tx.ts b/yarn-project/aztec.js/src/contract/sent_tx.ts index 9c52a6418a7..92f96a96967 100644 --- a/yarn-project/aztec.js/src/contract/sent_tx.ts +++ b/yarn-project/aztec.js/src/contract/sent_tx.ts @@ -19,7 +19,7 @@ export type WaitOpts = { debug?: boolean; }; -const DefaultWaitOpts: WaitOpts = { +export const DefaultWaitOpts: WaitOpts = { timeout: 60, interval: 1, waitForNotesSync: true, diff --git a/yarn-project/end-to-end/src/e2e_2_pxes.test.ts b/yarn-project/end-to-end/src/e2e_2_pxes.test.ts index cd3e221a237..660d89292b6 100644 --- a/yarn-project/end-to-end/src/e2e_2_pxes.test.ts +++ b/yarn-project/end-to-end/src/e2e_2_pxes.test.ts @@ -6,11 +6,13 @@ import { EthAddress, ExtendedNote, Fr, + GrumpkinScalar, Note, PXE, TxStatus, Wallet, computeMessageSecretHash, + getUnsafeSchnorrAccount, retryUntil, } from '@aztec/aztec.js'; import { ChildContract, TokenContract } from '@aztec/noir-contracts/types'; @@ -248,4 +250,22 @@ describe('e2e_2_pxes', () => { // Check that user B balance is 0 on server A await expectTokenBalance(walletA, completeTokenAddress.address, userB.address, 0n, checkIfSynchronized); }); + + it('permits migrating an account from one PXE to another', async () => { + const privateKey = GrumpkinScalar.random(); + const account = getUnsafeSchnorrAccount(pxeA, privateKey, Fr.random()); + const completeAddress = await account.getCompleteAddress(); + const wallet = await account.waitDeploy(); + + await expect(wallet.isAccountStateSynchronized(completeAddress.address)).resolves.toBe(true); + const accountOnB = getUnsafeSchnorrAccount(pxeB, privateKey, completeAddress); + const walletOnB = await accountOnB.getWallet(); + + // need to register first otherwise the new PXE won't know about the account + await expect(walletOnB.isAccountStateSynchronized(completeAddress.address)).rejects.toThrow(); + + await accountOnB.register(); + // registering should wait for the account to be synchronized + await expect(walletOnB.isAccountStateSynchronized(completeAddress.address)).resolves.toBe(true); + }); }); diff --git a/yarn-project/pxe/src/synchronizer/synchronizer.ts b/yarn-project/pxe/src/synchronizer/synchronizer.ts index 56a42dc377f..bc3ad8c03d6 100644 --- a/yarn-project/pxe/src/synchronizer/synchronizer.ts +++ b/yarn-project/pxe/src/synchronizer/synchronizer.ts @@ -240,7 +240,8 @@ export class Synchronizer { * @returns A promise that resolves once the account is added to the Synchronizer. */ public addAccount(publicKey: PublicKey, keyStore: KeyStore, startingBlock: number) { - const processor = this.noteProcessors.find(x => x.publicKey.equals(publicKey)); + const predicate = (x: NoteProcessor) => x.publicKey.equals(publicKey); + const processor = this.noteProcessors.find(predicate) ?? this.noteProcessorsToCatchUp.find(predicate); if (processor) return; this.noteProcessorsToCatchUp.push(new NoteProcessor(publicKey, keyStore, this.db, this.node, startingBlock)); @@ -259,7 +260,8 @@ export class Synchronizer { if (!completeAddress) { throw new Error(`Checking if account is synched is not possible for ${account} because it is not registered.`); } - const processor = this.noteProcessors.find(x => x.publicKey.equals(completeAddress.publicKey)); + const findByPublicKey = (x: NoteProcessor) => x.publicKey.equals(completeAddress.publicKey); + const processor = this.noteProcessors.find(findByPublicKey) ?? this.noteProcessorsToCatchUp.find(findByPublicKey); if (!processor) { throw new Error( `Checking if account is synched is not possible for ${account} because it is only registered as a recipient.`,