Fix/bug78 prevent foreign token withdrawal

contracts/bene_contract/contract_v2.es

@@ -136,8 +136,17 @@
         OUTPUTS(0).tokens.size == 1 || OUTPUTS(0).tokens.size == 2 || OUTPUTS(0).tokens.size == 3
       }
 
+      // Ensure replication output does not contain foreign tokens (only APT, PFT, and base token if non-ERG)
+      val noForeignTokens = {
+        val allowed = OUTPUTS(0).tokens.filter({ (token: (Coll[Byte], Long)) =>
+          val id = token._1
+          if (isERGBase) { id == selfId || id == pftTokenId } else { id == selfId || id == pftTokenId || id == baseTokenId }
+        })
+        allowed.size == OUTPUTS(0).tokens.size
+      }
+
       // Verify that the output box is a valid copy of the input box
-      sameId && sameBlockLimit && sameMinimumSold && sameExchangeRate && sameConstants && sameProjectContent && sameScript && noAddsOtherTokens
+      sameId && sameBlockLimit && sameMinimumSold && sameExchangeRate && sameConstants && sameProjectContent && sameScript && noAddsOtherTokens && noForeignTokens
     }
   }
 
@@ -161,6 +170,21 @@
       selfAmount == outAmount
     }
   }
+  val NonCoreTokensRemainConstant = !isReplicationBoxPresent || {
+    val selfOther = SELF.tokens.filter({ (token: (Coll[Byte], Long)) =>
+      val id = token._1
+      if (isERGBase) { id != selfId && id != pftTokenId } else { id != selfId && id != pftTokenId && id != baseTokenId }
+    })
+    val outOther = OUTPUTS(0).tokens.filter({ (token: (Coll[Byte], Long)) =>
+      val id = token._1
+      if (isERGBase) { id != selfId && id != pftTokenId } else { id != selfId && id != pftTokenId && id != baseTokenId }
+    })
+    val badMatches = selfOther.filter({ (token: (Coll[Byte], Long)) =>
+      val matches = outOther.filter({ (t: (Coll[Byte], Long)) => t._1 == token._1 && t._2 == token._2 })
+      matches.size == 0
+    })
+    badMatches.size == 0 && selfOther.size == outOther.size
+  }
   val soldCounterRemainsConstant = !isReplicationBoxPresent || (selfSoldCounter == OUTPUTS(0).R6[Coll[Long]].get(0))
   val refundCounterRemainsConstant = !isReplicationBoxPresent || (selfRefundCounter == OUTPUTS(0).R6[Coll[Long]].get(1))
   val auxiliarExchangeCounterRemainsConstant = !isReplicationBoxPresent || (selfAuxiliarExchangeCounter == OUTPUTS(0).R6[Coll[Long]].get(2))
@@ -444,7 +468,8 @@
         auxiliarExchangeCounterRemainsConstant,   
         // mantainValue,                           // Needs to extract value from the contract
         APTokenRemainsConstant,                    // There is no need to modify the auxiliar token, so it must be constant
-        ProofFundingTokenRemainsConstant           // There is no need to modify the proof of funding token, so it must be constant
+        ProofFundingTokenRemainsConstant,          // There is no need to modify the proof of funding token, so it must be constant
+        NonCoreTokensRemainConstant                // Prevents withdrawing unrelated tokens
       ))
 
       sigmaProp(allOf(Coll(

tests/contracts/withdraw_funds_unrelated_tokens.test.ts

@@ -0,0 +1,157 @@
+import { describe, it, expect, beforeEach } from "vitest";
+import { Box, OutputBuilder, TransactionBuilder, RECOMMENDED_MIN_FEE_VALUE, SAFE_MIN_BOX_VALUE } from "@fleet-sdk/core";
+import { SByte, SColl, SLong } from "@fleet-sdk/serializer";
+import { stringToBytes } from "@scure/base";
+import { compile } from "@fleet-sdk/compiler";
+import { setupBeneTestContext, ERG_BASE_TOKEN, ERG_BASE_TOKEN_NAME, type BeneTestContext, USD_BASE_TOKEN, USD_BASE_TOKEN_NAME, createR4 } from "./bene_contract_helpers";
+
+// This test demonstrates a bug in contract_v2.es: during withdrawFunds, the contract allows moving arbitrary extra tokens
+// that were previously injected into the SELF box, as long as APT and PFT remain constant and the base-token amounts are correct.
+// Expected behavior: Only the base token (ERG or configured token) should be extracted; any unrelated tokens on the contract
+// must remain on the replicated SELF box.
+// Current behavior: The contract does not restrict removal of unrelated tokens in withdrawFunds, enabling owner to sweep them.
+// This test is intentionally written to expect failure (result === false). It currently passes under the buggy contract,
+// therefore the test will FAIL, proving the issue.
+
+const baseModes = [
+  { name: "USD Token Mode", token: USD_BASE_TOKEN, tokenName: USD_BASE_TOKEN_NAME },
+  { name: "ERG Mode", token: ERG_BASE_TOKEN, tokenName: ERG_BASE_TOKEN_NAME },
+];
+
+describe.each(baseModes)("Bene Contract v1.2 - Withdraw funds should NOT move unrelated tokens (%s)", (mode) => {
+  let ctx: BeneTestContext;  // Test environment
+  let projectBox: Box;  // Contract box
+  let soldTokens: bigint;
+  let collectedFunds: bigint;
+
+  describe("Withdraw funds while attempting to sweep an unrelated token", () => {
+    beforeEach(() => {
+      // Initialize context
+      ctx = setupBeneTestContext(mode.token, mode.tokenName);
+
+      // Owner funded for tx fees
+      ctx.projectOwner.addBalance({ nanoergs: 10_000_000_000n });
+
+      // Prepare a project box with MINIMUM REACHED
+      soldTokens = ctx.minimumTokensSold;
+      collectedFunds = soldTokens * ctx.exchangeRate;
+
+      // Inject a foreign token into the contract box (not baseTokenId, not pftTokenId)
+      const foreignTokenId = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"; // 64-hex fake token id
+
+      let assets = [
+        { tokenId: ctx.projectNftId, amount: 1n + ctx.totalPFTokens - soldTokens },
+        { tokenId: foreignTokenId, amount: 100n }, // unrelated token present on the contract
+      ];
+
+      let value = SAFE_MIN_BOX_VALUE;
+
+      if (!ctx.isErgMode) {
+        assets.push({ tokenId: ctx.baseTokenId, amount: collectedFunds });
+      } else {
+        value += collectedFunds;
+        collectedFunds = value; // include SAFE_MIN_BOX_VALUE
+      }
+
+      ctx.beneContract.addUTxOs({
+        value: value,
+        ergoTree: ctx.beneErgoTree.toHex(),
+        assets: assets,
+        creationHeight: ctx.mockChain.height - 100,
+        additionalRegisters: {
+          R4: createR4(ctx),
+          R5: SLong(ctx.minimumTokensSold).toHex(),
+          R6: SColl(SLong, [soldTokens, 0n, ctx.totalPFTokens]).toHex(),
+          R7: SLong(ctx.exchangeRate).toHex(),
+          R8: ctx.constants.toHex(),
+          R9: SColl(SByte, stringToBytes("utf8", "{}")).toHex(),
+        },
+      });
+
+      projectBox = ctx.beneContract.utxos.toArray()[0];
+    });
+
+    it("should fail when withdraw attempts to move a foreign token from the contract", () => {
+      // Compute expected split
+      const withdrawAmount = ctx.isErgMode ? (collectedFunds - SAFE_MIN_BOX_VALUE) / 2n : collectedFunds / 2n;
+      const devFeeAmount = (withdrawAmount * BigInt(ctx.devFeePercentage)) / 100n;
+      const projectAmount = withdrawAmount - devFeeAmount;
+      const devFeeContract = compile(`{ sigmaProp(true) }`);
+
+      // Build outputs with partial withdraw (replicate SELF) and sweep the foreign token to owner
+      let remainingValue: bigint;
+      let remainingAssets: { tokenId: string; amount: bigint }[];
+      let projectValue: bigint;
+      let projectAssets: { tokenId: string; amount: bigint }[];
+      let devFeeValue: bigint;
+      let devFeeAssets: { tokenId: string; amount: bigint }[];
+
+      const foreignTokenIdOnSelf = projectBox.assets.find(t => t.tokenId !== ctx.projectNftId && t.tokenId !== ctx.pftTokenId && t.tokenId !== ctx.baseTokenId)?.tokenId!;
+      const foreignTokenAmountOnSelf = projectBox.assets.find(t => t.tokenId === foreignTokenIdOnSelf)?.amount ?? 0n;
+
+      if (ctx.isErgMode) {
+        // ERG mode: decrease ERG on SELF, move ERG to project/dev, and also move foreign token to project
+        remainingValue = collectedFunds - withdrawAmount; // keep remainder on self
+        remainingAssets = [
+          { tokenId: ctx.projectNftId, amount: projectBox.assets[0].amount }, // keep APT
+          // NOTE: we DO NOT keep the foreign token on SELF; we will try to move it out
+        ];
+
+        projectValue = projectAmount;
+        projectAssets = [];
+
+        devFeeValue = devFeeAmount;
+        devFeeAssets = [];
+      } else {
+        // Token mode: keep half base tokens on SELF, move half to project/dev, and also move foreign token to project
+        remainingValue = SAFE_MIN_BOX_VALUE;
+        remainingAssets = [
+          { tokenId: ctx.projectNftId, amount: projectBox.assets[0].amount },
+          { tokenId: ctx.baseTokenId, amount: collectedFunds - withdrawAmount },
+          // NOTE: we DO NOT keep the foreign token on SELF; we will try to move it out
+        ];
+
+        projectValue = SAFE_MIN_BOX_VALUE;
+        projectAssets = [
+          { tokenId: ctx.baseTokenId, amount: projectAmount },
+          // BUG: Also extracting unrelated token to owner
+          { tokenId: foreignTokenIdOnSelf, amount: foreignTokenAmountOnSelf },
+        ];
+
+        devFeeValue = SAFE_MIN_BOX_VALUE;
+        devFeeAssets = [
+          { tokenId: ctx.baseTokenId, amount: devFeeAmount },
+        ];
+      }
+
+      const transaction = new TransactionBuilder(ctx.mockChain.height)
+        .from([projectBox, ...ctx.projectOwner.utxos.toArray()])
+        .to([
+          // Output 0: Replicated SELF with updated funds (but foreign token removed)
+          new OutputBuilder(remainingValue, ctx.beneErgoTree)
+            .addTokens(remainingAssets)
+            .setAdditionalRegisters({
+              R4: projectBox.additionalRegisters.R4,
+              R5: SLong(ctx.minimumTokensSold).toHex(),
+              R6: SColl(SLong, [soldTokens, 0n, ctx.totalPFTokens]).toHex(),
+              R7: SLong(ctx.exchangeRate).toHex(),
+              R8: projectBox.additionalRegisters.R8,
+              R9: projectBox.additionalRegisters.R9,
+            }),
+          // Output 1: Project
+          new OutputBuilder(projectValue, ctx.projectOwner.address).addTokens(projectAssets),
+          // Output 2: Dev fee
+          new OutputBuilder(devFeeValue, devFeeContract).addTokens(devFeeAssets),
+        ])
+        .sendChangeTo(ctx.projectOwner.address)
+        .payFee(RECOMMENDED_MIN_FEE_VALUE)
+        .build();
+
+      const result = ctx.mockChain.execute(transaction, { signers: [ctx.projectOwner], throw: false });
+
+      // EXPECTED: The contract should reject removing unrelated tokens from SELF during withdraw
+      // ACTUAL: The contract currently allows it (APT and PFT remain constant; base token amounts are correct)
+      expect(result).toBe(false);
+    });
+  });
+});

tests/mocks/wallet-svelte-component.ts

@@ -0,0 +1,75 @@
+// Minimal Node-friendly mock for wallet-svelte-component used by backend tests
+
+// Simple writable store implementation compatible with basic subscribe/set usage
+export type Unsubscriber = () => void;
+
+export interface Writable<T> {
+  subscribe(run: (value: T) => void): Unsubscriber;
+  set(value: T): void;
+  update(updater: (value: T) => T): void;
+}
+
+function writable<T>(initial: T): Writable<T> {
+  let value = initial;
+  const subscribers = new Set<(v: T) => void>();
+
+  function subscribe(run: (v: T) => void): Unsubscriber {
+    subscribers.add(run);
+    // Emit current value on subscribe
+    try { run(value); } catch {}
+    return () => subscribers.delete(run);
+  }
+
+  function set(v: T) {
+    value = v;
+    for (const s of subscribers) {
+      try { s(value); } catch {}
+    }
+  }
+
+  function update(updater: (v: T) => T) {
+    set(updater(value));
+  }
+
+  return { subscribe, set, update };
+}
+
+// Stores exposed by the real library
+export const walletConnected: Writable<boolean> = writable(false);
+export const walletAddress: Writable<string> = writable("");
+export const explorerUri: Writable<string | null> = writable(null);
+
+// Wallet manager mock
+export const walletManager = {
+  async connectWallet(_name?: string): Promise<void> {
+    walletConnected.set(true);
+  },
+  async disconnect(): Promise<void> {
+    walletConnected.set(false);
+  },
+};
+
+// Node-side helpers used by actions (provide harmless stubs)
+export async function getCurrentHeight(): Promise<number> {
+  return 0;
+}
+
+export async function getChangeAddress(): Promise<string> {
+  // Dummy P2PK address (not actually used in unit tests)
+  return "9fMockChangeAddress000000000000000000000000000000000000000000000000000";
+}
+
+export async function getUtxos(_amount?: string): Promise<any[]> {
+  // Return empty list for tests that don't require wallet UTXOs
+  return [];
+}
+
+export async function signTransaction<T>(unsigned: T): Promise<T> {
+  // Echo back unsigned transaction for tests that don't submit
+  return unsigned;
+}
+
+export async function submitTransaction(_signed: any): Promise<string> {
+  // Return a dummy transaction id
+  return "mock-tx-id";
+}

vitest.config.ts

@@ -11,6 +11,8 @@ export default defineConfig({
   resolve: {
     alias: {
       $lib: path.resolve(__dirname, './src/lib'),
+      'wallet-svelte-component': path.resolve(__dirname, './tests/mocks/wallet-svelte-component.ts'),
+      'wallet-svelte-component/dist/wallet/wallet-manager': path.resolve(__dirname, './tests/mocks/wallet-svelte-component.ts'),
     },
   },
 });