From d42fd9fdd38046f998b93eab35ccf302837d7c29 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=CE=BA=CE=B1=CF=83=CF=83=CE=AC=CE=BD=CE=B4=CF=81=CE=B1=2Ee?= =?UTF-8?q?th?= <0xDADA@protonmail.com> Date: Tue, 1 Sep 2026 16:47:51 +0200 Subject: [PATCH] fix: shield --from flag supports stealth addys --- CHANGELOG.md | 2 + README.md | 3 +- src/commands/shield.ts | 68 ++++++++------------------ tests/shield-txs.test.ts | 102 ++++++++++++++++++++++++++++++++++++++- 4 files changed, 126 insertions(+), 49 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b888169..f0b836c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,6 +20,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). - fix `balances --verbose` to no longer sync privacy protocols twice unnecessarily. - `shield` and `unshield` keep a live progress timer during protocol sync (including Railgun WASM), matching `balances`. +- the `--from` flag on `shield` correctly supports `` or `sN` (stealth address by index) ## [0.0.4] — 2026-08-25 @@ -37,3 +38,4 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). ### Fixed - Unshielding to a custom / ephemeral recipient (`--to` an address that is not a stored HD account) no longer fails or mis-routes funds. + diff --git a/README.md b/README.md index ddc685c..c3147ed 100644 --- a/README.md +++ b/README.md @@ -418,7 +418,7 @@ Move funds from a **public** account into a private protocol. | `--protocol ` | Required unless `DEFAULT_PRIVACY_PROTOCOL` is set to one of those values. | | `--wallet ` | Wallet. | | `--password ` | Unlock password. | -| `--from ` | Sender public account (address or HD index). | +| `--from ` | Sender public account address, HD index, or stealth selector (`s0`). | | `--from-priv` | With `--broadcast`: derive private key by index from mnemonic if account not yet in stored public list. | | `--token ` | Token (default: `eth`). | | `--amount-wei ` | Amount in base units. | @@ -446,6 +446,7 @@ When a shield needs more than one on-chain call, the CLI uses EIP-7702 Simple770 ```bash kohaku shield --protocol tornado --wallet testWallet --from 0 --amount-formatted 0.1 --broadcast +kohaku shield --protocol tornado --wallet testWallet --from s0 --amount-formatted 0.1 --broadcast kohaku shield --protocol tornado --wallet testWallet --from 0 --amount-max --broadcast kohaku shield --protocol railgun --wallet testWallet --from 0 --token 0x1c7D4B196Cb0C7B01d743Fbc6116a902379C7238 --amount-formatted 10 --broadcast kohaku shield --protocol tornado --wallet testWallet --from 0 --amount-formatted 0.1 --without-tor diff --git a/src/commands/shield.ts b/src/commands/shield.ts index ba841e5..0dbc2ae 100644 --- a/src/commands/shield.ts +++ b/src/commands/shield.ts @@ -4,7 +4,6 @@ import chalk from "chalk"; import type { AssetAmount } from "@kohaku-eth/plugins"; import type { Command } from "commander"; import { formatUnits, getAddress, isAddress, parseUnits } from "viem"; -import { Mnemonic } from "derive-railgun-keys"; import { makeHost } from "../host/makeHost"; import { @@ -12,8 +11,10 @@ import { formatAccountSelector, formatPublicAccountBalanceLabel, listPublicAccountsWithBalance, + parseFromIndex, partitionShieldTxs, resolveShieldApprovalCalls, + resolveShieldSender, shieldTransactionConfirmMessage, summarizeMultiShieldPlan, toShieldTxs, @@ -42,7 +43,6 @@ import { } from "../utils/fee-preview.js"; import { resolveAddressOrName } from "../utils/resolve-name.js"; import { - addressFromPrivateKey, makeWalletClient, sendTransactionAndWait, simulateCallOrThrow, @@ -71,7 +71,6 @@ import { resolveWalletPassword, } from "../utils/wallets-util"; import { readSeedKeystore } from "../utils/mnemonic"; -import { makePublicAccountsStorage } from "../utils/public-accounts"; import { assertPpErc20TokenWhitelisted, createProtocolPlugin, @@ -135,13 +134,6 @@ function etherscanTxUrl(chainId: bigint, txHash: string): string { return `https://${host}/tx/${txHash}`; } -function parseFromIndex(fromValue: string): number | null { - if (!/^\d+$/.test(fromValue)) return null; - const parsed = Number(fromValue); - if (!Number.isInteger(parsed) || parsed < 0) return null; - return parsed; -} - function findAccountWithBalance( fromValue: string, accounts: PublicAccountWithBalance[] @@ -290,7 +282,10 @@ export function registerShieldCommand(program: Command): void { ) .option("--wallet ", cliOptions.walletPickList) .option("--password ", cliOptions.password) - .option("--from ", "Public sender address or public-account index") + .option( + "--from ", + "Public sender address, HD index, or stealth selector (s0)" + ) .option( "--from-priv", "With --broadcast: derive --from index from mnemonic when missing from public accounts (not required for dry-run)" @@ -455,8 +450,6 @@ export function registerShieldCommand(program: Command): void { return; } - const publicStorage = makePublicAccountsStorage(walletDir, mnemonic, password); - const withBalances = await listPublicAccountsWithBalance( rpcUrl, walletDir, @@ -676,40 +669,21 @@ export function registerShieldCommand(program: Command): void { } } - const fromIndex = parseFromIndex(fromValue); - let senderPrivateKey: string | undefined; let senderAddress: string; - if (fromIndex !== null) { - const account = publicStorage.getAccount(fromIndex); - if (account) { - senderPrivateKey = account.priv; - senderAddress = account.address; - } else if (opts.fromPriv || dryRun) { - senderPrivateKey = Mnemonic.to0xPrivateKeyByIndex(mnemonic, fromIndex); - senderAddress = addressFromPrivateKey(senderPrivateKey); - } else { - cliError( - `Public account index ${fromIndex} not found. Use --from-priv with --broadcast to derive from mnemonic, or omit --broadcast for a dry-run.` - ); - return; - } - } else if (isAddress(fromValue)) { - senderAddress = getAddress(fromValue); - const match = publicStorage - .getAccounts() - .find((x) => x.address.toLowerCase() === senderAddress.toLowerCase()); - if (match) { - senderPrivateKey = match.priv; - } else if (dryRun) { - senderPrivateKey = undefined; - } else { - cliError( - `Address ${senderAddress} is not in this wallet's public accounts. Use --broadcast with --from-priv and an index, or omit --broadcast to preview txs for this address.` - ); - return; - } - } else { - cliError("--from must be either a valid address or a non-negative index."); + let senderPrivateKey: string | undefined; + try { + const resolved = resolveShieldSender({ + fromValue, + walletDir, + mnemonic, + password, + dryRun, + allowDeriveFromMnemonic: !!opts.fromPriv, + }); + senderAddress = resolved.senderAddress; + senderPrivateKey = resolved.senderPrivateKey; + } catch (e) { + cliErrorFromCaught(e); return; } @@ -1005,7 +979,7 @@ export function registerShieldCommand(program: Command): void { if (!senderPrivateKey) { cliError( - "Cannot sign: no private key for this --from (use a saved public account or --from-priv with --broadcast)." + "Cannot sign: no private key for this --from (use a saved public/stealth account or --from-priv with --broadcast)." ); return; } diff --git a/tests/shield-txs.test.ts b/tests/shield-txs.test.ts index 505bfe8..d5a2379 100644 --- a/tests/shield-txs.test.ts +++ b/tests/shield-txs.test.ts @@ -1,15 +1,22 @@ import assert from "node:assert/strict"; +import { mkdtempSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; import { describe, it } from "node:test"; -import { encodeFunctionData, getAddress } from "viem"; +import { encodeFunctionData, getAddress, type Hex } from "viem"; +import { privateKeyToAccount } from "viem/accounts"; import { parseFromIndex, partitionShieldTxs, + resolveShieldSender, toShieldTxs, tryDecodeErc20Approve, type ShieldCall, } from "../src/lib/shield-flow.js"; +import { makeStealthAccountsStorage } from "../src/lib/stealth/storage.js"; import { ERC20_ABI } from "../src/utils/tokens-util.js"; +import { addressFromPrivateKey } from "../src/utils/viem-tx.js"; const TOKEN = "0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48"; const POOL_A = getAddress("0x12d66f87a04a9e220743712ce6d9bb1b5616b8fc"); @@ -32,6 +39,21 @@ function depositCall(to: string, value = 0n): ShieldCall { return { to, data: "0xdead", value }; } +const MNEMONIC = + "test test test test test test test test test test test junk"; +const STEALTH_PRIV = + "0xaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" as Hex; +const STEALTH_ADDRESS = privateKeyToAccount(STEALTH_PRIV).address; + +function withWalletDir(fn: (walletDir: string) => void): void { + const dir = mkdtempSync(join(tmpdir(), "kohaku-shield-from-")); + try { + fn(dir); + } finally { + rmSync(dir, { recursive: true, force: true }); + } +} + describe("parseFromIndex", () => { it("parses a non-negative decimal HD index", () => { assert.equal(parseFromIndex("0"), 0); @@ -50,6 +72,84 @@ describe("parseFromIndex", () => { }); }); +describe("resolveShieldSender stealth --from", () => { + const senderOpts = (walletDir: string, fromValue: string, dryRun = false) => ({ + fromValue, + walletDir, + mnemonic: MNEMONIC, + password: "pw", + dryRun, + allowDeriveFromMnemonic: false, + }); + + it("resolves --from s0 and stealth:0 to the stored stealth key", () => { + withWalletDir((walletDir) => { + makeStealthAccountsStorage(walletDir, "pw").upsertAccount({ + address: STEALTH_ADDRESS, + priv: STEALTH_PRIV, + ephemeralPublicKey: + "0x02aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + schemeId: 1, + lastUpdated: 1, + ethBalance: "0", + erc20Balances: {}, + }); + + for (const fromValue of ["s0", "S0", "stealth:0"]) { + const resolved = resolveShieldSender(senderOpts(walletDir, fromValue)); + assert.equal(resolved.senderAddress, getAddress(STEALTH_ADDRESS)); + assert.equal(resolved.senderPrivateKey, STEALTH_PRIV); + } + }); + }); + + it("resolves --from to the stored stealth key", () => { + withWalletDir((walletDir) => { + makeStealthAccountsStorage(walletDir, "pw").upsertAccount({ + address: STEALTH_ADDRESS, + priv: STEALTH_PRIV, + ephemeralPublicKey: + "0x02aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + schemeId: 1, + lastUpdated: 1, + ethBalance: "0", + erc20Balances: {}, + }); + + for (const fromValue of [STEALTH_ADDRESS, STEALTH_ADDRESS.toLowerCase()]) { + const resolved = resolveShieldSender( + senderOpts(walletDir, fromValue, false) + ); + assert.equal(resolved.senderAddress, getAddress(STEALTH_ADDRESS)); + assert.equal(resolved.senderPrivateKey, STEALTH_PRIV); + } + }); + }); + + it("rejects a missing stealth selector", () => { + withWalletDir((walletDir) => { + assert.throws( + () => resolveShieldSender(senderOpts(walletDir, "s0")), + /Stealth account s0 not found/ + ); + }); + }); + + it("still derives an HD index on dry-run when the public account is missing", () => { + withWalletDir((walletDir) => { + const resolved = resolveShieldSender(senderOpts(walletDir, "0", true)); + assert.equal( + resolved.senderAddress, + addressFromPrivateKey( + // index 0 of the well-known test mnemonic + "0xac0974bec39a17e36ba4a6b4d238ff944bacb478cbed5efcae784d7bf4f2ff80" + ) + ); + assert.ok(resolved.senderPrivateKey); + }); + }); +}); + describe("tryDecodeErc20Approve", () => { it("decodes spender and amount from approve calldata", () => { const data = encodeFunctionData({