Merge pull request #943 from ava-labs/fix/cap-tx-complexity

fix: cap tx complexity

Author: rictorlome

src/fixtures/pvm.ts

@@ -389,7 +389,7 @@ export const disableL1ValidatorTxBytes = () =>
   concatBytes(baseTxbytes(), idBytes(), bytesForInt(10), inputBytes());
 
 export const feeState = (): FeeState => ({
-  capacity: 1n,
+  capacity: 999_999n,
   excess: 1n,
   price: 1n,
   timestamp: new Date().toISOString(),

src/vms/pvm/etna-builder/spend-reducers/fixtures/reducers.ts

@@ -52,7 +52,7 @@ export const getSpendHelper = ({
 > = {}) => {
   return new SpendHelper({
     changeOutputs: [],
-    gasPrice: feeState.price,
+    feeState,
     initialComplexity,
     inputs: [],
     shouldConsolidateOutputs,

src/vms/pvm/etna-builder/spend-reducers/index.ts

@@ -2,5 +2,6 @@ export { handleFeeAndChange } from './handleFeeAndChange';
 export { useSpendableLockedUTXOs } from './useSpendableLockedUTXOs';
 export { useUnlockedUTXOs } from './useUnlockedUTXOs';
 export { verifyAssetsConsumed } from './verifyAssetsConsumed';
+export { verifyGasUsage } from './verifyGasUsage';
 
 export type * from './types';

src/vms/pvm/etna-builder/spend-reducers/verifyAssetsConsumed.test.ts

@@ -22,7 +22,7 @@ describe('verifyAssetsConsumed', () => {
     // Mock the verifyAssetsConsumed method to throw an error
     // Testing for this function can be found in the spendHelper.test.ts file
     spendHelper.verifyAssetsConsumed = vi.fn(() => {
-      throw new Error('Test error');
+      return new Error('Test error');
     });
 
     expect(() =>

src/vms/pvm/etna-builder/spend-reducers/verifyGasUsage.test.ts

@@ -0,0 +1,32 @@
+import { describe, expect, test, vi } from 'vitest';
+import { testContext } from '../../../../fixtures/context';
+import { getInitialReducerState, getSpendHelper } from './fixtures/reducers';
+import { verifyGasUsage } from './verifyGasUsage';
+
+describe('verifyGasUsage', () => {
+  test('returns original state if gas is under the threshold', () => {
+    const initialState = getInitialReducerState();
+    const spendHelper = getSpendHelper();
+    const spy = vi.spyOn(spendHelper, 'verifyGasUsage');
+
+    const state = verifyGasUsage(initialState, spendHelper, testContext);
+
+    expect(state).toBe(initialState);
+    expect(spy).toHaveBeenCalledTimes(1);
+  });
+
+  test('throws an error if gas is over the threshold', () => {
+    const initialState = getInitialReducerState();
+    const spendHelper = getSpendHelper();
+
+    // Mock the verifyGasUsage method to throw an error
+    // Testing for this function can be found in the spendHelper.test.ts file
+    spendHelper.verifyGasUsage = vi.fn(() => {
+      return new Error('Test error');
+    });
+
+    expect(() =>
+      verifyGasUsage(initialState, spendHelper, testContext),
+    ).toThrow('Test error');
+  });
+});

src/vms/pvm/etna-builder/spend-reducers/verifyGasUsage.ts

@@ -0,0 +1,16 @@
+import type { SpendReducerFunction } from './types';
+
+/**
+ * Verify that gas usage is within limits.
+ *
+ * Calls the spendHelper's verifyGasUsage method.
+ */
+export const verifyGasUsage: SpendReducerFunction = (state, spendHelper) => {
+  const verifyError = spendHelper.verifyGasUsage();
+
+  if (verifyError) {
+    throw verifyError;
+  }
+
+  return state;
+};

src/vms/pvm/etna-builder/spend.test.ts

@@ -4,6 +4,7 @@ import { Address, OutputOwners } from '../../../serializable';
 import { createDimensions } from '../../common/fees/dimensions';
 import {
   verifyAssetsConsumed,
+  verifyGasUsage,
   type SpendReducerFunction,
   type SpendReducerState,
   handleFeeAndChange,
@@ -14,6 +15,7 @@ import { feeState as testFeeState } from '../../../fixtures/pvm';
 import { bech32ToBytes } from '../../../utils';
 
 vi.mock('./spend-reducers', () => ({
+  verifyGasUsage: vi.fn<SpendReducerFunction>((state) => state),
   verifyAssetsConsumed: vi.fn<SpendReducerFunction>((state) => state),
   handleFeeAndChange: vi.fn<SpendReducerFunction>((state) => state),
 }));
@@ -51,6 +53,7 @@ describe('./src/vms/pvm/etna-builder/spend.test.ts', () => {
 
     expect(testReducer).toHaveBeenCalledTimes(1);
     expect(verifyAssetsConsumed).toHaveBeenCalledTimes(1);
+    expect(verifyGasUsage).toHaveBeenCalledTimes(1);
     expect(handleFeeAndChange).toHaveBeenCalledTimes(1);
   });
 

src/vms/pvm/etna-builder/spend.ts

@@ -9,7 +9,11 @@ import type { Dimensions } from '../../common/fees/dimensions';
 import type { Context } from '../../context';
 import type { FeeState } from '../models';
 import type { SpendReducerFunction, SpendReducerState } from './spend-reducers';
-import { handleFeeAndChange, verifyAssetsConsumed } from './spend-reducers';
+import {
+  handleFeeAndChange,
+  verifyAssetsConsumed,
+  verifyGasUsage,
+} from './spend-reducers';
 import { SpendHelper } from './spendHelper';
 
 type SpendResult = Readonly<{
@@ -118,11 +122,9 @@ export const spend = (
         fromAddresses.map((address) => address.toBytes()),
       );
 
-    const gasPrice: bigint = feeState.price;
-
     const spendHelper = new SpendHelper({
       changeOutputs: [],
-      gasPrice,
+      feeState,
       initialComplexity,
       inputs: [],
       shouldConsolidateOutputs,
@@ -147,6 +149,7 @@ export const spend = (
       ...spendReducers,
       verifyAssetsConsumed,
       handleFeeAndChange,
+      verifyGasUsage, // This should happen after change is added
       // Consolidation and sorting happens in the SpendHelper.
     ];
 

src/vms/pvm/etna-builder/spendHelper.test.ts

@@ -6,7 +6,8 @@ import {
 import { describe, test, expect } from 'vitest';
 
 import { id } from '../../../fixtures/common';
-import { stakeableLockOut } from '../../../fixtures/pvm';
+import type { FeeState } from '../models';
+import { stakeableLockOut, feeState } from '../../../fixtures/pvm';
 import { TransferableOutput } from '../../../serializable';
 import { isTransferOut } from '../../../utils';
 import type { Dimensions } from '../../common/fees/dimensions';
@@ -20,6 +21,7 @@ import { SpendHelper } from './spendHelper';
 import { getInputComplexity, getOutputComplexity } from '../txs/fee';
 
 const DEFAULT_GAS_PRICE = 3n;
+const DEFAULT_FEE_STATE: FeeState = { ...feeState(), price: DEFAULT_GAS_PRICE };
 
 const DEFAULT_WEIGHTS = createDimensions({
   bandwidth: 1,
@@ -30,7 +32,7 @@ const DEFAULT_WEIGHTS = createDimensions({
 
 const DEFAULT_PROPS: SpendHelperProps = {
   changeOutputs: [],
-  gasPrice: DEFAULT_GAS_PRICE,
+  feeState: DEFAULT_FEE_STATE,
   initialComplexity: createDimensions({
     bandwidth: 1,
     dbRead: 1,
@@ -372,6 +374,37 @@ describe('src/vms/pvm/etna-builder/spendHelper', () => {
       );
     });
   });
+  describe('SpendHelper.verifyGasUsage', () => {
+    test('returns null when gas is under capacity', () => {
+      const spendHelper = new SpendHelper({
+        ...DEFAULT_PROPS,
+      });
+
+      const changeOutput = transferableOutput();
+
+      spendHelper.addChangeOutput(changeOutput);
+
+      expect(spendHelper.verifyGasUsage()).toBe(null);
+    });
+
+    test('returns an error when gas is over capacity', () => {
+      const spendHelper = new SpendHelper({
+        ...DEFAULT_PROPS,
+        feeState: {
+          ...DEFAULT_FEE_STATE,
+          capacity: 0n,
+        },
+      });
+
+      const changeOutput = transferableOutput();
+
+      spendHelper.addChangeOutput(changeOutput);
+
+      expect(spendHelper.verifyGasUsage()).toEqual(
+        new Error('Gas usage of transaction (113) exceeds capacity (0)'),
+      );
+    });
+  });
 
   test('no consolidated outputs when `shouldConsolidateOutputs` is `false`', () => {
     const spendHelper = new SpendHelper(DEFAULT_PROPS);

src/vms/pvm/etna-builder/spendHelper.ts

@@ -11,11 +11,12 @@ import {
   dimensionsToGas,
 } from '../../common/fees/dimensions';
 import { consolidateOutputs } from '../../utils/consolidateOutputs';
+import type { FeeState } from '../models';
 import { getInputComplexity, getOutputComplexity } from '../txs/fee';
 
 export interface SpendHelperProps {
   changeOutputs: readonly TransferableOutput[];
-  gasPrice: bigint;
+  feeState: FeeState;
   initialComplexity: Dimensions;
   inputs: readonly TransferableInput[];
   shouldConsolidateOutputs: boolean;
@@ -32,7 +33,7 @@ export interface SpendHelperProps {
  * @class
  */
 export class SpendHelper {
-  private readonly gasPrice: bigint;
+  private readonly feeState: FeeState;
   private readonly initialComplexity: Dimensions;
   private readonly shouldConsolidateOutputs: boolean;
   private readonly toBurn: Map<string, bigint>;
@@ -47,7 +48,7 @@ export class SpendHelper {
 
   constructor({
     changeOutputs,
-    gasPrice,
+    feeState,
     initialComplexity,
     inputs,
     shouldConsolidateOutputs,
@@ -56,7 +57,7 @@ export class SpendHelper {
     toStake,
     weights,
   }: SpendHelperProps) {
-    this.gasPrice = gasPrice;
+    this.feeState = feeState;
     this.initialComplexity = initialComplexity;
     this.shouldConsolidateOutputs = shouldConsolidateOutputs;
     this.toBurn = toBurn;
@@ -217,13 +218,13 @@ export class SpendHelper {
   }
 
   /**
-   * Calculates the fee for the SpendHelper based on its complexity and gas price.
+   * Calculates the gas usage for the SpendHelper based on its complexity and the weights.
    * Provide an empty change output as a parameter to calculate the fee as if the change output was already added.
    *
    * @param {TransferableOutput} additionalOutput - The change output that has not yet been added to the SpendHelper.
-   * @returns {bigint} The fee for the SpendHelper.
+   * @returns {bigint} The gas usage for the SpendHelper.
    */
-  calculateFee(additionalOutput?: TransferableOutput): bigint {
+  private calculateGas(additionalOutput?: TransferableOutput): bigint {
     this.consolidateOutputs();
 
     const gas = dimensionsToGas(
@@ -233,7 +234,22 @@ export class SpendHelper {
       this.weights,
     );
 
-    return gas * this.gasPrice;
+    return gas;
+  }
+
+  /**
+   * Calculates the fee for the SpendHelper based on its complexity and gas price.
+   * Provide an empty change output as a parameter to calculate the fee as if the change output was already added.
+   *
+   * @param {TransferableOutput} additionalOutput - The change output that has not yet been added to the SpendHelper.
+   * @returns {bigint} The fee for the SpendHelper.
+   */
+  calculateFee(additionalOutput?: TransferableOutput): bigint {
+    const gas = this.calculateGas(additionalOutput);
+
+    const gasPrice = this.feeState.price;
+
+    return gas * gasPrice;
   }
 
   /**
@@ -281,6 +297,22 @@ export class SpendHelper {
     return null;
   }
 
+  /**
+   * Verifies that gas usage does not exceed the fee state maximum.
+   *
+   * @returns {Error | null} An error if gas usage exceeds maximum, null otherwise.
+   */
+  verifyGasUsage(): Error | null {
+    const gas = this.calculateGas();
+    if (this.feeState.capacity < gas) {
+      return new Error(
+        `Gas usage of transaction (${gas.toString()}) exceeds capacity (${this.feeState.capacity.toString()})`,
+      );
+    }
+
+    return null;
+  }
+
   /**
    * Gets the inputs, outputs, and UTXOs for the SpendHelper.
    *