tapchannel: update sweepContracts for HTLC sweeps

In this commit, we add some intermediate functions and types needed to
properly handle HTLC sweeps.

We now use the preimage info added in a prior commit to make sure that
once we learn of the preimage (after vPkt creation, but before publish),
we'll properly insert it in the correct location.

In sweepContracts, when we go to create the new anchor change output,
we'll make sure to omit any second level transactions, as they already
have a destination location specified.

Direct spends are spends directly from the commitment transaction. If we
don't have any direct spends, then we don't need to make a change addr,
as second level spends are already bound to an anchor output.

Author: Roasbeef

tapchannel/aux_sweeper.go

@@ -432,6 +432,61 @@ func (a *AuxSweeper) signSweepVpackets(vPackets []*tappsbt.VPacket,
 	return nil
 }
 
+// vPktsWithInput couples a vPkt along with the input that contained it.
+type vPktsWithInput struct {
+	// btcInput is the Bitcoin that the vPkt will the spending from (on the
+	// TAP layer).
+	btcInput input.Input
+
+	// vPkts is the set of vPacket that will be used to spend the input.
+	vPkts []*tappsbt.VPacket
+
+	// tapSigDesc houses the information we'll need to re-sign the vPackets
+	// above. Note that this is only set if this is a second level packet.
+	tapSigDesc lfn.Option[cmsg.TapscriptSigDesc]
+}
+
+// sweepVpkts contains the set of vPkts needed for sweeping an output. Most
+// outputs will only have the first level specified. The second level is needed
+// for HTLC outputs on our local commitment transaction.
+type sweepVpkts struct {
+	// firstLevel houses vPackets that are used to sweep outputs directly
+	// from the commitment transaction.
+	firstLevel []vPktsWithInput
+
+	// secondLevel is used to sweep outputs that are created by second level
+	// HTLC transactions.
+	secondLevel []vPktsWithInput
+}
+
+// isEmpty returns true if the sweepVpkts is empty.
+func (s sweepVpkts) isEmpty() bool {
+	return len(s.firstLevel) == 0 && len(s.secondLevel) == 0
+}
+
+// firstLevelPkts returns a slice of the first level pkts.
+func (s sweepVpkts) firstLevelPkts() []*tappsbt.VPacket {
+	return fn.FlatMap(
+		s.firstLevel, func(v vPktsWithInput) []*tappsbt.VPacket {
+			return v.vPkts
+		},
+	)
+}
+
+// secondLevelPkts returns a slice of the second level pkts.
+func (s sweepVpkts) secondLevelPkts() []*tappsbt.VPacket {
+	return fn.FlatMap(
+		s.secondLevel, func(v vPktsWithInput) []*tappsbt.VPacket {
+			return v.vPkts
+		},
+	)
+}
+
+// allPkts returns a slice of both the first and second level pkts.
+func (s sweepVpkts) allPkts() []*tappsbt.VPacket {
+	return append(s.firstLevelPkts(), s.secondLevelPkts()...)
+}
+
 // createAndSignSweepVpackets creates vPackets that sweep the funds from the
 // channel to the wallet, and then signs them as well.
 func (a *AuxSweeper) createAndSignSweepVpackets(
@@ -643,6 +698,8 @@ func remoteHtlcTimeoutSweepDesc(keyRing *lnwallet.CommitmentKeyRing,
 		return lfn.Err[tapscriptSweepDescs](err)
 	}
 
+	// TODO(roasbeef): use GenTaprootHtlcScript instead?
+
 	// Now that we have the script tree, we'll make the control block needed
 	// to spend it, but taking the revoked path.
 	ctrlBlock, err := htlcScriptTree.CtrlBlockForPath(
@@ -1813,8 +1870,8 @@ func newBlobWithWitnessInfo(i input.Input) lfn.Result[blobWithWitnessInfo] {
 		secondLevel  bool
 	)
 	switch i.WitnessType() {
-
 	// This is the case when we're sweeping the HTLC output on our local
+
 	// commitment transaction via a second level HTLC.
 	//
 	// The final witness stack is:
@@ -1874,40 +1931,118 @@ func newBlobWithWitnessInfo(i input.Input) lfn.Result[blobWithWitnessInfo] {
 	})
 }
 
+// prepVpkts decodes the set of vPkts, supplementing them as needed to ensure
+// all inputs can be swept properly.
+func prepVpkts(bRes lfn.Result[blobWithWitnessInfo],
+	secondLevel bool) (*vPktsWithInput, error) {
+
+	b, err := bRes.Unpack()
+	if err != nil {
+		return nil, err
+	}
+
+	var res cmsg.ContractResolution
+	err = res.Decode(bytes.NewReader(b.resolutionBlob))
+	if err != nil {
+		return nil, err
+	}
+
+	// For each vPacket, if we have a preimage to insert, then we'll we'll
+	// update the witness to insert the preimage at the correct index.
+	var tapSigDesc lfn.Option[cmsg.TapscriptSigDesc]
+	pkts := res.Vpkts1()
+	if secondLevel {
+		pkts = res.Vpkts2()
+		tapSigDesc = res.SigDescs()
+	}
+
+	b.preimageInfo.WhenSome(func(p preimageDesc) {
+		for _, pkt := range pkts {
+			newAsset := pkt.Outputs[0].Asset
+
+			prevWitness := newAsset.PrevWitnesses[0].TxWitness
+			prevWitness = slices.Insert(
+				prevWitness, p.witnessIndex,
+				p.preimage[:],
+			)
+			newAsset.UpdateTxWitness(0, prevWitness)
+		}
+	})
+
+	return &vPktsWithInput{
+		vPkts:      pkts,
+		btcInput:   b.input,
+		tapSigDesc: tapSigDesc,
+	}, nil
+}
+
 // extractInputVPackets extracts the vPackets from the inputs passed in. If
 // none of the inputs have any resolution blobs. Then an empty slice will be
 // returned.
-func extractInputVPackets(inputs []input.Input) lfn.Result[[]*tappsbt.VPacket] {
-	type returnType = []*tappsbt.VPacket
+func extractInputVPackets(inputs []input.Input) lfn.Result[sweepVpkts] {
+	type returnType = sweepVpkts
 
 	// Otherwise, we'll extract the set of resolution blobs from the inputs
 	// passed in.
 	relevantInputs := fn.Filter(inputs, func(i input.Input) bool {
 		return i.ResolutionBlob().IsSome()
 	})
-	resolutionBlobs := fn.Map(relevantInputs, func(i input.Input) tlv.Blob {
-		// We already know this has a blob from the filter above.
-		return i.ResolutionBlob().UnwrapOr(nil)
-	})
+	resolutionInfo := fn.Map(
+		relevantInputs,
+		func(i input.Input) lfn.Result[blobWithWitnessInfo] {
+			return newBlobWithWitnessInfo(i)
+		},
+	)
 
-	// With our set of resolution inputs extracted, we'll now decode them
-	// in the vPackets we'll use to generate the output to addr.
-	vPkts, err := fn.FlatMapErr(
-		resolutionBlobs,
-		func(b tlv.Blob) ([]*tappsbt.VPacket, error) {
-			var res cmsg.ContractResolution
-			if err := res.Decode(bytes.NewReader(b)); err != nil {
-				return nil, err
-			}
+	firstLevelSweeps := lfn.Filter(
+		func(info lfn.Result[blobWithWitnessInfo]) bool {
+			var secondLevel bool
+			info.WhenResult(func(i blobWithWitnessInfo) {
+				secondLevel = i.secondLevel
+			})
 
-			return res.Vpkts1(), nil
+			return !secondLevel
 		},
+		resolutionInfo,
 	)
-	if err != nil {
-		return lfn.Err[returnType](err)
+	secondLevelSweeps := lfn.Filter(
+		func(info lfn.Result[blobWithWitnessInfo]) bool {
+			var secondLevel bool
+			info.WhenResult(func(i blobWithWitnessInfo) {
+				secondLevel = i.secondLevel
+			})
+
+			return secondLevel
+		},
+		resolutionInfo,
+	)
+
+	// With our set of resolution inputs extracted, we'll now decode them in
+	// the vPackets we'll use to generate the output to addr.
+	var vPkts1 []vPktsWithInput
+	for _, bRes := range firstLevelSweeps {
+		vpkt, err := prepVpkts(bRes, false)
+		if err != nil {
+			return lfn.Err[sweepVpkts](err)
+		}
+
+		vPkts1 = append(vPkts1, *vpkt)
 	}
 
-	return lfn.Ok(vPkts)
+	var vPkts2 []vPktsWithInput
+	for _, bRes := range secondLevelSweeps {
+		vpkt, err := prepVpkts(bRes, true)
+		if err != nil {
+			return lfn.Err[sweepVpkts](err)
+		}
+
+		vPkts2 = append(vPkts2, *vpkt)
+	}
+
+	return lfn.Ok(sweepVpkts{
+		firstLevel:  vPkts1,
+		secondLevel: vPkts2,
+	})
 }
 
 // sweepContracts takes a set of inputs, and the change address we'd use to
@@ -1931,13 +2066,25 @@ func (a *AuxSweeper) sweepContracts(inputs []input.Input,
 
 	// Now that we know we have a relevant input set, extract all the
 	// vPackets from the inputs.
-	vPkts, err := extractInputVPackets(inputs).Unpack()
+	sPkts, err := extractInputVPackets(inputs).Unpack()
 	if err != nil {
 		return lfn.Err[returnType](err)
 	}
 
 	log.Infof("Generating anchor output for vpkts=%v",
-		limitSpewer.Sdump(vPkts))
+		limitSpewer.Sdump(sPkts))
+
+	// If this is a sweep from the local commitment transaction. Then we'll
+	// have both the first and second level sweeps. However for the first
+	// sweep, it's a broadcast of a pre-signed transaction, so we don't need
+	// an anchor output for those.
+	directPkts := sPkts.directSpendPkts()
+
+	// If there're no direct level vPkts, then we can just return a nil
+	// error as we don't have a real sweep output to create.
+	if len(directPkts) == 0 {
+		return lfn.Err[sweep.SweepOutput](nil)
+	}
 
 	// At this point, now that we're about to generate a new output, we'll
 	// need an internal key, so we can update all the vPkts.
@@ -1950,17 +2097,32 @@ func (a *AuxSweeper) sweepContracts(inputs []input.Input,
 	if err != nil {
 		return lfn.Err[returnType](err)
 	}
-	for idx := range vPkts {
-		for _, vOut := range vPkts[idx].Outputs {
+	for idx := range directPkts {
+		for _, vOut := range directPkts[idx].Outputs {
 			vOut.SetAnchorInternalKey(
 				internalKey, a.cfg.ChainParams.HDCoinType,
 			)
 		}
 	}
 
+	// For any second level outputs we're sweeping, we'll need to sign for
+	// it, as now we know the txid of the sweeping transaction. We'll do
+	// this again when we register for the final broadcast, we we need to
+	// sign the right prevIDs.
+	for _, sweepSet := range sPkts.secondLevel {
+		for _, vPkt := range sweepSet.vPkts {
+			for _, vIn := range vPkt.Inputs {
+				vIn.PrevID.OutPoint = sweepSet.btcInput.OutPoint()
+			}
+			for _, vOut := range vPkt.Outputs {
+				vOut.Asset.PrevWitnesses[0].PrevID.OutPoint = sweepSet.btcInput.OutPoint()
+			}
+		}
+	}
+
 	// Now that we have our set of resolutions, we'll make a new commitment
 	// out of all the vPackets contained.
-	outCommitments, err := tapsend.CreateOutputCommitments(vPkts)
+	outCommitments, err := tapsend.CreateOutputCommitments(directPkts)
 	if err != nil {
 		return lfn.Errf[returnType]("unable to create "+
 			"output commitments: %w", err)
@@ -2040,7 +2202,7 @@ func (a *AuxSweeper) registerAndBroadcastSweep(req *sweep.BumpRequest,
 
 	// If we don't have any vPackets that had our resolution data in them,
 	// then we can exit early.
-	if len(vPkts) == 0 {
+	if len(vPkts.firstLevel) == 0 && len(vPkts.secondLevel) == 0 {
 		log.Infof("Sweep request had no vPkts, exiting")
 		return nil
 	}
@@ -2063,16 +2225,16 @@ func (a *AuxSweeper) registerAndBroadcastSweep(req *sweep.BumpRequest,
 
 	// We'll also use the passed in context to set the anchor key again for
 	// all the vOuts.
-	for idx := range vPkts {
-		for _, vOut := range vPkts[idx].Outputs {
+	for idx := range vPkts.firstLevelPkts() {
+		for _, vOut := range vPkts.firstLevelPkts()[idx].Outputs {
 			vOut.SetAnchorInternalKey(
 				internalKey, a.cfg.ChainParams.HDCoinType,
 			)
 		}
 	}
 
 	// Now that we have our vPkts, we'll re-create the output commitments.
-	outCommitments, err := tapsend.CreateOutputCommitments(vPkts)
+	outCommitments, err := tapsend.CreateOutputCommitments(vPkts.allPkts())
 	if err != nil {
 		return fmt.Errorf("unable to create output "+
 			"commitments: %w", err)
@@ -2094,15 +2256,16 @@ func (a *AuxSweeper) registerAndBroadcastSweep(req *sweep.BumpRequest,
 	//
 	// TODO(roasbeef): base off allocations? then can serialize, then
 	// re-use the logic
-	for idx := range vPkts {
-		vPkt := vPkts[idx]
+	allVpkts := vPkts.allPkts()
+	for idx := range allVpkts {
+		vPkt := allVpkts[idx]
 		for outIdx := range vPkt.Outputs {
 			exclusionCreator := sweepExclusionProofGen(
 				changeInternalKey,
 			)
 
 			proofSuffix, err := tapsend.CreateProofSuffixCustom(
-				sweepTx, vPkt, outCommitments, outIdx, vPkts,
+				sweepTx, vPkt, outCommitments, outIdx, allVpkts,
 				exclusionCreator,
 			)
 			if err != nil {
@@ -2122,7 +2285,7 @@ func (a *AuxSweeper) registerAndBroadcastSweep(req *sweep.BumpRequest,
 	// We pass false for the last arg as we already updated our suffix
 	// proofs here.
 	return shipChannelTxn(
-		a.cfg.TxSender, sweepTx, outCommitments, vPkts, int64(fee),
+		a.cfg.TxSender, sweepTx, outCommitments, allVpkts, int64(fee),
 	)
 }
 

tapchannelmsg/records.go

@@ -2201,6 +2201,11 @@ func (c *ContractResolution) Decode(r io.Reader) error {
 	return nil
 }
 
+// SigDescs returns the list of tapscriptSigDescs.
+func (c *ContractResolution) SigDescs() lfn.Option[TapscriptSigDesc] {
+	return c.secondLevelSigDescs.ValOpt()
+}
+
 // Vpkts1 returns the set of first level Vpkts.
 func (c *ContractResolution) Vpkts1() []*tappsbt.VPacket {
 	return c.firstLevelSweepVpkts.Val.Pkts