feat: add batchUnstakingBuilder and withdrawUnbondedBuilder to support unstaking in polyx

modules/sdk-coin-polyx/src/lib/batchUnstakingBuilder.ts

@@ -0,0 +1,153 @@
+import { BaseCoin as CoinConfig } from '@bitgo/statics';
+import { methods } from '@substrate/txwrapper-polkadot';
+import { UnsignedTransaction, DecodedSigningPayload, DecodedSignedTx } from '@substrate/txwrapper-core';
+import { InvalidTransactionError, TransactionType } from '@bitgo/sdk-core';
+import BigNumber from 'bignumber.js';
+import { TransactionBuilder, Transaction } from '@bitgo/abstract-substrate';
+import { BatchArgs } from './iface';
+import { BatchUnstakingTransactionSchema } from './txnSchema';
+
+export class BatchUnstakingBuilder extends TransactionBuilder {
+  protected _amount: string;
+
+  constructor(_coinConfig: Readonly<CoinConfig>) {
+    super(_coinConfig);
+  }
+
+  /**
+   * Unbond tokens and chill (stop nominating validators)
+   *
+   * @returns {UnsignedTransaction} an unsigned Polyx transaction
+   */
+  protected buildTransaction(): UnsignedTransaction {
+    const baseTxInfo = this.createBaseTxInfo();
+
+    const chillCall = methods.staking.chill({}, baseTxInfo.baseTxInfo, baseTxInfo.options);
+
+    const unbondCall = methods.staking.unbond(
+      {
+        value: this._amount,
+      },
+      baseTxInfo.baseTxInfo,
+      baseTxInfo.options
+    );
+
+    // Create batch all transaction (atomic execution)
+    return methods.utility.batchAll(
+      {
+        calls: [chillCall.method, unbondCall.method],
+      },
+      baseTxInfo.baseTxInfo,
+      baseTxInfo.options
+    );
+  }
+
+  protected get transactionType(): TransactionType {
+    return TransactionType.Batch;
+  }
+
+  /**
+   * The amount to unstake.
+   *
+   * @param {string} amount
+   * @returns {BatchUnstakingBuilder} This unstake builder.
+   */
+  amount(amount: string): this {
+    this.validateValue(new BigNumber(amount));
+    this._amount = amount;
+    return this;
+  }
+
+  /**
+   * Get the amount to unstake
+   */
+  getAmount(): string {
+    return this._amount;
+  }
+
+  /** @inheritdoc */
+  validateDecodedTransaction(decodedTxn: DecodedSigningPayload | DecodedSignedTx): void {
+    const methodName = decodedTxn.method?.name as string;
+
+    if (methodName === 'utility.batchAll') {
+      const txMethod = decodedTxn.method.args as unknown as BatchArgs;
+      const calls = txMethod.calls;
+
+      if (calls.length !== 2) {
+        throw new InvalidTransactionError(
+          `Invalid batch unstaking transaction: expected 2 calls but got ${calls.length}`
+        );
+      }
+
+      // Check that first call is chill
+      if (calls[0].method !== 'staking.chill') {
+        throw new InvalidTransactionError(
+          `Invalid batch unstaking transaction: first call should be staking.chill but got ${calls[0].method}`
+        );
+      }
+
+      // Check that second call is unbond
+      if (calls[1].method !== 'staking.unbond') {
+        throw new InvalidTransactionError(
+          `Invalid batch unstaking transaction: second call should be staking.unbond but got ${calls[1].method}`
+        );
+      }
+
+      // Validate unbond amount
+      const unbondArgs = calls[1].args as { value: string };
+      const validationResult = BatchUnstakingTransactionSchema.validate({
+        value: unbondArgs.value,
+      });
+
+      if (validationResult.error) {
+        throw new InvalidTransactionError(`Invalid batch unstaking transaction: ${validationResult.error.message}`);
+      }
+    } else {
+      throw new InvalidTransactionError(`Invalid transaction type: ${methodName}. Expected utility.batchAll`);
+    }
+  }
+
+  /** @inheritdoc */
+  protected fromImplementation(rawTransaction: string): Transaction {
+    const tx = super.fromImplementation(rawTransaction);
+
+    if (this._method && (this._method.name as string) === 'utility.batchAll') {
+      const txMethod = this._method.args as unknown as BatchArgs;
+      const calls = txMethod.calls;
+
+      if (calls && calls.length === 2 && calls[1].method === 'staking.unbond') {
+        const unbondArgs = calls[1].args as { value: string };
+        this.amount(unbondArgs.value);
+      }
+    } else {
+      throw new InvalidTransactionError(`Invalid Transaction Type: ${this._method?.name}. Expected utility.batchAll`);
+    }
+
+    return tx;
+  }
+
+  /** @inheritdoc */
+  validateTransaction(_: Transaction): void {
+    super.validateTransaction(_);
+    this.validateFields(this._amount);
+  }
+
+  private validateFields(value: string): void {
+    const validationResult = BatchUnstakingTransactionSchema.validate({
+      value,
+    });
+
+    if (validationResult.error) {
+      throw new InvalidTransactionError(
+        `Batch Unstaking Builder Transaction validation failed: ${validationResult.error.message}`
+      );
+    }
+  }
+
+  /**
+   * Validates fields for testing
+   */
+  testValidateFields(): void {
+    this.validateFields(this._amount);
+  }
+}

modules/sdk-coin-polyx/src/lib/iface.ts

@@ -20,3 +20,14 @@ export interface TxMethod extends Omit<Interface.TxMethod, 'args'> {
 export interface DecodedTx extends Omit<DecodedUnsignedTx, 'method'> {
   method: TxMethod;
 }
+
+export interface WithdrawUnbondedArgs extends Args {
+  numSlashingSpans: number;
+}
+
+export interface BatchArgs {
+  calls: {
+    method: string;
+    args: Record<string, unknown>;
+  }[];
+}

modules/sdk-coin-polyx/src/lib/index.ts

@@ -11,4 +11,6 @@ export {
 export { TransactionBuilderFactory } from './transactionBuilderFactory';
 export { TransferBuilder } from './transferBuilder';
 export { RegisterDidWithCDDBuilder } from './registerDidWithCDDBuilder';
+export { BatchUnstakingBuilder } from './batchUnstakingBuilder';
+export { WithdrawUnbondedBuilder } from './withdrawUnbondedBuilder';
 export { Utils, default as utils } from './utils';

modules/sdk-coin-polyx/src/lib/transactionBuilderFactory.ts

@@ -3,6 +3,8 @@ import { BaseCoin as CoinConfig } from '@bitgo/statics';
 import { decode } from '@substrate/txwrapper-polkadot';
 import { TransferBuilder } from './transferBuilder';
 import { RegisterDidWithCDDBuilder } from './registerDidWithCDDBuilder';
+import { BatchUnstakingBuilder } from './batchUnstakingBuilder';
+import { WithdrawUnbondedBuilder } from './withdrawUnbondedBuilder';
 import utils from './utils';
 import { Interface, SingletonRegistry, TransactionBuilder } from './';
 import { TxMethod } from './iface';
@@ -27,6 +29,14 @@ export class TransactionBuilderFactory extends BaseTransactionBuilderFactory {
     return new RegisterDidWithCDDBuilder(this._coinConfig).material(this._material);
   }
 
+  getBatchUnstakingBuilder(): BatchUnstakingBuilder {
+    return new BatchUnstakingBuilder(this._coinConfig).material(this._material);
+  }
+
+  getWithdrawUnbondedBuilder(): WithdrawUnbondedBuilder {
+    return new WithdrawUnbondedBuilder(this._coinConfig).material(this._material);
+  }
+
   getWalletInitializationBuilder(): void {
     throw new NotImplementedError(`walletInitialization for ${this._coinConfig.name} not implemented`);
   }
@@ -54,8 +64,21 @@ export class TransactionBuilderFactory extends BaseTransactionBuilderFactory {
       return this.getTransferBuilder();
     } else if (methodName === Interface.MethodNames.RegisterDidWithCDD) {
       return this.getRegisterDidWithCDDBuilder();
-    } else {
-      throw new Error('Transaction cannot be parsed or has an unsupported transaction type');
+    } else if (methodName === 'utility.batchAll') {
+      const args = decodedTxn.method.args as { calls?: { method: string; args: Record<string, unknown> }[] };
+
+      if (
+        args.calls &&
+        args.calls.length === 2 &&
+        args.calls[0].method === 'staking.chill' &&
+        args.calls[1].method === 'staking.unbond'
+      ) {
+        return this.getBatchUnstakingBuilder();
+      }
+    } else if (methodName === 'staking.withdrawUnbonded') {
+      return this.getWithdrawUnbondedBuilder();
     }
+
+    throw new Error('Transaction cannot be parsed or has an unsupported transaction type');
   }
 }

modules/sdk-coin-polyx/src/lib/txnSchema.ts

@@ -8,3 +8,21 @@ export const RegisterDidWithCDDTransactionSchema = joi.object({
   secondaryKeys: joi.array().length(0).required(),
   expiry: joi.valid(null).required(),
 });
+
+export const BatchUnstakingTransactionSchema = {
+  validate: (value: { value: string }): joi.ValidationResult =>
+    joi
+      .object({
+        value: joi.string().required(),
+      })
+      .validate(value),
+};
+
+export const WithdrawUnbondedTransactionSchema = {
+  validate: (value: { slashingSpans: number }): joi.ValidationResult =>
+    joi
+      .object({
+        slashingSpans: joi.number().min(0).required(),
+      })
+      .validate(value),
+};

modules/sdk-coin-polyx/src/lib/withdrawUnbondedBuilder.ts

@@ -0,0 +1,109 @@
+import { BaseCoin as CoinConfig } from '@bitgo/statics';
+import { methods } from '@substrate/txwrapper-polkadot';
+import { UnsignedTransaction, DecodedSigningPayload, DecodedSignedTx } from '@substrate/txwrapper-core';
+import { InvalidTransactionError, TransactionType } from '@bitgo/sdk-core';
+import BigNumber from 'bignumber.js';
+import { TransactionBuilder, Transaction } from '@bitgo/abstract-substrate';
+import { WithdrawUnbondedTransactionSchema } from './txnSchema';
+import { WithdrawUnbondedArgs } from './iface';
+
+export class WithdrawUnbondedBuilder extends TransactionBuilder {
+  protected _slashingSpans = 0;
+
+  constructor(_coinConfig: Readonly<CoinConfig>) {
+    super(_coinConfig);
+  }
+
+  /**
+   * Withdraw unbonded tokens after the unbonding period has passed
+   *
+   * @returns {UnsignedTransaction} an unsigned Polyx transaction
+   */
+  protected buildTransaction(): UnsignedTransaction {
+    const baseTxInfo = this.createBaseTxInfo();
+
+    return methods.staking.withdrawUnbonded(
+      {
+        numSlashingSpans: this._slashingSpans,
+      },
+      baseTxInfo.baseTxInfo,
+      baseTxInfo.options
+    );
+  }
+
+  protected get transactionType(): TransactionType {
+    return TransactionType.StakingWithdraw;
+  }
+
+  /**
+   * The number of slashing spans, typically 0 for most users
+   *
+   * @param {number} slashingSpans
+   * @returns {WithdrawUnbondedBuilder} This withdrawUnbonded builder.
+   */
+  slashingSpans(slashingSpans: number): this {
+    this.validateValue(new BigNumber(slashingSpans));
+    this._slashingSpans = slashingSpans;
+    return this;
+  }
+
+  /**
+   * Get the slashing spans
+   */
+  getSlashingSpans(): number {
+    return this._slashingSpans;
+  }
+
+  /** @inheritdoc */
+  validateDecodedTransaction(decodedTxn: DecodedSigningPayload | DecodedSignedTx): void {
+    if (decodedTxn.method?.name === 'staking.withdrawUnbonded') {
+      const txMethod = decodedTxn.method.args as unknown as WithdrawUnbondedArgs;
+      const slashingSpans = txMethod.numSlashingSpans;
+      const validationResult = WithdrawUnbondedTransactionSchema.validate({ slashingSpans });
+
+      if (validationResult.error) {
+        throw new InvalidTransactionError(
+          `WithdrawUnbonded Transaction validation failed: ${validationResult.error.message}`
+        );
+      }
+    } else {
+      throw new InvalidTransactionError(
+        `Invalid transaction type: ${decodedTxn.method?.name}. Expected staking.withdrawUnbonded`
+      );
+    }
+  }
+
+  /** @inheritdoc */
+  protected fromImplementation(rawTransaction: string): Transaction {
+    const tx = super.fromImplementation(rawTransaction);
+
+    if (this._method && (this._method.name as string) === 'staking.withdrawUnbonded') {
+      const txMethod = this._method.args as unknown as WithdrawUnbondedArgs;
+      this.slashingSpans(txMethod.numSlashingSpans);
+    } else {
+      throw new InvalidTransactionError(
+        `Invalid Transaction Type: ${this._method?.name}. Expected staking.withdrawUnbonded`
+      );
+    }
+
+    return tx;
+  }
+
+  /** @inheritdoc */
+  validateTransaction(_: Transaction): void {
+    super.validateTransaction(_);
+    this.validateFields(this._slashingSpans);
+  }
+
+  private validateFields(slashingSpans: number): void {
+    const validationResult = WithdrawUnbondedTransactionSchema.validate({
+      slashingSpans,
+    });
+
+    if (validationResult.error) {
+      throw new InvalidTransactionError(
+        `WithdrawUnbonded Builder Transaction validation failed: ${validationResult.error.message}`
+      );
+    }
+  }
+}

modules/sdk-coin-polyx/test/resources/index.ts

@@ -50,6 +50,16 @@ export const rawTx = {
     unsigned:
       '0x90071460b685d82b315b70d7c7604f990a05395eab09d5e75bae5d2c519ca1b01e25e500004503040090d76a00070000002ace05e703aa50b48c0ccccfc8b424f7aab9a1e2c424ed12e45d20b1e8ffd0d6cbd4f0bb74e13c8c4da973b1a15c3df61ae3b82677b024ffa60faf7799d5ed4b',
   },
+  unstake: {
+    signed:
+      '0xcd018400bec110eab4d327d3b2b6bb68e888654a474694d3935ce35bd3926e4bc7ebd538011a740e63a85858c9fa99ba381ce3b9c12db872c0d976948df9d5206f35642c78a8c25a2f927b569a163985dcb7e27e63fe2faa926371e79a070703095607b787d502180029020811061102034353c5b3',
+    unsigned: '0x340429020811061102034353c5b3',
+  },
+  withdrawUnbonded: {
+    signed:
+      '0xb5018400bec110eab4d327d3b2b6bb68e888654a474694d3935ce35bd3926e4bc7ebd53801a67640e1f61a3881a6fa3d093e09149f00a75747f47facb497689c6bb2f71d49b91ebebe12ccc2febba86b6af869c979053b811f33ea8aba48938aff48b56488a5012000110300000000',
+    unsigned: '0x1c04110300000000',
+  },
 };
 
 export const { txVersion, specVersion, genesisHash, chainName, specName } = Networks.test.polyx;