Skip to content

Commit

Permalink
feat(utxo-lib): build transaction without nonWitnessUtxo when specified
Browse files Browse the repository at this point in the history 
TICKET: BTC-606
  • Loading branch information
@davidkaplanbitgo
davidkaplanbitgo committed Nov 8, 2023
1 parent 1a75e88 commit 5ec3759
Show file tree
Hide file tree
Showing 4 changed files with 109 additions and 25 deletions.
4 changes: 2 additions & 2 deletions modules/utxo-lib/src/bitgo/wallet/Psbt.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ import { UtxoTransaction } from '../UtxoTransaction';
import { createOutputScript2of3, getLeafHash, scriptTypeForChain, toXOnlyPublicKey } from '../outputScripts';
import { DerivedWalletKeys, RootWalletKeys } from './WalletKeys';
import { toPrevOutputWithPrevTx } from '../Unspent';
import { createPsbtFromTransaction } from '../transaction';
import { createPsbtFromHex, createPsbtFromTransaction } from '../transaction';
import { isWalletUnspent, WalletUnspent } from './Unspent';

import {
Expand Down Expand Up @@ -556,7 +556,7 @@ export function extractP2msOnlyHalfSignedTx(psbt: UtxoPsbt): UtxoTransaction<big
* It is not BIP-174 compliant, so use it carefully.
*/
export function clonePsbtWithoutNonWitnessUtxo(psbt: UtxoPsbt): UtxoPsbt {
const newPsbt = psbt.clone();
const newPsbt = createPsbtFromHex(psbt.toHex(), psbt.network);
const txInputs = psbt.txInputs;

psbt.data.inputs.forEach((input, i) => {
Expand Down
38 changes: 30 additions & 8 deletions modules/utxo-lib/src/bitgo/wallet/Unspent.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -163,7 +163,8 @@ export function updateReplayProtectionUnspentToPsbt(
psbt: UtxoPsbt,
inputIndex: number,
u: Unspent<bigint>,
redeemScript?: Buffer
redeemScript?: Buffer,
customParams?: { skipNonWitnessUtxo?: boolean }
): void {
if (!psbtIncludesUnspentAtIndex(psbt, inputIndex, u.id)) {
throw new Error(`unspent does not correspond to psbt input`);
Expand All @@ -178,12 +179,15 @@ export function updateReplayProtectionUnspentToPsbt(
// with the previous transaction. Therefore, we can treat Zcash non-segwit transactions as Bitcoin
// segwit transactions
const isZcash = getMainnet(psbt.network) === networks.zcash;
if (!isUnspentWithPrevTx(u) && !isZcash) {
if (!isUnspentWithPrevTx(u) && !isZcash && !customParams?.skipNonWitnessUtxo) {
throw new Error('Error, require previous tx to add to PSBT');
}
if (isZcash && !input.witnessUtxo) {
const { script, value } = toPrevOutput(u, psbt.network);
psbt.updateInput(inputIndex, { witnessUtxo: { script, value } });
} else if (customParams?.skipNonWitnessUtxo) {
const { script, value } = toPrevOutput(u, psbt.network);
psbt.updateInput(inputIndex, { witnessUtxo: { script, value } });
} else if (!isZcash && !input.nonWitnessUtxo) {
psbt.updateInput(inputIndex, { nonWitnessUtxo: (u as UnspentWithPrevTx<bigint>).prevTx });
}
Expand All @@ -202,28 +206,38 @@ function addUnspentToPsbt(
});
}

export function addReplayProtectionUnspentToPsbt(psbt: UtxoPsbt, u: Unspent<bigint>, redeemScript: Buffer): void {
export function addReplayProtectionUnspentToPsbt(
psbt: UtxoPsbt,
u: Unspent<bigint>,
redeemScript: Buffer,
customParams?: { skipNonWitnessUtxo?: boolean }
): void {
addUnspentToPsbt(psbt, u.id);
updateReplayProtectionUnspentToPsbt(psbt, psbt.inputCount - 1, u, redeemScript);
updateReplayProtectionUnspentToPsbt(psbt, psbt.inputCount - 1, u, redeemScript, customParams);
}

/**
* Update the PSBT with the unspent data for the input at the given index if the data is not there already.
*
* If skipNonWitnessUtxo is true, then the nonWitnessUtxo will not be added for an input that requires it (e.g. non-segwit)
* and instead the witnessUtxo will be added
*
* @param psbt
* @param inputIndex
* @param u
* @param rootWalletKeys
* @param signer
* @param cosigner
* @param customParams
*/
export function updateWalletUnspentForPsbt(
psbt: UtxoPsbt,
inputIndex: number,
u: WalletUnspent<bigint>,
rootWalletKeys: RootWalletKeys,
signer: KeyName,
cosigner: KeyName
cosigner: KeyName,
customParams?: { skipNonWitnessUtxo?: boolean }
): void {
if (!psbtIncludesUnspentAtIndex(psbt, inputIndex, u.id)) {
throw new Error(`unspent does not correspond to psbt input`);
Expand All @@ -234,7 +248,7 @@ export function updateWalletUnspentForPsbt(
// with the previous transaction. Therefore, we can treat Zcash non-segwit transactions as Bitcoin
// segwit transactions
const isZcashOrSegwit = isSegwit(u.chain) || getMainnet(psbt.network) === networks.zcash;
if (isZcashOrSegwit && !input.witnessUtxo) {
if ((isZcashOrSegwit && !input.witnessUtxo) || customParams?.skipNonWitnessUtxo) {
const { script, value } = toPrevOutput(u, psbt.network);
psbt.updateInput(inputIndex, { witnessUtxo: { script, value } });
} else if (!isZcashOrSegwit) {
Expand Down Expand Up @@ -354,10 +368,18 @@ export function addWalletUnspentToPsbt(
rootWalletKeys: RootWalletKeys,
signer: KeyName,
cosigner: KeyName,
customParams?: { sequenceNumber?: number }
customParams?: { sequenceNumber?: number; skipNonWitnessUtxo?: boolean }
): void {
addUnspentToPsbt(psbt, u.id, {
sequenceNumber: customParams ? customParams.sequenceNumber : TX_INPUT_SEQUENCE_NUMBER_FINAL,
});
updateWalletUnspentForPsbt(psbt, psbt.inputCount - 1, u, rootWalletKeys, signer, cosigner);
updateWalletUnspentForPsbt(
psbt,
psbt.inputCount - 1,
u,
rootWalletKeys,
signer,
cosigner,
customParams ? { skipNonWitnessUtxo: customParams.skipNonWitnessUtxo } : {}
);
}
38 changes: 27 additions & 11 deletions modules/utxo-lib/src/testutil/psbt.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ import {
UtxoPsbt,
UtxoTransaction,
verifySignatureWithUnspent,
withUnsafeNonSegwit,
} from '../bitgo';
import { Network } from '../networks';
import { mockReplayProtectionUnspent, mockWalletUnspent } from './mock';
Expand Down Expand Up @@ -116,19 +117,32 @@ export function signPsbtInput(
params?: {
signers?: { signerName: KeyName; cosignerName?: KeyName };
deterministic?: boolean;
skipNonWitnessUtxo?: boolean;
}
): void {
const { signers, deterministic } = params ?? {};
function signPsbt(psbt: UtxoPsbt, signFunc: () => void, skipNonWitnessUtxo?: boolean) {
if (skipNonWitnessUtxo) {
withUnsafeNonSegwit(psbt, signFunc);
} else {
signFunc();
}
}

const { signers, deterministic, skipNonWitnessUtxo } = params ?? {};
const { signerName, cosignerName } = signers ? signers : getSigners(input.scriptType);
if (sign === 'halfsigned') {
if (input.scriptType === 'p2shP2pk') {
psbt.signInput(inputIndex, rootWalletKeys[signerName]);
signPsbt(psbt, () => psbt.signInput(inputIndex, rootWalletKeys[signerName]), skipNonWitnessUtxo);
} else {
psbt.signInputHD(inputIndex, rootWalletKeys[signerName]);
signPsbt(psbt, () => psbt.signInputHD(inputIndex, rootWalletKeys[signerName]), skipNonWitnessUtxo);
}
}
if (sign === 'fullsigned' && cosignerName && input.scriptType !== 'p2shP2pk') {
psbt.signInputHD(inputIndex, rootWalletKeys[cosignerName], { deterministic });
signPsbt(
psbt,
() => psbt.signInputHD(inputIndex, rootWalletKeys[cosignerName], { deterministic }),
skipNonWitnessUtxo
);
}
}

Expand All @@ -144,11 +158,12 @@ export function signAllPsbtInputs(
params?: {
signers?: { signerName: KeyName; cosignerName?: KeyName };
deterministic?: boolean;
skipNonWitnessUtxo?: boolean;
}
): void {
const { signers, deterministic } = params ?? {};
const { signers, deterministic, skipNonWitnessUtxo } = params ?? {};
inputs.forEach((input, inputIndex) => {
signPsbtInput(psbt, input, inputIndex, rootWalletKeys, sign, { signers, deterministic });
signPsbtInput(psbt, input, inputIndex, rootWalletKeys, sign, { signers, deterministic, skipNonWitnessUtxo });
});
}

Expand All @@ -164,9 +179,10 @@ export function constructPsbt(
params?: {
signers?: { signerName: KeyName; cosignerName?: KeyName };
deterministic?: boolean;
skipNonWitnessUtxo?: boolean;
}
): UtxoPsbt {
const { signers, deterministic } = params ?? {};
const { signers, deterministic, skipNonWitnessUtxo } = params ?? {};
const totalInputAmount = inputs.reduce((sum, input) => sum + input.value, BigInt(0));
const outputInputAmount = outputs.reduce((sum, output) => sum + output.value, BigInt(0));
assert(totalInputAmount >= outputInputAmount, 'total output can not exceed total input');
Expand All @@ -181,11 +197,11 @@ export function constructPsbt(
unspents.forEach((u, i) => {
const { signerName, cosignerName } = signers ? signers : getSigners(inputs[i].scriptType);
if (isWalletUnspent(u) && cosignerName) {
addWalletUnspentToPsbt(psbt, u, rootWalletKeys, signerName, cosignerName);
addWalletUnspentToPsbt(psbt, u, rootWalletKeys, signerName, cosignerName, { skipNonWitnessUtxo });
} else {
const { redeemScript } = createOutputScriptP2shP2pk(rootWalletKeys[signerName].publicKey);
assert(redeemScript);
addReplayProtectionUnspentToPsbt(psbt, u, redeemScript);
addReplayProtectionUnspentToPsbt(psbt, u, redeemScript, { skipNonWitnessUtxo });
}
});

Expand All @@ -211,10 +227,10 @@ export function constructPsbt(
psbt.setAllInputsMusig2NonceHD(rootWalletKeys['user']);
psbt.setAllInputsMusig2NonceHD(rootWalletKeys['bitgo'], { deterministic });

signAllPsbtInputs(psbt, inputs, rootWalletKeys, 'halfsigned', { signers });
signAllPsbtInputs(psbt, inputs, rootWalletKeys, 'halfsigned', { signers, skipNonWitnessUtxo });

if (sign === 'fullsigned') {
signAllPsbtInputs(psbt, inputs, rootWalletKeys, sign, { signers, deterministic });
signAllPsbtInputs(psbt, inputs, rootWalletKeys, sign, { signers, deterministic, skipNonWitnessUtxo });
}

return psbt;
Expand Down
54 changes: 50 additions & 4 deletions modules/utxo-lib/test/bitgo/psbt/Psbt.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -190,7 +190,12 @@ function runExtractP2msOnlyHalfSignedTxTest(network: Network, inputs: Input[], o
});
}

function runBuildSignSendFlowTest(network: Network, inputs: Input[], outputs: Output[]) {
function runBuildSignSendFlowTest(
network: Network,
inputs: Input[],
outputs: Output[],
{ skipNonWitnessUtxo = false } = {}
) {
const coin = getNetworkName(network);

function assertValidate(psbt: UtxoPsbt) {
Expand All @@ -204,14 +209,24 @@ function runBuildSignSendFlowTest(network: Network, inputs: Input[], outputs: Ou
}

describe(`Build, sign & send flow for ${coin}`, function () {
it(`success for ${coin}`, function () {
const psbt = testutil.constructPsbt(inputs, outputs, network, rootWalletKeys, 'unsigned', {
/**
* Skip adding nonWitnessUtxos to psbts
* ------------------------------------
* In the instance that we want to doing a bulk sweep, for network and client performance reasons we are substituting
* the nonWitnessUtxo for p2sh and p2shP2pk inputs with a witnessUtxo. We need the witnessUtxo so that we can half
* sign the transaction locally with the user key. When we send the half signed to BitGo, the PSBT will be properly
* populated such that the non-segwit inputs have the nonWitnessUtxo. This means when we send it to BitGo we should
* remove the witnessUtxo so that it just has the partialSig and redeemScript.
*/
it(`success for ${coin}${skipNonWitnessUtxo ? ' without nonWitnessUtxo for p2sh' : ''}`, function () {
const parentPsbt = testutil.constructPsbt(inputs, outputs, network, rootWalletKeys, 'unsigned', {
signers: {
signerName: 'user',
cosignerName: 'bitgo',
},
});

let psbt = skipNonWitnessUtxo ? clonePsbtWithoutNonWitnessUtxo(parentPsbt) : parentPsbt;
addXpubsToPsbt(psbt, rootWalletKeys);
psbt.setAllInputsMusig2NonceHD(rootWalletKeys['user']);

Expand All @@ -231,6 +246,7 @@ function runBuildSignSendFlowTest(network: Network, inputs: Input[], outputs: Ou
signerName: 'user',
cosignerName: 'bitgo',
},
skipNonWitnessUtxo,
});

psbtWithoutPrevTx = clonePsbtWithoutNonWitnessUtxo(psbt);
Expand All @@ -253,6 +269,10 @@ function runBuildSignSendFlowTest(network: Network, inputs: Input[], outputs: Ou

psbtFromHsm = createPsbtFromHex(hexAtHsm, network);
deleteWitnessUtxoForNonSegwitInputs(psbtFromHsm);

if (skipNonWitnessUtxo) {
psbt = parentPsbt;
}
psbt.combine(psbtFromHsm);

assertValidate(psbt);
Expand All @@ -261,6 +281,28 @@ function runBuildSignSendFlowTest(network: Network, inputs: Input[], outputs: Ou
});
}

function runBuildPsbtWithSDK(network: Network, inputs: Input[], outputs: Output[]) {
const coin = getNetworkName(network);
it(`check that building a PSBT while skipping nonWitnessUtxo works - ${coin}`, async function () {
const psbtWithNonWitness = testutil.constructPsbt(inputs, outputs, network, rootWalletKeys, 'unsigned', {
signers: {
signerName: 'user',
cosignerName: 'bitgo',
},
});
const psbtWithoutNonWitness = testutil.constructPsbt(inputs, outputs, network, rootWalletKeys, 'unsigned', {
signers: {
signerName: 'user',
cosignerName: 'bitgo',
},
skipNonWitnessUtxo: true,
});

const clonedPsbt = clonePsbtWithoutNonWitnessUtxo(psbtWithNonWitness);
assert.deepStrictEqual(psbtWithoutNonWitness.toHex(), clonedPsbt.toHex());
});
}

getNetworkList()
.filter((v) => isMainnet(v) && v !== networks.bitcoinsv)
.forEach((network) => {
Expand All @@ -274,7 +316,11 @@ getNetworkList()
isSupportedScriptType(network, input.scriptType === 'taprootKeyPathSpend' ? 'p2trMusig2' : input.scriptType)
);
const supportedPsbtOutputs = psbtOutputs.filter((output) => isSupportedScriptType(network, output.scriptType));
runBuildSignSendFlowTest(network, supportedPsbtInputs, supportedPsbtOutputs);
[false, true].forEach((skipNonWitnessUtxo) =>
runBuildSignSendFlowTest(network, supportedPsbtInputs, supportedPsbtOutputs, { skipNonWitnessUtxo })
);

runBuildPsbtWithSDK(network, supportedPsbtInputs, supportedPsbtOutputs);
});

describe('isTransactionWithKeyPathSpendInput', function () {
Expand Down

0 comments on commit 5ec3759

Please sign in to comment.