Fix multisig partial spends

contracts/package.json

@@ -21,7 +21,8 @@
   "homepage": "https://docs.openzeppelin.com/contracts-compact/",
   "type": "module",
   "imports": {
-    "#test-utils/address.js": "./test-utils/address.js"
+    "#test-utils/address.js": "./test-utils/address.js",
+    "#test-utils/zswap.js": "./test-utils/zswap.js"
   },
   "scripts": {
     "compact": "compact-compiler --exclude '*/archive/*'",

contracts/src/multisig/ForwarderPrivate.compact

@@ -97,8 +97,12 @@ module ForwarderPrivate {
    * @description Spends a previously-deposited shielded coin to the parent,
    * a coin public key. The caller proves knowledge of `(parent, opSecret)`
    * matching the stored commitment; the coin is sent as a plain shielded note,
-   * so the parent stays hidden on-chain. If `coin.value` exceeds `value`, the
-   * change is re-emitted back to the contract for future drains.
+   * so the parent stays hidden on-chain. If `coin.value` exceeds `value`,
+   * `sendShielded` routes the change back to this contract as a self-owned
+   * output; `_drain` returns it untouched via `result.change`. The
+   * implementing contract chooses what to do with it: leave it dwelling at
+   * the contract for a future drain, or spend it onward to a different
+   * recipient in the same transaction via `sendImmediateShielded`.
    *
    * The parent is a `ZswapCoinPublicKey`, never a `ContractAddress`: a shielded
    * send to a contract publishes that contract's address in cleartext on the
@@ -160,14 +164,14 @@ module ForwarderPrivate {
       disclose(value)
     );
 
-    if (disclose(result.change.is_some)) {
-      sendImmediateShielded(
-        disclose(result.change.value),
-        Utils_selfAsRecipient(),
-        disclose(result.change.value.value)
-      );
-    }
-
+    // `sendShielded` already emits any change as a fresh output owned by this
+    // contract (dwelling at the contract address for future drains). `_drain`
+    // returns it untouched: it must not re-spend the change here and then hand
+    // that same (now-nullified) coin back via `result.change` — that is the
+    // double spend a node rejects on the next drain. A caller is free to spend
+    // the returned change onward itself (e.g. route it to another recipient);
+    // the rule is only that the change must never be re-spent AND still treated
+    // as held.
     return result;
   }
 

contracts/src/multisig/ShieldedTreasury.compact

@@ -85,8 +85,12 @@ module ShieldedTreasury {
    * @description Sends shielded tokens from the treasury.
    *
    * Looks up the stored coin by color, verifies sufficient value,
-   * and executes the shielded send. If the send produces change,
-   * it is sent back to the contract via `sendImmediateShielded`.
+   * and executes the shielded send. `sendShielded` already routes any
+   * change back to this contract as a self-owned output. Unlike the
+   * stateless variant, this treasury RETAINS the change: it is recorded in
+   * the UTXO map for the next spend. The returned `result.change` therefore
+   * belongs to the treasury and must not be re-spent by the caller (doing
+   * so would double-spend the recorded coin).
    *
    * @notice Access control is NOT enforced here.
    * The consuming contract must gate this behind its own
@@ -119,14 +123,13 @@ module ShieldedTreasury {
 
     const result = sendShielded(coin, disclose(recipient), disclose(amount));
 
+    // `sendShielded` spends the stored coin (revealing its nullifier) and,
+    // when there is change, emits it as a fresh output owned by this
+    // contract. That change coin is already spendable. Do NOT re-spend it
+    // with `sendImmediateShielded`, which would reveal its nullifier in this
+    // same transaction and leave the recorded coin permanently double-spent.
     if (disclose(result.change.is_some)) {
-      const changeCoin = result.change.value;
-      sendImmediateShielded(
-        changeCoin,
-        Utils_selfAsRecipient(),
-        changeCoin.value
-      );
-      _coins.insertCoin(disclose(color), changeCoin, Utils_selfAsRecipient());
+      _coins.insertCoin(disclose(color), result.change.value, Utils_selfAsRecipient());
     } else {
       _coins.remove(disclose(color));
     }

contracts/src/multisig/ShieldedTreasuryStateless.compact

@@ -20,7 +20,6 @@ pragma language_version >= 0.23.0;
  */
 module ShieldedTreasuryStateless {
   import CompactStandardLibrary;
-  import { selfAsRecipient } from "../utils/Utils" prefix Utils_;
 
   // ─── Deposit ────────────────────────────────────────────────────
 
@@ -38,8 +37,14 @@ module ShieldedTreasuryStateless {
    *
    * The coin is supplied by the caller rather than looked up from
    * on-chain state; the token color is derived from `coin.color`.
-   * Executes the shielded send and, if change is produced, returns it
-   * to the contract via `sendImmediateShielded`. No balance or
+   * Executes the shielded send and returns the result untouched.
+   * `sendShielded` already routes any change back to this contract as a
+   * self-owned output, so `result.change` is a live, spendable coin. The
+   * implementing contract chooses what to do with it: persist it and pass
+   * it back as the `coin` for a later spend, or spend it onward to a
+   * different recipient in the same transaction via `sendImmediateShielded`.
+   * (Re-spending change is only meaningful when it goes somewhere other than
+   * self — `sendShielded` already routes it to self.) No balance or
    * received/sent accounting is tracked on the ledger.
    *
    * @notice Access control is NOT enforced here.
@@ -63,14 +68,13 @@ module ShieldedTreasuryStateless {
   ): ShieldedSendResult {
     const result = sendShielded(disclose(coin), disclose(recipient), disclose(amount));
 
-    if (disclose(result.change.is_some)) {
-      sendImmediateShielded(
-        disclose(result.change.value),
-        Utils_selfAsRecipient(),
-        disclose(result.change.value.value)
-      );
-    }
-
+    // `sendShielded` already emits any change as a fresh, self-owned output.
+    // `_send` returns it untouched: it must not re-spend the change here and
+    // then hand that same (now-nullified) coin back via `result.change`. That
+    // is the double spend a node rejects on the next spend. The implementing
+    // contract decides what to do with the returned change: persist it and
+    // spend it later, or route it onward itself. The rule is only that the
+    // change must never be re-spent AND still treated as held.
     return result;
   }
 }

contracts/src/multisig/test/ForwarderPrivate.test.ts

@@ -1,6 +1,12 @@
 import fc from 'fast-check';
 import { beforeEach, describe, expect, it } from 'vitest';
 import * as utils from '#test-utils/address.js';
+import {
+  bytesToHex,
+  isNonceSpent,
+  zswapDelta,
+  zswapSnapshot,
+} from '#test-utils/zswap.js';
 import { MockForwarderPrivateSimulator } from './simulators/MockForwarderPrivateSimulator.js';
 
 // The drain parent is a `ZswapCoinPublicKey` (coin public key only). A contract
@@ -282,6 +288,149 @@ describe('ForwarderPrivate module', () => {
     });
   });
 
+  // `sendShielded` routes a partial drain's change back to the contract as a
+  // self-owned output, handed to the caller via `result.change`. That coin must
+  // stay spendable: if `_drain` also re-spent it with `sendImmediateShielded`,
+  // that would reveal its nullifier in the same transaction, making the returned
+  // coin a double spend the node rejects on the next drain. The dry simulator
+  // does not enforce nullifiers, so these tests read the recorded Zswap I/O: the
+  // change coin's nonce must not appear among the spent inputs.
+  describe('drain — change coin is spendable (no double spend)', () => {
+    // A non-zero deploy address so the change output (routed to self for future
+    // drains) carries a recognizable address rather than the zero
+    // `dummyContractAddress()` default.
+    const FORWARDER_ADDRESS = '3e'.repeat(32);
+    let mock: MockForwarderPrivateSimulator;
+
+    beforeEach(async () => {
+      mock = await MockForwarderPrivateSimulator.create(
+        commitment(PARENT_BYTES, OP_SECRET),
+        true,
+        { contractAddress: FORWARDER_ADDRESS },
+      );
+      await mock.deposit(makeCoin(COLOR, AMOUNT));
+    });
+
+    it('should send the note to the parent and route the change back to itself on a partial drain', async () => {
+      const snap = zswapSnapshot(mock);
+      const result = await mock.drain(
+        makeQualifiedCoin(COLOR, AMOUNT, 0n),
+        key(PARENT_BYTES),
+        OP_SECRET,
+        400n,
+      );
+      const { inputs, outputs } = zswapDelta(mock, snap);
+
+      // One coin consumed: the drained coin (nonce 0, full value).
+      expect(inputs).toHaveLength(1);
+      expect(inputs[0].value).toBe(AMOUNT);
+      expect(inputs[0].color).toStrictEqual(COLOR);
+      expect(inputs[0].nonce).toStrictEqual(new Uint8Array(32).fill(0));
+
+      // Two coins produced: the note to the parent key (`left` arm) and the
+      // change back to this contract (`right`/self arm).
+      expect(outputs).toHaveLength(2);
+      const toParent = outputs.filter((o) => o.recipient.is_left);
+      const toSelf = outputs.filter((o) => !o.recipient.is_left);
+      expect(toParent).toHaveLength(1);
+      expect(toSelf).toHaveLength(1);
+
+      // Note: 400 of COLOR to the parent key; equals result.sent.
+      expect(toParent[0].coinInfo.value).toBe(400n);
+      expect(toParent[0].coinInfo).toStrictEqual(result.sent);
+      expect(toParent[0].recipient.left.bytes).toStrictEqual(PARENT_BYTES);
+
+      // Change: the remainder back to THIS contract's address (for future
+      // drains), identical to the returned change coin, and NOT spent in this
+      // same tx.
+      expect(result.change.is_some).toBe(true);
+      expect(toSelf[0].coinInfo.value).toBe(AMOUNT - 400n);
+      expect(toSelf[0].coinInfo).toStrictEqual(result.change.value);
+      expect(bytesToHex(toSelf[0].recipient.right.bytes)).toBe(
+        FORWARDER_ADDRESS,
+      );
+      expect(isNonceSpent(inputs, result.change.value.nonce)).toBe(false);
+    });
+
+    it('should spend the coin and produce only the note when draining in full', async () => {
+      const snap = zswapSnapshot(mock);
+      const result = await mock.drain(
+        makeQualifiedCoin(COLOR, AMOUNT, 0n),
+        key(PARENT_BYTES),
+        OP_SECRET,
+        AMOUNT,
+      );
+      const { inputs, outputs } = zswapDelta(mock, snap);
+
+      // No change: one input (the drained coin), one output (the note to parent).
+      expect(result.change.is_some).toBe(false);
+      expect(inputs).toHaveLength(1);
+      expect(inputs[0].value).toBe(AMOUNT);
+      expect(outputs).toHaveLength(1);
+      expect(outputs[0].recipient.is_left).toBe(true);
+      expect(outputs[0].recipient.left.bytes).toStrictEqual(PARENT_BYTES);
+      expect(outputs[0].coinInfo.value).toBe(AMOUNT);
+      expect(outputs[0].coinInfo).toStrictEqual(result.sent);
+    });
+  });
+
+  // Tests that `_drain` (inner call in the mock `drainAndRouteChange`) returns a
+  // live, unspent change coin, so an implementing contract can spend it onward to
+  // a different recipient in the same tx (the only reason to re-spend change,
+  // since `sendShielded` already routes it to self). This would be impossible if
+  // `_drain` handed back a coin it had already spent.
+  describe('drain — implementing contract routes the change onward', () => {
+    const CHANGE_DEST = utils.createEitherTestUser('CHANGE_DEST');
+    let mock: MockForwarderPrivateSimulator;
+
+    beforeEach(async () => {
+      mock = await freshMock(PARENT_BYTES);
+    });
+
+    it('should let the caller send the change to a different recipient using the drain result', async () => {
+      const snap = zswapSnapshot(mock);
+      const routed = await mock.drainAndRouteChange(
+        makeQualifiedCoin(COLOR, AMOUNT, 0n),
+        key(PARENT_BYTES),
+        OP_SECRET,
+        400n,
+        CHANGE_DEST,
+      );
+      const { inputs, outputs } = zswapDelta(mock, snap);
+
+      // The onward send delivered the whole change to `changeRecipient`, so
+      // nothing is left over.
+      expect(routed.change.is_some).toBe(false);
+      expect(routed.sent.value).toBe(AMOUNT - 400n);
+
+      // Two coins are consumed: the drained coin, then the change coin (spent
+      // to route it onward — spending it here is correct, unlike the keep-change
+      // path).
+      expect(inputs).toHaveLength(2);
+
+      // The note still went to the parent for the drained `value`...
+      const toParent = outputs.filter(
+        (o) =>
+          o.recipient.is_left &&
+          bytesToHex(o.recipient.left.bytes) === bytesToHex(PARENT_BYTES),
+      );
+      expect(toParent).toHaveLength(1);
+      expect(toParent[0].coinInfo.value).toBe(400n);
+
+      // ...and the change was routed onward to `changeRecipient`, matching the
+      // returned coin.
+      const toChangeDest = outputs.filter(
+        (o) =>
+          o.recipient.is_left &&
+          bytesToHex(o.recipient.left.bytes) ===
+            bytesToHex(CHANGE_DEST.left.bytes),
+      );
+      expect(toChangeDest).toHaveLength(1);
+      expect(toChangeDest[0].coinInfo.value).toBe(AMOUNT - 400n);
+      expect(toChangeDest[0].coinInfo).toStrictEqual(routed.sent);
+    });
+  });
+
   describe('property: change arithmetic', () => {
     it('should preserve change.value == coin.value - drain.value on partial drain', async () => {
       await fc.assert(