CORE-2193 Clean up post StarkEx upgrade (#1855)

Co-authored-by: Alex Karolis <[email protected]>

packages/x-provider/src/signable-actions/withdrawal-actions/completeERC20Withdrawal.ts

@@ -1,11 +1,19 @@
+// Note that this file contains withdrawal functions that are shared
+// by both ERC20 and ETH in completeERC20WithdrawalAction and completeEthWithdrawalAction
 import { Signer } from '@ethersproject/abstract-signer';
-import { ERC20Token, StarkSigner } from '@imtbl/x-client';
+import { TransactionResponse } from '@ethersproject/providers';
+import {
+  Contracts,
+  ERC20Token,
+  ImmutableXConfiguration,
+  StarkSigner,
+  signRegisterEthAddress,
+} from '@imtbl/x-client';
 import { isRegisteredOnChain } from '../registration';
 import { getEncodeAssetInfo } from './getEncodeAssetInfo';
 import { validateChain } from '../helpers';
 import { ProviderConfiguration } from '../../config';
 import { getWithdrawalBalances } from './getWithdrawalBalance';
-import { executeRegisterAndWithdrawAllFungible, executeWithdrawAllFungible } from './completeEthWithdrawal';
 
 type CompleteERC20WithdrawalWorkflowParams = {
   ethSigner: Signer;
@@ -17,6 +25,80 @@ type CompleteERC20WithdrawalWorkflowParams = {
 
 const ERC20TokenType = 'ERC20';
 
+export async function executeRegisterAndWithdrawAllFungible(
+  ethSigner: Signer,
+  starkSigner: StarkSigner,
+  starkPublicKey: string,
+  assetType: string,
+  config: ImmutableXConfiguration,
+): Promise<TransactionResponse> {
+  const etherKey = await ethSigner.getAddress();
+
+  const starkSignature = await signRegisterEthAddress(
+    starkSigner,
+    etherKey,
+    starkPublicKey,
+  );
+
+  // we use registration v4 contract as a wrapper for the core contract
+  // so that v3 and v4 withdrawals, AND on-chain registration can be executed in a single transaction
+  const contract = Contracts.RegistrationV4.connect(
+    config.ethConfiguration.registrationV4ContractAddress || config.ethConfiguration.registrationContractAddress,
+    ethSigner,
+  );
+
+  const populatedTransaction = await contract.populateTransaction.registerAndWithdrawAll(
+    etherKey,
+    starkPublicKey,
+    starkSignature,
+    assetType,
+  );
+
+  return ethSigner.sendTransaction(populatedTransaction);
+}
+
+export async function executeWithdrawAllFungible(
+  ethSigner: Signer,
+  starkPublicKey: string,
+  assetType: string,
+  config: ImmutableXConfiguration,
+): Promise<TransactionResponse> {
+  // we use registration v4 contract as a wrapper for the core contract
+  // so that v3 and v4 withdrawals can be executed in a single transaction
+  // (if there are pending withdrawable funds for both)
+  const contract = Contracts.RegistrationV4.connect(
+    config.ethConfiguration.registrationV4ContractAddress || config.ethConfiguration.registrationContractAddress,
+    ethSigner,
+  );
+
+  const populatedTransaction = await contract.populateTransaction.withdrawAll(
+    await ethSigner.getAddress(),
+    starkPublicKey,
+    assetType,
+  );
+
+  return ethSigner.sendTransaction(populatedTransaction);
+}
+
+export async function executeWithdrawFungible(
+  ethSigner: Signer,
+  starkPublicKey: string,
+  assetType: string,
+  config: ImmutableXConfiguration,
+): Promise<TransactionResponse> {
+  const contract = Contracts.CoreV4.connect(
+    config.ethConfiguration.coreContractAddress,
+    ethSigner,
+  );
+
+  const populatedTransaction = await contract.populateTransaction.withdraw(
+    await ethSigner.getAddress(),
+    assetType,
+  );
+
+  return ethSigner.sendTransaction(populatedTransaction);
+}
+
 // equivilant to Core SDK completeERC20WithdrawalV1Workflow
 // in src/workflows/withdrawal/completeERC20Withdrawal.ts
 export async function completeERC20WithdrawalAction({
@@ -42,28 +124,31 @@ export async function completeERC20WithdrawalAction({
     config.immutableXConfig,
   );
 
-  if (v3Balance.isZero() && v4Balance.isZero()) {
-    throw new Error('No balance to withdraw');
-  }
-
-  const isRegistered = await isRegisteredOnChain(
-    starkPublicKey,
-    ethSigner,
-    config,
-  );
-
   const assetType = await getEncodeAssetInfo('asset', ERC20TokenType, config.immutableXConfig, {
     token_address: token.tokenAddress,
   });
 
-  if (isRegistered) {
-    return executeWithdrawAllFungible(ethSigner, starkPublicKey, assetType.asset_type, config.immutableXConfig);
+  if (!v3Balance.isZero() && !v3Balance.isNegative()) {
+    const isRegistered = await isRegisteredOnChain(
+      starkPublicKey,
+      ethSigner,
+      config,
+    );
+    if (isRegistered) {
+      return executeWithdrawAllFungible(ethSigner, starkPublicKey, assetType.asset_type, config.immutableXConfig);
+    }
+    return executeRegisterAndWithdrawAllFungible(
+      ethSigner,
+      starkSigner,
+      starkPublicKey,
+      assetType.asset_type,
+      config.immutableXConfig,
+    );
   }
-  return executeRegisterAndWithdrawAllFungible(
-    ethSigner,
-    starkSigner,
-    starkPublicKey,
-    assetType.asset_type,
-    config.immutableXConfig,
-  );
+
+  if (!v4Balance.isZero() && !v4Balance.isNegative()) {
+    return executeWithdrawFungible(ethSigner, starkPublicKey, assetType.asset_type, config.immutableXConfig);
+  }
+
+  throw new Error('No balance to withdraw');
 }

packages/x-provider/src/signable-actions/withdrawal-actions/completeERC721Withdrawal.test.ts

@@ -2,6 +2,7 @@ import { Contracts } from '@imtbl/x-client';
 import { imx } from '@imtbl/generated-clients';
 import * as encUtils from 'enc-utils';
 import { TransactionResponse } from '@ethersproject/providers';
+import { BigNumber } from '@ethersproject/bignumber';
 import { getEncodeAssetInfo } from './getEncodeAssetInfo';
 import {
   getSignableRegistrationOnchain,
@@ -24,6 +25,7 @@ jest.mock('@imtbl/generated-clients');
 
 async function act(): Promise<TransactionResponse> {
   const signers = await generateSigners(privateKey1);
+  const mintsApi = new imx.MintsApi(testConfig.immutableXConfig.apiConfiguration);
   return await completeERC721WithdrawalAction({
     ethSigner: signers.ethSigner,
     starkSigner: signers.starkSigner,
@@ -34,11 +36,10 @@ async function act(): Promise<TransactionResponse> {
       tokenId: '23',
       tokenAddress: '0x23cv1',
     },
-  });
+  }, mintsApi);
 }
 
-// TODO fix MintsApi mocking so that getMintableTokenDetailsByClientTokenId does not return undefined
-describe.skip('completeERC721Withdrawal action', () => {
+describe('completeERC721Withdrawal action', () => {
   describe('when ERC721 is mintable', () => {
     const mintableErc721Token: imx.MintableTokenDetails = {
       token_id: '23',
@@ -68,7 +69,8 @@ describe.skip('completeERC721Withdrawal action', () => {
       });
     });
     it('should complete ERC721 withdrawal with on-chain registered user', async () => {
-      (Contracts.Core.connect as jest.Mock).mockReturnValue({
+      (Contracts.CoreV4.connect as jest.Mock).mockReturnValue({
+        getWithdrawalBalance: jest.fn().mockReturnValue(BigNumber.from('1')),
         populateTransaction: {
           withdrawAndMint: jest.fn().mockResolvedValue(transactionResponse),
         },
@@ -80,9 +82,12 @@ describe.skip('completeERC721Withdrawal action', () => {
       await expect(response).toEqual(transactionResponse);
     });
     it('should complete ERC721 withdrawal with unregistered user', async () => {
-      (Contracts.Registration.connect as jest.Mock).mockReturnValue({
+      (Contracts.CoreV4.connect as jest.Mock).mockReturnValue({
+        getWithdrawalBalance: jest.fn().mockReturnValue(BigNumber.from('1')),
+      });
+      (Contracts.RegistrationV4.connect as jest.Mock).mockReturnValue({
         populateTransaction: {
-          regsiterAndWithdrawAndMint: jest
+          registerWithdrawAndMint: jest
             .fn()
             .mockResolvedValue(transactionResponse),
         },
@@ -114,10 +119,18 @@ describe.skip('completeERC721Withdrawal action', () => {
         getMintableTokenDetailsByClientTokenId: jest
           .fn()
           .mockRejectedValue(error),
+        getMint: jest.fn(),
+        listMints: jest.fn(),
+        mintTokens: jest.fn(),
+        basePath: jest.fn(),
+        axios: jest.fn(),
+        configuration: jest.fn(),
       });
     });
+
     it('should complete ERC721 withdrawal with on-chain registered user', async () => {
-      (Contracts.Core.connect as jest.Mock).mockReturnValue({
+      (Contracts.CoreV4.connect as jest.Mock).mockReturnValue({
+        getWithdrawalBalance: jest.fn().mockReturnValue(BigNumber.from('1')),
         populateTransaction: {
           withdrawNft: jest.fn().mockResolvedValue(transactionResponse),
         },
@@ -126,8 +139,12 @@ describe.skip('completeERC721Withdrawal action', () => {
       const response = await act();
       await expect(response).toEqual(transactionResponse);
     });
+
     it('should complete ERC721 withdrawal with unregistered user', async () => {
-      (Contracts.Registration.connect as jest.Mock).mockReturnValue({
+      (Contracts.CoreV4.connect as jest.Mock).mockReturnValue({
+        getWithdrawalBalance: jest.fn().mockReturnValue(BigNumber.from('1')),
+      });
+      (Contracts.RegistrationV4.connect as jest.Mock).mockReturnValue({
         populateTransaction: {
           registerAndWithdrawNft: jest
             .fn()
@@ -162,6 +179,7 @@ describe.skip('completeERC721Withdrawal action', () => {
 
     it('should throw error', async () => {
       const signers = await generateSigners(privateKey1);
+      const mintsApi = new imx.MintsApi(testConfig.immutableXConfig.apiConfiguration);
       await expect(
         completeERC721WithdrawalAction({
           ethSigner: signers.ethSigner,
@@ -173,7 +191,7 @@ describe.skip('completeERC721Withdrawal action', () => {
             tokenId: '23',
             tokenAddress: '0x23cv1',
           },
-        }),
+        }, mintsApi),
       ).rejects.toThrowError();
     });
   });

packages/x-provider/src/signable-actions/withdrawal-actions/completeERC721Withdrawal.ts

@@ -1,5 +1,4 @@
 import { Signer } from '@ethersproject/abstract-signer';
-import { imx } from '@imtbl/generated-clients';
 import {
   Contracts,
   ERC721Token,
@@ -14,7 +13,7 @@ import { ProviderConfiguration } from '../../config';
 import { getEncodeAssetInfo } from './getEncodeAssetInfo';
 import { isRegisteredOnChain } from '../registration';
 import { validateChain } from '../helpers';
-import { getWithdrawalBalances } from './getWithdrawalBalance';
+import { getWithdrawalBalancesERC721 } from './getWithdrawalBalance';
 
 interface MintableERC721Withdrawal {
   type: 'ERC721';
@@ -253,53 +252,48 @@ export async function completeERC721WithdrawalAction({
   starkPublicKey,
   token,
   config,
-}: CompleteERC721WithdrawalActionParams) {
+}: CompleteERC721WithdrawalActionParams, mintsApi: MintsApi) {
   await validateChain(ethSigner, config.immutableXConfig);
-
-  const mintsApi = new imx.MintsApi(config.immutableXConfig.apiConfiguration);
-
+  const ethAddress = await ethSigner.getAddress();
   const {
     v3Balance,
     v4Balance,
-  } = await getWithdrawalBalances(
+  } = await getWithdrawalBalancesERC721(
     ethSigner,
     starkPublicKey,
-    await ethSigner.getAddress(),
+    ethAddress,
     {
       type: ERC721TokenType,
       tokenAddress: token.tokenAddress,
       tokenId: token.tokenId,
     },
     config.immutableXConfig,
+    mintsApi,
   );
-
-  if (v3Balance.isZero() && v4Balance.isZero()) {
-    throw new Error('No balance to withdraw');
+  if (!v3Balance.isZero() && !v3Balance.isNegative()) {
+    const isRegistered = await isRegisteredOnChain(
+      starkPublicKey,
+      ethSigner,
+      config,
+    );
+    // if the user is already registered on-chain, we can withdraw using stark key as the owner key
+    if (isRegistered) {
+      return completeERC721Withdrawal(mintsApi, ethSigner, starkPublicKey, token, config.immutableXConfig);
+    }
+    // if not registered on-chain, we need to register the user on-chain using stark public key as the owner key
+    return completeERC721RegisterAndWithdrawal(
+      mintsApi,
+      ethSigner,
+      starkSigner,
+      token,
+      config.immutableXConfig,
+    );
   }
 
-  const ethAddress = await ethSigner.getAddress();
-
   // if v4 balance is NOT zero, the withdrawal was prepared using eth address (using v2/withdrawals API)
-  if (!v4Balance.isZero()) {
+  if (!v4Balance.isZero() && !v4Balance.isNegative()) {
     return completeERC721Withdrawal(mintsApi, ethSigner, ethAddress, token, config.immutableXConfig);
   }
 
-  const isRegistered = await isRegisteredOnChain(
-    starkPublicKey,
-    ethSigner,
-    config,
-  );
-
-  // if the user is already registered on-chain, we can withdraw using stark key as the owner key
-  if (isRegistered) {
-    return completeERC721Withdrawal(mintsApi, ethSigner, starkPublicKey, token, config.immutableXConfig);
-  }
-  // if not registered on-chain, we need to register the user on-chain using stark public key as the owner key
-  return completeERC721RegisterAndWithdrawal(
-    mintsApi,
-    ethSigner,
-    starkSigner,
-    token,
-    config.immutableXConfig,
-  );
+  throw new Error('No balance to withdraw');
 }