Merge pull request #784 from starius/sweepbatcher-custom-musig2

sweepbatcher: add option WithCustomSignMuSig2

Author: starius

sweepbatcher/sweep_batch.go

@@ -137,6 +137,12 @@ type batchConfig struct {
 	// the caller has to update it in the source of SweepInfo (interface
 	// SweepFetcher) and re-add the sweep by calling AddSweep.
 	noBumping bool
+
+	// customMuSig2Signer is a custom signer. If it is set, it is used to
+	// create musig2 signatures instead of musig2SignSweep and signerClient.
+	// 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
 }
 
 // rbfCache stores data related to our last fee bump.
@@ -503,9 +509,12 @@ func (b *batch) Run(ctx context.Context) error {
 		close(b.finished)
 	}()
 
-	if b.muSig2SignSweep == nil {
+	if b.muSig2SignSweep == nil && b.cfg.customMuSig2Signer == nil {
 		return fmt.Errorf("no musig2 signer available")
 	}
+	if b.muSig2SignSweep != nil && b.cfg.customMuSig2Signer != nil {
+		return fmt.Errorf("both musig2 signers provided")
+	}
 
 	blockChan, blockErrChan, err :=
 		b.chainNotifier.RegisterBlockEpochNtfn(runCtx)
@@ -727,10 +736,11 @@ func (b *batch) publishBatch(ctx context.Context) (btcutil.Amount, error) {
 
 	weightEstimate.AddP2TROutput()
 
-	fee = b.rbfCache.FeeRate.FeeForWeight(weightEstimate.Weight())
+	weight := weightEstimate.Weight()
+	feeForWeight := b.rbfCache.FeeRate.FeeForWeight(weight)
 
 	// Clamp the calculated fee to the max allowed fee amount for the batch.
-	fee = clampBatchFee(fee, batchAmt)
+	fee = clampBatchFee(feeForWeight, batchAmt)
 
 	// Add the batch transaction output, which excludes the fees paid to
 	// miners.
@@ -761,8 +771,9 @@ func (b *batch) publishBatch(ctx context.Context) (btcutil.Amount, error) {
 	}
 
 	b.log.Infof("attempting to publish non-coop tx=%v with feerate=%v, "+
-		"totalfee=%v, sweeps=%d, destAddr=%s", batchTx.TxHash(),
-		b.rbfCache.FeeRate, fee, len(batchTx.TxIn), address)
+		"weight=%v, feeForWeight=%v, fee=%v, sweeps=%d, destAddr=%s",
+		batchTx.TxHash(), b.rbfCache.FeeRate, weight, feeForWeight, fee,
+		len(batchTx.TxIn), address)
 
 	b.debugLogTx("serialized non-coop sweep", batchTx)
 
@@ -857,10 +868,11 @@ func (b *batch) publishBatchCoop(ctx context.Context) (btcutil.Amount,
 
 	weightEstimate.AddP2TROutput()
 
-	fee = b.rbfCache.FeeRate.FeeForWeight(weightEstimate.Weight())
+	weight := weightEstimate.Weight()
+	feeForWeight := b.rbfCache.FeeRate.FeeForWeight(weight)
 
 	// Clamp the calculated fee to the max allowed fee amount for the batch.
-	fee = clampBatchFee(fee, batchAmt)
+	fee = clampBatchFee(feeForWeight, batchAmt)
 
 	// Add the batch transaction output, which excludes the fees paid to
 	// miners.
@@ -896,19 +908,18 @@ func (b *batch) publishBatchCoop(ctx context.Context) (btcutil.Amount,
 		return fee, err, false
 	}
 
-	prevOutputFetcher := txscript.NewMultiPrevOutFetcher(prevOuts)
-
 	// Attempt to cooperatively sign the batch tx with the server.
 	err = b.coopSignBatchTx(
-		ctx, packet, sweeps, prevOutputFetcher, prevOuts, psbtBuf,
+		ctx, packet, sweeps, prevOuts, psbtBuf.Bytes(),
 	)
 	if err != nil {
 		return fee, err, false
 	}
 
 	b.log.Infof("attempting to publish coop tx=%v with feerate=%v, "+
-		"totalfee=%v, sweeps=%d, destAddr=%s", batchTx.TxHash(),
-		b.rbfCache.FeeRate, fee, len(batchTx.TxIn), address)
+		"weight=%v, feeForWeight=%v, fee=%v, sweeps=%d, destAddr=%s",
+		batchTx.TxHash(), b.rbfCache.FeeRate, weight, feeForWeight, fee,
+		len(batchTx.TxIn), address)
 
 	b.debugLogTx("serialized coop sweep", batchTx)
 
@@ -942,119 +953,84 @@ func (b *batch) debugLogTx(msg string, tx *wire.MsgTx) {
 // coopSignBatchTx collects the necessary signatures from the server in order
 // to cooperatively sweep the funds.
 func (b *batch) coopSignBatchTx(ctx context.Context, packet *psbt.Packet,
-	sweeps []sweep, prevOutputFetcher *txscript.MultiPrevOutFetcher,
-	prevOuts map[wire.OutPoint]*wire.TxOut, psbtBuf bytes.Buffer) error {
+	sweeps []sweep, prevOuts map[wire.OutPoint]*wire.TxOut,
+	psbt []byte) error {
 
 	for i, sweep := range sweeps {
 		sweep := sweep
 
-		sigHashes := txscript.NewTxSigHashes(
-			packet.UnsignedTx, prevOutputFetcher,
-		)
-
-		sigHash, err := txscript.CalcTaprootSignatureHash(
-			sigHashes, txscript.SigHashDefault, packet.UnsignedTx,
-			i, prevOutputFetcher,
+		finalSig, err := b.musig2sign(
+			ctx, i, sweep, packet.UnsignedTx, prevOuts, psbt,
 		)
 		if err != nil {
 			return err
 		}
 
-		var (
-			signers       [][]byte
-			muSig2Version input.MuSig2Version
-		)
-
-		// Depending on the MuSig2 version we either pass 32 byte
-		// Schnorr public keys or normal 33 byte public keys.
-		if sweep.protocolVersion >= loopdb.ProtocolVersionMuSig2 {
-			muSig2Version = input.MuSig2Version100RC2
-			signers = [][]byte{
-				sweep.htlcKeys.SenderInternalPubKey[:],
-				sweep.htlcKeys.ReceiverInternalPubKey[:],
-			}
-		} else {
-			muSig2Version = input.MuSig2Version040
-			signers = [][]byte{
-				sweep.htlcKeys.SenderInternalPubKey[1:],
-				sweep.htlcKeys.ReceiverInternalPubKey[1:],
-			}
+		packet.UnsignedTx.TxIn[i].Witness = wire.TxWitness{
+			finalSig,
 		}
+	}
 
-		htlcScript, ok := sweep.htlc.HtlcScript.(*swap.HtlcScriptV3)
-		if !ok {
-			return fmt.Errorf("invalid htlc script version")
-		}
+	return nil
+}
 
-		// Now we're creating a local MuSig2 session using the receiver
-		// key's key locator and the htlc's root hash.
-		musig2SessionInfo, err := b.signerClient.MuSig2CreateSession(
-			ctx, muSig2Version,
-			&sweep.htlcKeys.ClientScriptKeyLocator, signers,
-			lndclient.MuSig2TaprootTweakOpt(
-				htlcScript.RootHash[:], false,
-			),
-		)
-		if err != nil {
-			return err
-		}
+// musig2sign signs one sweep using musig2.
+func (b *batch) musig2sign(ctx context.Context, inputIndex int, sweep sweep,
+	unsignedTx *wire.MsgTx, prevOuts map[wire.OutPoint]*wire.TxOut,
+	psbt []byte) ([]byte, error) {
 
-		// With the session active, we can now send the server our
-		// public nonce and the sig hash, so that it can create it's own
-		// MuSig2 session and return the server side nonce and partial
-		// signature.
-		serverNonce, serverSig, err := b.muSig2SignSweep(
-			ctx, sweep.protocolVersion, sweep.swapHash,
-			sweep.swapInvoicePaymentAddr,
-			musig2SessionInfo.PublicNonce[:], psbtBuf.Bytes(),
-			prevOuts,
-		)
-		if err != nil {
-			return err
-		}
+	prevOutputFetcher := txscript.NewMultiPrevOutFetcher(prevOuts)
 
-		var serverPublicNonce [musig2.PubNonceSize]byte
-		copy(serverPublicNonce[:], serverNonce)
+	sigHashes := txscript.NewTxSigHashes(unsignedTx, prevOutputFetcher)
 
-		// Register the server's nonce before attempting to create our
-		// partial signature.
-		haveAllNonces, err := b.signerClient.MuSig2RegisterNonces(
-			ctx, musig2SessionInfo.SessionID,
-			[][musig2.PubNonceSize]byte{serverPublicNonce},
-		)
-		if err != nil {
-			return err
-		}
+	sigHash, err := txscript.CalcTaprootSignatureHash(
+		sigHashes, txscript.SigHashDefault, unsignedTx, inputIndex,
+		prevOutputFetcher,
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	var (
+		signers       [][]byte
+		muSig2Version input.MuSig2Version
+	)
 
-		// Sanity check that we have all the nonces.
-		if !haveAllNonces {
-			return fmt.Errorf("invalid MuSig2 session: " +
-				"nonces missing")
+	// Depending on the MuSig2 version we either pass 32 byte
+	// Schnorr public keys or normal 33 byte public keys.
+	if sweep.protocolVersion >= loopdb.ProtocolVersionMuSig2 {
+		muSig2Version = input.MuSig2Version100RC2
+		signers = [][]byte{
+			sweep.htlcKeys.SenderInternalPubKey[:],
+			sweep.htlcKeys.ReceiverInternalPubKey[:],
+		}
+	} else {
+		muSig2Version = input.MuSig2Version040
+		signers = [][]byte{
+			sweep.htlcKeys.SenderInternalPubKey[1:],
+			sweep.htlcKeys.ReceiverInternalPubKey[1:],
 		}
+	}
 
-		var digest [32]byte
-		copy(digest[:], sigHash)
+	htlcScript, ok := sweep.htlc.HtlcScript.(*swap.HtlcScriptV3)
+	if !ok {
+		return nil, fmt.Errorf("invalid htlc script version")
+	}
 
-		// Since our MuSig2 session has all nonces, we can now create
-		// the local partial signature by signing the sig hash.
-		_, err = b.signerClient.MuSig2Sign(
-			ctx, musig2SessionInfo.SessionID, digest, false,
-		)
-		if err != nil {
-			return err
-		}
+	var digest [32]byte
+	copy(digest[:], sigHash)
 
-		// Now combine the partial signatures to use the final combined
-		// signature in the sweep transaction's witness.
-		haveAllSigs, finalSig, err := b.signerClient.MuSig2CombineSig(
-			ctx, musig2SessionInfo.SessionID, [][]byte{serverSig},
+	// If a custom signer is installed, use it instead of b.signerClient
+	// and b.muSig2SignSweep.
+	if b.cfg.customMuSig2Signer != nil {
+		// Produce a signature.
+		finalSig, err := b.cfg.customMuSig2Signer(
+			ctx, muSig2Version, sweep.swapHash,
+			htlcScript.RootHash, digest,
 		)
 		if err != nil {
-			return err
-		}
-
-		if !haveAllSigs {
-			return fmt.Errorf("failed to combine signatures")
+			return nil, fmt.Errorf("customMuSig2Signer failed: %w",
+				err)
 		}
 
 		// To be sure that we're good, parse and validate that the
@@ -1064,15 +1040,88 @@ func (b *batch) coopSignBatchTx(ctx context.Context, packet *psbt.Packet,
 			htlcScript.TaprootKey, sigHash, finalSig,
 		)
 		if err != nil {
-			return err
+			return nil, fmt.Errorf("verifySchnorrSig failed: %w",
+				err)
 		}
 
-		packet.UnsignedTx.TxIn[i].Witness = wire.TxWitness{
-			finalSig,
-		}
+		return finalSig, nil
 	}
 
-	return nil
+	// Now we're creating a local MuSig2 session using the receiver key's
+	// key locator and the htlc's root hash.
+	keyLocator := &sweep.htlcKeys.ClientScriptKeyLocator
+	musig2SessionInfo, err := b.signerClient.MuSig2CreateSession(
+		ctx, muSig2Version, keyLocator, signers,
+		lndclient.MuSig2TaprootTweakOpt(htlcScript.RootHash[:], false),
+	)
+	if err != nil {
+		return nil, fmt.Errorf("signerClient.MuSig2CreateSession "+
+			"failed: %w", err)
+	}
+
+	// With the session active, we can now send the server our
+	// public nonce and the sig hash, so that it can create it's own
+	// MuSig2 session and return the server side nonce and partial
+	// signature.
+	serverNonce, serverSig, err := b.muSig2SignSweep(
+		ctx, sweep.protocolVersion, sweep.swapHash,
+		sweep.swapInvoicePaymentAddr,
+		musig2SessionInfo.PublicNonce[:], psbt, prevOuts,
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	var serverPublicNonce [musig2.PubNonceSize]byte
+	copy(serverPublicNonce[:], serverNonce)
+
+	// Register the server's nonce before attempting to create our
+	// partial signature.
+	haveAllNonces, err := b.signerClient.MuSig2RegisterNonces(
+		ctx, musig2SessionInfo.SessionID,
+		[][musig2.PubNonceSize]byte{serverPublicNonce},
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	// Sanity check that we have all the nonces.
+	if !haveAllNonces {
+		return nil, fmt.Errorf("invalid MuSig2 session: " +
+			"nonces missing")
+	}
+
+	// Since our MuSig2 session has all nonces, we can now create
+	// the local partial signature by signing the sig hash.
+	_, err = b.signerClient.MuSig2Sign(
+		ctx, musig2SessionInfo.SessionID, digest, false,
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	// Now combine the partial signatures to use the final combined
+	// signature in the sweep transaction's witness.
+	haveAllSigs, finalSig, err := b.signerClient.MuSig2CombineSig(
+		ctx, musig2SessionInfo.SessionID, [][]byte{serverSig},
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	if !haveAllSigs {
+		return nil, fmt.Errorf("failed to combine signatures")
+	}
+
+	// To be sure that we're good, parse and validate that the
+	// combined signature is indeed valid for the sig hash and the
+	// internal pubkey.
+	err = b.verifySchnorrSig(htlcScript.TaprootKey, sigHash, finalSig)
+	if err != nil {
+		return nil, err
+	}
+
+	return finalSig, nil
 }
 
 // updateRbfRate updates the fee rate we should use for the new batch