wallet: add label transaction function

Add function which allows you to retrospectively label a transaction,
optionally overwriting the existing label.

wallet/wallet.go

@@ -65,6 +65,16 @@ var (
 	// down.
 	ErrWalletShuttingDown = errors.New("wallet shutting down")
 
+	// ErrUnknownTransaction is returned when an attempt is made to label
+	// a transaction that is not known to the wallet.
+	ErrUnknownTransaction = errors.New("cannot label transaction not " +
+		"known to wallet")
+
+	// ErrTxLabelExists is returned when a transaction already has a label
+	// and an attempt has been made to label it without setting overwrite
+	// to true.
+	ErrTxLabelExists = errors.New("transaction already labelled")
+
 	// Namespace bucket keys.
 	waddrmgrNamespaceKey = []byte("waddrmgr")
 	wtxmgrNamespaceKey   = []byte("wtxmgr")
@@ -1591,6 +1601,58 @@ func (w *Wallet) PubKeyForAddress(a btcutil.Address) (*btcec.PublicKey, error) {
 	return pubKey, err
 }
 
+// LabelTransaction adds a label to the transaction with the hash provided. The
+// call will fail if the label is too long, or if the transaction already has
+// a label and the overwrite boolean is not set.
+func (w *Wallet) LabelTransaction(hash chainhash.Hash, label string,
+	overwrite bool) error {
+
+	// Check that the transaction is known to the wallet, and fail if it is
+	// unknown. If the transaction is known, check whether it already has
+	// a label.
+	err := walletdb.View(w.db, func(tx walletdb.ReadTx) error {
+		txmgrNs := tx.ReadBucket(wtxmgrNamespaceKey)
+
+		dbTx, err := w.TxStore.TxDetails(txmgrNs, &hash)
+		if err != nil {
+			return err
+		}
+
+		// If the transaction looked up is nil, it was not found. We
+		// do not allow labelling of unknown transactions so we fail.
+		if dbTx == nil {
+			return ErrUnknownTransaction
+		}
+
+		_, err = wtxmgr.FetchTxLabel(txmgrNs, hash)
+		return err
+	})
+
+	switch err {
+	// If no labels have been written yet, we can silence the error.
+	// Likewise if there is no label, we do not need to do any overwrite
+	// checks.
+	case wtxmgr.ErrNoLabelBucket:
+	case wtxmgr.ErrTxLabelNotFound:
+
+	// If we successfully looked up a label, fail if the overwrite param
+	// is not set.
+	case nil:
+		if !overwrite {
+			return ErrTxLabelExists
+		}
+
+	// In another unrelated error occurred, return it.
+	default:
+		return err
+	}
+
+	return walletdb.Update(w.db, func(tx walletdb.ReadWriteTx) error {
+		txmgrNs := tx.ReadWriteBucket(wtxmgrNamespaceKey)
+		return w.TxStore.PutTxLabel(txmgrNs, hash, label)
+	})
+}
+
 // PrivKeyForAddress looks up the associated private key for a P2PKH or P2PK
 // address.
 func (w *Wallet) PrivKeyForAddress(a btcutil.Address) (*btcec.PrivateKey, error) {

wallet/wallet_test.go

@@ -1,8 +1,20 @@
 package wallet
 
 import (
+	"encoding/hex"
 	"testing"
 	"time"
+
+	"github.com/btcsuite/btcwallet/walletdb"
+	"github.com/btcsuite/btcwallet/wtxmgr"
+
+	"github.com/btcsuite/btcutil"
+)
+
+var (
+	TstSerializedTx, _ = hex.DecodeString("010000000114d9ff358894c486b4ae11c2a8cf7851b1df64c53d2e511278eff17c22fb7373000000008c493046022100995447baec31ee9f6d4ec0e05cb2a44f6b817a99d5f6de167d1c75354a946410022100c9ffc23b64d770b0e01e7ff4d25fbc2f1ca8091053078a247905c39fce3760b601410458b8e267add3c1e374cf40f1de02b59213a82e1d84c2b94096e22e2f09387009c96debe1d0bcb2356ffdcf65d2a83d4b34e72c62eccd8490dbf2110167783b2bffffffff0280969800000000001976a914479ed307831d0ac19ebc5f63de7d5f1a430ddb9d88ac38bfaa00000000001976a914dadf9e3484f28b385ddeaa6c575c0c0d18e9788a88ac00000000")
+	TstTx, _           = btcutil.NewTxFromBytes(TstSerializedTx)
+	TstTxHash          = TstTx.Hash()
 )
 
 // TestLocateBirthdayBlock ensures we can properly map a block in the chain to a
@@ -83,3 +95,110 @@ func TestLocateBirthdayBlock(t *testing.T) {
 		}
 	}
 }
+
+// TestLabelTransaction tests labelling of transactions with invalid labels,
+// and failure to label a transaction when it already has a label.
+func TestLabelTransaction(t *testing.T) {
+	tests := []struct {
+		name string
+
+		// Whether the transaction should be known to the wallet.
+		txKnown bool
+
+		// Whether the test should write an existing label to disk.
+		existingLabel bool
+
+		// The overwrite parameter to call label transaction with.
+		overwrite bool
+
+		// The error we expect to be returned.
+		expectedErr error
+	}{
+		{
+			name:          "existing label, not overwrite",
+			txKnown:       true,
+			existingLabel: true,
+			overwrite:     false,
+			expectedErr:   ErrTxLabelExists,
+		},
+		{
+			name:          "existing label, overwritten",
+			txKnown:       true,
+			existingLabel: true,
+			overwrite:     true,
+			expectedErr:   nil,
+		},
+		{
+			name:          "no prexisting label, ok",
+			txKnown:       true,
+			existingLabel: false,
+			overwrite:     false,
+			expectedErr:   nil,
+		},
+		{
+			name:          "transaction unknown",
+			txKnown:       false,
+			existingLabel: false,
+			overwrite:     false,
+			expectedErr:   ErrUnknownTransaction,
+		},
+	}
+
+	for _, test := range tests {
+		test := test
+
+		t.Run(test.name, func(t *testing.T) {
+			t.Parallel()
+
+			w, cleanup := testWallet(t)
+			defer cleanup()
+
+			// If the transaction should be known to the store, we
+			// write txdetail to disk.
+			if test.txKnown {
+				rec, err := wtxmgr.NewTxRecord(
+					TstSerializedTx, time.Now(),
+				)
+				if err != nil {
+					t.Fatal(err)
+				}
+
+				err = walletdb.Update(w.db,
+					func(tx walletdb.ReadWriteTx) error {
+
+						ns := tx.ReadWriteBucket(
+							wtxmgrNamespaceKey,
+						)
+
+						return w.TxStore.InsertTx(
+							ns, rec, nil,
+						)
+					})
+				if err != nil {
+					t.Fatalf("could not insert tx: %v", err)
+				}
+			}
+
+			// If we want to setup an existing label for the purpose
+			// of the test, write one to disk.
+			if test.existingLabel {
+				err := w.LabelTransaction(
+					*TstTxHash, "existing label", false,
+				)
+				if err != nil {
+					t.Fatalf("could not write label: %v",
+						err)
+				}
+			}
+
+			newLabel := "new label"
+			err := w.LabelTransaction(
+				*TstTxHash, newLabel, test.overwrite,
+			)
+			if err != test.expectedErr {
+				t.Fatalf("expected: %v, got: %v",
+					test.expectedErr, err)
+			}
+		})
+	}
+}