Repository: coinbase/rosetta-cli
Pull request #77 (2020-07-25 01:18:44)

Coin Storage (automated UTXO tracking)
Related PR: #68 

To support automated testing of UTXO-based blockchains, we need to provide coin tracking support.

### Changes
This PR introduces `CoinStorage` as a `BlockWorker` that automatically tracks coins that put either `utxo_created` or `utxo_spent` in the `Operation.Metadata`.


Patch:
diff --git a/internal/storage/block_storage.go b/internal/storage/block_storage.go
index c1f5b2a..64ee203 100644
--- a/internal/storage/block_storage.go
+++ b/internal/storage/block_storage.go
@@ -411,33 +411,26 @@ func (b *BlockStorage) storeTransactionHash(
 		return err
 	}
 
+	var blocks map[string]int64
 	if !exists {
-		m := make(map[string]int64)
-		m[blockIdentifier.Hash] = blockIdentifier.Index
-		encodedResult, err := encode(m)
-		if err != nil {
-			return fmt.Errorf("%w: unable to encode transaction data", err)
+		blocks = make(map[string]int64)
+	} else {
+		if err := decode(val, &blocks); err != nil {
+			return fmt.Errorf("%w: could not decode transaction hash contents", err)
 		}
 
-		return transaction.Set(ctx, hashKey, encodedResult)
-	}
-
-	var blocks map[string]int64
-	if err := decode(val, &blocks); err != nil {
-		return fmt.Errorf("%w: could not decode transaction hash contents", err)
-	}
-
-	if _, exists := blocks[blockIdentifier.Hash]; exists {
-		return fmt.Errorf(
-			"%w: duplicate transaction %s found in block %s:%d",
-			ErrDuplicateTransactionHash,
-			transactionIdentifier.Hash,
-			blockIdentifier.Hash,
-			blockIdentifier.Index,
-		)
+		if _, exists := blocks[blockIdentifier.Hash]; exists {
+			return fmt.Errorf(
+				"%w: duplicate transaction %s found in block %s:%d",
+				ErrDuplicateTransactionHash,
+				transactionIdentifier.Hash,
+				blockIdentifier.Hash,
+				blockIdentifier.Index,
+			)
+		}
 	}
-
 	blocks[blockIdentifier.Hash] = blockIdentifier.Index
+
 	encodedResult, err := encode(blocks)
 	if err != nil {
 		return fmt.Errorf("%w: unable to encode transaction data", err)
diff --git a/internal/storage/coin_storage.go b/internal/storage/coin_storage.go
new file mode 100644
index 0000000..c0af020
--- /dev/null
+++ b/internal/storage/coin_storage.go
@@ -0,0 +1,324 @@
+// Copyright 2020 Coinbase, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package storage
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/coinbase/rosetta-sdk-go/types"
+)
+
+const (
+	coinNamespace        = "coinNamespace"
+	coinAccountNamespace = "coinAccountNamespace"
+
+	// For UTXOs to be recognized by CoinStorage, they must contain
+	// coinCreated or coinSpent in the Operation.Metadata with a value
+	// of the unique identifier of the coin. In Bitcoin, this unique
+	// identifier would be the outpoint (tx_hash:index).
+	coinCreated = "utxo_created"
+	coinSpent   = "utxo_spent"
+)
+
+var _ BlockWorker = (*CoinStorage)(nil)
+
+// CoinStorage implements storage methods for storing
+// UTXOs.
+type CoinStorage struct {
+	db Database
+}
+
+// NewCoinStorage returns a new CoinStorage.
+func NewCoinStorage(
+	db Database,
+) *CoinStorage {
+	return &CoinStorage{
+		db: db,
+	}
+}
+
+// Coin represents some spendable output (typically
+// referred to as a UTXO).
+type Coin struct {
+	Identifier  string             `json:"identifier"` // uses "utxo_created" or "utxo_spent"
+	Transaction *types.Transaction `json:"transaction"`
+	Operation   *types.Operation   `json:"operation"`
+}
+
+func getCoinKey(identifier string) []byte {
+	return []byte(fmt.Sprintf("%s/%s", coinNamespace, identifier))
+}
+
+func getCoinAccountKey(accountIdentifier *types.AccountIdentifier) []byte {
+	return []byte(fmt.Sprintf("%s/%s", coinAccountNamespace, types.Hash(accountIdentifier)))
+}
+
+func getAndDecodeCoin(
+	ctx context.Context,
+	transaction DatabaseTransaction,
+	coinIdentifier string,
+) (bool, *Coin, error) {
+	exists, val, err := transaction.Get(ctx, getCoinKey(coinIdentifier))
+	if err != nil {
+		return false, nil, fmt.Errorf("%w: unable to query for coin", err)
+	}
+
+	if !exists { // this could occur if coin was created before we started syncing
+		return false, nil, nil
+	}
+
+	var coin Coin
+	if err := decode(val, &coin); err != nil {
+		return false, nil, fmt.Errorf("%w: unable to decode coin", err)
+	}
+
+	return true, &coin, nil
+}
+
+func (c *CoinStorage) tryAddingCoin(
+	ctx context.Context,
+	transaction DatabaseTransaction,
+	blockTransaction *types.Transaction,
+	operation *types.Operation,
+	identiferKey string,
+) error {
+	rawIdentifier, ok := operation.Metadata[identiferKey]
+	if ok {
+		coinIdentifier, ok := rawIdentifier.(string)
+		if !ok {
+			return fmt.Errorf("unable to parse created coin %v", rawIdentifier)
+		}
+
+		newCoin := &Coin{
+			Identifier:  coinIdentifier,
+			Transaction: blockTransaction,
+			Operation:   operation,
+		}
+
+		encodedResult, err := encode(newCoin)
+		if err != nil {
+			return fmt.Errorf("%w: unable to encode coin data", err)
+		}
+
+		if err := transaction.Set(ctx, getCoinKey(coinIdentifier), encodedResult); err != nil {
+			return fmt.Errorf("%w: unable to store coin", err)
+		}
+
+		accountExists, coins, err := getAndDecodeCoins(ctx, transaction, operation.Account)
+		if err != nil {
+			return fmt.Errorf("%w: unable to query coin account", err)
+		}
+
+		if !accountExists {
+			coins = map[string]struct{}{}
+		}
+
+		if _, exists := coins[coinIdentifier]; exists {
+			return fmt.Errorf(
+				"coin %s already exists in account %s",
+				coinIdentifier,
+				types.PrettyPrintStruct(operation.Account),
+			)
+		}
+
+		coins[coinIdentifier] = struct{}{}
+
+		if err := encodeAndSetCoins(ctx, transaction, operation.Account, coins); err != nil {
+			return fmt.Errorf("%w: unable to set coin account", err)
+		}
+	}
+
+	return nil
+}
+
+func encodeAndSetCoins(
+	ctx context.Context,
+	transaction DatabaseTransaction,
+	accountIdentifier *types.AccountIdentifier,
+	coins map[string]struct{},
+) error {
+	encodedResult, err := encode(coins)
+	if err != nil {
+		return fmt.Errorf("%w: unable to encode coins", err)
+	}
+
+	if err := transaction.Set(ctx, getCoinAccountKey(accountIdentifier), encodedResult); err != nil {
+		return fmt.Errorf("%w: unable to set coin account", err)
+	}
+
+	return nil
+}
+
+func getAndDecodeCoins(
+	ctx context.Context,
+	transaction DatabaseTransaction,
+	accountIdentifier *types.AccountIdentifier,
+) (bool, map[string]struct{}, error) {
+	accountExists, val, err := transaction.Get(ctx, getCoinAccountKey(accountIdentifier))
+	if err != nil {
+		return false, nil, fmt.Errorf("%w: unable to query coin account", err)
+	}
+
+	if !accountExists {
+		return false, nil, nil
+	}
+
+	var coins map[string]struct{}
+	if err := decode(val, &coins); err != nil {
+		return false, nil, fmt.Errorf("%w: unable to decode coin account", err)
+	}
+
+	return true, coins, nil
+}
+
+func (c *CoinStorage) tryRemovingCoin(
+	ctx context.Context,
+	transaction DatabaseTransaction,
+	operation *types.Operation,
+	identiferKey string,
+) error {
+	rawIdentifier, ok := operation.Metadata[identiferKey]
+	if ok {
+		coinIdentifier, ok := rawIdentifier.(string)
+		if !ok {
+			return fmt.Errorf("unable to parse spent coin %v", rawIdentifier)
+		}
+
+		exists, _, err := transaction.Get(ctx, getCoinKey(coinIdentifier))
+		if err != nil {
+			return fmt.Errorf("%w: unable to query for coin", err)
+		}
+
+		if !exists { // this could occur if coin was created before we started syncing
+			return nil
+		}
+
+		if err := transaction.Delete(ctx, getCoinKey(coinIdentifier)); err != nil {
+			return fmt.Errorf("%w: unable to delete coin", err)
+		}
+
+		accountExists, coins, err := getAndDecodeCoins(ctx, transaction, operation.Account)
+		if err != nil {
+			return fmt.Errorf("%w: unable to query coin account", err)
+		}
+
+		if !accountExists {
+			return fmt.Errorf("%w: unable to find owner of coin", err)
+		}
+
+		if _, exists := coins[coinIdentifier]; !exists {
+			return fmt.Errorf(
+				"unable to find coin %s in account %s",
+				coinIdentifier,
+				types.PrettyPrintStruct(operation.Account),
+			)
+		}
+
+		delete(coins, coinIdentifier)
+
+		if err := encodeAndSetCoins(ctx, transaction, operation.Account, coins); err != nil {
+			return fmt.Errorf("%w: unable to set coin account", err)
+		}
+	}
+
+	return nil
+}
+
+// AddingBlock is called by BlockStorage when adding a block.
+func (c *CoinStorage) AddingBlock(
+	ctx context.Context,
+	block *types.Block,
+	transaction DatabaseTransaction,
+) (CommitWorker, error) {
+	for _, txn := range block.Transactions {
+		for _, operation := range txn.Operations {
+			if operation.Amount == nil {
+				continue
+			}
+
+			if err := c.tryAddingCoin(ctx, transaction, txn, operation, coinCreated); err != nil {
+				return nil, fmt.Errorf("%w: unable to add coin", err)
+			}
+
+			if err := c.tryRemovingCoin(ctx, transaction, operation, coinSpent); err != nil {
+				return nil, fmt.Errorf("%w: unable to remove coin", err)
+			}
+		}
+	}
+
+	return nil, nil
+}
+
+// RemovingBlock is called by BlockStorage when removing a block.
+func (c *CoinStorage) RemovingBlock(
+	ctx context.Context,
+	block *types.Block,
+	transaction DatabaseTransaction,
+) (CommitWorker, error) {
+	for _, txn := range block.Transactions {
+		for _, operation := range txn.Operations {
+			if operation.Amount == nil {
+				continue
+			}
+
+			// We add spent coins and remove created coins during a re-org (opposite of
+			// AddingBlock).
+			if err := c.tryAddingCoin(ctx, transaction, txn, operation, coinSpent); err != nil {
+				return nil, fmt.Errorf("%w: unable to add coin", err)
+			}
+
+			if err := c.tryRemovingCoin(ctx, transaction, operation, coinCreated); err != nil {
+				return nil, fmt.Errorf("%w: unable to remove coin", err)
+			}
+		}
+	}
+
+	return nil, nil
+}
+
+// GetCoins returns all unspent coins for a provided *types.AccountIdentifier.
+func (c *CoinStorage) GetCoins(
+	ctx context.Context,
+	accountIdentifier *types.AccountIdentifier,
+) ([]*Coin, error) {
+	transaction := c.db.NewDatabaseTransaction(ctx, false)
+	defer transaction.Discard(ctx)
+
+	accountExists, coins, err := getAndDecodeCoins(ctx, transaction, accountIdentifier)
+	if err != nil {
+		return nil, fmt.Errorf("%w: unable to query account identifier", err)
+	}
+
+	if !accountExists {
+		return []*Coin{}, nil
+	}
+
+	coinArr := []*Coin{}
+	for coinIdentifier := range coins {
+		exists, coin, err := getAndDecodeCoin(ctx, transaction, coinIdentifier)
+		if err != nil {
+			return nil, fmt.Errorf("%w: unable to query coin", err)
+		}
+
+		if !exists {
+			return nil, fmt.Errorf("%w: unable to get coin %s", err, coinIdentifier)
+		}
+
+		coinArr = append(coinArr, coin)
+	}
+
+	return coinArr, nil
+}


Test patch:
diff --git a/internal/storage/coin_storage_test.go b/internal/storage/coin_storage_test.go
new file mode 100644
index 0000000..aef2102
--- /dev/null
+++ b/internal/storage/coin_storage_test.go
@@ -0,0 +1,249 @@
+// Copyright 2020 Coinbase, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package storage
+
+import (
+	"context"
+	"testing"
+
+	"github.com/coinbase/rosetta-cli/internal/utils"
+
+	"github.com/coinbase/rosetta-sdk-go/types"
+	"github.com/stretchr/testify/assert"
+)
+
+var (
+	account = &types.AccountIdentifier{
+		Address: "blah",
+	}
+
+	account2 = &types.AccountIdentifier{
+		Address: "blah2",
+	}
+
+	account3 = &types.AccountIdentifier{
+		Address: "blah",
+		SubAccount: &types.SubAccountIdentifier{
+			Address: "extra account",
+		},
+	}
+
+	accountCoins = []*Coin{
+		{
+			Identifier:  "coin1",
+			Transaction: coinBlock.Transactions[0],
+			Operation:   coinBlock.Transactions[0].Operations[0],
+		},
+	}
+
+	account2Coins = []*Coin{
+		{
+			Identifier:  "coin2",
+			Transaction: coinBlock.Transactions[0],
+			Operation:   coinBlock.Transactions[0].Operations[1],
+		},
+	}
+
+	account3Coins = []*Coin{
+		{
+			Identifier:  "coin3",
+			Transaction: coinBlock3.Transactions[0],
+			Operation:   coinBlock3.Transactions[0].Operations[0],
+		},
+		{
+			Identifier:  "coin4",
+			Transaction: coinBlock3.Transactions[1],
+			Operation:   coinBlock3.Transactions[1].Operations[0],
+		},
+	}
+
+	coinBlock = &types.Block{
+		Transactions: []*types.Transaction{
+			{
+				Operations: []*types.Operation{
+					{
+						Account: account,
+						Amount: &types.Amount{
+							Value: "10",
+						},
+						Metadata: map[string]interface{}{
+							coinCreated: "coin1",
+						},
+					},
+					{
+						Account: account2,
+						Amount: &types.Amount{
+							Value: "15",
+						},
+						Metadata: map[string]interface{}{
+							coinSpent: "coin2",
+						},
+					},
+				},
+			},
+		},
+	}
+
+	coinBlock2 = &types.Block{
+		Transactions: []*types.Transaction{
+			{
+				Operations: []*types.Operation{
+					{
+						Account: account,
+						Amount: &types.Amount{
+							Value: "-10",
+						},
+						Metadata: map[string]interface{}{
+							coinSpent: "coin1",
+						},
+					},
+				},
+			},
+		},
+	}
+
+	coinBlock3 = &types.Block{
+		Transactions: []*types.Transaction{
+			{
+				Operations: []*types.Operation{
+					{
+						Account: account3,
+						Amount: &types.Amount{
+							Value: "4",
+						},
+						Metadata: map[string]interface{}{
+							coinCreated: "coin3",
+						},
+					},
+				},
+			},
+			{
+				Operations: []*types.Operation{
+					{
+						Account: account3,
+						Amount: &types.Amount{
+							Value: "6",
+						},
+						Metadata: map[string]interface{}{
+							coinCreated: "coin4",
+						},
+					},
+				},
+			},
+		},
+	}
+)
+
+func TestCoinStorage(t *testing.T) {
+	ctx := context.Background()
+
+	newDir, err := utils.CreateTempDir()
+	assert.NoError(t, err)
+	defer utils.RemoveTempDir(newDir)
+
+	database, err := NewBadgerStorage(ctx, newDir)
+	assert.NoError(t, err)
+	defer database.Close(ctx)
+
+	c := NewCoinStorage(database)
+
+	t.Run("get coins of unset account", func(t *testing.T) {
+		coins, err := c.GetCoins(ctx, account)
+		assert.NoError(t, err)
+		assert.Equal(t, []*Coin{}, coins)
+	})
+
+	t.Run("add block", func(t *testing.T) {
+		tx := c.db.NewDatabaseTransaction(ctx, true)
+		commitFunc, err := c.AddingBlock(ctx, coinBlock, tx)
+		assert.Nil(t, commitFunc)
+		assert.NoError(t, err)
+		assert.NoError(t, tx.Commit(ctx))
+
+		coins, err := c.GetCoins(ctx, account)
+		assert.NoError(t, err)
+		assert.Equal(t, accountCoins, coins)
+	})
+
+	t.Run("add duplicate coin", func(t *testing.T) {
+		tx := c.db.NewDatabaseTransaction(ctx, true)
+		commitFunc, err := c.AddingBlock(ctx, coinBlock, tx)
+		assert.Nil(t, commitFunc)
+		assert.Error(t, err)
+		tx.Discard(ctx)
+
+		coins, err := c.GetCoins(ctx, account)
+		assert.NoError(t, err)
+		assert.Equal(t, accountCoins, coins)
+	})
+
+	t.Run("remove block", func(t *testing.T) {
+		tx := c.db.NewDatabaseTransaction(ctx, true)
+		commitFunc, err := c.RemovingBlock(ctx, coinBlock, tx)
+		assert.Nil(t, commitFunc)
+		assert.NoError(t, err)
+		assert.NoError(t, tx.Commit(ctx))
+
+		coins, err := c.GetCoins(ctx, account)
+		assert.NoError(t, err)
+		assert.Equal(t, []*Coin{}, coins)
+
+		coins, err = c.GetCoins(ctx, account2)
+		assert.NoError(t, err)
+		assert.Equal(t, account2Coins, coins)
+	})
+
+	t.Run("spend coin", func(t *testing.T) {
+		tx := c.db.NewDatabaseTransaction(ctx, true)
+		commitFunc, err := c.AddingBlock(ctx, coinBlock, tx)
+		assert.Nil(t, commitFunc)
+		assert.NoError(t, err)
+		assert.NoError(t, tx.Commit(ctx))
+
+		coins, err := c.GetCoins(ctx, account)
+		assert.NoError(t, err)
+		assert.Equal(t, accountCoins, coins)
+
+		tx = c.db.NewDatabaseTransaction(ctx, true)
+		commitFunc, err = c.AddingBlock(ctx, coinBlock2, tx)
+		assert.Nil(t, commitFunc)
+		assert.NoError(t, err)
+		assert.NoError(t, tx.Commit(ctx))
+
+		coins, err = c.GetCoins(ctx, account)
+		assert.NoError(t, err)
+		assert.Equal(t, []*Coin{}, coins)
+
+		coins, err = c.GetCoins(ctx, account2)
+		assert.NoError(t, err)
+		assert.Equal(t, []*Coin{}, coins)
+	})
+
+	t.Run("add block with multiple outputs for 1 account", func(t *testing.T) {
+		tx := c.db.NewDatabaseTransaction(ctx, true)
+		commitFunc, err := c.AddingBlock(ctx, coinBlock3, tx)
+		assert.Nil(t, commitFunc)
+		assert.NoError(t, err)
+		assert.NoError(t, tx.Commit(ctx))
+
+		coins, err := c.GetCoins(ctx, account)
+		assert.NoError(t, err)
+		assert.Equal(t, []*Coin{}, coins)
+
+		coins, err = c.GetCoins(ctx, account3)
+		assert.NoError(t, err)
+		assert.ElementsMatch(t, account3Coins, coins)
+	})
+}