storage: Preserve counter across Encrypt() calls

storage/encryption/encryption.go

@@ -31,12 +31,14 @@ type Key []byte
 type Encryption interface {
 	Encrypt(data []byte) ([]byte, error)
 	Decrypt(data []byte) ([]byte, error)
+	Reset()
 }
 
 type encryption struct {
 	key      Key              // the encryption key (hashSize bytes long)
 	keyLen   int              // length of the key = length of blockcipher block
 	padding  int              // encryption will pad the data upto this if > 0
+	index    int              // counter index
 	initCtr  uint32           // initial counter used for counter mode blockcipher
 	hashFunc func() hash.Hash // hasher constructor function
 }
@@ -79,18 +81,24 @@ func (e *encryption) Decrypt(data []byte) ([]byte, error) {
 	return out, nil
 }
 
-//
+// Reset resets the counter. It is only safe to call after an encryption operation is completed
+// After Reset is called, the Encryption object can be re-used for other data
+func (e *encryption) Reset() {
+	e.index = 0
+}
+
+// split up input into keylength segments and encrypt sequentially
 func (e *encryption) transform(in, out []byte) {
 	inLength := len(in)
 	wg := sync.WaitGroup{}
 	wg.Add((inLength-1)/e.keyLen + 1)
 	for i := 0; i < inLength; i += e.keyLen {
 		l := min(e.keyLen, inLength-i)
-		// call transformations per segment (asyncronously)
 		go func(i int, x, y []byte) {
 			defer wg.Done()
 			e.Transcrypt(i, x, y)
-		}(i/e.keyLen, in[i:i+l], out[i:i+l])
+		}(e.index, in[i:i+l], out[i:i+l])
+		e.index++
 	}
 	// pad the rest if out is longer
 	pad(out[inLength:])

storage/encryption/encryption_test.go

@@ -18,6 +18,7 @@ package encryption
 
 import (
 	"bytes"
+	crand "crypto/rand"
 	"testing"
 
 	"github.com/ethereum/go-ethereum/common"
@@ -37,6 +38,7 @@ func init() {
 	if err != nil {
 		panic(err.Error())
 	}
+	testutil.Init()
 }
 
 func TestEncryptDataLongerThanPadding(t *testing.T) {
@@ -132,6 +134,7 @@ func testEncryptDecryptIsIdentity(t *testing.T, padding int, initCtr uint32, dat
 		t.Fatalf("Expected no error got %v", err)
 	}
 
+	enc.Reset()
 	decrypted, err := enc.Decrypt(encrypted)
 	if err != nil {
 		t.Fatalf("Expected no error got %v", err)
@@ -149,3 +152,42 @@ func testEncryptDecryptIsIdentity(t *testing.T, padding int, initCtr uint32, dat
 		t.Fatalf("Expected decrypted %v got %v", common.Bytes2Hex(data), common.Bytes2Hex(decrypted))
 	}
 }
+
+// TestEncryptSectioned tests that the cipherText is the same regardless of size of data input buffer
+func TestEncryptSectioned(t *testing.T) {
+	data := make([]byte, 4096)
+	c, err := crand.Read(data)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if c < 4096 {
+		t.Fatalf("short read %d", c)
+	}
+
+	key := make([]byte, KeyLength)
+	c, err = crand.Read(key)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if c < KeyLength {
+		t.Fatalf("short read %d", c)
+	}
+
+	enc := New(key, 0, uint32(42), sha3.NewLegacyKeccak256)
+	whole, err := enc.Encrypt(data)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	enc.Reset()
+	for i := 0; i < 4096; i += KeyLength {
+		cipher, err := enc.Encrypt(data[i : i+KeyLength])
+		if err != nil {
+			t.Fatal(err)
+		}
+		wholeSection := whole[i : i+KeyLength]
+		if !bytes.Equal(cipher, wholeSection) {
+			t.Fatalf("index %d, expected %x, got %x", i/KeyLength, wholeSection, cipher)
+		}
+	}
+}