Implement rsmt tree wrapper for nmt

p2p/ipld/nmt_wrapper.go

@@ -0,0 +1,91 @@
+package ipld
+
+import (
+	"crypto/sha256"
+
+	"github.com/lazyledger/lazyledger-core/types"
+	"github.com/lazyledger/nmt"
+	"github.com/lazyledger/nmt/namespace"
+	"github.com/lazyledger/rsmt2d"
+)
+
+// Fulfills the rsmt2d.Tree interface and rsmt2d.TreeConstructorFn function
+var _ rsmt2d.TreeConstructorFn = ErasuredNamespacedMerkleTree{}.Constructor
+var _ rsmt2d.Tree = &ErasuredNamespacedMerkleTree{}
+
+// ErasuredNamespacedMerkleTree wraps NamespaceMerkleTree to conform to the
+// rsmt2d.Tree interface while catering specifically to erasure data. For the
+// first half of the tree, it uses the first DefaultNamespaceIDLen number of
+// bytes of the data pushed to determine the namespace. For the second half, it
+// uses the parity namespace ID
+type ErasuredNamespacedMerkleTree struct {
+	squareSize uint64
+	pushCount  uint64
+	options    []nmt.Option
+	tree       *nmt.NamespacedMerkleTree
+}
+
+// NewErasuredNamespacedMerkleTree issues a new ErasuredNamespacedMerkleTree
+func NewErasuredNamespacedMerkleTree(squareSize uint64, setters ...nmt.Option) ErasuredNamespacedMerkleTree {
+	return ErasuredNamespacedMerkleTree{squareSize: squareSize, options: setters}
+}
+
+// Constructor acts as the rsmt2d.TreeConstructorFn for
+// ErasuredNamespacedMerkleTree
+func (w ErasuredNamespacedMerkleTree) Constructor() rsmt2d.Tree {
+	w.tree = nmt.New(sha256.New(), w.options...)
+	return &w
+}
+
+// Push adds the provided data to the underlying NamespaceMerkleTree, and
+// automatically uses the first DefaultNamespaceIDLen number of bytes as the
+// namespace unless the data pushed to the second half of the tree. Fulfills the
+// rsmt.Tree interface. NOTE: panics if there's an error pushing to underlying
+// NamespaceMerkleTree or if the tree size is exceeded
+func (w *ErasuredNamespacedMerkleTree) Push(data []byte) {
+	// determine the namespace based on where in the tree we're pushing
+	nsID := make(namespace.ID, types.NamespaceSize)
+
+	switch {
+	// panic if the tree size is exceeded
+	case w.pushCount > 2*w.squareSize:
+		panic("tree size exceeded")
+
+	// if the namespace is included in the data, use that ns
+	case w.pushCount+1 <= w.squareSize/2:
+		copy(nsID, data[:types.NamespaceSize])
+
+	// if the data is erasure data use the parity ns
+	default:
+		copy(nsID, types.ParitySharesNamespaceID)
+	}
+
+	// push to the underlying tree
+	err := w.tree.Push(nsID, data)
+	// panic on error
+	if err != nil {
+		panic(err)
+	}
+
+	w.pushCount++
+}
+
+// Prove fulfills the rsmt.Tree interface by generating and returning a single
+// leaf proof using the underlying NamespacedMerkleTree. NOTE: panics if the
+// underlying NamespaceMerkleTree errors.
+func (w *ErasuredNamespacedMerkleTree) Prove(
+	idx int,
+) (merkleRoot []byte, proofSet [][]byte, proofIndex uint64, numLeaves uint64) {
+	proof, err := w.tree.Prove(idx)
+	if err != nil {
+		panic(err)
+	}
+	nodes := proof.Nodes()
+	return w.Root(), nodes, uint64(proof.Start()), uint64(len(nodes))
+}
+
+// Root fulfills the rsmt.Tree interface by generating and returning the
+// underlying NamespaceMerkleTree Root.
+func (w *ErasuredNamespacedMerkleTree) Root() []byte {
+	return w.tree.Root().Bytes()
+}

p2p/ipld/nmt_wrapper_test.go

@@ -0,0 +1,120 @@
+package ipld
+
+import (
+	"crypto/sha256"
+	"testing"
+
+	"github.com/lazyledger/lazyledger-core/types"
+	"github.com/lazyledger/nmt"
+	"github.com/lazyledger/rsmt2d"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestPushErasuredNamespacedMerkleTree(t *testing.T) {
+	testCases := []struct {
+		name       string
+		squareSize int
+	}{
+		{"extendedSquareSize = 16", 8},
+		{"extendedSquareSize = 256", 128},
+	}
+	for _, tc := range testCases {
+		tc := tc
+		n := NewErasuredNamespacedMerkleTree(uint64(tc.squareSize))
+		tree := n.Constructor()
+
+		// push test data to the tree
+		for _, d := range generateErasuredData(t, tc.squareSize) {
+			// push will panic if there's an error
+			tree.Push(d)
+		}
+	}
+}
+
+func TestRootErasuredNamespacedMerkleTree(t *testing.T) {
+	// check that the root is different from a standard nmt tree this should be
+	// the case, because the ErasuredNamespacedMerkleTree should add namespaces
+	// to the second half of the tree
+	size := 16
+	data := generateRandNamespacedRawData(size, types.NamespaceSize, AdjustedMessageSize)
+	n := NewErasuredNamespacedMerkleTree(uint64(16))
+	tree := n.Constructor()
+	nmtTree := nmt.New(sha256.New())
+
+	for _, d := range data {
+		tree.Push(d)
+		err := nmtTree.Push(d[:types.NamespaceSize], d[types.NamespaceSize:])
+		if err != nil {
+			t.Error(err)
+		}
+	}
+
+	assert.NotEqual(t, nmtTree.Root().Bytes(), tree.Root())
+}
+
+func TestErasureNamespacedMerkleTreePanics(t *testing.T) {
+	testCases := []struct {
+		name  string
+		pFucn assert.PanicTestFunc
+	}{
+		{
+			"push over square size",
+			assert.PanicTestFunc(
+				func() {
+					data := generateErasuredData(t, 16)
+					n := NewErasuredNamespacedMerkleTree(uint64(15))
+					tree := n.Constructor()
+					for _, d := range data {
+						tree.Push(d)
+					}
+				}),
+		},
+		{
+			"push in incorrect lexigraphic order",
+			assert.PanicTestFunc(
+				func() {
+					data := generateErasuredData(t, 16)
+					n := NewErasuredNamespacedMerkleTree(uint64(16))
+					tree := n.Constructor()
+					for i := len(data) - 1; i > 0; i-- {
+						tree.Push(data[i])
+					}
+				},
+			),
+		},
+		{
+			"Prove non existent leaf",
+			assert.PanicTestFunc(
+				func() {
+					size := 16
+					data := generateErasuredData(t, size)
+					n := NewErasuredNamespacedMerkleTree(uint64(size))
+					tree := n.Constructor()
+					for _, d := range data {
+						tree.Push(d)
+					}
+					tree.Prove(size + 100)
+				},
+			),
+		},
+	}
+	for _, tc := range testCases {
+		assert.Panics(t, tc.pFucn)
+
+	}
+}
+
+// generateErasuredData produces a slice that is twice as long as it erasures
+// the data
+func generateErasuredData(t *testing.T, numLeaves int) [][]byte {
+	raw := generateRandNamespacedRawData(
+		numLeaves,
+		types.NamespaceSize,
+		AdjustedMessageSize,
+	)
+	erasuredData, err := rsmt2d.Encode(raw, rsmt2d.RSGF8)
+	if err != nil {
+		t.Error(err)
+	}
+	return append(raw, erasuredData...)
+}

p2p/ipld/read.go

@@ -8,6 +8,11 @@ import (
 	"github.com/ipfs/go-cid"
 	coreiface "github.com/ipfs/interface-go-ipfs-core"
 	"github.com/ipfs/interface-go-ipfs-core/path"
+	"github.com/lazyledger/lazyledger-core/types"
+)
+
+const (
+	AdjustedMessageSize = types.ShareSize - types.NamespaceSize
 )
 
 // /////////////////////////////////////