From 07bad81f491115a6e76681a1808c6dc578366af9 Mon Sep 17 00:00:00 2001
From: rene <41963722+renaynay@users.noreply.github.com>
Date: Tue, 26 Apr 2022 12:29:50 +0100
Subject: [PATCH] chore: update uses of rsmt2d

---
 ipld/read.go             |  8 ++++++--
 ipld/read_test.go        | 13 +++++++++++--
 service/header/header.go |  5 ++++-
 3 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/ipld/read.go b/ipld/read.go
index 86d3decd96..db012ddb1b 100644
--- a/ipld/read.go
+++ b/ipld/read.go
@@ -42,12 +42,16 @@ func RetrieveData(
 	}
 	batchAdder := NewNmtNodeAdder(parentCtx, ipld.NewBatch(parentCtx, dag, ipld.MaxSizeBatchOption(batchSize(edsWidth))))
 	tree := wrapper.NewErasuredNamespacedMerkleTree(uint64(edsWidth)/2, nmt.NodeVisitor(batchAdder.Visit))
-	extended, err := rsmt2d.RepairExtendedDataSquare(dah.RowsRoots, dah.ColumnRoots, dataSquare, codec, tree.Constructor)
+	err = rsmt2d.RepairExtendedDataSquare(dah.RowsRoots, dah.ColumnRoots, dataSquare, codec, tree.Constructor)
 	if err != nil {
 		return nil, err
 	}
 
-	return extended, batchAdder.Commit()
+	eds, err := rsmt2d.ImportExtendedDataSquare(dataSquare, codec, tree.Constructor)
+	if err != nil {
+		return nil, err
+	}
+	return eds, batchAdder.Commit()
 }
 
 type quadrant struct {
diff --git a/ipld/read_test.go b/ipld/read_test.go
index 2f77ef4b23..d2bf0732bb 100644
--- a/ipld/read_test.go
+++ b/ipld/read_test.go
@@ -92,11 +92,13 @@ func TestBlockRecovery(t *testing.T) {
 
 			flat := flatten(eds)
 
+			dataSquare := removeRandShares(flat, tc.d)
+
 			// recover a partially complete square
-			reds, err := rsmt2d.RepairExtendedDataSquare(
+			err = rsmt2d.RepairExtendedDataSquare(
 				rowRoots,
 				colRoots,
-				removeRandShares(flat, tc.d),
+				dataSquare,
 				rsmt2d.NewRSGF8Codec(),
 				recoverTree.Constructor,
 			)
@@ -108,6 +110,13 @@ func TestBlockRecovery(t *testing.T) {
 			}
 			assert.NoError(t, err)
 
+			reds, err := rsmt2d.ImportExtendedDataSquare(dataSquare, rsmt2d.NewRSGF8Codec(), recoverTree.Constructor)
+			// if an error is expected during attempted EDS reparation,
+			// then an error should be expected during import
+			if tc.expectErr {
+				require.Error(t, err)
+			}
+			assert.NoError(t, err)
 			// check that the squares are equal
 			assert.Equal(t, flatten(eds), flatten(reds))
 		})
diff --git a/service/header/header.go b/service/header/header.go
index e19606ad43..7aac61c0cc 100644
--- a/service/header/header.go
+++ b/service/header/header.go
@@ -45,7 +45,10 @@ func MakeExtendedHeader(
 ) (*ExtendedHeader, error) {
 	var dah DataAvailabilityHeader
 	if len(b.Txs) > 0 {
-		namespacedShares, _ := b.Data.ComputeShares()
+		namespacedShares, _, err := b.Data.ComputeShares(b.Data.OriginalSquareSize)
+		if err != nil {
+			return nil, err
+		}
 		extended, err := ipld.PutData(ctx, namespacedShares.RawShares(), dag)
 		if err != nil {
 			return nil, err