Fix ReconstructSome crash (#303)

Fixes #302

When selectively reconstructing shards, all data shards must be reconstructed for the parity to be reconstructible.

It seems like it would be possible to do all reconstruction in one pass, which would be faster with the generated code.

Will explore that separately.

Author: klauspost

reedsolomon.go

@@ -1411,10 +1411,10 @@ func (r *reedSolomon) ReconstructData(shards [][]byte) error {
 // As the reconstructed shard set may contain missing parity shards,
 // calling the Verify function is likely to fail.
 func (r *reedSolomon) ReconstructSome(shards [][]byte, required []bool) error {
-	if len(required) == r.totalShards {
-		return r.reconstruct(shards, false, required)
+	if len(required) != r.dataShards && len(required) != r.totalShards {
+		return ErrInvalidInput
 	}
-	return r.reconstruct(shards, true, required)
+	return r.reconstruct(shards, len(required) == r.dataShards, required)
 }
 
 // reconstruct will recreate the missing data totalShards, and unless
@@ -1442,14 +1442,20 @@ func (r *reedSolomon) reconstruct(shards [][]byte, dataOnly bool, required []boo
 	numberPresent := 0
 	dataPresent := 0
 	missingRequired := 0
+	needAllData := false
 	for i := 0; i < r.totalShards; i++ {
 		if len(shards[i]) != 0 {
 			numberPresent++
 			if i < r.dataShards {
 				dataPresent++
 			}
-		} else if required != nil && required[i] {
-			missingRequired++
+		} else if required != nil {
+			if !dataOnly && i >= r.dataShards {
+				needAllData = true
+			}
+			if required[i] {
+				missingRequired++
+			}
 		}
 	}
 	if numberPresent == r.totalShards || dataOnly && dataPresent == r.dataShards ||
@@ -1526,23 +1532,24 @@ func (r *reedSolomon) reconstruct(shards [][]byte, dataOnly bool, required []boo
 	// The input to the coding is all of the shards we actually
 	// have, and the output is the missing data shards.  The computation
 	// is done using the special decode matrix we just built.
+	outputCount := 0
 	outputs := make([][]byte, r.parityShards)
 	matrixRows := make([][]byte, r.parityShards)
-	outputCount := 0
-
-	for iShard := 0; iShard < r.dataShards; iShard++ {
-		if len(shards[iShard]) == 0 && (required == nil || required[iShard]) {
-			if cap(shards[iShard]) >= shardSize {
-				shards[iShard] = shards[iShard][0:shardSize]
-			} else {
-				shards[iShard] = AllocAligned(1, shardSize)[0]
+	if dataPresent != r.dataShards {
+		for iShard := 0; iShard < r.dataShards; iShard++ {
+			if len(shards[iShard]) == 0 && (needAllData || required == nil || required[iShard]) {
+				if cap(shards[iShard]) >= shardSize {
+					shards[iShard] = shards[iShard][0:shardSize]
+				} else {
+					shards[iShard] = AllocAligned(1, shardSize)[0]
+				}
+				outputs[outputCount] = shards[iShard]
+				matrixRows[outputCount] = dataDecodeMatrix[iShard]
+				outputCount++
 			}
-			outputs[outputCount] = shards[iShard]
-			matrixRows[outputCount] = dataDecodeMatrix[iShard]
-			outputCount++
 		}
+		r.codeSomeShards(matrixRows, subShards, outputs[:outputCount], shardSize)
 	}
-	r.codeSomeShards(matrixRows, subShards, outputs[:outputCount], shardSize)
 
 	if dataOnly {
 		// Exit out early if we are only interested in the data shards
@@ -1568,7 +1575,9 @@ func (r *reedSolomon) reconstruct(shards [][]byte, dataOnly bool, required []boo
 			outputCount++
 		}
 	}
-	r.codeSomeShards(matrixRows, shards[:r.dataShards], outputs[:outputCount], shardSize)
+	if outputCount > 0 {
+		r.codeSomeShards(matrixRows, shards[:r.dataShards], outputs[:outputCount], shardSize)
+	}
 	return nil
 }
 

reedsolomon_test.go

@@ -2240,3 +2240,110 @@ func TestReentrant(t *testing.T) {
 		}
 	}
 }
+
+func TestRSReconstructSome(t *testing.T) {
+	const (
+		dataShards   = 3
+		parityShards = 2
+		sectorSize   = 1 << 20 // 1 MiB
+	)
+	stripedSplit := func(data []byte, dataShards [][]byte) {
+		const leafSize = 64
+		buf := bytes.NewBuffer(data)
+		for off := 0; buf.Len() > 0; off += leafSize {
+			for _, shard := range dataShards {
+				copy(shard[off:], buf.Next(leafSize))
+			}
+		}
+	}
+
+	// prepare random data
+	buf := make([]byte, sectorSize*dataShards)
+	rand.Read(buf)
+
+	// prepare shards
+	shards := make([][]byte, dataShards+parityShards)
+	for i := range shards {
+		shards[i] = make([]byte, sectorSize)
+	}
+
+	for i, o := range testOpts() {
+		toFill := []int{0, 1, 2, 3, 4}
+		t.Run(fmt.Sprintf("options %d", i), func(t *testing.T) {
+			// Shuffle the indices to fill
+			rand.Shuffle(len(toFill), func(i, j int) {
+				toFill[i], toFill[j] = toFill[j], toFill[i]
+			})
+			stripedSplit(buf, shards[:dataShards])
+			// encode shards
+			enc, err := New(dataShards, parityShards, o...)
+			if err != nil {
+				t.Fatal(err)
+			}
+			err = enc.Encode(shards)
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			for _, required := range [][]bool{
+				{false, false, false, false, true},
+				{false, false, false, true, true},
+				{false, false, true, false, true},
+				{false, false, true, true, true},
+				{false, true, false, false, true},
+				{false, true, false, true, true},
+				{false, true, true, false, true},
+				{false, true, true, true, true},
+				{true, false, false, false, true},
+				{true, false, false, true, true},
+				{true, false, true, false, true},
+				{true, false, true, true, true},
+				{true, true, false, false, true},
+				{true, true, false, true, true},
+				{true, true, true, false, true},
+				{true, true, true, true, true},
+
+				{false, false, false, false, false},
+				{false, false, false, true, false},
+				{false, false, true, false, false},
+				{false, false, true, true, false},
+				{false, true, false, false, false},
+				{false, true, false, true, false},
+				{false, true, true, false, false},
+				{false, true, true, true, false},
+				{true, false, false, false, false},
+				{true, false, false, true, false},
+				{true, false, true, false, false},
+				{true, false, true, true, false},
+				{true, true, false, false, false},
+				{true, true, false, true, false},
+				{true, true, true, false, false},
+				{true, true, true, true, false},
+			} {
+				downloaded := make([][]byte, len(shards))
+				downloaded[toFill[0]] = append([]byte{}, shards[toFill[0]]...)
+				downloaded[toFill[1]] = append([]byte{}, shards[toFill[1]]...)
+				downloaded[toFill[2]] = append([]byte{}, shards[toFill[2]]...)
+				err = enc.ReconstructSome(downloaded, required)
+				if err != nil {
+					t.Fatal(err)
+				}
+				for i, shard := range downloaded {
+					if shard == nil && required[i] {
+						t.Errorf("shard idx %d should not be nil", i)
+					}
+					if !required[i] && shard == nil {
+						// Ignore, if it wasn't reconstructed
+						continue
+					}
+					if len(shard) != sectorSize {
+						t.Errorf("shard idx %d has wrong size: %d, want %d", i, len(shard), sectorSize)
+					}
+					if !bytes.Equal(shard, shards[i]) {
+						t.Errorf("shard idx %d does not match original", i)
+					}
+				}
+			}
+		})
+	}
+}