From b9e473a5ebd53031da47f9bd491c51e3b5d7b681 Mon Sep 17 00:00:00 2001
From: Kaloyan Raev <kaloyan-raev@users.noreply.github.com>
Date: Thu, 18 Oct 2018 17:21:08 +0300
Subject: [PATCH] Encapsulate key encryption in the Key and EncryptedPrivateKey
 types (#499)

---
 pkg/encryption/encryption.go | 25 +++++++++++++++++++++++--
 pkg/storage/streams/store.go | 11 ++++-------
 2 files changed, 27 insertions(+), 9 deletions(-)

diff --git a/pkg/encryption/encryption.go b/pkg/encryption/encryption.go
index 83292dc25..09f73309e 100644
--- a/pkg/encryption/encryption.go
+++ b/pkg/encryption/encryption.go
@@ -28,8 +28,13 @@ const (
 type Key [KeySize]byte
 
 // Bytes returns the key as a byte array pointer
-func (key *Key) Bytes() *[KeySize]byte {
-	return (*[KeySize]byte)(key)
+func (k *Key) Bytes() *[KeySize]byte {
+	return (*[KeySize]byte)(k)
+}
+
+// Encrypt encrypts the current key with the given key and nonce
+func (k *Key) Encrypt(cipher Cipher, key *Key, nonce *Nonce) (EncryptedPrivateKey, error) {
+	return cipher.Encrypt(k[:], key, nonce)
 }
 
 // Nonce represents the largest nonce used by any encryption protocol
@@ -60,6 +65,22 @@ func (nonce *AESGCMNonce) Bytes() *[AESGCMNonceSize]byte {
 	return (*[AESGCMNonceSize]byte)(nonce)
 }
 
+// EncryptedPrivateKey is a private key that has been encrypted
+type EncryptedPrivateKey []byte
+
+// Decrypt decrypts the current key with the given key and nonce
+func (k EncryptedPrivateKey) Decrypt(cipher Cipher, key *Key, nonce *Nonce) (*Key, error) {
+	plainData, err := cipher.Decrypt(k, key, nonce)
+	if err != nil {
+		return nil, err
+	}
+
+	var decryptedKey Key
+	copy(decryptedKey[:], plainData)
+
+	return &decryptedKey, nil
+}
+
 // Encrypt encrypts byte data with a key and nonce. The cipher data is returned
 // The type of encryption to use can be modified with encType
 func (cipher Cipher) Encrypt(data []byte, key *Key, nonce *Nonce) (cipherData []byte, err error) {
diff --git a/pkg/storage/streams/store.go b/pkg/storage/streams/store.go
index 2401d0615..adba917a2 100644
--- a/pkg/storage/streams/store.go
+++ b/pkg/storage/streams/store.go
@@ -155,7 +155,7 @@ func (s *streamStore) Put(ctx context.Context, path paths.Path, data io.Reader,
 			return Meta{}, err
 		}
 
-		encryptedKey, err := cipher.Encrypt(contentKey[:], (*encryption.Key)(derivedKey), &keyNonce)
+		encryptedKey, err := contentKey.Encrypt(cipher, (*encryption.Key)(derivedKey), &keyNonce)
 		if err != nil {
 			return Meta{}, err
 		}
@@ -637,7 +637,7 @@ func (s *streamStore) cancelHandler(ctx context.Context, totalSegments int64, pa
 	}
 }
 
-func getEncryptedKeyAndNonce(m *pb.SegmentMeta) ([]byte, *encryption.Nonce) {
+func getEncryptedKeyAndNonce(m *pb.SegmentMeta) (encryption.EncryptedPrivateKey, *encryption.Nonce) {
 	if m == nil {
 		return nil, nil
 	}
@@ -662,14 +662,11 @@ func decryptStreamInfo(ctx context.Context, item segments.Meta, path paths.Path,
 
 	cipher := encryption.Cipher(streamMeta.EncryptionType)
 	encryptedKey, keyNonce := getEncryptedKeyAndNonce(streamMeta.LastSegmentMeta)
-	e, err := cipher.Decrypt(encryptedKey, (*encryption.Key)(derivedKey), keyNonce)
+	contentKey, err := encryptedKey.Decrypt(cipher, (*encryption.Key)(derivedKey), keyNonce)
 	if err != nil {
 		return nil, err
 	}
 
-	var contentKey encryption.Key
-	copy(contentKey[:], e)
-
 	// decrypt metadata with the content encryption key and zero nonce
-	return cipher.Decrypt(streamMeta.EncryptedStreamInfo, &contentKey, &encryption.Nonce{})
+	return cipher.Decrypt(streamMeta.EncryptedStreamInfo, contentKey, &encryption.Nonce{})
 }