From 2fb1f55ec405346bbabb1701c48c2f3d5e5a8672 Mon Sep 17 00:00:00 2001
From: Gregor <gregor.mitscha-baude@gmx.at>
Date: Tue, 11 Jun 2024 11:37:48 +0200
Subject: [PATCH] fix concurrent transfer test in a more principled way

---
 FungibleToken.ts                   | 23 ++++++++++++++++-------
 FungibleTokenAdmin.ts              | 20 +++++++++++++-------
 examples/concurrent-transfer.eg.ts |  3 +--
 3 files changed, 30 insertions(+), 16 deletions(-)

diff --git a/FungibleToken.ts b/FungibleToken.ts
index 790a174..d752794 100644
--- a/FungibleToken.ts
+++ b/FungibleToken.ts
@@ -76,13 +76,20 @@ export class FungibleToken extends TokenContract implements FungibleTokenLike {
     this.actionState.set(Reducer.initialActionState)
   }
 
-  public getAdminContract(): FungibleTokenAdminBase {
-    return (new FungibleToken.adminContract(this.admin.getAndRequireEquals()))
+  public async getAdminContract(): Promise<FungibleTokenAdminBase> {
+    const admin = await Provable.witnessAsync(PublicKey, async () => {
+      let pk = await this.admin.fetch()
+      assert(pk !== undefined, "could not fetch admin contract key")
+      return pk
+    })
+    this.admin.requireEquals(admin)
+    return (new FungibleToken.adminContract(admin))
   }
 
   @method
   async setAdmin(admin: PublicKey) {
-    const canChangeAdmin = await this.getAdminContract().canChangeAdmin(admin)
+    const adminContract = await this.getAdminContract()
+    const canChangeAdmin = await adminContract.canChangeAdmin(admin)
     canChangeAdmin.assertTrue()
     this.admin.set(admin)
     this.emitEvent("SetAdmin", admin)
@@ -92,8 +99,8 @@ export class FungibleToken extends TokenContract implements FungibleTokenLike {
   async mint(recipient: PublicKey, amount: UInt64) {
     this.paused.getAndRequireEquals().assertFalse()
     const accountUpdate = this.internal.mint({ address: recipient, amount })
-    const canMint = await this.getAdminContract()
-      .canMint(accountUpdate)
+    const adminContract = await this.getAdminContract()
+    const canMint = await adminContract.canMint(accountUpdate)
     canMint.assertTrue()
     this.approve(accountUpdate)
     this.emitEvent("Mint", new MintEvent({ recipient, amount }))
@@ -112,14 +119,16 @@ export class FungibleToken extends TokenContract implements FungibleTokenLike {
 
   @method
   async pause() {
-    const canPause = await this.getAdminContract().canPause()
+    const adminContract = await this.getAdminContract()
+    const canPause = await adminContract.canPause()
     canPause.assertTrue()
     this.paused.set(Bool(true))
   }
 
   @method
   async resume() {
-    const canResume = await this.getAdminContract().canResume()
+    const adminContract = await this.getAdminContract()
+    const canResume = await adminContract.canResume()
     canResume.assertTrue()
     this.paused.set(Bool(false))
   }
diff --git a/FungibleTokenAdmin.ts b/FungibleTokenAdmin.ts
index cb1fdf2..a00953b 100644
--- a/FungibleTokenAdmin.ts
+++ b/FungibleTokenAdmin.ts
@@ -1,13 +1,14 @@
 import {
   AccountUpdate,
+  assert,
   Bool,
   DeployArgs,
   method,
+  Provable,
   PublicKey,
   SmartContract,
   State,
   state,
-  UInt64,
 } from "o1js"
 
 export type FungibleTokenAdminBase = SmartContract & {
@@ -38,32 +39,37 @@ export class FungibleTokenAdmin extends SmartContract implements FungibleTokenAd
     this.adminPublicKey.set(props.adminPublicKey)
   }
 
-  private ensureAdminSignature() {
-    const admin = this.adminPublicKey.getAndRequireEquals()
+  private async ensureAdminSignature() {
+    const admin = await Provable.witnessAsync(PublicKey, async () => {
+      let pk = await this.adminPublicKey.fetch()
+      assert(pk !== undefined, "could not fetch admin public key")
+      return pk
+    })
+    this.adminPublicKey.requireEquals(admin)
     return AccountUpdate.createSigned(admin)
   }
 
   @method.returns(Bool)
   public async canMint(_accountUpdate: AccountUpdate) {
-    this.ensureAdminSignature()
+    await this.ensureAdminSignature()
     return Bool(true)
   }
 
   @method.returns(Bool)
   public async canChangeAdmin(_admin: PublicKey) {
-    this.ensureAdminSignature()
+    await this.ensureAdminSignature()
     return Bool(true)
   }
 
   @method.returns(Bool)
   public async canPause(): Promise<Bool> {
-    this.ensureAdminSignature()
+    await this.ensureAdminSignature()
     return Bool(true)
   }
 
   @method.returns(Bool)
   public async canResume(): Promise<Bool> {
-    this.ensureAdminSignature()
+    await this.ensureAdminSignature()
     return Bool(true)
   }
 }
diff --git a/examples/concurrent-transfer.eg.ts b/examples/concurrent-transfer.eg.ts
index f2f8255..bf82c7f 100644
--- a/examples/concurrent-transfer.eg.ts
+++ b/examples/concurrent-transfer.eg.ts
@@ -1,4 +1,4 @@
-import { AccountUpdate, fetchAccount, Mina, PrivateKey, PublicKey, UInt64, UInt8 } from "o1js"
+import { AccountUpdate, Mina, PrivateKey, PublicKey, UInt64, UInt8 } from "o1js"
 import { FungibleToken, FungibleTokenAdmin } from "../index.js"
 
 const url = "https://proxy.devnet.minaexplorer.com/graphql"
@@ -100,7 +100,6 @@ const deployTxResult = await deployTx.send().then((v) => v.wait())
 console.log("Deploy tx:", deployTxResult.hash)
 
 console.log("Minting new tokens to Alexa.")
-await fetchAccount({ publicKey: admin.publicKey }) // hack to ensure the admin account is fetched
 const mintTx = await Mina.transaction({
   sender: feepayer.publicKey,
   fee,