Merge pull request #1468 from lightninglabs/send-asset-label

Add Transfer Label Support for Asset Sends

Author: guggero

itest/addrs_test.go

@@ -3,6 +3,8 @@ package itest
 import (
 	"bytes"
 	"context"
+	"fmt"
+	"time"
 
 	"github.com/btcsuite/btcd/btcec/v2/schnorr"
 	"github.com/btcsuite/btcd/wire"
@@ -964,25 +966,40 @@ func sendAsset(t *harnessTest, sender *tapdHarness,
 
 	require.NotEmpty(t.t, options.sendAssetRequest.TapAddrs)
 
-	// We need the first address's scriptkey to subscribe to events.
-	firstAddr, err := address.DecodeAddress(
-		options.sendAssetRequest.TapAddrs[0], &address.RegressionNetTap,
-	)
-	require.NoError(t.t, err)
-	scriptKey := firstAddr.ScriptKey.SerializeCompressed()
+	// Assign a default transfer label using a Unix timestamp if none is
+	// provided. This will be used in filtering send events.
+	if options.sendAssetRequest.Label == "" {
+		options.sendAssetRequest.Label = fmt.Sprintf(
+			"%d", time.Now().UnixNano(),
+		)
+	}
 
+	// Construct send event stream.
 	ctxc, streamCancel := context.WithCancel(ctxb)
 	stream, err := sender.SubscribeSendEvents(
 		ctxc, &taprpc.SubscribeSendEventsRequest{
-			FilterScriptKey: scriptKey,
+			FilterLabel: options.sendAssetRequest.Label,
 		},
 	)
 	require.NoError(t.t, err)
+
+	// Formulate a subscription handler for the send event stream.
 	sub := &EventSubscription[*taprpc.SendEvent]{
 		ClientEventStream: stream,
 		Cancel:            streamCancel,
+
+		// Use the filter callback to ensure that the server side filter
+		// is working as expected.
+		ShouldNotify: func(e *taprpc.SendEvent) (bool, error) {
+			require.Equal(
+				t.t, e.TransferLabel,
+				options.sendAssetRequest.Label,
+			)
+			return true, nil
+		},
 	}
 
+	// Kick off the send asset request.
 	resp, err := sender.SendAsset(ctxt, &options.sendAssetRequest)
 	if options.errText != "" {
 		require.ErrorContains(t.t, err, options.errText)
@@ -992,8 +1009,12 @@ func sendAsset(t *harnessTest, sender *tapdHarness,
 	require.NoError(t.t, err)
 
 	// We'll get events up to the point where we broadcast the transaction.
+	//
+	// Don't specify a target script key here, we are already filtering
+	// events by the label.
+	var targetScriptKey []byte = nil
 	AssertSendEvents(
-		t.t, scriptKey, sub,
+		t.t, targetScriptKey, sub,
 		tapfreighter.SendStateVirtualCommitmentSelect,
 		tapfreighter.SendStateBroadcast,
 	)

itest/assertions.go

@@ -961,6 +961,36 @@ func AssertReceiveEvents(t *testing.T, addr *taprpc.Addr,
 	}
 }
 
+// makeFilterSendEventScriptKey returns a filter function that checks if the
+// given script key is present in the send event. If it is, the event is
+// included in the stream.
+func makeFilterSendEventScriptKey(
+	scriptKey btcec.PublicKey) func(*taprpc.SendEvent) (bool, error) {
+
+	return func(event *taprpc.SendEvent) (bool, error) {
+		for _, vPacketBytes := range event.VirtualPackets {
+			vPkt, err := tappsbt.Decode(vPacketBytes)
+			if err != nil {
+				return false, err
+			}
+
+			for _, vOut := range vPkt.Outputs {
+				if vOut.ScriptKey.PubKey == nil {
+					continue
+				}
+
+				// This packet sends to the filtered script key,
+				// we want to include this event.
+				if vOut.ScriptKey.PubKey.IsEqual(&scriptKey) {
+					return true, nil
+				}
+			}
+		}
+
+		return false, nil
+	}
+}
+
 // AssertSendEventsComplete makes sure the two remaining events for the given
 // script key are sent on the stream.
 func AssertSendEventsComplete(t *testing.T, scriptKey []byte,
@@ -974,39 +1004,27 @@ func AssertSendEventsComplete(t *testing.T, scriptKey []byte,
 
 // AssertSendEvents makes sure all events with incremental status are sent
 // on the stream for the given script key.
-func AssertSendEvents(t *testing.T, scriptKey []byte,
+func AssertSendEvents(t *testing.T, targetScriptKey []byte,
 	stream *EventSubscription[*taprpc.SendEvent], from,
 	to tapfreighter.SendState) {
 
 	success := make(chan struct{})
 	timeout := time.After(defaultWaitTimeout)
 	startStatus := from
 
-	targetScriptKey, err := btcec.ParsePubKey(scriptKey)
-	require.NoError(t, err)
-
-	sendsToKey := func(e *taprpc.SendEvent) bool {
-		for _, vPacketBytes := range e.VirtualPackets {
-			vPkt, err := tappsbt.Decode(vPacketBytes)
-			require.NoError(t, err)
-
-			for _, vOut := range vPkt.Outputs {
-				if vOut.ScriptKey.PubKey == nil {
-					continue
-				}
-
-				// This packet sends to the filtered script key,
-				// we want to include this event.
-				if vOut.ScriptKey.PubKey.IsEqual(
-					targetScriptKey,
-				) {
+	// By default, if the target script key is not given we will accept all
+	// send events.
+	sendsToKey := func(*taprpc.SendEvent) (bool, error) {
+		return true, nil
+	}
 
-					return true
-				}
-			}
-		}
+	// If a target script key is given, we will only accept send events that
+	// relate to that key.
+	if targetScriptKey != nil {
+		targetScriptKey, err := btcec.ParsePubKey(targetScriptKey)
+		require.NoError(t, err)
 
-		return false
+		sendsToKey = makeFilterSendEventScriptKey(*targetScriptKey)
 	}
 
 	// To make sure we don't forever hang on receiving on the stream, we'll
@@ -1026,7 +1044,11 @@ func AssertSendEvents(t *testing.T, scriptKey []byte,
 		event, err := stream.Recv()
 		require.NoError(t, err, "receiving send event")
 
-		require.True(t, sendsToKey(event))
+		// Ensure that the event complies with the target script key
+		// check.
+		res, err := sendsToKey(event)
+		require.NoError(t, err)
+		require.True(t, res)
 
 		// Check the event's error field for unexpected errors. Perform
 		// this check before verifying the expected send state, as

itest/utils.go

@@ -3,6 +3,7 @@ package itest
 import (
 	"bytes"
 	"context"
+	"io"
 	"testing"
 	"time"
 
@@ -61,6 +62,51 @@ type ClientEventStream[T any] interface {
 type EventSubscription[T any] struct {
 	ClientEventStream[T]
 	Cancel context.CancelFunc
+
+	// ShouldNotify is an optional filter predicate function that can be
+	// used to filter events received from the client stream.
+	//
+	// If set, it will be called for each event received from the stream. If
+	// it returns true, the event is returned. If it returns false, the
+	// event is ignored and the next event is received from the stream.
+	ShouldNotify func(T) (bool, error)
+}
+
+// Recv receives an event from the client stream. If a filter is set, it will
+// check if the event matches the filter. If it does, it returns the event.
+// If not, it continues receiving events until it finds one that matches the
+// filter.
+func (e *EventSubscription[T]) Recv() (T, error) {
+	var zero T
+
+	// If no filter predicate is set, we can just return the event.
+	if e.ShouldNotify == nil {
+		return e.ClientEventStream.Recv()
+	}
+
+	// If a filter is set, we need to check if the event matches the
+	// filter. If it does, we return the event. If not, we continue
+	// receiving events until we find one that matches the filter.
+	for {
+		event, err := e.ClientEventStream.Recv()
+		if err != nil {
+			if err == io.EOF {
+				// Handle end of stream.
+				return zero, err
+			}
+
+			return zero, err
+		}
+
+		match, err := e.ShouldNotify(event)
+		if err != nil {
+			return zero, err
+		}
+
+		if match {
+			return event, nil
+		}
+	}
 }
 
 // CopyRequest is a helper function to copy a request so that we can modify it.

rpcserver.go

@@ -3375,7 +3375,7 @@ func (r *rpcServer) SendAsset(ctx context.Context,
 	}
 
 	resp, err := r.cfg.ChainPorter.RequestShipment(
-		tapfreighter.NewAddressParcel(feeRate, tapAddrs...),
+		tapfreighter.NewAddressParcel(feeRate, req.Label, tapAddrs...),
 	)
 	if err != nil {
 		return nil, err
@@ -3662,6 +3662,7 @@ func marshalOutboundParcel(
 		AnchorTxBlockHeight: parcel.AnchorTxBlockHeight,
 		Inputs:              rpcInputs,
 		Outputs:             rpcOutputs,
+		Label:               parcel.Label,
 	}, nil
 }
 
@@ -3846,23 +3847,69 @@ func (r *rpcServer) SubscribeReceiveEvents(
 	)
 }
 
+// shouldNotifyAssetSendEvent returns true if the given AssetSendEvent matches
+// all specified filter criteria. Unset filter values are ignored. If set, all
+// filters must match (logical AND).
+func shouldNotifyAssetSendEvent(event tapfreighter.AssetSendEvent,
+	targetScriptKey fn.Option[btcec.PublicKey], targetLabel string) bool {
+
+	// By default, if no filter values are provided, match on event.
+	match := true
+
+	// Filter by label if specified.
+	if targetLabel != "" {
+		match = match && (event.TransferLabel == targetLabel)
+	}
+
+	// Filter by target script key if specified.
+	if targetScriptKey.IsSome() {
+		// If script key is specified but there are no virtual packets,
+		// early return as a match is impossible.
+		if len(event.VirtualPackets) == 0 {
+			return false
+		}
+
+		scriptKey := targetScriptKey.UnwrapToPtr()
+		found := false
+
+		for _, vPacket := range event.VirtualPackets {
+			if found {
+				break
+			}
+
+			for _, vOut := range vPacket.Outputs {
+				if vOut.ScriptKey.PubKey == nil {
+					continue
+				}
+				if vOut.ScriptKey.PubKey.IsEqual(scriptKey) {
+					found = true
+					break
+				}
+			}
+		}
+
+		match = match && found
+	}
+
+	return match
+}
+
 // SubscribeSendEvents registers a subscription to the event notification
 // stream which relates to the asset sending process.
 func (r *rpcServer) SubscribeSendEvents(req *taprpc.SubscribeSendEventsRequest,
 	ntfnStream sendEventStream) error {
 
-	var (
-		targetScriptKey *btcec.PublicKey
-		err             error
-	)
+	var targetScriptKey fn.Option[btcec.PublicKey]
 	if len(req.FilterScriptKey) > 0 {
-		targetScriptKey, err = btcec.ParsePubKey(req.FilterScriptKey)
+		scriptKey, err := btcec.ParsePubKey(req.FilterScriptKey)
 		if err != nil {
 			return fmt.Errorf("error parsing script key: %w", err)
 		}
+
+		targetScriptKey = fn.MaybeSome(scriptKey)
 	}
 
-	filter := func(event fn.Event) (bool, error) {
+	shouldNotify := func(event fn.Event) (bool, error) {
 		var e *tapfreighter.AssetSendEvent
 		switch typedEvent := event.(type) {
 		case *tapfreighter.AssetSendEvent:
@@ -3878,44 +3925,14 @@ func (r *rpcServer) SubscribeSendEvents(req *taprpc.SubscribeSendEventsRequest,
 				event)
 		}
 
-		// If we're not filtering on a specific script key, we return
-		// all events.
-		if targetScriptKey == nil {
-			return true, nil
-		}
-
-		// We can only match the script key on the active virtual
-		// packets, so if there are none, then this isn't an event we're
-		// interested in.
-		if len(e.VirtualPackets) == 0 {
-			return false, nil
-		}
-
-		for _, vPacket := range e.VirtualPackets {
-			for _, vOut := range vPacket.Outputs {
-				if vOut.ScriptKey.PubKey == nil {
-					continue
-				}
-
-				// This packet sends to the filtered script key,
-				// we want to include this event.
-				if vOut.ScriptKey.PubKey.IsEqual(
-					targetScriptKey,
-				) {
-
-					return true, nil
-				}
-			}
-		}
-
-		// No packets with outputs going to the filtered script key
-		// found.
-		return false, nil
+		return shouldNotifyAssetSendEvent(
+			*e, targetScriptKey, req.FilterLabel,
+		), nil
 	}
 
 	return handleEvents[bool, *taprpc.SendEvent](
-		r.cfg.ChainPorter, ntfnStream, marshalSendEvent, filter, r.quit,
-		false,
+		r.cfg.ChainPorter, ntfnStream, marshalSendEvent, shouldNotify,
+		r.quit, false,
 	)
 }
 
@@ -4162,6 +4179,7 @@ func marshalSendEvent(event fn.Event) (*taprpc.SendEvent, error) {
 		SendState:             e.SendState.String(),
 		VirtualPackets:        make([][]byte, len(e.VirtualPackets)),
 		PassiveVirtualPackets: make([][]byte, len(e.PassivePackets)),
+		TransferLabel:         e.TransferLabel,
 	}
 
 	if e.Error != nil {

tapdb/assets_store.go

@@ -2442,6 +2442,7 @@ func (a *AssetStore) LogPendingParcel(ctx context.Context,
 			HeightHint:       int32(spend.AnchorTxHeightHint),
 			AnchorTxid:       newAnchorTXID[:],
 			TransferTimeUnix: spend.TransferTime,
+			Label:            sqlStr(spend.Label),
 		})
 		if err != nil {
 			return fmt.Errorf("unable to insert asset transfer: "+
@@ -3562,6 +3563,7 @@ func (a *AssetStore) QueryParcels(ctx context.Context,
 				ChainFees:          dbAnchorTx.ChainFees,
 				Inputs:             inputs,
 				Outputs:            outputs,
+				Label:              dbT.Label.String,
 			}
 
 			// Set the block height if the anchor is marked as