o1js
Version:
TypeScript framework for zk-SNARKs and zkApps
471 lines (416 loc) • 16.1 kB
text/typescript
import { expect } from 'expect';
import { AccountUpdate, Mina, Permissions, PrivateKey, UInt64 } from 'o1js';
import { getProfiler } from '../../utils/profiler.js';
import { TokenContract, addresses, createDex, keys, tokenIds } from './dex.js';
let proofsEnabled = false;
console.log('starting upgradeability tests');
await upgradeabilityTests({
withVesting: false,
});
console.log('all upgradeability tests were successful! 🎉');
console.log('starting atomic actions tests');
await atomicActionsTest({
withVesting: false,
});
console.log('all atomic actions tests were successful!');
async function atomicActionsTest({ withVesting }: { withVesting: boolean }) {
const DexProfiler = getProfiler('DEX profiler atomic actions');
DexProfiler.start('DEX test flow');
let Local = await Mina.LocalBlockchain({
proofsEnabled,
enforceTransactionLimits: false,
});
Mina.setActiveInstance(Local);
let [feePayer] = Local.testAccounts;
let tx, balances;
let options = withVesting ? { lockedLiquiditySlots: 2 } : undefined;
let { Dex, DexTokenHolder, getTokenBalances } = createDex(options);
// analyze methods for quick error feedback
await DexTokenHolder.analyzeMethods();
await Dex.analyzeMethods();
if (proofsEnabled) {
// compile & deploy all zkApps
console.log('compile (dex token holder)...');
await DexTokenHolder.compile();
console.log('compile (dex main contract)...');
await Dex.compile();
}
let tokenX = new TokenContract(addresses.tokenX);
let tokenY = new TokenContract(addresses.tokenY);
let dex = new Dex(addresses.dex);
let dexTokenHolderX = new DexTokenHolder(addresses.dex, tokenIds.X);
let dexTokenHolderY = new DexTokenHolder(addresses.dex, tokenIds.Y);
console.log('deploy & init token contracts...');
tx = await Mina.transaction(feePayer, async () => {
await tokenX.deploy();
await tokenY.deploy();
// pay fees for creating 2 token contract accounts, and fund them so each can create 2 accounts themselves
const accountFee = Mina.getNetworkConstants().accountCreationFee;
let feePayerUpdate = AccountUpdate.fundNewAccount(feePayer, 2);
feePayerUpdate.send({ to: tokenX.self, amount: accountFee.mul(2) });
feePayerUpdate.send({ to: tokenY.self, amount: accountFee.mul(2) });
});
await tx.prove();
tx.sign([feePayer.key, keys.tokenX, keys.tokenY]);
await tx.send();
balances = getTokenBalances();
console.log(
'Token contract tokens (X, Y):',
balances.tokenContract.X,
balances.tokenContract.Y
);
/**
* # Atomic Actions 1
*
* Preconditions:
* - SC deployed with a non-writable delegate field (impossible permission)
*
* Actions:
* - User sets permissions for delegate to be edible
* - User changes delegate field to something
*
* Expected:
* - delegate now holds a new address
*/
console.log('deploy dex contracts...');
tx = await Mina.transaction(feePayer, async () => {
// pay fees for creating 3 dex accounts
AccountUpdate.fundNewAccount(feePayer, 3);
await dex.deploy();
await dexTokenHolderX.deploy();
await tokenX.approveAccountUpdate(dexTokenHolderX.self);
await dexTokenHolderY.deploy();
await tokenY.approveAccountUpdate(dexTokenHolderY.self);
console.log('manipulating setDelegate field to impossible...');
// setting the setDelegate permission field to impossible
let dexAccount = AccountUpdate.create(addresses.dex);
dexAccount.account.permissions.set({
...Permissions.initial(),
setDelegate: Permissions.impossible(),
});
dexAccount.requireSignature();
});
await tx.prove();
tx.sign([feePayer.key, keys.dex]);
await tx.send();
console.log(
'trying to change delegate (setDelegate=impossible, should fail)'
);
tx = await Mina.transaction(feePayer, async () => {
// setting the delegate field to something, although permissions forbid it
let dexAccount = AccountUpdate.create(addresses.dex);
dexAccount.account.delegate.set(PrivateKey.random().toPublicKey());
dexAccount.requireSignature();
});
await tx.prove();
await expect(tx.sign([feePayer.key, keys.dex]).send()).rejects.toThrow(
/Cannot update field 'delegate'/
);
console.log('changing delegate permission back to normal');
tx = await Mina.transaction(feePayer, async () => {
let dexAccount = AccountUpdate.create(addresses.dex);
dexAccount.account.permissions.set({
...Permissions.initial(),
setDelegate: Permissions.proofOrSignature(),
});
dexAccount.requireSignature();
});
await tx.prove();
await tx.sign([feePayer.key, keys.dex]).send();
console.log('changing delegate field to a new address');
let newDelegate = PrivateKey.random().toPublicKey();
tx = await Mina.transaction(feePayer, async () => {
let dexAccount = AccountUpdate.create(addresses.dex);
dexAccount.account.delegate.set(newDelegate);
dexAccount.requireSignature();
});
await tx.prove();
await tx.sign([feePayer.key, keys.dex]).send();
Mina.getAccount(addresses.dex).delegate?.assertEquals(newDelegate);
/**
* # Atomic Actions 2
*
* Preconditions:
* - Similar to happy path (SC deployed), but changing order of transaction and updates (within one transaction)
*
* Actions:
* - Include multiple actions of different order in the same transaction
*
* Expected:
* - tx fails if the order of actions is not valid
*/
console.log(
'changing permission to impossible and then trying to change delegate field - in one transaction'
);
tx = await Mina.transaction(feePayer, async () => {
// changing the permission to impossible and then trying to change the delegate field
let permissionUpdate = AccountUpdate.create(addresses.dex);
permissionUpdate.account.permissions.set({
...Permissions.initial(),
setDelegate: Permissions.impossible(),
});
permissionUpdate.requireSignature();
let fieldUpdate = AccountUpdate.create(addresses.dex);
fieldUpdate.account.delegate.set(PrivateKey.random().toPublicKey());
fieldUpdate.requireSignature();
});
await tx.prove();
await expect(tx.sign([feePayer.key, keys.dex]).send()).rejects.toThrow(
/Cannot update field 'delegate'/
);
/**
* # Atomic Actions 3
*
* Preconditions:
* - Similar to happy path (SC deployed) and the test before, but editing fields and changing permissions will be successful
*
* Actions:
* - Include multiple actions in one transaction
* edit field -> change permission back -> successful state transition
*
* Expected:
* - tx is successful and state updated
*/
console.log('creating multiple valid account updates in one transaction');
newDelegate = PrivateKey.random().toPublicKey();
tx = await Mina.transaction(feePayer, async () => {
// changing field
let fieldUpdate = AccountUpdate.create(addresses.dex);
fieldUpdate.account.delegate.set(newDelegate);
fieldUpdate.requireSignature();
// changing permissions back to impossible
let permissionUpdate2 = AccountUpdate.create(addresses.dex);
permissionUpdate2.account.permissions.set({
...Permissions.initial(),
setDelegate: Permissions.impossible(),
});
permissionUpdate2.requireSignature();
});
await tx.prove();
await tx.sign([feePayer.key, keys.dex]).send();
Mina.getAccount(addresses.dex).delegate?.assertEquals(newDelegate);
DexProfiler.stop().store();
}
async function upgradeabilityTests({ withVesting }: { withVesting: boolean }) {
const DexProfiler = getProfiler('DEX profiler upgradeability tests');
DexProfiler.start('DEX test flow');
let Local = await Mina.LocalBlockchain({
proofsEnabled: proofsEnabled,
enforceTransactionLimits: false,
});
Mina.setActiveInstance(Local);
let [feePayer] = Local.testAccounts;
let tx, balances, oldBalances;
let options = withVesting ? { lockedLiquiditySlots: 2 } : undefined;
let {
Dex,
DexTokenHolder,
ModifiedDexTokenHolder,
ModifiedDex,
getTokenBalances,
} = createDex(options);
// analyze methods for quick error feedback
await DexTokenHolder.analyzeMethods();
await Dex.analyzeMethods();
// compile & deploy all zkApps
console.log('compile (token contract)...');
await TokenContract.compile();
console.log('compile (dex token holder)...');
await DexTokenHolder.compile();
console.log('compile (dex main contract)...');
await Dex.compile();
let tokenX = new TokenContract(addresses.tokenX);
let tokenY = new TokenContract(addresses.tokenY);
let dex = new Dex(addresses.dex);
let dexTokenHolderX = new DexTokenHolder(addresses.dex, tokenIds.X);
let dexTokenHolderY = new DexTokenHolder(addresses.dex, tokenIds.Y);
console.log('deploy & init token contracts...');
tx = await Mina.transaction(feePayer, async () => {
await tokenX.deploy();
await tokenY.deploy();
// pay fees for creating 2 token contract accounts, and fund them so each can create 2 accounts themselves
const accountFee = Mina.getNetworkConstants().accountCreationFee;
let feePayerUpdate = AccountUpdate.fundNewAccount(feePayer, 2);
feePayerUpdate.send({ to: tokenX.self, amount: accountFee.mul(2) });
feePayerUpdate.send({ to: tokenY.self, amount: accountFee.mul(2) });
});
await tx.prove();
tx.sign([feePayer.key, keys.tokenX, keys.tokenY]);
await tx.send();
balances = getTokenBalances();
console.log(
'Token contract tokens (X, Y):',
balances.tokenContract.X,
balances.tokenContract.Y
);
/**
* # Upgradeability 1 - Happy Path
*
* Preconditions:
* - Initially the SC was deployed with the ability to be upgradable
*
* Actions:
* - deploy valid SC with upgradeable permissions
* - compile and upgrade new contract with different source code
* - set permissions to be non upgradeable
*
* Expected:
* - upgrade is successful
* - updates to the source code and methods are recognized
* (e.g. new formula applies to the swap)
*/
console.log('deploy dex contracts...');
tx = await Mina.transaction(feePayer, async () => {
// pay fees for creating 3 dex accounts
AccountUpdate.fundNewAccount(feePayer, 3);
await dex.deploy();
await dexTokenHolderX.deploy();
await tokenX.approveAccountUpdate(dexTokenHolderX.self);
await dexTokenHolderY.deploy();
await tokenY.approveAccountUpdate(dexTokenHolderY.self);
});
await tx.prove();
tx.sign([feePayer.key, keys.dex]);
await tx.send();
console.log('transfer tokens to user');
tx = await Mina.transaction(
{
sender: feePayer,
fee: Mina.getNetworkConstants().accountCreationFee.mul(1),
},
async () => {
let au = AccountUpdate.createSigned(feePayer);
au.balance.subInPlace(
Mina.getNetworkConstants().accountCreationFee.mul(4)
);
au.send({ to: addresses.user, amount: 20e9 }); // give users MINA to pay fees
au.send({ to: addresses.user2, amount: 20e9 });
// transfer to fee payer so they can provide initial liquidity
await tokenX.transfer(addresses.tokenX, feePayer, 10_000);
await tokenY.transfer(addresses.tokenY, feePayer, 10_000);
// mint tokens to the user (this is additional to the tokens minted at the beginning, so we can overflow the balance
await tokenX.init2();
await tokenY.init2();
}
);
await tx.prove();
tx.sign([feePayer.key, keys.tokenX, keys.tokenY]);
await tx.send();
[ ] = [balances, getTokenBalances()];
console.log('User tokens (X, Y):', balances.user.X, balances.user.Y);
console.log('User MINA:', balances.user.MINA);
console.log(
'deploying an upgraded DexTokenHolder contract (adjusted swap method) and Dex contract'
);
console.log('compiling modified DexTokenHolder contract...');
await ModifiedDexTokenHolder.compile();
console.log('compiling modified Dex contract...');
await ModifiedDex.compile();
let modifiedDex = new ModifiedDex(addresses.dex);
let modifiedDexTokenHolderX = new ModifiedDexTokenHolder(
addresses.dex,
tokenIds.X
);
let modifiedDexTokenHolderY = new ModifiedDexTokenHolder(
addresses.dex,
tokenIds.Y
);
tx = await Mina.transaction(feePayer, async () => {
await modifiedDex.deploy();
await modifiedDexTokenHolderX.deploy();
await tokenX.approveAccountUpdate(modifiedDexTokenHolderX.self);
await modifiedDexTokenHolderY.deploy();
await tokenY.approveAccountUpdate(modifiedDexTokenHolderY.self);
});
await tx.prove();
tx.sign([feePayer.key, keys.dex]);
await tx.send();
// Making sure that both token holder accounts have been updated with the new modified verification key
expect(
Mina.getAccount(addresses.dex, tokenX.deriveTokenId()).zkapp
?.verificationKey?.data
).toEqual(ModifiedDexTokenHolder._verificationKey?.data);
expect(
Mina.getAccount(addresses.dex, tokenY.deriveTokenId()).zkapp
?.verificationKey?.data
).toEqual(ModifiedDexTokenHolder._verificationKey?.data);
// this is important; we have to re-enable proof production (and verification) to make sure the proofs are valid against the newly deployed VK
Local.setProofsEnabled(true);
console.log('supply liquidity -- base');
tx = await Mina.transaction(
{
sender: feePayer,
fee: Mina.getNetworkConstants().accountCreationFee.mul(1),
},
async () => {
AccountUpdate.fundNewAccount(feePayer);
await modifiedDex.supplyLiquidityBase(
UInt64.from(10_000),
UInt64.from(10_000)
);
}
);
await tx.prove();
tx.sign([feePayer.key]);
await tx.send();
[ ] = [balances, getTokenBalances()];
console.log('DEX liquidity (X, Y):', balances.dex.X, balances.dex.Y);
let USER_DX = 10n;
console.log('swap 10 X for Y');
tx = await Mina.transaction(addresses.user, async () => {
await modifiedDex.swapX(UInt64.from(USER_DX));
});
await tx.prove();
await tx.sign([keys.user]).send();
[ ] = [balances, getTokenBalances()];
console.log('User tokens (X, Y):', balances.user.X, balances.user.Y);
// according to the newly modified formula `dy = y.mul(dx).div(x.add(dx)).add(15)`;
// the user should have a balance of 1_000_024 Y
expect(oldBalances.user.Y).toEqual(1_000_000n);
expect(balances.user.Y).toEqual(1_000_024n);
/**
* # Upgradeability 2 and 3 - Upgrading forbidden, previous methods should still be valid
*
* Preconditions:
* - DEX contract deployed, but upgrading forbidden
*
* Actions:
* - setVerificationKey permission to impossible
* - try upgrading contract
*
* Expected:
* - tx fails and contract cannot be upgraded
* - methods from the previous contract should still be valid
*/
console.log('changing upgrade permissions to impossible');
tx = await Mina.transaction(feePayer, async () => {
// pay fees for creating 3 dex accounts
let update = AccountUpdate.createSigned(addresses.dex);
update.account.permissions.set({
...Permissions.initial(),
setVerificationKey:
Permissions.VerificationKey.impossibleDuringCurrentVersion(),
});
});
await tx.prove();
tx.sign([feePayer.key, keys.dex]);
await tx.send();
console.log('trying to upgrade contract - should fail');
tx = await Mina.transaction(feePayer, async () => {
await modifiedDex.deploy(); // cannot deploy new VK because its forbidden
});
await tx.prove();
await expect(tx.sign([feePayer.key, keys.dex]).send()).rejects.toThrow(
/Cannot update field 'verificationKey'/
);
console.log('trying to invoke modified swap method');
// method should still be valid since the upgrade was forbidden
USER_DX = 10n;
console.log('swap 10 X for Y');
tx = await Mina.transaction(addresses.user, async () => {
await modifiedDex.swapX(UInt64.from(USER_DX));
});
await tx.prove();
await tx.sign([keys.user]).send();
DexProfiler.stop().store();
}