Adding unfreeze and withdraw for tron unstaking

modules/sdk-coin-trx/src/lib/enum.ts

@@ -14,6 +14,14 @@ export enum ContractType {
    * This is a smart contract type.
    */
   TriggerSmartContract,
+  /**
+   * This is the contract for unfreezing balances
+   */
+  UnfreezeBalanceV2,
+  /**
+   * This is the contract for withdrawing expired unfrozen balances
+   */
+  WithdrawExpireUnfreeze,
 }
 
 export enum PermissionType {

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

@@ -42,7 +42,12 @@ export interface RawData {
   ref_block_hash: string;
   fee_limit?: number;
   contractType?: ContractType;
-  contract: TransferContract[] | AccountPermissionUpdateContract[] | TriggerSmartContract[];
+  contract:
+    | TransferContract[]
+    | AccountPermissionUpdateContract[]
+    | TriggerSmartContract[]
+    | UnfreezeBalanceV2Contract[]
+    | WithdrawExpireUnfreezeContract[];
 }
 
 export interface Value {
@@ -117,3 +122,51 @@ export interface AccountInfo {
   active_permission: [{ keys: [PermissionKey] }];
   trc20: [Record<string, string>];
 }
+
+/**
+ * Unfreeze transaction value fields
+ */
+export interface UnfreezeBalanceValueFields {
+  resource: string;
+  unfreeze_balance: number;
+  owner_address: string;
+}
+
+/**
+ * Unfreeze balance contract value interface
+ */
+export interface UnfreezeBalanceValue {
+  type_url?: string;
+  value: UnfreezeBalanceValueFields;
+}
+
+/**
+ * Unfreeze balance v2 contract interface
+ */
+export interface UnfreezeBalanceV2Contract {
+  parameter: UnfreezeBalanceValue;
+  type?: string;
+}
+
+/**
+ * Withdraw transaction value fields
+ */
+export interface WithdrawBalanceValueFields {
+  owner_address: string;
+}
+
+/**
+ * Withdraw balance contract value interface
+ */
+export interface WithdrawBalanceValue {
+  type_url?: string;
+  value: WithdrawBalanceValueFields;
+}
+
+/**
+ * Withdraw expire unfreeze contract interface
+ */
+export interface WithdrawExpireUnfreezeContract {
+  parameter: WithdrawBalanceValue;
+  type?: string;
+}

modules/sdk-coin-trx/src/lib/resourceTypes.ts

@@ -0,0 +1,7 @@
+/**
+ * Valid resource types for Tron freezing and unfreezing
+ */
+export enum TronResource {
+  BANDWIDTH = 'BANDWIDTH',
+  ENERGY = 'ENERGY',
+}

modules/sdk-coin-trx/src/lib/transaction.ts

@@ -17,7 +17,15 @@ import {
   tokenMainnetContractAddresses,
   tokenTestnetContractAddresses,
 } from './utils';
-import { ContractEntry, RawData, TransactionReceipt, TransferContract, TriggerSmartContract } from './iface';
+import {
+  ContractEntry,
+  RawData,
+  TransactionReceipt,
+  TransferContract,
+  TriggerSmartContract,
+  UnfreezeBalanceV2Contract,
+  WithdrawExpireUnfreezeContract,
+} from './iface';
 
 /**
  * Tron transaction model.
@@ -126,6 +134,30 @@ export class Transaction extends BaseTransaction {
           value: '0',
         };
         break;
+      case ContractType.UnfreezeBalanceV2:
+        this._type = TransactionType.StakingUnlock;
+        const unfreezeValues = (rawData.contract[0] as UnfreezeBalanceV2Contract).parameter.value;
+        output = {
+          address: unfreezeValues.owner_address,
+          value: unfreezeValues.unfreeze_balance.toString(),
+        };
+        input = {
+          address: unfreezeValues.owner_address,
+          value: unfreezeValues.unfreeze_balance.toString(),
+        };
+        break;
+      case ContractType.WithdrawExpireUnfreeze:
+        this._type = TransactionType.StakingWithdraw;
+        const withdrawValues = (rawData.contract[0] as WithdrawExpireUnfreezeContract).parameter.value;
+        output = {
+          address: withdrawValues.owner_address,
+          value: '0', // no value field
+        };
+        input = {
+          address: withdrawValues.owner_address,
+          value: '0',
+        };
+        break;
       default:
         throw new ParseTransactionError('Unsupported contract type');
     }

modules/sdk-coin-trx/src/lib/unfreezeBuilder.ts

@@ -0,0 +1,94 @@
+import { InvalidTransactionError, TransactionType } from '@bitgo/sdk-core';
+import { TransactionBuilder } from './transactionBuilder';
+import { TronResource } from './resourceTypes';
+
+export class UnfreezeBuilder extends TransactionBuilder {
+  /** @inheritdoc */
+  protected get transactionType(): TransactionType {
+    return TransactionType.StakingUnlock;
+  }
+
+  /** Override to initialize this builder from a raw transaction */
+  initBuilder(rawTransaction: string | any): void {
+    this.transaction = this.fromImplementation(rawTransaction);
+    // Explicitly set the transaction type after initialization
+    this.transaction.setTransactionType(this.transactionType);
+  }
+
+  validateTransaction(transaction: any): void {
+    if (transaction && typeof transaction.toJson === 'function') {
+      super.validateTransaction(transaction);
+      // Get the raw transaction data from the Transaction object
+      const rawTx = transaction.toJson();
+      this.validateUnfreezeTransaction(rawTx);
+    } else {
+      // If it's already a raw transaction object, validate it directly
+      this.validateUnfreezeTransaction(transaction);
+    }
+  }
+
+  /**
+   * Validates if the transaction is a valid unfreeze transaction
+   * @param transaction The transaction to validate
+   * @throws {InvalidTransactionError} when the transaction is invalid
+   */
+  private validateUnfreezeTransaction(transaction: any): void {
+    if (
+      !transaction ||
+      !transaction.raw_data ||
+      !transaction.raw_data.contract ||
+      transaction.raw_data.contract.length === 0
+    ) {
+      throw new InvalidTransactionError('Invalid transaction: missing or empty contract array');
+    }
+
+    const contract = transaction.raw_data.contract[0];
+
+    // Validate contract type
+    if (contract.type !== 'UnfreezeBalanceV2Contract') {
+      throw new InvalidTransactionError(
+        `Invalid unfreeze transaction: expected contract type UnfreezeBalanceV2Contract but got ${contract.type}`
+      );
+    }
+
+    // Validate parameter value
+    if (!contract.parameter || !contract.parameter.value) {
+      throw new InvalidTransactionError('Invalid unfreeze transaction: missing parameter value');
+    }
+
+    const value = contract.parameter.value;
+
+    // Validate resource
+    if (!Object.values(TronResource).includes(value.resource)) {
+      throw new InvalidTransactionError(
+        `Invalid unfreeze transaction: resource must be ${Object.values(TronResource).join(' or ')}, got ${
+          value.resource
+        }`
+      );
+    }
+
+    // Validate unfreeze_balance
+    if (!value.unfreeze_balance || value.unfreeze_balance <= 0) {
+      throw new InvalidTransactionError('Invalid unfreeze transaction: unfreeze_balance must be positive');
+    }
+
+    // Validate owner_address
+    if (!value.owner_address || typeof value.owner_address !== 'string' || value.owner_address.length === 0) {
+      throw new InvalidTransactionError('Invalid unfreeze transaction: missing or invalid owner_address');
+    }
+  }
+
+  /**
+   * Check if the transaction is a valid unfreeze transaction
+   * @param transaction Transaction to check
+   * @returns True if the transaction is a valid unfreeze transaction
+   */
+  canSign(transaction: any): boolean {
+    try {
+      this.validateUnfreezeTransaction(transaction);
+      return true;
+    } catch (e) {
+      return false;
+    }
+  }
+}

modules/sdk-coin-trx/src/lib/utils.ts

@@ -90,21 +90,21 @@ export function getHexAddressFromBase58Address(base58: string): string {
   // pulled from: https://github.com/TRON-US/tronweb/blob/dcb8efa36a5ebb65c4dab3626e90256a453f3b0d/src/utils/help.js#L17
   // but they don't surface this call in index.js
   const bytes = tronweb.utils.crypto.decodeBase58Address(base58);
-  return getHexAddressFromByteArray(bytes);
+  return getHexAddressFromByteArray(bytes as any);
 }
 
 /**
  * @param privateKey
  */
 export function getPubKeyFromPriKey(privateKey: TronBinaryLike): ByteArray {
-  return tronweb.utils.crypto.getPubKeyFromPriKey(privateKey);
+  return tronweb.utils.crypto.getPubKeyFromPriKey(privateKey as any);
 }
 
 /**
  * @param privateKey
  */
 export function getAddressFromPriKey(privateKey: TronBinaryLike): ByteArray {
-  return tronweb.utils.crypto.getAddressFromPriKey(privateKey);
+  return tronweb.utils.crypto.getAddressFromPriKey(privateKey as any);
 }
 
 /**
@@ -127,7 +127,7 @@ export function getBase58AddressFromHex(hex: string): string {
  * @param transaction
  */
 export function signTransaction(privateKey: string | ByteArray, transaction: TransactionReceipt): TransactionReceipt {
-  return tronweb.utils.crypto.signTransaction(privateKey, transaction);
+  return tronweb.utils.crypto.signTransaction(privateKey, transaction) as unknown as TransactionReceipt;
 }
 
 /**
@@ -136,14 +136,14 @@ export function signTransaction(privateKey: string | ByteArray, transaction: Tra
  * @param useTronHeader
  */
 export function signString(message: string, privateKey: string | ByteArray, useTronHeader = true): string {
-  return tronweb.Trx.signString(message, privateKey, useTronHeader);
+  return tronweb.Trx.signString(message, privateKey as any, useTronHeader);
 }
 
 /**
  * @param pubBytes
  */
 export function getRawAddressFromPubKey(pubBytes: TronBinaryLike): ByteArray {
-  return tronweb.utils.crypto.computeAddress(pubBytes);
+  return tronweb.utils.crypto.computeAddress(pubBytes as any);
 }
 
 /**
@@ -175,6 +175,14 @@ export function decodeTransaction(hexString: string): RawData {
       contractType = ContractType.TriggerSmartContract;
       contract = exports.decodeTriggerSmartContract(rawTransaction.contracts[0].parameter.value);
       break;
+    case 'type.googleapis.com/protocol.WithdrawExpireUnfreezeContract':
+      contract = decodeWithdrawExpireUnfreezeContract(rawTransaction.contracts[0].parameter.value);
+      contractType = ContractType.WithdrawExpireUnfreeze;
+      break;
+    case 'type.googleapis.com/protocol.UnfreezeBalanceV2Contract':
+      contract = decodeUnfreezeBalanceV2Contract(rawTransaction.contracts[0].parameter.value);
+      contractType = ContractType.UnfreezeBalanceV2;
+      break;
     default:
       throw new UtilsError('Unsupported contract type');
   }
@@ -360,6 +368,97 @@ export function decodeAccountPermissionUpdateContract(base64: string): AccountPe
   };
 }
 
+/**
+ * Deserialize the segment of the txHex corresponding with unfreeze balance contract
+ *
+ * @param {string} base64 - The base64 encoded contract data
+ * @returns {Array} - Array containing the decoded unfreeze contract
+ */
+export function decodeUnfreezeBalanceV2Contract(base64: string): any[] {
+  let unfreezeContract;
+  try {
+    unfreezeContract = protocol.UnfreezeBalanceContract.decode(Buffer.from(base64, 'base64')).toJSON();
+  } catch (e) {
+    throw new UtilsError('There was an error decoding the unfreeze contract in the transaction.');
+  }
+
+  if (!unfreezeContract.ownerAddress) {
+    throw new UtilsError('Owner address does not exist in this unfreeze contract.');
+  }
+
+  if (!unfreezeContract.resource) {
+    throw new UtilsError('Resource type does not exist in this unfreeze contract.');
+  }
+
+  if (!unfreezeContract.hasOwnProperty('unfrozenBalance')) {
+    throw new UtilsError('Unfreeze balance does not exist in this unfreeze contract.');
+  }
+
+  // deserialize attributes
+  const owner_address = getBase58AddressFromByteArray(
+    getByteArrayFromHexAddress(Buffer.from(unfreezeContract.ownerAddress, 'base64').toString('hex'))
+  );
+
+  // Convert ResourceCode enum value to string resource name
+  const resourceValue = unfreezeContract.resource;
+  let resource: string;
+  if (resourceValue === protocol.ResourceCode.BANDWIDTH) {
+    resource = 'BANDWIDTH';
+  } else if (resourceValue === protocol.ResourceCode.ENERGY) {
+    resource = 'ENERGY';
+  } else {
+    throw new UtilsError(`Unknown resource type: ${resourceValue}`);
+  }
+
+  const unfreeze_balance = unfreezeContract.unfrozenBalance;
+
+  return [
+    {
+      parameter: {
+        value: {
+          resource,
+          unfreeze_balance: Number(unfreeze_balance),
+          owner_address,
+        },
+      },
+    },
+  ];
+}
+
+/**
+ * Deserialize the segment of the txHex corresponding with withdraw expire unfreeze contract
+ *
+ * @param {string} base64 - The base64 encoded contract data
+ * @returns {Array} - Array containing the decoded withdraw contract
+ */
+export function decodeWithdrawExpireUnfreezeContract(base64: string): any[] {
+  let withdrawContract;
+  try {
+    withdrawContract = protocol.WithdrawBalanceContract.decode(Buffer.from(base64, 'base64')).toJSON();
+  } catch (e) {
+    throw new UtilsError('There was an error decoding the withdraw contract in the transaction.');
+  }
+
+  if (!withdrawContract.ownerAddress) {
+    throw new UtilsError('Owner address does not exist in this withdraw contract.');
+  }
+
+  // deserialize attributes
+  const owner_address = getBase58AddressFromByteArray(
+    getByteArrayFromHexAddress(Buffer.from(withdrawContract.ownerAddress, 'base64').toString('hex'))
+  );
+
+  return [
+    {
+      parameter: {
+        value: {
+          owner_address,
+        },
+      },
+    },
+  ];
+}
+
 /**
  * @param raw
  */