Repository: celestiaorg/celestia-node
Pull request #737 (2022-05-20 17:00:25)

deps: bump rsmt2d to 0.5.0
I tried updating to this version before but didn't like that we needed to keep both flattened and imported squares, so that code became even bulkier for the next currently unpushed iteration of Retriever, so I decided to wait until the requested rsmt2d API changes to land before upgrading(cc @adlerjohn), but now we have to update as core updated already. 


Patch:
diff --git a/go.mod b/go.mod
index 5dc45c96..efc57408 100644
--- a/go.mod
+++ b/go.mod
@@ -9,7 +9,7 @@ require (
 	github.com/celestiaorg/celestia-app v0.4.0
 	github.com/celestiaorg/go-libp2p-messenger v0.1.0
 	github.com/celestiaorg/nmt v0.8.0
-	github.com/celestiaorg/rsmt2d v0.4.0
+	github.com/celestiaorg/rsmt2d v0.5.0
 	github.com/cosmos/cosmos-sdk v0.44.3
 	github.com/dgraph-io/badger/v2 v2.2007.4
 	github.com/gammazero/workerpool v1.1.2
diff --git a/go.sum b/go.sum
index 3cd5b300..1c244a40 100644
--- a/go.sum
+++ b/go.sum
@@ -202,8 +202,9 @@ github.com/celestiaorg/merkletree v0.0.0-20210714075610-a84dc3ddbbe4/go.mod h1:f
 github.com/celestiaorg/nmt v0.8.0 h1:wtX7GRouLbmBe+ffnc8+cOg2UbWteM+Y1imZuZ/EeqU=
 github.com/celestiaorg/nmt v0.8.0/go.mod h1:3bqzTj8xKj0DgQUpOgZzoxvtNkC3MS/hTbQ6dn8SIa0=
 github.com/celestiaorg/rsmt2d v0.3.1/go.mod h1:2Frw4GEYUnVu6Mvlo+CUzuC2/8wn+zLwVVtp+muN6vg=
-github.com/celestiaorg/rsmt2d v0.4.0 h1:mpICbeXYKE1yNW8VBzgBrcnzuMKzutpROH1g2xI7Ls4=
 github.com/celestiaorg/rsmt2d v0.4.0/go.mod h1:EZ+O2KdCq8xI7WFwjATLdhtMdrdClmAs2w7zENDr010=
+github.com/celestiaorg/rsmt2d v0.5.0 h1:Wa0uNZUXl8lIMJnSunjoD65ktqBedXZD0z2ZU3xKYYw=
+github.com/celestiaorg/rsmt2d v0.5.0/go.mod h1:EZ+O2KdCq8xI7WFwjATLdhtMdrdClmAs2w7zENDr010=
 github.com/cenkalti/backoff v2.2.1+incompatible h1:tNowT99t7UNflLxfYYSlKYsBpXdEet03Pg2g16Swow4=
 github.com/cenkalti/backoff v2.2.1+incompatible/go.mod h1:90ReRw6GdpyfrHakVjL/QHaoyV4aDUVVkXQJJJ3NXXM=
 github.com/cenkalti/backoff/v4 v4.1.1/go.mod h1:scbssz8iZGpm3xbr14ovlUdkxfGXNInqkPWOWmG2CLw=
diff --git a/ipld/retriever.go b/ipld/retriever.go
index 2177542f..a8de6715 100644
--- a/ipld/retriever.go
+++ b/ipld/retriever.go
@@ -50,22 +50,22 @@ func NewRetriever(dag format.DAGService) *Retriever {
 // RetrieveQuadrantTimeout or unrecoverable.
 func (r *Retriever) Retrieve(ctx context.Context, dah *da.DataAvailabilityHeader) (*rsmt2d.ExtendedDataSquare, error) {
 	log.Debugw("retrieving data square", "data_hash", hex.EncodeToString(dah.Hash()), "size", len(dah.RowsRoots))
-	ses := r.newSession(ctx, dah)
-	for _, qs := range newQuadrants(dah) {
-		for _, q := range qs {
-			eds, err := ses.retrieve(ctx, q)
-			if err == nil {
-				return eds, nil
-			}
-			var (
-				errRow *rsmt2d.ErrByzantineRow
-				errCol *rsmt2d.ErrByzantineCol
-			)
-			if errors.As(err, &errRow) || errors.As(err, &errCol) {
-				return nil, NewErrByzantine(ctx, r.dag, dah, errRow, errCol)
-			}
-			log.Warnw("not enough shares to reconstruct data square, requesting more...", "err", err)
+	ses, err := r.newSession(ctx, dah)
+	if err != nil {
+		return nil, err
+	}
+
+	for _, q := range newQuadrants(dah) {
+		eds, err := ses.retrieve(ctx, q)
+		if err == nil {
+			return eds, nil
 		}
+
+		var errByz *rsmt2d.ErrByzantineData
+		if errors.As(err, &errByz) {
+			return nil, NewErrByzantine(ctx, r.dag, dah, errByz)
+		}
+		log.Warnw("not enough shares to reconstruct data square, requesting more...", "err", err)
 		// retry quadrants until we can reconstruct the EDS or error out
 	}
 
@@ -79,15 +79,16 @@ type retrieverSession struct {
 	treeFn rsmt2d.TreeConstructorFn
 	codec  rsmt2d.Codec
 
-	dah      *da.DataAvailabilityHeader
-	squareLk sync.RWMutex
-	square   [][]byte
+	dah            *da.DataAvailabilityHeader
+	squareLk       sync.RWMutex
+	square         [][]byte
+	squareImported *rsmt2d.ExtendedDataSquare
 }
 
-func (r *Retriever) newSession(ctx context.Context, dah *da.DataAvailabilityHeader) *retrieverSession {
+func (r *Retriever) newSession(ctx context.Context, dah *da.DataAvailabilityHeader) (*retrieverSession, error) {
 	size := len(dah.RowsRoots)
 	adder := NewNmtNodeAdder(ctx, format.NewBatch(ctx, r.dag, format.MaxSizeBatchOption(batchSize(size))))
-	return &retrieverSession{
+	ses := &retrieverSession{
 		dag:   merkledag.NewSession(ctx, r.dag),
 		adder: adder,
 		treeFn: func() rsmt2d.Tree {
@@ -98,6 +99,14 @@ func (r *Retriever) newSession(ctx context.Context, dah *da.DataAvailabilityHead
 		dah:    dah,
 		square: make([][]byte, size*size),
 	}
+
+	square, err := rsmt2d.ImportExtendedDataSquare(ses.square, ses.codec, ses.treeFn)
+	if err != nil {
+		return nil, err
+	}
+
+	ses.squareImported = square
+	return ses, nil
 }
 
 func (rs *retrieverSession) retrieve(ctx context.Context, q *quadrant) (*rsmt2d.ExtendedDataSquare, error) {
@@ -117,13 +126,13 @@ func (rs *retrieverSession) retrieve(ctx context.Context, q *quadrant) (*rsmt2d.
 	// TODO: Avoid reimporting of the square which can potentially remove the requirement for the lock
 	rs.squareLk.Lock()
 	defer rs.squareLk.Unlock()
-	err := rsmt2d.RepairExtendedDataSquare(rs.dah.RowsRoots, rs.dah.ColumnRoots, rs.square, rs.codec, rs.treeFn)
+	err := rs.squareImported.Repair(rs.dah.RowsRoots, rs.dah.ColumnRoots, rs.codec, rs.treeFn)
 	if err != nil {
 		return nil, err
 	}
 
 	log.Infow("data square reconstructed", "data_hash", hex.EncodeToString(rs.dah.Hash()), "size", len(rs.dah.RowsRoots))
-	return rsmt2d.ImportExtendedDataSquare(rs.square, rs.codec, rs.treeFn)
+	return rs.squareImported, nil
 }
 
 func (rs *retrieverSession) request(ctx context.Context, q *quadrant) {
diff --git a/ipld/retriever_byzantine.go b/ipld/retriever_byzantine.go
index 1d93cc87..fb27d801 100644
--- a/ipld/retriever_byzantine.go
+++ b/ipld/retriever_byzantine.go
@@ -35,30 +35,17 @@ func NewErrByzantine(
 	ctx context.Context,
 	dag format.NodeGetter,
 	dah *da.DataAvailabilityHeader,
-	errRow *rsmt2d.ErrByzantineRow,
-	errCol *rsmt2d.ErrByzantineCol) *ErrByzantine {
-	var (
-		errShares [][]byte
-		root      []byte
-		index     uint8
-	)
-	isRow := false
-	if errRow != nil {
-		errShares = errRow.Shares
-		root = dah.RowsRoots[errRow.RowNumber]
-		index = uint8(errRow.RowNumber)
-		isRow = true
-	} else {
-		errShares = errCol.Shares
-		root = dah.ColumnRoots[errCol.ColNumber]
-		index = uint8(errCol.ColNumber)
-	}
-
+	errByz *rsmt2d.ErrByzantineData,
+) *ErrByzantine {
+	root := [][][]byte{
+		dah.RowsRoots,
+		dah.ColumnRoots,
+	}[errByz.Axis][errByz.Index]
 	sharesWithProof, err := GetProofsForShares(
 		ctx,
 		dag,
 		plugin.MustCidFromNamespacedSha256(root),
-		errShares,
+		errByz.Shares,
 	)
 	if err != nil {
 		// Fatal as rsmt2d proved that error is byzantine,
@@ -69,5 +56,9 @@ func NewErrByzantine(
 		log.Fatalw("getting proof for ErrByzantine", "err", err)
 	}
 
-	return &ErrByzantine{Index: index, Shares: sharesWithProof, IsRow: isRow}
+	return &ErrByzantine{
+		Index:  uint8(errByz.Index),
+		Shares: sharesWithProof,
+		IsRow:  errByz.Axis == rsmt2d.Row,
+	}
 }
diff --git a/ipld/retriever_quadrant.go b/ipld/retriever_quadrant.go
index 3eb2d4da..18dafd91 100644
--- a/ipld/retriever_quadrant.go
+++ b/ipld/retriever_quadrant.go
@@ -34,7 +34,7 @@ type quadrant struct {
 // newQuadrants constructs a slice of quadrants from DAHeader.
 // There are always 4 quadrants per each source (row and col), so 8 in total.
 // The ordering of quadrants is random.
-func newQuadrants(dah *da.DataAvailabilityHeader) [][]*quadrant {
+func newQuadrants(dah *da.DataAvailabilityHeader) []*quadrant {
 	// combine all the roots into one slice, so they can be easily accessible by index
 	daRoots := [][][]byte{
 		dah.RowsRoots,
@@ -66,12 +66,14 @@ func newQuadrants(dah *da.DataAvailabilityHeader) [][]*quadrant {
 				source: source,
 			}
 		}
-		// shuffle quadrants to be fetched in random order
-		rand.Shuffle(len(quadrants), func(i, j int) { quadrants[i], quadrants[j] = quadrants[j], quadrants[i] })
 	}
-	// randomize sources, s.t. downloading on the network happens from both rows and cols
-	rand.Shuffle(len(sources), func(i, j int) { sources[i], sources[j] = sources[j], sources[i] })
-	return sources
+	quadrants := make([]*quadrant, 0, numQuadrants*2)
+	for _, qs := range sources {
+		quadrants = append(quadrants, qs...)
+	}
+	// shuffle quadrants to be fetched in random order
+	rand.Shuffle(len(quadrants), func(i, j int) { quadrants[i], quadrants[j] = quadrants[j], quadrants[i] })
+	return quadrants
 }
 
 // index calculates index for a share in a data square slice flattened by rows.


Test patch:
diff --git a/ipld/get_test.go b/ipld/get_test.go
index 138c669f..f4d08cd3 100644
--- a/ipld/get_test.go
+++ b/ipld/get_test.go
@@ -93,14 +93,14 @@ func TestBlockRecovery(t *testing.T) {
 
 			// recover a partially complete square
 			rdata := removeRandShares(flat, tc.d)
-			err = rsmt2d.RepairExtendedDataSquare(
-				rowRoots,
-				colRoots,
+			eds, err = rsmt2d.ImportExtendedDataSquare(
 				rdata,
 				rsmt2d.NewRSGF8Codec(),
 				recoverTree.Constructor,
 			)
+			require.NoError(t, err)
 
+			err = eds.Repair(rowRoots, colRoots, rsmt2d.NewRSGF8Codec(), recoverTree.Constructor)
 			if tc.expectErr {
 				require.Error(t, err)
 				require.Contains(t, err.Error(), tc.errString)