service/share, ipld: adds ability to retrieve shares by namespace ID (#…

…170)

go.mod

@@ -7,7 +7,7 @@ replace github.com/ipfs/go-verifcid => github.com/celestiaorg/go-verifcid v0.0.1
 require (
 	github.com/BurntSushi/toml v0.4.1
 	github.com/celestiaorg/celestia-core v0.0.2-0.20210924001615-488ac31b4b3c
-	github.com/celestiaorg/nmt v0.7.0
+	github.com/celestiaorg/nmt v0.8.0
 	github.com/celestiaorg/rsmt2d v0.3.0
 	github.com/gogo/protobuf v1.3.2
 	github.com/hashicorp/golang-lru v0.5.4

ipld/errors.go

@@ -0,0 +1,5 @@
+package ipld
+
+import "fmt"
+
+var ErrNotFoundInRange = fmt.Errorf("namespaceID not found in range")

ipld/plugin/nmt.go

@@ -44,6 +44,11 @@ const (
 	namespaceSize = 8
 	// nmtHashSize is the size of a digest created by an NMT in bytes.
 	nmtHashSize = 2*namespaceSize + sha256.Size
+
+	// mhOverhead is the size of the prepended buffer of the CID encoding
+	// for NamespacedSha256. For more information, see:
+	// https://multiformats.io/multihash/#the-multihash-format
+	mhOverhead = 4
 )
 
 func init() {
@@ -406,3 +411,8 @@ func MustCidFromNamespacedSha256(hash []byte) cid.Cid {
 	}
 	return cidFromHash
 }
+
+// NamespacedSha256FromCID derives the Namespaced hash from the given CID.
+func NamespacedSha256FromCID(cid cid.Cid) []byte {
+	return cid.Hash()[mhOverhead:]
+}

ipld/plugin/nmt_test.go

@@ -0,0 +1,43 @@
+package plugin
+
+import (
+	"math"
+	"strconv"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/celestiaorg/celestia-core/pkg/da"
+	"github.com/celestiaorg/celestia-core/testutils"
+)
+
+// TestNamespaceFromCID checks that deriving the Namespaced hash from
+// the given CID works correctly.
+func TestNamespaceFromCID(t *testing.T) {
+	var tests = []struct {
+		randData [][]byte
+	}{
+		{randData: testutils.GenerateRandNamespacedRawData(4, namespaceSize, ShareSize)},
+		{randData: testutils.GenerateRandNamespacedRawData(16, 16, ShareSize)},
+		{randData: testutils.GenerateRandNamespacedRawData(4, 4, ShareSize)},
+		{randData: testutils.GenerateRandNamespacedRawData(4, namespaceSize, ShareSize/2)},
+	}
+
+	for i, tt := range tests {
+		t.Run(strconv.Itoa(i), func(t *testing.T) {
+			// create DAH from rand data
+			squareSize := uint64(math.Sqrt(float64(len(tt.randData))))
+			dah, err := da.NewDataAvailabilityHeader(squareSize, tt.randData)
+			require.NoError(t, err)
+			// check to make sure NamespacedHash is correctly derived from CID
+			for _, row := range dah.RowsRoots {
+				c, err := CidFromNamespacedSha256(row)
+				require.NoError(t, err)
+
+				got := NamespacedSha256FromCID(c)
+				assert.Equal(t, row, got)
+			}
+		})
+	}
+}

ipld/read.go

@@ -9,11 +9,12 @@ import (
 	"github.com/ipfs/go-cid"
 	ipld "github.com/ipfs/go-ipld-format"
 
-	"github.com/celestiaorg/rsmt2d"
-
 	"github.com/celestiaorg/celestia-core/pkg/da"
 	"github.com/celestiaorg/celestia-core/pkg/wrapper"
 	"github.com/celestiaorg/celestia-node/ipld/plugin"
+	"github.com/celestiaorg/nmt"
+	"github.com/celestiaorg/nmt/namespace"
+	"github.com/celestiaorg/rsmt2d"
 )
 
 var ErrRetrieveTimeout = errors.New("retrieve data timeout")
@@ -220,7 +221,7 @@ func GetLeafData(
 	return nd.RawData()[1:], nil
 }
 
-// GetLeafData fetches and returns the raw leaf.
+// GetLeaf fetches and returns the raw leaf.
 // It walks down the IPLD NMT tree until it finds the requested one.
 func GetLeaf(ctx context.Context, dag ipld.NodeGetter, root cid.Cid, leaf, total int) (ipld.Node, error) {
 	// request the node
@@ -247,3 +248,41 @@ func GetLeaf(ctx context.Context, dag ipld.NodeGetter, root cid.Cid, leaf, total
 	// recursively walk down through selected children
 	return GetLeaf(ctx, dag, root, leaf, total)
 }
+
+// GetLeavesByNamespace returns all the shares from the given DataAvailabilityHeader root
+// with the given namespace.ID.
+func GetLeavesByNamespace(
+	ctx context.Context,
+	dag ipld.NodeGetter,
+	root cid.Cid,
+	nID namespace.ID,
+) (out []ipld.Node, err error) {
+	rootH := plugin.NamespacedSha256FromCID(root)
+	if nID.Less(nmt.MinNamespace(rootH, nID.Size())) || !nID.LessOrEqual(nmt.MaxNamespace(rootH, nID.Size())) {
+		return nil, ErrNotFoundInRange
+	}
+	// request the node
+	nd, err := dag.Get(ctx, root)
+	if err != nil {
+		return
+	}
+	// check links
+	lnks := nd.Links()
+	if len(lnks) == 1 {
+		// if there is one link, then this is a leaf node, so just return it
+		out = append(out, nd)
+		return
+	}
+	// if there are some links, then traverse them
+	for _, lnk := range nd.Links() {
+		nds, err := GetLeavesByNamespace(ctx, dag, lnk.Cid, nID)
+		if err != nil {
+			if err == ErrNotFoundInRange {
+				// There is always right and left child and it is ok if one of them does not have a required nID.
+				continue
+			}
+		}
+		out = append(nds, out...)
+	}
+	return out, err
+}

ipld/read_test.go

@@ -6,6 +6,7 @@ import (
 	"crypto/sha256"
 	"math"
 	"math/rand"
+	"strconv"
 	"testing"
 	"time"
 
@@ -15,12 +16,14 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
-	"github.com/celestiaorg/nmt"
-	"github.com/celestiaorg/rsmt2d"
-
+	"github.com/celestiaorg/celestia-core/pkg/da"
 	"github.com/celestiaorg/celestia-core/pkg/wrapper"
+	"github.com/celestiaorg/celestia-core/testutils"
 	"github.com/celestiaorg/celestia-node/ipld/plugin"
 	"github.com/celestiaorg/celestia-node/service/header"
+	"github.com/celestiaorg/nmt"
+	"github.com/celestiaorg/nmt/namespace"
+	"github.com/celestiaorg/rsmt2d"
 )
 
 func TestGetLeafData(t *testing.T) {
@@ -125,7 +128,7 @@ func TestRetrieveBlockData(t *testing.T) {
 
 		tc := tc
 		t.Run(tc.name, func(t *testing.T) {
-			// // generate EDS
+			// generate EDS
 			eds := generateRandEDS(t, tc.squareSize)
 
 			shares := ExtractODSShares(eds)
@@ -231,3 +234,66 @@ func removeRandShares(data [][]byte, d int) [][]byte {
 	}
 	return data
 }
+
+func TestGetLeavesByNamespace(t *testing.T) {
+	var tests = []struct {
+		rawData [][]byte
+	}{
+		{rawData: testutils.GenerateRandNamespacedRawData(16, NamespaceSize, plugin.ShareSize)},
+		{rawData: testutils.GenerateRandNamespacedRawData(16, NamespaceSize, 8)},
+		{rawData: testutils.GenerateRandNamespacedRawData(4, NamespaceSize, plugin.ShareSize)},
+		{rawData: testutils.GenerateRandNamespacedRawData(16, NamespaceSize, 8)},
+	}
+
+	for i, tt := range tests {
+		t.Run(strconv.Itoa(i), func(t *testing.T) {
+			squareSize := uint64(math.Sqrt(float64(len(tt.rawData))))
+
+			// choose random nID from rand shares
+			expected := tt.rawData[len(tt.rawData)/2]
+			nID := expected[:8]
+
+			// change rawData to contain several shares with same nID
+			tt.rawData[(len(tt.rawData)/2)+1] = expected
+
+			// generate DAH
+			dah, err := da.NewDataAvailabilityHeader(squareSize, tt.rawData)
+			require.NoError(t, err)
+
+			// put raw data in DAG
+			dag := mdutils.Mock()
+			_, err = PutData(context.Background(), tt.rawData, dag)
+			require.NoError(t, err)
+
+			rowRootCIDs, err := rowRootsByNamespaceID(nID, &dah)
+			require.NoError(t, err)
+
+			for _, rowCID := range rowRootCIDs {
+				nodes, err := GetLeavesByNamespace(context.Background(), dag, rowCID, nID)
+				require.NoError(t, err)
+
+				for _, node := range nodes {
+					// TODO @renaynay: nID is prepended twice for some reason.
+					share := node.RawData()[1:]
+					assert.Equal(t, expected, share[8:])
+				}
+			}
+		})
+	}
+}
+
+// rowRootsByNamespaceID is a convenience method that finds the row root(s)
+// that contain the given namespace ID.
+func rowRootsByNamespaceID(nID namespace.ID, dah *da.DataAvailabilityHeader) ([]cid.Cid, error) {
+	roots := make([]cid.Cid, 0)
+	for _, row := range dah.RowsRoots {
+		// if nID exists within range of min -> max of row, return the row
+		if !nID.Less(nmt.MinNamespace(row, nID.Size())) && nID.LessOrEqual(nmt.MaxNamespace(row, nID.Size())) {
+			roots = append(roots, plugin.MustCidFromNamespacedSha256(row))
+		}
+	}
+	if len(roots) == 0 {
+		return nil, ErrNotFoundInRange
+	}
+	return roots, nil
+}

service/share/share.go

@@ -11,11 +11,10 @@ import (
 	"github.com/ipfs/go-merkledag"
 
 	"github.com/celestiaorg/celestia-core/pkg/da"
-
-	"github.com/celestiaorg/nmt/namespace"
-
 	"github.com/celestiaorg/celestia-node/ipld"
 	"github.com/celestiaorg/celestia-node/ipld/plugin"
+	"github.com/celestiaorg/nmt"
+	"github.com/celestiaorg/nmt/namespace"
 )
 
 var log = logging.Logger("share")
@@ -60,13 +59,14 @@ type Service interface {
 	// It also optimistically executes erasure coding recovery.
 	GetShares(context.Context, *Root) ([][]Share, error)
 
-	// GetSharesByNamespace loads all the Shares committed to the given DataAvailabilityHeader as a 1D array/slice.
+	// GetSharesByNamespace loads all the Shares of the given namespace.ID committed to the given
+	// DataAvailabilityHeader as a 1D array/slice.
 	GetSharesByNamespace(context.Context, *Root, namespace.ID) ([]Share, error)
 
-	// Starts the Service.
+	// Start starts the Service.
 	Start(context.Context) error
 
-	// Stops the Service.
+	// Stop stops the Service.
 	Stop(context.Context) error
 }
 
@@ -132,8 +132,29 @@ func (s *service) GetShares(context.Context, *Root) ([][]Share, error) {
 	panic("implement me")
 }
 
-func (s *service) GetSharesByNamespace(context.Context, *Root, namespace.ID) ([]Share, error) {
-	panic("implement me")
+func (s *service) GetSharesByNamespace(ctx context.Context, root *Root, nID namespace.ID) ([]Share, error) {
+	rowRootCIDs := make([]cid.Cid, 0)
+	for _, row := range root.RowsRoots {
+		if !nID.Less(nmt.MinNamespace(row, nID.Size())) && nID.LessOrEqual(nmt.MaxNamespace(row, nID.Size())) {
+			rowRootCIDs = append(rowRootCIDs, plugin.MustCidFromNamespacedSha256(row))
+		}
+	}
+	if len(rowRootCIDs) == 0 {
+		return nil, ipld.ErrNotFoundInRange
+	}
+
+	namespacedShares := make([]Share, 0)
+	for _, rootCID := range rowRootCIDs {
+		nodes, err := ipld.GetLeavesByNamespace(ctx, s.dag, rootCID, nID)
+		if err != nil {
+			return nil, err
+		}
+
+		for _, node := range nodes {
+			namespacedShares = append(namespacedShares, node.RawData()[1:])
+		}
+	}
+	return namespacedShares, nil
 }
 
 // translate transforms square coordinates into IPLD NMT tree path to a leaf node.

service/share/share_test.go

@@ -2,6 +2,7 @@ package share
 
 import (
 	"context"
+	"strconv"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -28,3 +29,30 @@ func TestGetShare(t *testing.T) {
 	err = serv.Stop(ctx)
 	require.NoError(t, err)
 }
+
+func TestService_GetSharesByNamespace(t *testing.T) {
+	var tests = []struct {
+		amountShares       int
+		expectedShareCount int
+	}{
+		{amountShares: 4, expectedShareCount: 1},
+		{amountShares: 16, expectedShareCount: 2},
+		{amountShares: 128, expectedShareCount: 1},
+	}
+
+	for i, tt := range tests {
+		t.Run(strconv.Itoa(i), func(t *testing.T) {
+			serv, root := RandServiceWithSquare(t, tt.amountShares)
+			randNID := root.RowsRoots[(len(root.RowsRoots)-1)/2][:8]
+			if tt.expectedShareCount > 1 {
+				// make it so that two rows have the same namespace ID
+				root.RowsRoots[(len(root.RowsRoots) / 2)] = root.RowsRoots[(len(root.RowsRoots)-1)/2]
+			}
+
+			shares, err := serv.GetSharesByNamespace(context.Background(), root, randNID)
+			require.NoError(t, err)
+			assert.Len(t, shares, tt.expectedShareCount)
+			assert.Equal(t, randNID, []byte(shares[0].NamespaceID()))
+		})
+	}
+}