Charge a Fee on TRON Transfers

const FEE_BPS = 50n; // 0.5%

const percentageFee = (amount: string) => ((BigInt(amount) * FEE_BPS) / 10_000n).toString();

import { TronWeb } from 'tronweb';

const TOKEN_CONTRACT = process.env.TOKEN_CONTRACT!; // the token users send, for example mainnet USDT
const FEE_WALLET = process.env.FEE_WALLET!;

type ExpectedTransfer = { owner: string; to: string; minAmount: string };

const verifyTransfer = async (txid: string, expected: ExpectedTransfer) => {
  const [tx, info] = await Promise.all([
    tronPost('/wallet/gettransactionbyid', { value: txid, visible: true }),
    waitForConfirmation(txid),
  ]);

  const contract = tx.raw_data?.contract?.[0];
  const value = contract?.parameter?.value;
  const data: string = value?.data ?? '';
  const checks = {
    type: contract?.type === 'TriggerSmartContract',
    owner: value?.owner_address === expected.owner,
    token: value?.contract_address === TOKEN_CONTRACT,
    selector: data.startsWith('a9059cbb'),
    recipient: data.length >= 136 && TronWeb.address.fromHex('41' + data.slice(32, 72)) === expected.to,
    amount: data.length >= 136 && BigInt('0x' + data.slice(72, 136)) >= BigInt(expected.minAmount),
  };
  const failed = Object.entries(checks).filter(([, ok]) => !ok).map(([name]) => name);
  return {
    failed,
    succeeded: info.receipt?.result === 'SUCCESS',
    energyUsed: info.receipt?.energy_usage ?? 0,
    confirmedAt: info.blockTimeStamp as number,
  };
};

// Use a database in production, so every server instance sees the same transaction IDs.
const usedTxids = new Set<string>();

export const verifyFeePaid = async (
  userAddress: string,
  transfer: { txid: string; recipient: string; amount: string },
  fee: { txid: string; amount: string },
) => {
  if (usedTxids.has(transfer.txid) || usedTxids.has(fee.txid)) {
    throw new Error('A new transaction ID is required');
  }
  const checked = [
    await verifyTransfer(transfer.txid, { owner: userAddress, to: transfer.recipient, minAmount: transfer.amount }),
    await verifyTransfer(fee.txid, { owner: userAddress, to: FEE_WALLET, minAmount: fee.amount }),
  ];
  if (checked.some((result) => result.failed.length > 0 || !result.succeeded)) {
    throw new Error('The transfer or the fee failed verification');
  }
  usedTxids.add(transfer.txid);
  usedTxids.add(fee.txid);
};

const FEE_MULTIPLIER = Number(process.env.FEE_MULTIPLIER ?? '0.9'); // share of the TRX burn you charge
const FIXED_FEE_BASE_UNITS = process.env.FIXED_FEE_BASE_UNITS;
const PENDING_FEE_TTL_MS = Number(process.env.PENDING_FEE_TTL_MS ?? 10 * 60 * 1000);
const BANDWIDTH_PER_TRANSFER = 400; // a TRC-20 transfer uses about 345, plus headroom

const encodeTransfer = (to: string, amount: string) =>
  TronWeb.address.toHex(to).replace(/^41/, '').padStart(64, '0') +
  BigInt(amount).toString(16).padStart(64, '0');

const estimateEnergy = async (owner: string, to: string, amount: string): Promise<number> => {
  const response = await tronPost('/wallet/triggerconstantcontract', {
    owner_address: owner,
    contract_address: TOKEN_CONTRACT,
    function_selector: 'transfer(address,uint256)',
    parameter: encodeTransfer(to, amount),
    visible: true,
  });
  if (!response.result?.result) {
    throw new Error(`Energy estimate failed: ${JSON.stringify(response.result)}`);
  }
  return response.energy_used;
};

// Replace with your own price source. Returns USDT per 1 TRX, for example 0.25.
const getTrxUsdtPrice = async (): Promise<number> => {
  throw new Error('Implement getTrxUsdtPrice');
};

const getChainParameter = async (key: string): Promise<number> => {
  const { chainParameter } = await tronPost('/wallet/getchainparameters', {});
  const entry = chainParameter.find((p: { key: string }) => p.key === key);
  if (!entry) throw new Error(`Chain parameter ${key} not found`);
  return entry.value;
};

const priceGasFee = async (userAddress: string, recipient: string, amount: string) => {
  const energyMain = await estimateEnergy(userAddress, recipient, amount);
  const energyFeeLeg = await estimateEnergy(userAddress, FEE_WALLET, '1');
  const energy = Math.ceil((energyMain + energyFeeLeg) * 1.1);
  const bandwidth = 2 * BANDWIDTH_PER_TRANSFER;

  if (FIXED_FEE_BASE_UNITS) {
    return { energy, bandwidth, feeAmount: FIXED_FEE_BASE_UNITS };
  }

  const burnSun =
    energy * (await getChainParameter('getEnergyFee')) + bandwidth * (await getChainParameter('getTransactionFee'));
  // SUN × (USDT per TRX) = USDT base units, because both have 6 decimals
  const feeAmount = Math.ceil(burnSun * FEE_MULTIPLIER * (await getTrxUsdtPrice())).toString();
  return { energy, bandwidth, feeAmount };
};

const wholeTransfers = (bandwidth: number) =>
  Math.floor(Math.max(bandwidth, 0) / BANDWIDTH_PER_TRANSFER) * BANDWIDTH_PER_TRANSFER;

const delegateEnergy = async (energyNeeded: number, userAddress: string) => {
  const user = await tronPost('/wallet/getaccountresource', { address: userAddress, visible: true });
  const deficit = energyNeeded - ((user.EnergyLimit ?? 0) - (user.EnergyUsed ?? 0));
  if (deficit <= 0) return;

  const { TotalEnergyLimit, TotalEnergyWeight } = await tronPost('/wallet/getaccountresource', {
    address: sponsorAddress,
    visible: true,
  });
  const sun = Math.ceil((deficit * TotalEnergyWeight) / TotalEnergyLimit) * 1_000_000;

  const { max_size: maxSun = 0 } = await tronPost('/wallet/getcandelegatedmaxsize', {
    owner_address: sponsorAddress,
    type: 1, // ENERGY
    visible: true,
  });
  if (maxSun < sun) {
    throw new Error('Not enough free stake: reclaim from other users or stake more TRX');
  }

  const delegateTransaction = await tronPost('/wallet/delegateresource', {
    owner_address: sponsorAddress,
    receiver_address: userAddress,
    balance: sun,
    resource: 'ENERGY',
    lock: false,
    visible: true,
  });
  const { txid } = await signAndBroadcast(delegateTransaction);
  await waitForConfirmation(txid);
};

const topUpBandwidth = async (bandwidthNeeded: number, userAddress: string) => {
  const user = await tronPost('/wallet/getaccountresource', { address: userAddress, visible: true });
  const covered =
    wholeTransfers((user.freeNetLimit ?? 0) - (user.freeNetUsed ?? 0)) +
    wholeTransfers((user.NetLimit ?? 0) - (user.NetUsed ?? 0));
  const deficit = bandwidthNeeded - covered;
  if (deficit <= 0) return;

  const { balance = 0 } = await tronPost('/wallet/getaccount', { address: userAddress, visible: true });
  const sun = deficit * (await getChainParameter('getTransactionFee')) - balance;
  if (sun <= 0) return;

  const transferTransaction = await tronPost('/wallet/createtransaction', {
    owner_address: sponsorAddress,
    to_address: userAddress,
    amount: sun,
    visible: true,
  });
  const { txid } = await signAndBroadcast(transferTransaction);
  await waitForConfirmation(txid);
};

// The recipe's step 6, waiting for confirmation so the stake is free again before the next delegation.
const reclaimGas = async (userAddress: string) => {
  const { delegatedResource = [] } = await tronPost('/wallet/getdelegatedresourcev2', {
    fromAddress: sponsorAddress,
    toAddress: userAddress,
    visible: true,
  });
  const delegatedSun = delegatedResource.reduce(
    (sum: number, resource: { frozen_balance_for_energy?: number }) => sum + (resource.frozen_balance_for_energy ?? 0),
    0,
  );
  if (delegatedSun === 0) return;

  const undelegateTransaction = await tronPost('/wallet/undelegateresource', {
    owner_address: sponsorAddress,
    receiver_address: userAddress,
    balance: delegatedSun,
    resource: 'ENERGY',
    visible: true,
  });
  const { txid } = await signAndBroadcast(undelegateTransaction);
  await waitForConfirmation(txid);
};

const getEnergyUsed = async (userAddress: string): Promise<number> => {
  const user = await tronPost('/wallet/getaccountresource', { address: userAddress, visible: true });
  return user.EnergyUsed ?? 0;
};

type PendingFee = {
  feeAmount: string;
  recipient: string;
  amount: string;
  energyUsedAtSponsor: number;
  gasUsed: boolean;
  createdAt: number;
  feeTxid?: string; // set once the fee is verified on-chain
};
// Use a database in production. Claim a wallet's entry with one atomic write, such as an
// insert under a unique key on userAddress, so concurrent requests can't both pass the check.
const pendingFees = new Map<string, PendingFee>();

export const getGasQuote = async (userAddress: string, recipient: string, amount: string) => {
  const { feeAmount } = await priceGasFee(userAddress, recipient, amount);
  return { feeAmount, feeRecipient: FEE_WALLET };
};

export const sponsorWithFee = async (userAddress: string, recipient: string, amount: string, maxFeeAmount: string) => {
  if (pendingFees.has(userAddress)) {
    throw new Error('A previous gas fee is unsettled');
  }
  const pending: PendingFee = {
    feeAmount: maxFeeAmount,
    recipient,
    amount,
    energyUsedAtSponsor: 0,
    gasUsed: false,
    createdAt: Date.now(),
  };
  pendingFees.set(userAddress, pending);

  try {
    const { energy, bandwidth, feeAmount } = await priceGasFee(userAddress, recipient, amount);
    if (BigInt(feeAmount) > BigInt(maxFeeAmount)) {
      throw new Error('The gas fee is higher than the quote. Request a new quote.');
    }
    pending.feeAmount = feeAmount;
    pending.energyUsedAtSponsor = await getEnergyUsed(userAddress);

    await delegateEnergy(energy, userAddress);
    await topUpBandwidth(bandwidth, userAddress);
    return { feeAmount, feeRecipient: FEE_WALLET };
  } catch (error) {
    await reclaimGas(userAddress);
    pendingFees.delete(userAddress);
    throw error;
  }
};

// A failed transfer clears the fee only if it is the only sponsored Energy the user spent.
const sponsoredGasPaidForAnother = async (userAddress: string, pending: PendingFee, energyUsed: number) =>
  (await getEnergyUsed(userAddress)) - pending.energyUsedAtSponsor > energyUsed;

export const payFee = async (userAddress: string, feeTxid: string) => {
  const pending = pendingFees.get(userAddress);
  if (!pending) throw new Error('No unsettled fee');
  if (pending.feeTxid) throw new Error('The fee is already paid');
  if (usedTxids.has(feeTxid)) throw new Error('A new fee transaction ID is required');

  const fee = await verifyTransfer(feeTxid, { owner: userAddress, to: FEE_WALLET, minAmount: pending.feeAmount });
  if (fee.confirmedAt < pending.createdAt) fee.failed.push('sentAfterSponsorship');
  if (fee.failed.length > 0) {
    throw new Error(`Fee transaction failed checks: ${fee.failed.join(', ')}`);
  }
  if (!fee.succeeded) {
    if (await sponsoredGasPaidForAnother(userAddress, pending, fee.energyUsed)) {
      throw new Error('The sponsored gas paid for another transaction. Pay the fee instead.');
    }
    await reclaimGas(userAddress);
    pendingFees.delete(userAddress);
    return { status: 'fee_transfer_failed' as const };
  }

  usedTxids.add(feeTxid);
  pending.feeTxid = feeTxid;
  return { status: 'fee_paid' as const };
};

// The fee is paid, so the transfer's result no longer decides whether it's owed.
export const settleFee = async (userAddress: string, mainTxid: string) => {
  const pending = pendingFees.get(userAddress);
  if (!pending) throw new Error('No unsettled fee');
  if (!pending.feeTxid) throw new Error('Pay the fee first');
  if (usedTxids.has(mainTxid)) throw new Error('A new main transaction ID is required');

  const main = await verifyTransfer(mainTxid, { owner: userAddress, to: pending.recipient, minAmount: pending.amount });
  if (main.confirmedAt < pending.createdAt) main.failed.push('sentAfterSponsorship');
  if (main.failed.length > 0) {
    throw new Error(`Main transaction failed checks: ${main.failed.join(', ')}`);
  }

  usedTxids.add(mainTxid);
  await reclaimGas(userAddress);
  pendingFees.delete(userAddress);
  return { status: main.succeeded ? ('paid' as const) : ('main_transfer_failed' as const) };
};

// Reclaiming moves the Energy the user spent back to the sponsor, so record the answer first.
const sponsoredGasWasUsed = async (userAddress: string, pending: PendingFee) => {
  if (!pending.gasUsed) {
    pending.gasUsed = (await getEnergyUsed(userAddress)) > pending.energyUsedAtSponsor;
  }
  return pending.gasUsed;
};

export const cancelSponsorship = async (userAddress: string) => {
  const pending = pendingFees.get(userAddress);
  if (!pending) return;
  // After the fee, cancelling only ends the sponsorship.
  if (!pending.feeTxid && (await sponsoredGasWasUsed(userAddress, pending))) {
    throw new Error('The sponsored gas was already used. Pay the fee instead.');
  }
  await reclaimGas(userAddress);
  pendingFees.delete(userAddress);
};

export const expireStaleSponsorships = async () => {
  for (const [userAddress, pending] of pendingFees) {
    if (Date.now() - pending.createdAt < PENDING_FEE_TTL_MS) continue;
    if (pending.feeTxid) {
      await reclaimGas(userAddress);
      pendingFees.delete(userAddress);
      continue;
    }
    if (pending.gasUsed) continue;
    const used = await sponsoredGasWasUsed(userAddress, pending);
    await reclaimGas(userAddress);
    if (!used) pendingFees.delete(userAddress);
  }
};

import { getDefaultClient, getPrimaryWalletAccount } from '@dynamic-labs-sdk/client';
import { getTronWeb, isTronWalletAccount } from '@dynamic-labs-sdk/tron';

type TronWebInstance = Awaited<ReturnType<typeof getTronWeb>>;

const postJson = async (path: string, body: Record<string, unknown>) => {
  const response = await fetch(path, {
    method: 'POST',
    headers: {
      'Content-Type': 'application/json',
      Authorization: `Bearer ${getDefaultClient().token}`,
    },
    body: JSON.stringify(body),
  });
  if (!response.ok) {
    throw new Error(`${path} failed: ${await response.text()}`);
  }
  return response.json();
};

const sendUsdt = async (tronWeb: TronWebInstance, from: string, to: string, amount: string) => {
  const { transaction } = await tronWeb.transactionBuilder.triggerSmartContract(
    USDT_CONTRACT,
    'transfer(address,uint256)',
    { feeLimit: 100_000_000 },
    [
      { type: 'address', value: to },
      { type: 'uint256', value: amount },
    ],
    from,
  );
  const signedTransaction = await tronWeb.trx.sign(transaction);
  const result = await tronWeb.trx.sendRawTransaction(signedTransaction);
  if (!result.result) {
    throw new Error(`Broadcast failed: ${JSON.stringify(result)}`);
  }
  return result.txid as string;
};

export const getGasQuote = async (recipient: string, amount: string) => {
  const walletAccount = getPrimaryWalletAccount();
  if (!walletAccount || !isTronWalletAccount(walletAccount)) {
    throw new Error('No TRON wallet');
  }
  return postJson('/api/gas-quote', { userAddress: walletAccount.address, recipient, amount }) as Promise<{
    feeAmount: string;
    feeRecipient: string;
  }>;
};

export const sendGaslessTransferWithFee = async (recipient: string, amount: string, quotedFeeAmount: string) => {
  const walletAccount = getPrimaryWalletAccount();
  if (!walletAccount || !isTronWalletAccount(walletAccount)) {
    return;
  }
  const userAddress = walletAccount.address;

  const { feeAmount, feeRecipient } = await postJson('/api/sponsor', {
    userAddress,
    recipient,
    amount,
    maxFeeAmount: quotedFeeAmount,
  });

  const tronWeb = await getTronWeb({ walletAccount });

  let feeTxid: string;
  try {
    feeTxid = await sendUsdt(tronWeb, userAddress, feeRecipient, feeAmount);
  } catch (error) {
    await postJson('/api/sponsor/cancel', { userAddress });
    throw error;
  }
  const paid = await postJson('/api/sponsor/fee', { userAddress, feeTxid });
  if (paid.status !== 'fee_paid') return paid;

  let mainTxid: string;
  try {
    mainTxid = await sendUsdt(tronWeb, userAddress, recipient, amount);
  } catch (error) {
    // The fee is paid, so cancelling only ends the sponsorship.
    await postJson('/api/sponsor/cancel', { userAddress });
    throw error;
  }
  return postJson('/api/sponsor/settle', { userAddress, mainTxid });
};