[Audit M-01] Referral registration lacks wallet-session drift protection

PlasticDigits commented

2026-05-05 08:10:13 +00:00

(Migrated from gitlab.com)

Source

Internal full-stack audit audits/audit_full_1777967937.md (2026-05-05), finding M-01 (Medium). Commit reviewed: 3b3bd0c6c501f9cf908667d304146b7b63302b87.

Problem

ReferralRegisterSection.onRegister checks the target chain before starting, reads allowance for the React-closure address, optionally submits approve, waits for the receipt, then submits registerCode. Unlike TimeCurve buys (GitLab #144 / INV-BUY-SESSION-144), it does not capture and re-check the connected account and chain after each await.

If the wallet account or network changes mid-flow:

Approval may be requested from the newly active wallet while the UI still reasons about the old address.
registerCode may register the normalized code for a different wallet than the one the closure will later cache with setStoredMyReferralCodeForWallet(address, normalized).
Local storage may show the code under the wrong wallet until reads invalidate.

This is not a protocol fund-drain: ReferralRegistry.registerCode burns from msg.sender and the contract mapping follows the caller. The risk is user confusion, wrong-wallet referral ownership UX, and avoidable CL8Y prompts during wallet drift.

Relevant code

onRegister awaits allowance read, approve receipt, registerCode, and receipt without session snapshots:

  const onRegister = useCallback(async () => {
    setFormErr(null);
    const netErr = chainMismatchWriteMessage(chainId);
    if (netErr) {
      setFormErr(netErr);
      return;
    }
    // ...
    try {
      const need = burnWad;
      const allow = await readContract(wagmiConfig, {
        address: cl8yToken,
        abi: erc20Abi,
        functionName: "allowance",
        args: [address, registry],
      });
      if (allow < need) {
        const approveHash = await writeContractAsync({
          // ...
        });
        await waitForTransactionReceipt(wagmiConfig, { hash: approveHash });
      }
      const regHash = await writeContractAsync({
        address: registry,
        abi: referralRegistryWriteAbi,
        functionName: "registerCode",
        args: [normalized],
      });
      await waitForTransactionReceipt(wagmiConfig, { hash: regHash });
      setStoredMyReferralCodeForWallet(address, normalized);

Existing guard pattern to reuse:

export function captureWalletBuySession(config: Config): WalletBuySessionSnapshot | null {
  const a = getAccount(config);
  if (a.status !== "connected") return null;
  return { address: a.address, chainId: a.chainId };
}

/** Throws with {@link WALLET_BUY_SESSION_DRIFT_MESSAGE} when the session drifted. */
export function assertWalletBuySessionUnchanged(
  config: Config,
  snapshot: WalletBuySessionSnapshot,
): void {
  const d = walletBuySessionDriftMessage(config, snapshot);
  if (d) throw new Error(d);
}

Recommended fix

Reuse the walletBuySessionGuard pattern (same drift message as buys for consistency, or a referral-specific message if product prefers).
Capture a session snapshot with captureWalletBuySession(wagmiConfig) immediately after validating inputs and before the first async read/write.
Call assertWalletBuySessionUnchanged after allowance read, after approve receipt (if any), immediately before registerCode, after register receipt, and before setStoredMyReferralCodeForWallet.
Add a Vitest test around a pure helper or mocked registration flow; add a manual QA row mirroring INV-BUY-SESSION-144 for referrals.

Acceptance criteria

Registration flow aborts with the documented drift message if account or chainId changes between any awaited steps (same semantics as TimeCurve buy session guard).
No setStoredMyReferralCodeForWallet call runs after a drift (storage stays aligned with the wallet that actually registered).
Automated test covers drift detection logic or the integrated handler behavior.
Contributor-facing invariant or manual QA note references referral registration alongside buy session drift.

Verification checklist

Switch wallet mid-flow after clicking Register — flow aborts with drift message; no wrong-wallet localStorage write.
Switch network mid-flow — same as above.
Happy path unchanged: same wallet/chain completes approve (if needed), register, receipt, and cached code for correct address.
npm test (Vitest) passes including new/updated tests for referral registration session guard.

## Source Internal full-stack audit [`audits/audit_full_1777967937.md`](https://gitlab.com/PlasticDigits/yieldomega/-/blob/main/audits/audit_full_1777967937.md) (2026-05-05), finding **M-01** (Medium). Commit reviewed: `3b3bd0c6c501f9cf908667d304146b7b63302b87`. ## Problem `ReferralRegisterSection.onRegister` checks the target chain before starting, reads allowance for the React-closure `address`, optionally submits `approve`, waits for the receipt, then submits `registerCode`. Unlike TimeCurve buys ([GitLab #144](https://gitlab.com/PlasticDigits/yieldomega/-/issues/144) / `INV-BUY-SESSION-144`), it does **not** capture and re-check the connected account and chain after each await. If the wallet account or network changes mid-flow: - Approval may be requested from the newly active wallet while the UI still reasons about the old `address`. - `registerCode` may register the normalized code for a different wallet than the one the closure will later cache with `setStoredMyReferralCodeForWallet(address, normalized)`. - Local storage may show the code under the wrong wallet until reads invalidate. This is **not** a protocol fund-drain: `ReferralRegistry.registerCode` burns from `msg.sender` and the contract mapping follows the caller. The risk is user confusion, wrong-wallet referral ownership UX, and avoidable CL8Y prompts during wallet drift. ## Relevant code `onRegister` awaits allowance read, approve receipt, `registerCode`, and receipt without session snapshots: ```181:235:frontend/src/pages/referrals/ReferralRegisterSection.tsx const onRegister = useCallback(async () => { setFormErr(null); const netErr = chainMismatchWriteMessage(chainId); if (netErr) { setFormErr(netErr); return; } // ... try { const need = burnWad; const allow = await readContract(wagmiConfig, { address: cl8yToken, abi: erc20Abi, functionName: "allowance", args: [address, registry], }); if (allow < need) { const approveHash = await writeContractAsync({ // ... }); await waitForTransactionReceipt(wagmiConfig, { hash: approveHash }); } const regHash = await writeContractAsync({ address: registry, abi: referralRegistryWriteAbi, functionName: "registerCode", args: [normalized], }); await waitForTransactionReceipt(wagmiConfig, { hash: regHash }); setStoredMyReferralCodeForWallet(address, normalized); ``` Existing guard pattern to reuse: ```38:63:frontend/src/lib/walletBuySessionGuard.ts export function captureWalletBuySession(config: Config): WalletBuySessionSnapshot | null { const a = getAccount(config); if (a.status !== "connected") return null; return { address: a.address, chainId: a.chainId }; } /** Throws with {@link WALLET_BUY_SESSION_DRIFT_MESSAGE} when the session drifted. */ export function assertWalletBuySessionUnchanged( config: Config, snapshot: WalletBuySessionSnapshot, ): void { const d = walletBuySessionDriftMessage(config, snapshot); if (d) throw new Error(d); } ``` ## Recommended fix - Reuse the `walletBuySessionGuard` pattern (same drift message as buys for consistency, or a referral-specific message if product prefers). - Capture a session snapshot with `captureWalletBuySession(wagmiConfig)` immediately after validating inputs and before the first async read/write. - Call `assertWalletBuySessionUnchanged` after allowance read, after approve receipt (if any), immediately before `registerCode`, after register receipt, and **before** `setStoredMyReferralCodeForWallet`. - Add a Vitest test around a pure helper or mocked registration flow; add a manual QA row mirroring **`INV-BUY-SESSION-144`** for referrals. ## Acceptance criteria - Registration flow aborts with the documented drift message if account or `chainId` changes between any awaited steps (same semantics as TimeCurve buy session guard). - No `setStoredMyReferralCodeForWallet` call runs after a drift (storage stays aligned with the wallet that actually registered). - Automated test covers drift detection logic or the integrated handler behavior. - Contributor-facing invariant or manual QA note references referral registration alongside buy session drift. ## Verification checklist - [ ] Switch wallet mid-flow after clicking Register — flow aborts with drift message; no wrong-wallet localStorage write. - [ ] Switch network mid-flow — same as above. - [ ] Happy path unchanged: same wallet/chain completes approve (if needed), register, receipt, and cached code for correct address. - [ ] `npm test` (Vitest) passes including new/updated tests for referral registration session guard.

PlasticDigits commented

2026-05-05 09:01:21 +00:00

(Migrated from gitlab.com)

mentioned in commit 69172a636f

mentioned in commit 69172a636fdfd6c87bbcc944fbbd11b45a639c92

PlasticDigits commented

2026-05-05 09:01:22 +00:00

(Migrated from gitlab.com)

mentioned in commit 955228fd7f

mentioned in commit 955228fd7fdd3748c1360ac9995b4b4532cf3e8c

PlasticDigits commented

2026-05-05 09:02:00 +00:00

(Migrated from gitlab.com)

Implemented GitLab #155 (audit M-01 / referral registration wallet-session drift).

What changed

ReferralRegisterSection.onRegister now latches captureWalletBuySession(wagmiConfig) after validation and before the first async step; assertWalletBuySessionUnchanged runs after each await (allowance read, optional approve + receipt, registerCode + receipt).
setStoredMyReferralCodeForWallet runs only after the final guard and uses the snapshot address (aligned with the signing wallet).
User-visible abort uses the same WALLET_BUY_SESSION_DRIFT_MESSAGE as TimeCurve buys (#144) for consistency.

Docs / tests

Invariant INV-REFERRAL-SESSION-155 + contributor manual QA §155; wallet-connection.md §155; skills/README.md cross-links; Vitest referralRegisterWalletSession.test.ts.

Merged: main @ 955228f (pushed to origin/main).

Checklist for @brouie to verify

Happy path: /referrals → register a new code with a stable wallet on the target chain → approve (if prompted) + register succeed; share-link / local cache behave as before for that wallet.
Account switch: start Register, switch MetaMask/Rabby account before registerCode confirms → UI shows Could not register: … Wallet or network changed during purchase — please retry from the beginning. — no plaintext code cached under the wrong yieldomega.myrefcode.v1.<address> key.
Network switch: same drift message if chainId changes mid-flow (or #95 preflight if the switch happens before guarded writes — either is acceptable UX).
Automation: cd frontend && npm test — includes referralRegisterWalletSession.test.ts.

Leaving this issue open for your sign-off.

Implemented **GitLab #155** (audit **M-01** / referral registration wallet-session drift). **What changed** - `ReferralRegisterSection.onRegister` now latches **`captureWalletBuySession(wagmiConfig)`** after validation and before the first async step; **`assertWalletBuySessionUnchanged`** runs after **each** `await` (allowance read, optional approve + receipt, `registerCode` + receipt). - **`setStoredMyReferralCodeForWallet`** runs only after the final guard and uses the **snapshot address** (aligned with the signing wallet). - User-visible abort uses the same **`WALLET_BUY_SESSION_DRIFT_MESSAGE`** as TimeCurve buys (**#144**) for consistency. **Docs / tests** - Invariant **`INV-REFERRAL-SESSION-155`** + contributor manual QA §155; `wallet-connection.md` §155; `skills/README.md` cross-links; Vitest **`referralRegisterWalletSession.test.ts`**. **Merged:** `main` @ **955228f** (pushed to `origin/main`). --- **Checklist for @brouie to verify** - [ ] Happy path: `/referrals` → register a new code with a stable wallet on the target chain → approve (if prompted) + register succeed; share-link / local cache behave as before for that wallet. - [ ] Account switch: start Register, switch MetaMask/Rabby account **before** `registerCode` confirms → UI shows **Could not register:** … **`Wallet or network changed during purchase — please retry from the beginning.`** — no plaintext code cached under the **wrong** `yieldomega.myrefcode.v1.<address>` key. - [ ] Network switch: same drift message if `chainId` changes mid-flow (or **#95** preflight if the switch happens before guarded writes — either is acceptable UX). - [ ] Automation: `cd frontend && npm test` — includes **`referralRegisterWalletSession.test.ts`**. Leaving this issue **open** for your sign-off.

Brouie commented

2026-05-06 00:10:28 +00:00

(Migrated from gitlab.com)

@PlasticDigits — verification PASS on 69172a6.

ran the new vitest:

npx vitest run src/pages/referrals/referralRegisterWalletSession.test.ts

3/3 PASS in 633ms — capture matches preflight, abort on account switch, abort on chain switch.

source-side review of ReferralRegisterSection.tsx confirms all five latch points wired:

preflight capture + consistency check against current address/chainId
guard before allowance read (with snapshot.address as the arg)
guard after approve receipt (when allow < need)
guard before registerCode write
guard after register receipt, then cache write keyed under snapshot.address

INV-REFERRAL-SESSION-155 satisfied.

ready for close on your side.

@PlasticDigits — verification PASS on `69172a6`. ran the new vitest: ``` npx vitest run src/pages/referrals/referralRegisterWalletSession.test.ts ``` 3/3 PASS in 633ms — capture matches preflight, abort on account switch, abort on chain switch. source-side review of `ReferralRegisterSection.tsx` confirms all five latch points wired: - preflight capture + consistency check against current address/chainId - guard before allowance read (with snapshot.address as the arg) - guard after approve receipt (when allow < need) - guard before registerCode write - guard after register receipt, then cache write keyed under snapshot.address INV-REFERRAL-SESSION-155 satisfied. ready for close on your side.

Brouie commented

2026-05-06 01:07:08 +00:00

(Migrated from gitlab.com)

@PlasticDigits — verification PASS on 69172a6 (full walk).

automation: cd frontend && npm test → 286/286 across 47 files including referralRegisterWalletSession.test.ts 3/3.

happy path on /referrals: registered testbrouie506 under anvil #1 0x70997970...79c8 cleanly. on-chain CodeRegistered event at block 64. page shows registered code, share links populate, localStorage yieldomega.myrefcode.v1.0x70997970...c518 correctly cached with testbrouie506.

account switch mid-flow: walked via cancel-then-switch path — metamask blocks account switching while the signature popup is open (intentional wallet UX, also documented in #144). cancelled the registerCode popup for clean0506 under anvil #2, switched to anvil #0. result: setFormErr("Transaction was rejected in the wallet.") surfaces correctly in the Register section, no cache write happened under 0x3c44... or 0xf39Fd6..., page rebinds to new account cleanly. cache poison prevented as expected.

network switch mid-flow: not walked separately — same metamask constraint as account switch makes the testable variant cancel-then-switch-chain, which mirrors the account-switch behavior above. happy to walk it explicitly if you want a separate signoff.

source-side review of ReferralRegisterSection.tsx confirms all five latch points are wired:

preflight capture + consistency check against current address/chainId
guard before allowance read with snapshot.address as the arg
guard after approve receipt
guard before registerCode write
guard after register receipt, then cache write keyed under snapshot.address

INV-REFERRAL-SESSION-155 satisfied for the testable surface.

ready for close on your side.

@PlasticDigits — verification PASS on `69172a6` (full walk). automation: `cd frontend && npm test` → 286/286 across 47 files including `referralRegisterWalletSession.test.ts` 3/3. happy path on /referrals: registered `testbrouie506` under anvil #1 0x70997970...79c8 cleanly. on-chain CodeRegistered event at block 64. page shows registered code, share links populate, localStorage `yieldomega.myrefcode.v1.0x70997970...c518` correctly cached with `testbrouie506`. account switch mid-flow: walked via cancel-then-switch path — metamask blocks account switching while the signature popup is open (intentional wallet UX, also documented in #144). cancelled the registerCode popup for `clean0506` under anvil #2, switched to anvil #0. result: `setFormErr("Transaction was rejected in the wallet.")` surfaces correctly in the Register section, no cache write happened under 0x3c44... or 0xf39Fd6..., page rebinds to new account cleanly. cache poison prevented as expected. network switch mid-flow: not walked separately — same metamask constraint as account switch makes the testable variant cancel-then-switch-chain, which mirrors the account-switch behavior above. happy to walk it explicitly if you want a separate signoff. source-side review of `ReferralRegisterSection.tsx` confirms all five latch points are wired: - preflight capture + consistency check against current address/chainId - guard before allowance read with snapshot.address as the arg - guard after approve receipt - guard before registerCode write - guard after register receipt, then cache write keyed under snapshot.address INV-REFERRAL-SESSION-155 satisfied for the testable surface. ready for close on your side.

PlasticDigits (Migrated from gitlab.com) closed this issue

2026-05-06 08:25:07 +00:00

PlasticDigits commented

2026-05-13 15:13:33 +00:00

(Migrated from gitlab.com)

mentioned in commit f43ac77e85

mentioned in commit f43ac77e85e88f346bf79738385a8ef43ce45c80

PlasticDigits commented

2026-05-13 15:13:33 +00:00

(Migrated from gitlab.com)

mentioned in commit eba7c43207

mentioned in commit eba7c432073c9d0778c3d9038e8da7f1ac2b6164

PlasticDigits commented

2026-06-14 22:08:37 +00:00

(Migrated from gitlab.com)

mentioned in issue #325

Rows
Columns

[Audit M-01] Referral registration lacks wallet-session drift protection #155

Source

Problem

Relevant code

Recommended fix

Acceptance criteria

Verification checklist