diff --git a/sweepbatcher/sweep_batch.go b/sweepbatcher/sweep_batch.go index f78996c..549d8f8 100644 --- a/sweepbatcher/sweep_batch.go +++ b/sweepbatcher/sweep_batch.go @@ -7,9 +7,11 @@ import ( "errors" "fmt" "math" + "strings" "sync" "time" + "github.com/btcsuite/btcd/blockchain" "github.com/btcsuite/btcd/btcec/v2" "github.com/btcsuite/btcd/btcec/v2/schnorr/musig2" "github.com/btcsuite/btcd/btcutil" @@ -107,6 +109,11 @@ type sweep struct { // has to be spent using preimage. This is only used in fee estimations // when selecting a batch for the sweep to minimize fees. nonCoopHint bool + + // coopFailed is set, if we have tried to spend the sweep cooperatively, + // but it failed. We try to spend a sweep cooperatively only once. This + // status is not persisted in the DB. + coopFailed bool } // batchState is the state of the batch. @@ -160,6 +167,16 @@ type batchConfig struct { // Note that musig2SignSweep must be nil in this case, however signer // client must still be provided, as it is used for non-coop spendings. customMuSig2Signer SignMuSig2 + + // mixedBatch instructs sweepbatcher to create mixed batches with regard + // to cooperativeness. Such a batch can include sweeps signed both + // cooperatively and non-cooperatively. If cooperative signing fails for + // a sweep, transaction is updated to sign that sweep non-cooperatively + // and another round of cooperative signing runs on the remaining + // sweeps. The remaining sweeps are signed in non-cooperative (more + // expensive) way. If the whole procedure fails for whatever reason, the + // batch is signed non-cooperatively (the fallback). + mixedBatch bool } // rbfCache stores data related to our last fee bump. @@ -421,13 +438,18 @@ func (b *batch) addSweep(ctx context.Context, sweep *sweep) (bool, error) { // Before we run through the acceptance checks, let's just see if this // sweep is already in our batch. In that case, just update the sweep. - _, ok := b.sweeps[sweep.swapHash] + oldSweep, ok := b.sweeps[sweep.swapHash] if ok { + // Preserve coopFailed value not to forget about cooperative + // spending failure in this sweep. + tmp := *sweep + tmp.coopFailed = oldSweep.coopFailed + // If the sweep was resumed from storage, and the swap requested // to sweep again, a new sweep notifier will be created by the // swap. By re-assigning to the batch's sweep we make sure that // everything, including the notifier, is up to date. - b.sweeps[sweep.swapHash] = *sweep + b.sweeps[sweep.swapHash] = tmp // If this is the primary sweep, we also need to update the // batch's confirmation target and fee rate. @@ -724,7 +746,7 @@ func (b *batch) publish(ctx context.Context) error { var ( err error fee btcutil.Amount - coopSuccess bool + signSuccess bool ) if len(b.sweeps) == 0 { @@ -739,12 +761,19 @@ func (b *batch) publish(ctx context.Context) error { return err } - fee, err, coopSuccess = b.publishBatchCoop(ctx) - if err != nil { - b.log.Warnf("co-op publish error: %v", err) + if b.cfg.mixedBatch { + fee, err, signSuccess = b.publishMixedBatch(ctx) + if err != nil { + b.log.Warnf("Mixed batch publish error: %v", err) + } + } else { + fee, err, signSuccess = b.publishBatchCoop(ctx) + if err != nil { + b.log.Warnf("co-op publish error: %v", err) + } } - if !coopSuccess { + if !signSuccess { fee, err = b.publishBatch(ctx) } if err != nil { @@ -1091,6 +1120,361 @@ func (b *batch) publishBatchCoop(ctx context.Context) (btcutil.Amount, return fee, nil, true } +// constructUnsignedTx creates unsigned tx from the sweeps, paying to the addr. +// It also returns absolute fee (from weight and clamped). +func (b *batch) constructUnsignedTx(sweeps []sweep, + address btcutil.Address) (*wire.MsgTx, lntypes.WeightUnit, + btcutil.Amount, btcutil.Amount, error) { + + // Sanity check, there should be at least 1 sweep in this batch. + if len(sweeps) == 0 { + return nil, 0, 0, 0, fmt.Errorf("no sweeps in batch") + } + + // Create the batch transaction. + batchTx := &wire.MsgTx{ + Version: 2, + LockTime: uint32(b.currentHeight), + } + + // Add transaction inputs and estimate its weight. + var weightEstimate input.TxWeightEstimator + for _, sweep := range sweeps { + if sweep.nonCoopHint || sweep.coopFailed { + // Non-cooperative sweep. + batchTx.AddTxIn(&wire.TxIn{ + PreviousOutPoint: sweep.outpoint, + Sequence: sweep.htlc.SuccessSequence(), + }) + + err := sweep.htlcSuccessEstimator(&weightEstimate) + if err != nil { + return nil, 0, 0, 0, fmt.Errorf("sweep."+ + "htlcSuccessEstimator failed: %w", err) + } + } else { + // Cooperative sweep. + batchTx.AddTxIn(&wire.TxIn{ + PreviousOutPoint: sweep.outpoint, + }) + + weightEstimate.AddTaprootKeySpendInput( + txscript.SigHashDefault, + ) + } + } + + // Convert the destination address to pkScript. + batchPkScript, err := txscript.PayToAddrScript(address) + if err != nil { + return nil, 0, 0, 0, fmt.Errorf("txscript.PayToAddrScript "+ + "failed: %w", err) + } + + // Add the output to weight estimates. + err = sweeppkg.AddOutputEstimate(&weightEstimate, address) + if err != nil { + return nil, 0, 0, 0, fmt.Errorf("sweep.AddOutputEstimate "+ + "failed: %w", err) + } + + // Keep track of the total amount this batch is sweeping back. + batchAmt := btcutil.Amount(0) + for _, sweep := range sweeps { + batchAmt += sweep.value + } + + // Find weight and fee. + weight := weightEstimate.Weight() + feeForWeight := b.rbfCache.FeeRate.FeeForWeight(weight) + + // Clamp the calculated fee to the max allowed fee amount for the batch. + fee := clampBatchFee(feeForWeight, batchAmt) + + // Add the batch transaction output, which excludes the fees paid to + // miners. + batchTx.AddTxOut(&wire.TxOut{ + PkScript: batchPkScript, + Value: int64(batchAmt - fee), + }) + + return batchTx, weight, feeForWeight, fee, nil +} + +// publishMixedBatch constructs and publishes a batch transaction that can +// include sweeps spent both cooperatively and non-cooperatively. If a sweep is +// marked with nonCoopHint or coopFailed flags, it is spent non-cooperatively. +// If a cooperative sweep fails to sign cooperatively, the whole transaction +// is re-signed again, with this sweep signing non-cooperatively. This process +// is optimized, trying to detect all non-cooperative sweeps in one round. The +// function returns the absolute fee. The last result of the function indicates +// if signing succeeded. +func (b *batch) publishMixedBatch(ctx context.Context) (btcutil.Amount, error, + bool) { + + // Sanity check, there should be at least 1 sweep in this batch. + if len(b.sweeps) == 0 { + return 0, fmt.Errorf("no sweeps in batch"), false + } + + // Append this sweep to an array of sweeps. This is needed to keep the + // order of sweeps stored, as iterating the sweeps map does not + // guarantee same order. + sweeps := make([]sweep, 0, len(b.sweeps)) + for _, sweep := range b.sweeps { + sweeps = append(sweeps, sweep) + } + + // Determine if an external address is used. + addrOverride := false + for _, sweep := range sweeps { + if sweep.isExternalAddr { + addrOverride = true + } + } + + // Find destination address. + var address btcutil.Address + if addrOverride { + // Sanity check, there should be exactly 1 sweep in this batch. + if len(sweeps) != 1 { + return 0, fmt.Errorf("external address sweep batched " + + "with other sweeps"), false + } + + address = sweeps[0].destAddr + } else { + var err error + address, err = b.getBatchDestAddr(ctx) + if err != nil { + return 0, err, false + } + } + + // Each iteration of this loop is one attempt to sign the transaction + // cooperatively. We try cooperative signing only for the sweeps not + // known in advance to be non-cooperative (nonCoopHint) and not failed + // to sign cooperatively in previous rounds (coopFailed). If any of them + // fails, the sweep is excluded from all following rounds and another + // round is attempted. Otherwise the cycle completes and we sign the + // remaining sweeps non-cooperatively. + var ( + tx *wire.MsgTx + weight lntypes.WeightUnit + feeForWeight btcutil.Amount + fee btcutil.Amount + coopInputs int + ) + for attempt := 1; ; attempt++ { + b.log.Infof("Attempt %d of collecting cooperative signatures.", + attempt) + + // Construct unsigned batch transaction. + var err error + tx, weight, feeForWeight, fee, err = b.constructUnsignedTx( + sweeps, address, + ) + if err != nil { + return 0, fmt.Errorf("failed to construct tx: %w", err), + false + } + + // Create PSBT and prevOutsMap. + psbtBytes, prevOutsMap, err := b.createPsbt(tx, sweeps) + if err != nil { + return 0, fmt.Errorf("createPsbt failed: %w", err), + false + } + + // Keep track if any new sweep failed to sign cooperatively. + newCoopFailures := false + + // Try to sign all cooperative sweeps first. + coopInputs = 0 + for i, sweep := range sweeps { + // Skip non-cooperative sweeps. + if sweep.nonCoopHint || sweep.coopFailed { + continue + } + + // Try to sign the sweep cooperatively. + finalSig, err := b.musig2sign( + ctx, i, sweep, tx, prevOutsMap, psbtBytes, + ) + if err != nil { + b.log.Infof("cooperative signing failed for "+ + "sweep %x: %v", sweep.swapHash[:6], err) + + // Set coopFailed flag for this sweep in all the + // places we store the sweep. + sweep.coopFailed = true + sweeps[i] = sweep + b.sweeps[sweep.swapHash] = sweep + + // Update newCoopFailures to know if we need + // another attempt of cooperative signing. + newCoopFailures = true + } else { + // Put the signature to witness of the input. + tx.TxIn[i].Witness = wire.TxWitness{finalSig} + coopInputs++ + } + } + + // If there was any failure of cooperative signing, we need to + // update weight estimates (since non-cooperative signing has + // larger witness) and hence update the whole transaction and + // all the signatures. Otherwise we complete cooperative part. + if !newCoopFailures { + break + } + } + + // Calculate the expected number of non-cooperative sweeps. + nonCoopInputs := len(sweeps) - coopInputs + + // Now sign the remaining sweeps' inputs non-cooperatively. + // For that, first collect sign descriptors for the signatures. + // Also collect prevOuts for all inputs. + signDescs := make([]*lndclient.SignDescriptor, 0, nonCoopInputs) + prevOutsList := make([]*wire.TxOut, 0, len(sweeps)) + for i, sweep := range sweeps { + // Create and store the previous outpoint for this sweep. + prevOut := &wire.TxOut{ + Value: int64(sweep.value), + PkScript: sweep.htlc.PkScript, + } + prevOutsList = append(prevOutsList, prevOut) + + // Skip cooperative sweeps. + if !sweep.nonCoopHint && !sweep.coopFailed { + continue + } + + key, err := btcec.ParsePubKey( + sweep.htlcKeys.ReceiverScriptKey[:], + ) + if err != nil { + return 0, fmt.Errorf("btcec.ParsePubKey failed: %w", + err), false + } + + // Create and store the sign descriptor for this sweep. + signDesc := lndclient.SignDescriptor{ + WitnessScript: sweep.htlc.SuccessScript(), + Output: prevOut, + HashType: sweep.htlc.SigHash(), + InputIndex: i, + KeyDesc: keychain.KeyDescriptor{ + PubKey: key, + }, + } + + if sweep.htlc.Version == swap.HtlcV3 { + signDesc.SignMethod = input.TaprootScriptSpendSignMethod + } + + signDescs = append(signDescs, &signDesc) + } + + // Sanity checks. + if len(signDescs) != nonCoopInputs { + // This must not happen by construction. + return 0, fmt.Errorf("unexpected size of signDescs: %d != %d", + len(signDescs), nonCoopInputs), false + } + if len(prevOutsList) != len(sweeps) { + // This must not happen by construction. + return 0, fmt.Errorf("unexpected size of prevOutsList: "+ + "%d != %d", len(prevOutsList), len(sweeps)), false + } + + var rawSigs [][]byte + if nonCoopInputs > 0 { + // Produce the signatures for our inputs using sign descriptors. + var err error + rawSigs, err = b.signerClient.SignOutputRaw( + ctx, tx, signDescs, prevOutsList, + ) + if err != nil { + return 0, fmt.Errorf("signerClient.SignOutputRaw "+ + "failed: %w", err), false + } + } + + // Sanity checks. + if len(rawSigs) != nonCoopInputs { + // This must not happen by construction. + return 0, fmt.Errorf("unexpected size of rawSigs: %d != %d", + len(rawSigs), nonCoopInputs), false + } + + // Generate success witnesses for non-cooperative sweeps. + sigIndex := 0 + for i, sweep := range sweeps { + // Skip cooperative sweeps. + if !sweep.nonCoopHint && !sweep.coopFailed { + continue + } + + witness, err := sweep.htlc.GenSuccessWitness( + rawSigs[sigIndex], sweep.preimage, + ) + if err != nil { + return 0, fmt.Errorf("sweep.htlc.GenSuccessWitness "+ + "failed: %w", err), false + } + sigIndex++ + + // Add the success witness to our batch transaction's inputs. + tx.TxIn[i].Witness = witness + } + + // Log transaction's details. + var coopHexs, nonCoopHexs []string + for _, sweep := range sweeps { + swapHex := fmt.Sprintf("%x", sweep.swapHash[:6]) + if sweep.nonCoopHint || sweep.coopFailed { + nonCoopHexs = append(nonCoopHexs, swapHex) + } else { + coopHexs = append(coopHexs, swapHex) + } + } + txHash := tx.TxHash() + b.log.Infof("attempting to publish mixed tx=%v with feerate=%v, "+ + "weight=%v, feeForWeight=%v, fee=%v, sweeps=%d, "+ + "%d cooperative: (%s) and %d non-cooperative (%s), destAddr=%s", + txHash, b.rbfCache.FeeRate, weight, feeForWeight, fee, + len(tx.TxIn), coopInputs, strings.Join(coopHexs, ", "), + nonCoopInputs, strings.Join(nonCoopHexs, ", "), address) + + b.debugLogTx("serialized mixed batch", tx) + + // Make sure tx weight matches the expected value. + realWeight := lntypes.WeightUnit( + blockchain.GetTransactionWeight(btcutil.NewTx(tx)), + ) + if realWeight != weight { + b.log.Warnf("actual weight of tx %v is %v, estimated as %d", + txHash, realWeight, weight) + } + + // Publish the transaction. + err := b.wallet.PublishTransaction( + ctx, tx, b.cfg.txLabeler(b.id), + ) + if err != nil { + return 0, fmt.Errorf("publishing tx failed: %w", err), true + } + + // Store the batch transaction's txid and pkScript, for monitoring + // purposes. + b.batchTxid = &txHash + b.batchPkScript = tx.TxOut[0].PkScript + + return fee, nil, true +} + func (b *batch) debugLogTx(msg string, tx *wire.MsgTx) { // Serialize the transaction and convert to hex string. buf := bytes.NewBuffer(make([]byte, 0, tx.SerializeSize())) diff --git a/sweepbatcher/sweep_batcher.go b/sweepbatcher/sweep_batcher.go index fb6f266..6165b50 100644 --- a/sweepbatcher/sweep_batcher.go +++ b/sweepbatcher/sweep_batcher.go @@ -286,6 +286,16 @@ type Batcher struct { // Note that musig2SignSweep must be nil in this case, however signer // client must still be provided, as it is used for non-coop spendings. customMuSig2Signer SignMuSig2 + + // mixedBatch instructs sweepbatcher to create mixed batches with regard + // to cooperativeness. Such a batch can include sweeps signed both + // cooperatively and non-cooperatively. If cooperative signing fails for + // a sweep, transaction is updated to sign that sweep non-cooperatively + // and another round of cooperative signing runs on the remaining + // sweeps. The remaining sweeps are signed in non-cooperative (more + // expensive) way. If the whole procedure fails for whatever reason, the + // batch is signed non-cooperatively (the fallback). + mixedBatch bool } // BatcherConfig holds batcher configuration. @@ -321,6 +331,16 @@ type BatcherConfig struct { // Note that musig2SignSweep must be nil in this case, however signer // client must still be provided, as it is used for non-coop spendings. customMuSig2Signer SignMuSig2 + + // mixedBatch instructs sweepbatcher to create mixed batches with regard + // to cooperativeness. Such a batch can include sweeps signed both + // cooperatively and non-cooperatively. If cooperative signing fails for + // a sweep, transaction is updated to sign that sweep non-cooperatively + // and another round of cooperative signing runs on the remaining + // sweeps. The remaining sweeps are signed in non-cooperative (more + // expensive) way. If the whole procedure fails for whatever reason, the + // batch is signed non-cooperatively (the fallback). + mixedBatch bool } // BatcherOption configures batcher behaviour. @@ -386,6 +406,20 @@ func WithCustomSignMuSig2(customMuSig2Signer SignMuSig2) BatcherOption { } } +// WithMixedBatch instructs sweepbatcher to create mixed batches with +// regard to cooperativeness. Such a batch can include both sweeps signed +// both cooperatively and non-cooperatively. If cooperative signing fails +// for a sweep, transaction is updated to sign that sweep non-cooperatively +// and another round of cooperative signing runs on the remaining sweeps. +// The remaining sweeps are signed in non-cooperative (more expensive) way. +// If the whole procedure fails for whatever reason, the batch is signed +// non-cooperatively (the fallback). +func WithMixedBatch() BatcherOption { + return func(cfg *BatcherConfig) { + cfg.mixedBatch = true + } +} + // NewBatcher creates a new Batcher instance. func NewBatcher(wallet lndclient.WalletKitClient, chainNotifier lndclient.ChainNotifierClient, @@ -433,6 +467,7 @@ func NewBatcher(wallet lndclient.WalletKitClient, customFeeRate: cfg.customFeeRate, txLabeler: cfg.txLabeler, customMuSig2Signer: cfg.customMuSig2Signer, + mixedBatch: cfg.mixedBatch, } } @@ -1050,6 +1085,7 @@ func (b *Batcher) newBatchConfig(maxTimeoutDistance int32) batchConfig { txLabeler: b.txLabeler, customMuSig2Signer: b.customMuSig2Signer, clock: b.clock, + mixedBatch: b.mixedBatch, } } diff --git a/sweepbatcher/sweep_batcher_test.go b/sweepbatcher/sweep_batcher_test.go index ef00170..e48100e 100644 --- a/sweepbatcher/sweep_batcher_test.go +++ b/sweepbatcher/sweep_batcher_test.go @@ -2880,6 +2880,552 @@ func testCustomSignMuSig2(t *testing.T, store testStore, checkBatcherError(t, runErr) } +// testWithMixedBatch tests mixed batches construction. It also tests +// non-cooperative sweeping (using a preimage). Sweeps are added one by one. +func testWithMixedBatch(t *testing.T, store testStore, + batcherStore testBatcherStore) { + + defer test.Guard(t)() + + lnd := test.NewMockLnd() + ctx, cancel := context.WithCancel(context.Background()) + + // Extract payment address from the invoice. + swapPaymentAddr, err := utils.ObtainSwapPaymentAddr( + swapInvoice, lnd.ChainParams, + ) + require.NoError(t, err) + + // Use sweepFetcher to provide NonCoopHint for swapHash1. + sweepFetcher := &sweepFetcherMock{ + store: map[lntypes.Hash]*SweepInfo{}, + } + + // Create 3 sweeps: + // 1. known in advance to be non-cooperative, + // 2. fails cosigning during an attempt, + // 3. co-signs successfully. + + // Create 3 preimages, for 3 sweeps. + var preimages = []lntypes.Preimage{ + {1}, + {2}, + {3}, + } + + // Swap hashes must match the preimages, for non-cooperative spending + // path to work. + var swapHashes = []lntypes.Hash{ + preimages[0].Hash(), + preimages[1].Hash(), + preimages[2].Hash(), + } + + // Create muSig2SignSweep working only for 3rd swapHash. + muSig2SignSweep := func(ctx context.Context, + protocolVersion loopdb.ProtocolVersion, swapHash lntypes.Hash, + paymentAddr [32]byte, nonce []byte, sweepTxPsbt []byte, + prevoutMap map[wire.OutPoint]*wire.TxOut) ( + []byte, []byte, error) { + + if swapHash == swapHashes[2] { + return nil, nil, nil + } else { + return nil, nil, fmt.Errorf("test error") + } + } + + // Use mixed batches. + batcher := NewBatcher(lnd.WalletKit, lnd.ChainNotifier, lnd.Signer, + muSig2SignSweep, testVerifySchnorrSig, lnd.ChainParams, + batcherStore, sweepFetcher, WithMixedBatch()) + + var wg sync.WaitGroup + wg.Add(1) + + var runErr error + go func() { + defer wg.Done() + runErr = batcher.Run(ctx) + }() + + // Wait for the batcher to be initialized. + <-batcher.initDone + + // Expected weights for transaction having 1, 2, and 3 sweeps. + wantWeights := []lntypes.WeightUnit{559, 952, 1182} + + // Two non-cooperative sweeps, one cooperative. + wantWitnessSizes := []int{4, 4, 1} + + // Create 3 swaps and 3 sweeps. + for i, swapHash := range swapHashes { + // Publish a block to trigger republishing. + err = lnd.NotifyHeight(601 + int32(i)) + require.NoError(t, err) + + // Put a swap into store to satisfy SQL constraints. + swap := &loopdb.LoopOutContract{ + SwapContract: loopdb.SwapContract{ + CltvExpiry: 111, + AmountRequested: 1_000_000, + ProtocolVersion: loopdb.ProtocolVersionMuSig2, + HtlcKeys: htlcKeys, + + // Make preimage unique to pass SQL constraints. + Preimage: preimages[i], + }, + + DestAddr: destAddr, + SwapInvoice: swapInvoice, + SweepConfTarget: 111, + } + + require.NoError(t, store.CreateLoopOut(ctx, swapHash, swap)) + store.AssertLoopOutStored() + + // Add SweepInfo to sweepFetcher. + htlc, err := utils.GetHtlc( + swapHash, &swap.SwapContract, lnd.ChainParams, + ) + require.NoError(t, err) + + sweepInfo := &SweepInfo{ + Preimage: preimages[i], + ConfTarget: 123, + Timeout: 111, + SwapInvoicePaymentAddr: *swapPaymentAddr, + ProtocolVersion: loopdb.ProtocolVersionMuSig2, + HTLCKeys: htlcKeys, + HTLC: *htlc, + HTLCSuccessEstimator: htlc.AddSuccessToEstimator, + DestAddr: destAddr, + } + // The first sweep is known in advance to be non-cooperative. + if i == 0 { + sweepInfo.NonCoopHint = true + } + sweepFetcher.store[swapHash] = sweepInfo + + // Create sweep request. + sweepReq := SweepRequest{ + SwapHash: swapHash, + Value: 1_000_000, + Outpoint: wire.OutPoint{ + Hash: chainhash.Hash{1, 1}, + Index: 1, + }, + Notifier: &dummyNotifier, + } + require.NoError(t, batcher.AddSweep(&sweepReq)) + + if i == 0 { + // Since a batch was created we check that it registered + // for its primary sweep's spend. + <-lnd.RegisterSpendChannel + } + + // Expect mockSigner.SignOutputRaw call to sign non-cooperative + // sweeps. + <-lnd.SignOutputRawChannel + + // A transaction is published. + tx := <-lnd.TxPublishChannel + require.Equal(t, i+1, len(tx.TxIn)) + + // Check types of inputs. + var witnessSizes []int + for _, txIn := range tx.TxIn { + witnessSizes = append(witnessSizes, len(txIn.Witness)) + } + // The order of inputs is not deterministic, because they + // are stored in map. + require.ElementsMatch(t, wantWitnessSizes[:i+1], witnessSizes) + + // Calculate expected values. + feeRate := test.DefaultMockFee + for range i { + // Bump fee the number of blocks passed. + feeRate += defaultFeeRateStep + } + amt := btcutil.Amount(1_000_000 * (i + 1)) + weight := wantWeights[i] + expectedFee := feeRate.FeeForWeight(weight) + + // Check weight. + gotWeight := lntypes.WeightUnit( + blockchain.GetTransactionWeight(btcutil.NewTx(tx)), + ) + require.Equal(t, weight, gotWeight, "weights don't match") + + // Check fee. + out := btcutil.Amount(tx.TxOut[0].Value) + gotFee := amt - out + require.Equal(t, expectedFee, gotFee, "fees don't match") + + // Check fee rate. + gotFeeRate := chainfee.NewSatPerKWeight(gotFee, gotWeight) + require.Equal(t, feeRate, gotFeeRate, "fee rates don't match") + } + + // Make sure we have stored the batch. + batches, err := batcherStore.FetchUnconfirmedSweepBatches(ctx) + require.NoError(t, err) + require.Len(t, batches, 1) + + // Now make the batcher quit by canceling the context. + cancel() + wg.Wait() + + // Make sure the batcher exited without an error. + checkBatcherError(t, runErr) +} + +// testWithMixedBatchCustom tests mixed batches construction, custom scenario. +// All sweeps are added at once. +func testWithMixedBatchCustom(t *testing.T, store testStore, + batcherStore testBatcherStore, preimages []lntypes.Preimage, + muSig2SignSweep MuSig2SignSweep, nonCoopHints []bool, + expectSignOutputRawChannel bool, wantWeight lntypes.WeightUnit, + wantWitnessSizes []int) { + + defer test.Guard(t)() + + lnd := test.NewMockLnd() + ctx, cancel := context.WithCancel(context.Background()) + + // Extract payment address from the invoice. + swapPaymentAddr, err := utils.ObtainSwapPaymentAddr( + swapInvoice, lnd.ChainParams, + ) + require.NoError(t, err) + + // Use sweepFetcher to provide NonCoopHint for swapHash1. + sweepFetcher := &sweepFetcherMock{ + store: map[lntypes.Hash]*SweepInfo{}, + } + + // Swap hashes must match the preimages, for non-cooperative spending + // path to work. + swapHashes := make([]lntypes.Hash, len(preimages)) + for i, preimage := range preimages { + swapHashes[i] = preimage.Hash() + } + + // Use mixed batches. + batcher := NewBatcher(lnd.WalletKit, lnd.ChainNotifier, lnd.Signer, + muSig2SignSweep, testVerifySchnorrSig, lnd.ChainParams, + batcherStore, sweepFetcher, WithMixedBatch()) + + var wg sync.WaitGroup + wg.Add(1) + + var runErr error + go func() { + defer wg.Done() + runErr = batcher.Run(ctx) + }() + + // Wait for the batcher to be initialized. + <-batcher.initDone + + // Create swaps and sweeps. + for i, swapHash := range swapHashes { + // Put a swap into store to satisfy SQL constraints. + swap := &loopdb.LoopOutContract{ + SwapContract: loopdb.SwapContract{ + CltvExpiry: 111, + AmountRequested: 1_000_000, + ProtocolVersion: loopdb.ProtocolVersionMuSig2, + HtlcKeys: htlcKeys, + + // Make preimage unique to pass SQL constraints. + Preimage: preimages[i], + }, + + DestAddr: destAddr, + SwapInvoice: swapInvoice, + SweepConfTarget: 111, + } + + require.NoError(t, store.CreateLoopOut(ctx, swapHash, swap)) + store.AssertLoopOutStored() + + // Add SweepInfo to sweepFetcher. + htlc, err := utils.GetHtlc( + swapHash, &swap.SwapContract, lnd.ChainParams, + ) + require.NoError(t, err) + + sweepFetcher.store[swapHash] = &SweepInfo{ + Preimage: preimages[i], + NonCoopHint: nonCoopHints[i], + + ConfTarget: 123, + Timeout: 111, + SwapInvoicePaymentAddr: *swapPaymentAddr, + ProtocolVersion: loopdb.ProtocolVersionMuSig2, + HTLCKeys: htlcKeys, + HTLC: *htlc, + HTLCSuccessEstimator: htlc.AddSuccessToEstimator, + DestAddr: destAddr, + } + + // Create sweep request. + sweepReq := SweepRequest{ + SwapHash: swapHash, + Value: 1_000_000, + Outpoint: wire.OutPoint{ + Hash: chainhash.Hash{1, 1}, + Index: 1, + }, + Notifier: &dummyNotifier, + } + require.NoError(t, batcher.AddSweep(&sweepReq)) + + if i == 0 { + // Since a batch was created we check that it registered + // for its primary sweep's spend. + <-lnd.RegisterSpendChannel + } + } + + if expectSignOutputRawChannel { + // Expect mockSigner.SignOutputRaw call to sign non-cooperative + // sweeps. + <-lnd.SignOutputRawChannel + } + + // A transaction is published. + tx := <-lnd.TxPublishChannel + require.Equal(t, len(preimages), len(tx.TxIn)) + + // Check types of inputs. + var witnessSizes []int + for _, txIn := range tx.TxIn { + witnessSizes = append(witnessSizes, len(txIn.Witness)) + } + // The order of inputs is not deterministic, because they + // are stored in map. + require.ElementsMatch(t, wantWitnessSizes, witnessSizes) + + // Calculate expected values. + feeRate := test.DefaultMockFee + amt := btcutil.Amount(1_000_000 * len(preimages)) + expectedFee := feeRate.FeeForWeight(wantWeight) + + // Check weight. + gotWeight := lntypes.WeightUnit( + blockchain.GetTransactionWeight(btcutil.NewTx(tx)), + ) + require.Equal(t, wantWeight, gotWeight, "weights don't match") + + // Check fee. + out := btcutil.Amount(tx.TxOut[0].Value) + gotFee := amt - out + require.Equal(t, expectedFee, gotFee, "fees don't match") + + // Check fee rate. + gotFeeRate := chainfee.NewSatPerKWeight(gotFee, gotWeight) + require.Equal(t, feeRate, gotFeeRate, "fee rates don't match") + + // Make sure we have stored the batch. + batches, err := batcherStore.FetchUnconfirmedSweepBatches(ctx) + require.NoError(t, err) + require.Len(t, batches, 1) + + // Now make the batcher quit by canceling the context. + cancel() + wg.Wait() + + // Make sure the batcher exited without an error. + checkBatcherError(t, runErr) +} + +// testWithMixedBatchLarge tests mixed batches construction, many sweeps. +// All sweeps are added at once. +func testWithMixedBatchLarge(t *testing.T, store testStore, + batcherStore testBatcherStore) { + + // Create 9 sweeps. 3 groups of 3 sweeps. + // 1. known in advance to be non-cooperative, + // 2. fails cosigning during an attempt, + // 3. co-signs successfully. + var preimages = []lntypes.Preimage{ + {1}, {2}, {3}, + {4}, {5}, {6}, + {7}, {8}, {9}, + } + + // Create muSig2SignSweep. It fails all the sweeps, works only one time + // for swapHashes[2] and works any number of times for 5 and 8. This + // emulates client disconnect after first successful co-signing. + swapHash2Used := false + muSig2SignSweep := func(ctx context.Context, + protocolVersion loopdb.ProtocolVersion, swapHash lntypes.Hash, + paymentAddr [32]byte, nonce []byte, sweepTxPsbt []byte, + prevoutMap map[wire.OutPoint]*wire.TxOut) ( + []byte, []byte, error) { + + switch { + case swapHash == preimages[2].Hash(): + if swapHash2Used { + return nil, nil, fmt.Errorf("disconnected") + } else { + swapHash2Used = true + + return nil, nil, nil + } + + case swapHash == preimages[5].Hash(): + return nil, nil, nil + + case swapHash == preimages[8].Hash(): + return nil, nil, nil + + default: + return nil, nil, fmt.Errorf("test error") + } + } + + // The first sweep in a group is known in advance to be + // non-cooperative. + nonCoopHints := []bool{ + true, false, false, + true, false, false, + true, false, false, + } + + // Expect mockSigner.SignOutputRaw call to sign non-cooperative + // sweeps. + expectSignOutputRawChannel := true + + // Two non-cooperative sweeps, one cooperative. + wantWitnessSizes := []int{4, 4, 4, 4, 4, 1, 4, 4, 1} + + // Expected weight. + wantWeight := lntypes.WeightUnit(3377) + + testWithMixedBatchCustom(t, store, batcherStore, preimages, + muSig2SignSweep, nonCoopHints, expectSignOutputRawChannel, + wantWeight, wantWitnessSizes) +} + +// testWithMixedBatchCoopOnly tests mixed batches construction, +// All sweeps are added at once. All the sweeps are cooperative. +func testWithMixedBatchCoopOnly(t *testing.T, store testStore, + batcherStore testBatcherStore) { + + // Create 3 sweeps, all cooperative. + var preimages = []lntypes.Preimage{ + {1}, {2}, {3}, + } + + // Create muSig2SignSweep, working for all sweeps. + muSig2SignSweep := func(ctx context.Context, + protocolVersion loopdb.ProtocolVersion, swapHash lntypes.Hash, + paymentAddr [32]byte, nonce []byte, sweepTxPsbt []byte, + prevoutMap map[wire.OutPoint]*wire.TxOut) ( + []byte, []byte, error) { + + return nil, nil, nil + } + + // All the sweeps are cooperative. + nonCoopHints := []bool{false, false, false} + + // Do not expect a mockSigner.SignOutputRaw call, because there are no + // non-cooperative sweeps. + expectSignOutputRawChannel := false + + // Two non-cooperative sweeps, one cooperative. + wantWitnessSizes := []int{1, 1, 1} + + // Expected weight. + wantWeight := lntypes.WeightUnit(856) + + testWithMixedBatchCustom(t, store, batcherStore, preimages, + muSig2SignSweep, nonCoopHints, expectSignOutputRawChannel, + wantWeight, wantWitnessSizes) +} + +// testWithMixedBatchNonCoopHintOnly tests mixed batches construction, +// All sweeps are added at once. All the sweeps are known to be non-cooperative +// in advance. +func testWithMixedBatchNonCoopHintOnly(t *testing.T, store testStore, + batcherStore testBatcherStore) { + + // Create 3 sweeps, all known to be non-cooperative in advance. + var preimages = []lntypes.Preimage{ + {1}, {2}, {3}, + } + + // Create muSig2SignSweep, panicking for all sweeps. + muSig2SignSweep := func(ctx context.Context, + protocolVersion loopdb.ProtocolVersion, swapHash lntypes.Hash, + paymentAddr [32]byte, nonce []byte, sweepTxPsbt []byte, + prevoutMap map[wire.OutPoint]*wire.TxOut) ( + []byte, []byte, error) { + + panic("must not be called in this test") + } + + // All the sweeps are non-cooperative, this is known in advance. + nonCoopHints := []bool{true, true, true} + + // Expect mockSigner.SignOutputRaw call to sign non-cooperative + // sweeps. + expectSignOutputRawChannel := true + + // Two non-cooperative sweeps, one cooperative. + wantWitnessSizes := []int{4, 4, 4} + + // Expected weight. + wantWeight := lntypes.WeightUnit(1345) + + testWithMixedBatchCustom(t, store, batcherStore, preimages, + muSig2SignSweep, nonCoopHints, expectSignOutputRawChannel, + wantWeight, wantWitnessSizes) +} + +// testWithMixedBatchCoopFailedOnly tests mixed batches construction, +// All sweeps are added at once. All the sweeps fail co-signing. +func testWithMixedBatchCoopFailedOnly(t *testing.T, store testStore, + batcherStore testBatcherStore) { + + // Create 3 sweeps, all fail co-signing. + var preimages = []lntypes.Preimage{ + {1}, {2}, {3}, + } + + // Create muSig2SignSweep, failing any co-sign attempt. + muSig2SignSweep := func(ctx context.Context, + protocolVersion loopdb.ProtocolVersion, swapHash lntypes.Hash, + paymentAddr [32]byte, nonce []byte, sweepTxPsbt []byte, + prevoutMap map[wire.OutPoint]*wire.TxOut) ( + []byte, []byte, error) { + + return nil, nil, fmt.Errorf("test error") + } + + // All the sweeps are non-cooperative, but this is not known in advance. + nonCoopHints := []bool{false, false, false} + + // Expect mockSigner.SignOutputRaw call to sign non-cooperative + // sweeps. + expectSignOutputRawChannel := true + + // Two non-cooperative sweeps, one cooperative. + wantWitnessSizes := []int{4, 4, 4} + + // Expected weight. + wantWeight := lntypes.WeightUnit(1345) + + testWithMixedBatchCustom(t, store, batcherStore, preimages, + muSig2SignSweep, nonCoopHints, expectSignOutputRawChannel, + wantWeight, wantWitnessSizes) +} + // TestSweepBatcherBatchCreation tests that sweep requests enter the expected // batch based on their timeout distance. func TestSweepBatcherBatchCreation(t *testing.T) { @@ -2980,6 +3526,37 @@ func TestCustomSignMuSig2(t *testing.T) { runTests(t, testCustomSignMuSig2) } +// TestWithMixedBatch tests mixed batches construction. It also tests +// non-cooperative sweeping (using a preimage). Sweeps are added one by one. +func TestWithMixedBatch(t *testing.T) { + runTests(t, testWithMixedBatch) +} + +// TestWithMixedBatchLarge tests mixed batches construction, many sweeps. +// All sweeps are added at once. +func TestWithMixedBatchLarge(t *testing.T) { + runTests(t, testWithMixedBatchLarge) +} + +// TestWithMixedBatchCoopOnly tests mixed batches construction, +// All sweeps are added at once. All the sweeps are cooperative. +func TestWithMixedBatchCoopOnly(t *testing.T) { + runTests(t, testWithMixedBatchCoopOnly) +} + +// TestWithMixedBatchNonCoopHintOnly tests mixed batches construction, +// All sweeps are added at once. All the sweeps are known to be non-cooperative +// in advance. +func TestWithMixedBatchNonCoopHintOnly(t *testing.T) { + runTests(t, testWithMixedBatchNonCoopHintOnly) +} + +// TestWithMixedBatchCoopFailedOnly tests mixed batches construction, +// All sweeps are added at once. All the sweeps fail co-signing. +func TestWithMixedBatchCoopFailedOnly(t *testing.T) { + runTests(t, testWithMixedBatchCoopFailedOnly) +} + // testBatcherStore is BatcherStore used in tests. type testBatcherStore interface { BatcherStore