A function to extract calldata suffix from deeply nested EVM calls

trace.ts

import { createMemoryClient } from "tevm";
import type { Message } from "tevm/actions";
import { createAddress } from "tevm/address";
import {
  type Abi,
  type Address,
  bytesToHex,
  type Chain,
  decodeEventLog,
  decodeFunctionData,
  encodeFunctionData,
  getAddress,
  type Hex,
  hexToBigInt,
  hexToBytes,
  hexToString,
  isAddress,
  type PublicClient,
  type Transport,
} from "viem";

export type CallFilterResult = {
  to: Address;
  data: Hex;
  value: string;
  dataSuffix?: Hex;
  decodedDataSuffix?: string;
  logs?: {
    address: Address;
    topics: Hex[];
    data: Hex;
    decodedLog: {
      eventName?: string;
      args?: readonly unknown[];
    };
  }[];
};

export async function filterCalls({
  hash,
  publicClient,
  filter: _filter,
}: {
  hash: Hex;
  publicClient: PublicClient<Transport, Chain>;
  filter: Record<string, { abi: Abi }>;
}) {
  const filter = Object.entries(_filter).reduce((acc, [key, value]) => {
    acc[getAddress(key.toLowerCase())] = value;
    return acc;
  }, {} as Record<Address, { abi: Abi }>);

  const tx = await publicClient.getTransaction({
    hash,
  });

  const block = await publicClient.getBlock({
    blockNumber: tx.blockNumber,
  });

  const client = createMemoryClient({
    fork: {
      transport: publicClient,
      // Use previous block state to ensure we don't duplicate the same call
      blockTag: tx.blockNumber - 1n,
    },
    customPrecompiles: [
      // P256Verify Mock
      {
        address: createAddress("0x0000000000000000000000000000000000000100"),
        function: (_) => {
          return {
            executionGasUsed: hexToBigInt("0xd7a"),
            returnValue: hexToBytes(
              "0x0000000000000000000000000000000000000000000000000000000000000001"
            ),
          };
        },
      },
    ],
  });

  const matchedCalls: CallFilterResult[] = [];
  let currentMessage: Message;

  await client.tevmCall({
    to: tx.to!,
    data: tx.input,
    from: tx.from!,
    gas: tx.gas,
    gasPrice: tx.gasPrice,
    throwOnFail: false,
    blockOverrideSet: {
      number: tx.blockNumber,
      time: block.timestamp,
      baseFee: block.baseFeePerGas!,
    },
    onBeforeMessage: (message, next) => {
      currentMessage = message;
      next?.();
    },
    onAfterMessage: (result, next) => {
      const to = currentMessage?.to?.toString();

      if (!to || !isAddress(to)) return;
      const matchedCall = filter[getAddress(to)];

      if (matchedCall) {
        const calldata = bytesToHex(currentMessage.data);

        const filterResult: CallFilterResult = {
          to: currentMessage.to!.toString(),
          data: bytesToHex(currentMessage.data),
          value: currentMessage.value.toString(),
        };

        try {
          const dataSuffix = decodeDataSuffix({
            data: calldata,
            abi: matchedCall.abi,
          });

          filterResult.dataSuffix = dataSuffix;
          filterResult.decodedDataSuffix = hexToString(dataSuffix);
          matchedCalls.push(filterResult);

          if (result.execResult.logs) {
            const logs = result.execResult.logs.map((log) => {
              const address = bytesToHex(log[0]);
              const topics = log[1].map((topic) => bytesToHex(topic));
              const data = bytesToHex(log[2]);

              const decodedLog = decodeEventLog({
                abi: matchedCall.abi,
                data,
                topics: topics as any,
              });

              return { decodedLog, address, topics, data };
            });

            filterResult.logs = logs;
          }
        } catch (error) {}
      }
      next?.();
    },
  });

  return matchedCalls;
}

export function decodeDataSuffix({ data, abi }: { data: Hex; abi: Abi }) {
  const decodedFunctionCall = decodeFunctionData({
    abi,
    data,
  });

  const encodedFunctionCall = encodeFunctionData({
    abi,
    functionName: decodedFunctionCall.functionName,
    args: decodedFunctionCall.args,
  });

  return `0x${data.slice(encodedFunctionCall.length)}` as const;
}