[Utility] Use TreeStore as source of truth (#937)

ibc/ibc_msg_mempool_test.go

@@ -250,10 +250,39 @@ func TestHandleMessage_ErrorAlreadyInMempool(t *testing.T) {
 func TestHandleMessage_ErrorAlreadyCommitted(t *testing.T) {
 	// Prepare the environment
 	_, _, utilityMod, persistenceMod, _ := prepareEnvironment(t, 0, 0, 0, 0)
-	idxTx := prepareIndexedMessage(t, persistenceMod.GetTxIndexer())
+
+	privKey, err := crypto.GeneratePrivateKey()
+	require.NoError(t, err)
+	_, validPruneTx := preparePruneMessage(t, []byte("key"))
+	require.NoError(t, err)
+	err = validPruneTx.Sign(privKey)
+	require.NoError(t, err)
+	txProtoBytes, err := codec.GetCodec().Marshal(validPruneTx)
+	require.NoError(t, err)
+
+	idxTx := &coreTypes.IndexedTransaction{
+		Tx:            txProtoBytes,
+		Height:        0,
+		Index:         0,
+		ResultCode:    0,
+		Error:         "h5law",
+		SignerAddr:    "h5law",
+		RecipientAddr: "h5law",
+		MessageType:   "h5law",
+	}
+
+	// Index a test transaction
+	err = persistenceMod.GetTxIndexer().Index(idxTx)
+	require.NoError(t, err)
+
+	rwCtx, err := persistenceMod.NewRWContext(0)
+	require.NoError(t, err)
+	_, err = rwCtx.ComputeStateHash()
+	require.NoError(t, err)
+	rwCtx.Release()
 
 	// Error on having an indexed transaction
-	err := utilityMod.HandleTransaction(idxTx.Tx)
+	err = utilityMod.HandleTransaction(idxTx.Tx)
 	require.Error(t, err)
 	require.EqualError(t, err, coreTypes.ErrTransactionAlreadyCommitted().Error())
 }

persistence/block.go

@@ -1,33 +1,29 @@
 package persistence
 
 import (
+	"bytes"
 	"encoding/hex"
-	"errors"
 	"fmt"
 
-	"github.com/dgraph-io/badger/v3"
+	"github.com/pokt-network/pocket/persistence/trees"
 	"github.com/pokt-network/pocket/persistence/types"
 	"github.com/pokt-network/pocket/shared/codec"
 	coreTypes "github.com/pokt-network/pocket/shared/core/types"
 	"github.com/pokt-network/pocket/shared/crypto"
 	"google.golang.org/protobuf/types/known/timestamppb"
 )
 
-func (p *persistenceModule) TransactionExists(transactionHash string) (bool, error) {
-	hash, err := hex.DecodeString(transactionHash)
+func (p *persistenceModule) TransactionExists(txHash, txProtoBz []byte) (bool, error) {
+	exists, err := p.GetBus().GetTreeStore().Prove(trees.TransactionsTreeName, txHash, txProtoBz)
 	if err != nil {
 		return false, err
 	}
-	res, err := p.txIndexer.GetByHash(hash)
-	if res == nil {
-		// check for not found
-		if err != nil && errors.Is(err, badger.ErrKeyNotFound) {
-			return false, nil
-		}
-		return false, err
+	// exclusion proof verification
+	if bytes.Equal(txProtoBz, nil) && exists {
+		return false, nil
 	}
-
-	return true, nil
+	// inclusion proof verification
+	return exists, nil
 }
 
 func (p *PostgresContext) GetMinimumBlockHeight() (latestHeight uint64, err error) {

persistence/trees/trees.go

@@ -103,6 +103,24 @@ func (t *treeStore) GetTree(name string) ([]byte, kvstore.KVStore) {
 	return nil, nil
 }
 
+// Prove generates and verifies a proof against the tree name stored in the TreeStore
+// using the given key-value pair. If value == nil this will be an exclusion proof,
+// otherwise it will be an inclusion proof.
+func (t *treeStore) Prove(name string, key, value []byte) (bool, error) {
+	st, ok := t.merkleTrees[name]
+	if !ok {
+		return false, fmt.Errorf("tree not found: %s", name)
+	}
+	proof, err := st.tree.Prove(key)
+	if err != nil {
+		return false, fmt.Errorf("error generating proof (%s): %w", name, err)
+	}
+	if valid := smt.VerifyProof(proof, st.tree.Root(), key, value, st.tree.Spec()); !valid {
+		return false, nil
+	}
+	return true, nil
+}
+
 // GetTreeHashes returns a map of tree names to their root hashes for all
 // the trees tracked by the treestore, excluding the root tree
 func (t *treeStore) GetTreeHashes() map[string]string {

persistence/trees/trees_test.go

@@ -1,6 +1,13 @@
 package trees
 
-import "testing"
+import (
+	"fmt"
+	"testing"
+
+	"github.com/pokt-network/pocket/persistence/kvstore"
+	"github.com/pokt-network/smt"
+	"github.com/stretchr/testify/require"
+)
 
 // TECHDEBT(#836): Tests added in https://github.com/pokt-network/pocket/pull/836
 func TestTreeStore_Update(t *testing.T) {
@@ -22,3 +29,83 @@ func TestTreeStore_DebugClearAll(t *testing.T) {
 func TestTreeStore_GetTreeHashes(t *testing.T) {
 	t.Skip("TODO: Write test case for GetTreeHashes method") // context: https://github.com/pokt-network/pocket/pull/915#discussion_r1267313664
 }
+
+func TestTreeStore_Prove(t *testing.T) {
+	nodeStore := kvstore.NewMemKVStore()
+	tree := smt.NewSparseMerkleTree(nodeStore, smtTreeHasher)
+	testTree := &stateTree{
+		name:      "test",
+		tree:      tree,
+		nodeStore: nodeStore,
+	}
+
+	require.NoError(t, testTree.tree.Update([]byte("key"), []byte("value")))
+	require.NoError(t, testTree.tree.Commit())
+
+	treeStore := &treeStore{
+		merkleTrees: make(map[string]*stateTree, 1),
+	}
+	treeStore.merkleTrees["test"] = testTree
+
+	testCases := []struct {
+		name        string
+		treeName    string
+		key         []byte
+		value       []byte
+		valid       bool
+		expectedErr error
+	}{
+		{
+			name:        "valid inclusion proof: key and value in tree",
+			treeName:    "test",
+			key:         []byte("key"),
+			value:       []byte("value"),
+			valid:       true,
+			expectedErr: nil,
+		},
+		{
+			name:        "valid exclusion proof: key not in tree",
+			treeName:    "test",
+			key:         []byte("key2"),
+			value:       nil,
+			valid:       true,
+			expectedErr: nil,
+		},
+		{
+			name:        "invalid proof: tree not in store",
+			treeName:    "unstored tree",
+			key:         []byte("key"),
+			value:       []byte("value"),
+			valid:       false,
+			expectedErr: fmt.Errorf("tree not found: %s", "unstored tree"),
+		},
+		{
+			name:        "invalid inclusion proof: key in tree, wrong value",
+			treeName:    "test",
+			key:         []byte("key"),
+			value:       []byte("wrong value"),
+			valid:       false,
+			expectedErr: nil,
+		},
+		{
+			name:        "invalid exclusion proof: key in tree",
+			treeName:    "test",
+			key:         []byte("key"),
+			value:       nil,
+			valid:       false,
+			expectedErr: nil,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			valid, err := treeStore.Prove(tc.treeName, tc.key, tc.value)
+			require.Equal(t, valid, tc.valid)
+			if tc.expectedErr == nil {
+				require.NoError(t, err)
+				return
+			}
+			require.ErrorAs(t, err, &tc.expectedErr)
+		})
+	}
+}

shared/modules/persistence_module.go

@@ -33,7 +33,9 @@ type PersistenceModule interface {
 
 	// Indexer operations
 	GetTxIndexer() indexer.TxIndexer
-	TransactionExists(transactionHash string) (bool, error)
+
+	// TreeStore operations
+	TransactionExists(txHash, txProtoBz []byte) (bool, error)
 
 	// Debugging / development only
 	HandleDebugMessage(*messaging.DebugMessage) error

shared/modules/treestore_module.go

@@ -31,6 +31,9 @@ type TreeStoreModule interface {
 	Update(pgtx pgx.Tx, height uint64) (string, error)
 	// DebugClearAll completely clears the state of the trees. For debugging purposes only.
 	DebugClearAll() error
+	// Prove generates and verifies a proof against the tree with the matching name using the given
+	// key and value. If value == nil, it will verify non-membership of the key, otherwise membership.
+	Prove(treeName string, key, value []byte) (bool, error)
 	// GetTree returns the specified tree's root and nodeStore in order to be imported elsewhere
 	GetTree(name string) ([]byte, kvstore.KVStore)
 	// GetTreeHashes returns a map of tree names to their root hashes

utility/transaction.go

@@ -7,19 +7,21 @@ import (
 	"github.com/dgraph-io/badger/v3"
 	"github.com/pokt-network/pocket/shared/codec"
 	coreTypes "github.com/pokt-network/pocket/shared/core/types"
+	"github.com/pokt-network/pocket/shared/crypto"
 )
 
 // HandleTransaction implements the exposed functionality of the shared utilityModule interface.
 func (u *utilityModule) HandleTransaction(txProtoBytes []byte) error {
-	txHash := coreTypes.TxHash(txProtoBytes)
+	txHash := crypto.SHA3Hash(txProtoBytes)
 
 	// Is the tx already in the mempool (in memory)?
-	if u.mempool.Contains(txHash) {
+	if u.mempool.Contains(hex.EncodeToString(txHash)) {
 		return coreTypes.ErrDuplicateTransaction()
 	}
 
 	// Is the tx already committed & indexed (on disk)?
-	if txExists, err := u.GetBus().GetPersistenceModule().TransactionExists(txHash); err != nil {
+	txExists, err := u.GetBus().GetPersistenceModule().TransactionExists(txHash, txProtoBytes)
+	if err != nil {
 		return err
 	} else if txExists {
 		return coreTypes.ErrTransactionAlreadyCommitted()