diff --git a/noir-projects/noir-contracts/contracts/escrow_contract/src/main.nr b/noir-projects/noir-contracts/contracts/escrow_contract/src/main.nr index b6cab7aa13b7..cd6d891ad7e0 100644 --- a/noir-projects/noir-contracts/contracts/escrow_contract/src/main.nr +++ b/noir-projects/noir-contracts/contracts/escrow_contract/src/main.nr @@ -1,6 +1,8 @@ // Sample escrow contract that stores a balance of a private token on behalf of an owner. use dep::aztec::macros::aztec; +mod test; + #[aztec] pub contract Escrow { use dep::aztec::{ diff --git a/noir-projects/noir-contracts/contracts/escrow_contract/src/test.nr b/noir-projects/noir-contracts/contracts/escrow_contract/src/test.nr new file mode 100644 index 000000000000..ed3337a83caf --- /dev/null +++ b/noir-projects/noir-contracts/contracts/escrow_contract/src/test.nr @@ -0,0 +1,120 @@ +// TODO(ek): Clean up this test +use crate::Escrow; +use dep::token::Token; + +use aztec::{ + oracle::{execution::{get_block_number, get_contract_address}, storage::storage_read}, + prelude::AztecAddress, + protocol_types::storage::map::derive_storage_slot_in_map, + test::helpers::{cheatcodes, test_environment::TestEnvironment}, +}; + +pub unconstrained fn get_public_balance( + token_contract_address: AztecAddress, + address: AztecAddress, +) -> U128 { + let current_contract_address = get_contract_address(); + cheatcodes::set_contract_address(token_contract_address); + let block_number = get_block_number(); + + let balances_slot = Token::storage_layout().public_balances.slot; + let address_slot = derive_storage_slot_in_map(balances_slot, address); + let amount: U128 = storage_read(token_contract_address, address_slot, block_number); + cheatcodes::set_contract_address(current_contract_address); + amount +} + +pub unconstrained fn get_private_balance( + token_contract_address: AztecAddress, + address: AztecAddress, +) -> U128 { + let current_contract_address = get_contract_address(); + cheatcodes::set_contract_address(token_contract_address); + // Direct call to unconstrained + let amt = Token::balance_of_private(address); + cheatcodes::set_contract_address(current_contract_address); + amt +} + +global MINT_AMOUNT: U128 = U128::from_integer(200000); + +unconstrained fn deploy_contracts( + env: &mut TestEnvironment, + admin_and_owner: AztecAddress, +) -> (AztecAddress, AztecAddress) { + env.impersonate(admin_and_owner); + + // Deploy token contract + let donation_token_initializer_call_interface = Token::interface().constructor( + admin_and_owner, + "Token00000000000000000000000000", + "TKN0000000000000000000000000000", + 18, + ); + let donation_token_contract = env + .deploy("./@token_contract", "Token") + .with_public_void_initializer(donation_token_initializer_call_interface); + let token_contract_address = donation_token_contract.to_address(); + env.advance_block_by(1); + + // Deploy Escrow contract with public keys + let escrow_contract_initializer_call_interface = + Escrow::interface().constructor(admin_and_owner); + let escrow_contract = env + .deploy_with_public_keys("./@escrow_contract", "Escrow", 6969) + .with_private_initializer(escrow_contract_initializer_call_interface); + let escrow_contract_address = escrow_contract.to_address(); + + env.advance_block_by(1); + + Token::at(token_contract_address) + .mint_to_private(admin_and_owner, admin_and_owner, MINT_AMOUNT) + .call(&mut env.private()); + + env.advance_block_by(1); + + let private_balance_after_mint = get_private_balance(token_contract_address, admin_and_owner); + assert(private_balance_after_mint == MINT_AMOUNT); + + (token_contract_address, escrow_contract_address) +} + +#[test] +unconstrained fn main() { + let mut env = TestEnvironment::new(); + + let (account_1, account_2) = (env.create_account_contract(1), env.create_account_contract(2)); + + let (token_contract_address, escrow_contract_address) = deploy_contracts(&mut env, account_1); + + // We transfer tokens to the escrow contract + let TRANSFER_AMOUNT = U128::from_integer(20000); + Token::at(token_contract_address).transfer(escrow_contract_address, TRANSFER_AMOUNT).call( + &mut env.private(), + ); + env.advance_block_by(1); + + let balance_of_escrow_after_transfer = + get_private_balance(token_contract_address, escrow_contract_address); + + assert_eq(balance_of_escrow_after_transfer, TRANSFER_AMOUNT); + + // We then withdraw some escrowed funds to account_2 + let balance_of_account_2_before_withdrawal = + get_private_balance(token_contract_address, account_2); + assert(balance_of_account_2_before_withdrawal == U128::zero()); + + let WITHDRAWAL_AMOUNT = U128::from_integer(69); + Escrow::at(escrow_contract_address) + .withdraw(token_contract_address, WITHDRAWAL_AMOUNT, account_2) + .call(&mut env.private()); + env.advance_block_by(1); + + let balance_of_account_2_after_withdrawal = + get_private_balance(token_contract_address, account_2); + assert(balance_of_account_2_after_withdrawal == WITHDRAWAL_AMOUNT); + + let balance_of_escrow_after_withdrawal = + get_private_balance(token_contract_address, escrow_contract_address); + assert(balance_of_escrow_after_withdrawal == TRANSFER_AMOUNT - WITHDRAWAL_AMOUNT); +} diff --git a/yarn-project/txe/src/txe_service/txe_service.ts b/yarn-project/txe/src/txe_service/txe_service.ts index 76211e383a44..b9553bb3f385 100644 --- a/yarn-project/txe/src/txe_service/txe_service.ts +++ b/yarn-project/txe/src/txe_service/txe_service.ts @@ -95,12 +95,13 @@ export class TXEService { ]); if (!fromSingle(secret).equals(Fr.ZERO)) { - await this.createAccount(secret); + await this.addAccount(artifact, instance, secret); + } else { + await (this.typedOracle as TXE).addContractInstance(instance); + await (this.typedOracle as TXE).addContractArtifact(instance.contractClassId, artifact); + this.logger.debug(`Deployed ${artifact.name} at ${instance.address}`); } - this.logger.debug(`Deployed ${artifact.name} at ${instance.address}`); - await (this.typedOracle as TXE).addContractInstance(instance); - await (this.typedOracle as TXE).addContractArtifact(instance.contractClassId, artifact); return toForeignCallResult([ toArray([ instance.salt, @@ -137,6 +138,7 @@ export class TXEService { async createAccount(secret: ForeignCallSingle) { const keyStore = (this.typedOracle as TXE).getKeyStore(); const secretFr = fromSingle(secret); + // This is a footgun ! const completeAddress = await keyStore.addAccount(secretFr, secretFr); const accountStore = (this.typedOracle as TXE).getTXEDatabase(); await accountStore.setAccount(completeAddress.address, completeAddress);