diff --git a/wallet/wallet.go b/wallet/wallet.go index e38e6c0..cf2990f 100644 --- a/wallet/wallet.go +++ b/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) { diff --git a/wallet/wallet_test.go b/wallet/wallet_test.go index f70f43c..5be8d74 100644 --- a/wallet/wallet_test.go +++ b/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) + } + }) + } +}