Fix CalcEncryptedSize and CalcPieceSize (#1525)

pkg/eestream/encode.go

@@ -267,10 +267,8 @@ func (er *EncodedRanger) Range(ctx context.Context, offset, length int64) ([]io.
 // CalcPieceSize calculates what would be the piece size of the encoded data
 // after erasure coding data with dataSize using the given ErasureScheme.
 func CalcPieceSize(dataSize int64, scheme ErasureScheme) int64 {
-	stripes := dataSize / int64(scheme.StripeSize())
-	if dataSize%int64(scheme.StripeSize()) != 0 {
-		stripes++
-	}
+	stripeSize := int64(scheme.StripeSize())
+	stripes := (dataSize + uint32Size + stripeSize - 1) / stripeSize
 
 	encodedSize := stripes * int64(scheme.StripeSize())
 	pieceSize := encodedSize / int64(scheme.RequiredCount())

pkg/eestream/rs_test.go

@@ -16,10 +16,13 @@ import (
 	"time"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"github.com/vivint/infectious"
 	"github.com/zeebo/errs"
 
+	"storj.io/storj/internal/memory"
 	"storj.io/storj/internal/readcloser"
+	"storj.io/storj/internal/testcontext"
 	"storj.io/storj/pkg/encryption"
 	"storj.io/storj/pkg/ranger"
 	"storj.io/storj/pkg/storj"
@@ -640,3 +643,38 @@ func BenchmarkReedSolomonErasureScheme(b *testing.B) {
 		}
 	}
 }
+
+func TestCalcPieceSize(t *testing.T) {
+	ctx := testcontext.New(t)
+	defer ctx.Cleanup()
+
+	for i, dataSize := range []int64{
+		0,
+		1,
+		1*memory.KiB.Int64() - uint32Size,
+		1 * memory.KiB.Int64(),
+		32*memory.KiB.Int64() - uint32Size,
+		32 * memory.KiB.Int64(),
+		32*memory.KiB.Int64() + 100,
+	} {
+		errTag := fmt.Sprintf("%d. %+v", i, dataSize)
+
+		fc, err := infectious.NewFEC(2, 4)
+		require.NoError(t, err, errTag)
+		es := NewRSScheme(fc, 1*memory.KiB.Int())
+		rs, err := NewRedundancyStrategy(es, 0, 0)
+		require.NoError(t, err, errTag)
+
+		calculatedSize := CalcPieceSize(dataSize, es)
+
+		randReader := ioutil.NopCloser(io.LimitReader(rand.New(rand.NewSource(rand.Int63())), dataSize))
+		readers, err := EncodeReader(ctx, PadReader(randReader, es.StripeSize()), rs)
+		require.NoError(t, err, errTag)
+
+		for _, reader := range readers {
+			piece, err := ioutil.ReadAll(reader)
+			assert.NoError(t, err, errTag)
+			assert.EqualValues(t, calculatedSize, len(piece), errTag)
+		}
+	}
+}

pkg/encryption/encryption.go

@@ -10,8 +10,12 @@ import (
 	"storj.io/storj/pkg/storj"
 )
 
-// AESGCMNonceSize is the size of an AES-GCM nonce
-const AESGCMNonceSize = 12
+const (
+	// AESGCMNonceSize is the size of an AES-GCM nonce
+	AESGCMNonceSize = 12
+	// unit32Size is the number of bytes in the uint32 type
+	uint32Size = 4
+)
 
 // AESGCMNonce represents the nonce used by the AES-GCM protocol
 type AESGCMNonce [AESGCMNonceSize]byte
@@ -135,10 +139,8 @@ func CalcEncryptedSize(dataSize int64, scheme storj.EncryptionScheme) (int64, er
 		return 0, err
 	}
 
-	blocks := dataSize / int64(transformer.InBlockSize())
-	if dataSize%int64(transformer.InBlockSize()) != 0 {
-		blocks++
-	}
+	inBlockSize := int64(transformer.InBlockSize())
+	blocks := (dataSize + uint32Size + inBlockSize - 1) / inBlockSize
 
 	encryptedSize := blocks * int64(transformer.OutBlockSize())
 

pkg/encryption/encryption_test.go

@@ -0,0 +1,70 @@
+// Copyright (C) 2019 Storj Labs, Inc.
+// See LICENSE for copying information.
+
+package encryption_test
+
+import (
+	"fmt"
+	"io"
+	"io/ioutil"
+	"math/rand"
+	"testing"
+
+	"storj.io/storj/pkg/eestream"
+	"storj.io/storj/pkg/encryption"
+	"storj.io/storj/pkg/storj"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"storj.io/storj/internal/memory"
+	"storj.io/storj/internal/testcontext"
+)
+
+const (
+	uint32Size = 4
+)
+
+func TestCalcEncryptedSize(t *testing.T) {
+	ctx := testcontext.New(t)
+	defer ctx.Cleanup()
+
+	forAllCiphers(func(cipher storj.Cipher) {
+		for i, dataSize := range []int64{
+			0,
+			1,
+			1*memory.KiB.Int64() - uint32Size,
+			1 * memory.KiB.Int64(),
+			32*memory.KiB.Int64() - uint32Size,
+			32 * memory.KiB.Int64(),
+			32*memory.KiB.Int64() + 100,
+		} {
+			errTag := fmt.Sprintf("%d-%d. %+v", cipher, i, dataSize)
+
+			scheme := storj.EncryptionScheme{Cipher: cipher, BlockSize: 1 * memory.KiB.Int32()}
+
+			calculatedSize, err := encryption.CalcEncryptedSize(dataSize, scheme)
+			require.NoError(t, err, errTag)
+
+			encrypter, err := encryption.NewEncrypter(scheme.Cipher, new(storj.Key), new(storj.Nonce), int(scheme.BlockSize))
+			require.NoError(t, err, errTag)
+
+			randReader := ioutil.NopCloser(io.LimitReader(rand.New(rand.NewSource(rand.Int63())), dataSize))
+			reader := encryption.TransformReader(eestream.PadReader(randReader, encrypter.InBlockSize()), encrypter, 0)
+
+			cipherData, err := ioutil.ReadAll(reader)
+			assert.NoError(t, err, errTag)
+			assert.EqualValues(t, calculatedSize, len(cipherData), errTag)
+		}
+	})
+}
+
+func forAllCiphers(test func(cipher storj.Cipher)) {
+	for _, cipher := range []storj.Cipher{
+		storj.Unencrypted,
+		storj.AESGCM,
+		storj.SecretBox,
+	} {
+		test(cipher)
+	}
+}