add withdraw all

src/workflows/withdrawal/completeERC20Withdrawal.ts

@@ -5,9 +5,10 @@ import { ImmutableXConfiguration } from '../../config';
 import {
   Core,
   Core__factory,
+  CoreV4__factory,
   Registration,
   Registration__factory,
-  CoreV4__factory,
+  RegistrationV4__factory,
 } from '../../contracts';
 import { ERC20Token } from '../../types';
 import {
@@ -125,3 +126,30 @@ export async function completeERC20WithdrawalV2Workflow(
 
   return signer.sendTransaction(populatedTransaction);
 }
+
+export async function completeAllERC20WithdrawalWorkflow(
+  signer: Signer,
+  starkPublicKey: string,
+  token: ERC20Token,
+  encodingApi: EncodingApi,
+  config: ImmutableXConfiguration,
+): Promise<TransactionResponse> {
+  const assetType = await getEncodeAssetInfo('asset', 'ERC20', encodingApi, {
+    token_address: token.tokenAddress,
+  });
+
+  const registrationContract = RegistrationV4__factory.connect(
+    config.ethConfiguration.registrationContractAddress,
+    signer,
+  );
+
+  const ethAddress = await signer.getAddress();
+  const populatedTransaction =
+    await registrationContract.populateTransaction.withdrawAll(
+      ethAddress,
+      starkPublicKey,
+      assetType.asset_id,
+    );
+
+  return signer.sendTransaction(populatedTransaction);
+}

src/workflows/withdrawal/completeEthWithdrawal.ts

@@ -5,9 +5,10 @@ import { ImmutableXConfiguration } from '../../config';
 import {
   Core,
   Core__factory,
+  CoreV4__factory,
   Registration,
   Registration__factory,
-  CoreV4__factory,
+  RegistrationV4__factory,
 } from '../../contracts';
 import {
   getSignableRegistrationOnchain,
@@ -118,3 +119,27 @@ export async function completeEthWithdrawalV2Workflow(
 
   return signer.sendTransaction(populatedTransaction);
 }
+
+export async function completeAllEthWithdrawalWorkflow(
+  signer: Signer,
+  starkPublicKey: string,
+  encodingApi: EncodingApi,
+  config: ImmutableXConfiguration,
+): Promise<TransactionResponse> {
+  const assetType = await getEncodeAssetInfo('asset', 'ETH', encodingApi);
+
+  const registrationContract = RegistrationV4__factory.connect(
+    config.ethConfiguration.registrationContractAddress,
+    signer,
+  );
+
+  const ethAddress = await signer.getAddress();
+  const populatedTransaction =
+    await registrationContract.populateTransaction.withdrawAll(
+      ethAddress,
+      starkPublicKey,
+      assetType.asset_id,
+    );
+
+  return signer.sendTransaction(populatedTransaction);
+}

src/workflows/withdrawal/completeWithdrawal.ts

@@ -1,13 +1,15 @@
 import { Signer } from '@ethersproject/abstract-signer';
 import { TransactionResponse } from '@ethersproject/providers';
 import { AnyToken } from 'src/types';
+import { EncodingApi, MintsApi, UsersApi } from 'src/api';
+import { ImmutableXConfiguration } from 'src/config';
 import {
+  completeAllEthWithdrawalWorkflow,
   completeEthWithdrawalV1Workflow,
   completeEthWithdrawalV2Workflow,
 } from './completeEthWithdrawal';
-import { EncodingApi, MintsApi, UsersApi } from 'src/api';
-import { ImmutableXConfiguration } from 'src/config';
 import {
+  completeAllERC20WithdrawalWorkflow,
   completeERC20WithdrawalV1Workflow,
   completeERC20WithdrawalV2Workflow,
 } from './completeERC20Withdrawal';
@@ -85,3 +87,40 @@ export async function completeWithdrawalV2Workflow(
       );
   }
 }
+
+export async function completeAllWithdrawalWorkflow(
+  signer: Signer,
+  starkPublicKey: string,
+  token: AnyToken,
+  encodingApi: EncodingApi,
+  mintsApi: MintsApi,
+  config: ImmutableXConfiguration,
+): Promise<TransactionResponse> {
+  switch (token.type) {
+    case 'ETH':
+      return completeAllEthWithdrawalWorkflow(
+        signer,
+        starkPublicKey,
+        encodingApi,
+        config,
+      );
+    case 'ERC20':
+      return completeAllERC20WithdrawalWorkflow(
+        signer,
+        starkPublicKey,
+        token,
+        encodingApi,
+        config,
+      );
+    case 'ERC721':
+      // for ERC721, if the v3 balance > 0, then the v4 balance is 0
+      return completeERC721WithdrawalV2Workflow(
+        signer,
+        starkPublicKey,
+        token,
+        encodingApi,
+        mintsApi,
+        config,
+      );
+  }
+}

src/workflows/withdrawal/getWithdrawalBalance.ts

@@ -0,0 +1,149 @@
+import { Signer } from '@ethersproject/abstract-signer';
+import { EncodingApi, MintsApi } from '../../api';
+import { ImmutableXConfiguration } from '../../config';
+import { CoreV4__factory } from '../../contracts';
+import { AnyToken, ERC721Token } from '../../types';
+import { getEncodeAssetInfo } from './getEncodeAssetInfo';
+import { BigNumber } from 'ethers';
+
+async function getWithdrawalBalance(
+  signer: Signer,
+  ownerKey: string,
+  assetId: string,
+  config: ImmutableXConfiguration,
+) {
+  const coreContract = CoreV4__factory.connect(
+    config.ethConfiguration.coreContractAddress,
+    signer,
+  );
+  return coreContract.getWithdrawalBalance(ownerKey, assetId);
+}
+
+async function getETHWithdrawalBalance(
+  signer: Signer,
+  ownerKey: string,
+  encodingApi: EncodingApi,
+  config: ImmutableXConfiguration,
+): Promise<BigNumber> {
+  const assetType = await getEncodeAssetInfo('asset', 'ETH', encodingApi);
+  return await getWithdrawalBalance(
+    signer,
+    ownerKey,
+    assetType.asset_id,
+    config,
+  );
+}
+
+async function getERC20WithdrawalBalance(
+  signer: Signer,
+  ownerKey: string,
+  tokenAddress: string,
+  encodingApi: EncodingApi,
+  config: ImmutableXConfiguration,
+): Promise<BigNumber> {
+  const assetType = await getEncodeAssetInfo('asset', 'ERC20', encodingApi, {
+    token_address: tokenAddress,
+  });
+  return await getWithdrawalBalance(
+    signer,
+    ownerKey,
+    assetType.asset_id,
+    config,
+  );
+}
+
+async function getERC721WithdrawalBalance(
+  signer: Signer,
+  ownerKey: string,
+  token: ERC721Token,
+  encodingApi: EncodingApi,
+  mintsApi: MintsApi,
+  config: ImmutableXConfiguration,
+): Promise<BigNumber> {
+  const tokenAddress = token.tokenAddress;
+  const tokenId = token.tokenId;
+  return await mintsApi
+    .getMintableTokenDetailsByClientTokenId({
+      tokenAddress,
+      tokenId,
+    })
+    .then(async mintableToken => {
+      const assetType = await getEncodeAssetInfo(
+        'mintable-asset',
+        'ERC721',
+        encodingApi,
+        {
+          id: tokenId,
+          token_address: tokenAddress,
+          ...(mintableToken.data.blueprint && {
+            blueprint: mintableToken.data.blueprint,
+          }),
+        },
+      );
+      return await getWithdrawalBalance(
+        signer,
+        ownerKey,
+        assetType.asset_id,
+        config,
+      );
+    })
+    .catch(async error => {
+      if (error.response?.status === 404) {
+        // token is not a mintable ERC721 token
+        const assetType = await getEncodeAssetInfo(
+          'asset',
+          'ERC721',
+          encodingApi,
+          {
+            token_id: token.tokenId,
+            token_address: token.tokenAddress,
+          },
+        );
+        return await getWithdrawalBalance(
+          signer,
+          ownerKey,
+          assetType.asset_id,
+          config,
+        );
+      }
+      throw error; // unable to recover from any other kind of error
+    });
+}
+
+export async function getWithdrawalBalanceWorkflow(
+  signer: Signer,
+  ownerKey: string,
+  token: AnyToken,
+  encodingApi: EncodingApi,
+  mintsApi: MintsApi,
+  config: ImmutableXConfiguration,
+): Promise<BigNumber> {
+  switch (token.type) {
+    case 'ETH':
+      return await getETHWithdrawalBalance(
+        signer,
+        ownerKey,
+        encodingApi,
+        config,
+      );
+    case 'ERC20':
+      return await getERC20WithdrawalBalance(
+        signer,
+        ownerKey,
+        token.tokenAddress,
+        encodingApi,
+        config,
+      );
+    case 'ERC721':
+      return await getERC721WithdrawalBalance(
+        signer,
+        ownerKey,
+        token,
+        encodingApi,
+        mintsApi,
+        config,
+      );
+    default:
+      throw new Error('Unsupported token type');
+  }
+}

src/workflows/workflows.ts

@@ -68,9 +68,12 @@ import {
 } from './primarySales';
 import { TransactionResponse } from '@ethersproject/providers';
 import {
+  completeAllWithdrawalWorkflow,
   completeWithdrawalV1Workflow,
   completeWithdrawalV2Workflow,
 } from './withdrawal/completeWithdrawal';
+import { BigNumber } from 'ethers';
+import { getWithdrawalBalanceWorkflow } from './withdrawal/getWithdrawalBalance';
 
 export class Workflows {
   private readonly depositsApi: DepositsApi;
@@ -298,6 +301,7 @@ export class Workflows {
     const majorContractVersion = await this.parseMajorContractVersion(
       starkExContractInfo.data.version,
     );
+    const starkPublicKey = await walletConnection.starkSigner.getAddress();
 
     if (majorContractVersion === 3) {
       const starkPublicKey = await walletConnection.starkSigner.getAddress();
@@ -311,7 +315,46 @@ export class Workflows {
         this.config,
       );
     } else if (majorContractVersion >= 4) {
-      const ethAddress = await walletConnection.ethSigner.getAddress();
+      return this.completeWithdrawalAll(
+        walletConnection,
+        starkPublicKey,
+        token,
+      );
+    } else {
+      throw new Error(
+        `Invalid StarkEx contract version (${majorContractVersion}). Please try again later.`,
+      );
+    }
+  }
+
+  private async completeWithdrawalAll(
+    walletConnection: WalletConnection,
+    starkPublicKey: string,
+    token: AnyToken,
+  ): Promise<TransactionResponse> {
+    const ethAddress = await walletConnection.ethSigner.getAddress();
+    const { v3Balance, v4Balance } = await this.getWithdrawalBalances(
+      walletConnection.ethSigner,
+      starkPublicKey,
+      ethAddress,
+      token,
+    );
+
+    if (v3Balance.gt(0)) {
+      const isRegistered = await this.isRegisteredOnchain(walletConnection);
+      if (isRegistered) {
+        return completeAllWithdrawalWorkflow(
+          walletConnection.ethSigner,
+          starkPublicKey,
+          token,
+          this.encodingApi,
+          this.mintsApi,
+          this.config,
+        );
+      }
+      throw new Error('User unregistered');
+    }
+    if (v4Balance.gt(0)) {
       return completeWithdrawalV2Workflow(
         walletConnection.ethSigner,
         ethAddress,
@@ -320,11 +363,33 @@ export class Workflows {
         this.mintsApi,
         this.config,
       );
-    } else {
-      throw new Error(
-        `Invalid StarkEx contract version (${majorContractVersion}). Please try again later.`,
-      );
     }
+    throw new Error('Nothing to withdraw');
+  }
+
+  private async getWithdrawalBalances(
+    signer: Signer,
+    starkPublicKey: string,
+    ethAddress: string,
+    token: AnyToken,
+  ): Promise<{ v3Balance: BigNumber; v4Balance: BigNumber }> {
+    const v3Balance = await getWithdrawalBalanceWorkflow(
+      signer,
+      starkPublicKey,
+      token,
+      this.encodingApi,
+      this.mintsApi,
+      this.config,
+    );
+    const v4Balance = await getWithdrawalBalanceWorkflow(
+      signer,
+      ethAddress,
+      token,
+      this.encodingApi,
+      this.mintsApi,
+      this.config,
+    );
+    return { v3Balance, v4Balance };
   }
 
   public async createOrder(