From 29a5a92d13cad45794c6e42cb97260a9ab9900ab Mon Sep 17 00:00:00 2001
From: Gustav Simonsson <gustav.simonsson@gmail.com>
Date: Wed, 15 Apr 2015 13:24:12 +0200
Subject: [PATCH] Add key header to encrypted keys

* Add key header containing key version, kdf and kdf params
* Store key header as JSON in the key file
* Read in KDF params from key header
* Include key header in MAC calculation and MAC verification
---
 crypto/key.go                  | 23 +++++++++++++++----
 crypto/key_store_passphrase.go | 41 +++++++++++++++++++++++++++++++---
 2 files changed, 57 insertions(+), 7 deletions(-)

diff --git a/crypto/key.go b/crypto/key.go
index 5e1f3637ea..067a5a294a 100644
--- a/crypto/key.go
+++ b/crypto/key.go
@@ -48,17 +48,32 @@ type plainKeyJSON struct {
 	PrivateKey []byte
 }
 
+type encryptedKeyJSON struct {
+	Id      []byte
+	Address []byte
+	Crypto  cipherJSON
+}
+
 type cipherJSON struct {
 	MAC        []byte
 	Salt       []byte
 	IV         []byte
+	KeyHeader  keyHeaderJSON
 	CipherText []byte
 }
 
-type encryptedKeyJSON struct {
-	Id      []byte
-	Address []byte
-	Crypto  cipherJSON
+type keyHeaderJSON struct {
+	Version   string
+	Kdf       string
+	KdfParams scryptParamsJSON // TODO: make more generic?
+}
+
+type scryptParamsJSON struct {
+	N       int
+	R       int
+	P       int
+	DkLen   int
+	SaltLen int
 }
 
 func (k *Key) MarshalJSON() (j []byte, err error) {
diff --git a/crypto/key_store_passphrase.go b/crypto/key_store_passphrase.go
index 739483d9fd..00717b5d13 100644
--- a/crypto/key_store_passphrase.go
+++ b/crypto/key_store_passphrase.go
@@ -81,6 +81,8 @@ import (
 )
 
 const (
+	keyHeaderVersion = "1"
+	keyHeaderKDF     = "scrypt"
 	// 2^18 / 8 / 1 uses 256MB memory and approx 1s CPU time on a modern CPU.
 	scryptN     = 1 << 18
 	scryptr     = 8
@@ -140,12 +142,32 @@ func (ks keyStorePassphrase) StoreKey(key *Key, auth string) (err error) {
 	cipherText := make([]byte, len(toEncrypt))
 	AES128CBCEncrypter.CryptBlocks(cipherText, toEncrypt)
 
-	mac := Sha3(derivedKey[16:32], cipherText)
+	paramsJSON := scryptParamsJSON{
+		N:       scryptN,
+		R:       scryptr,
+		P:       scryptp,
+		DkLen:   scryptdkLen,
+		SaltLen: 32,
+	}
+
+	keyHeaderJSON := keyHeaderJSON{
+		Version:   keyHeaderVersion,
+		Kdf:       keyHeaderKDF,
+		KdfParams: paramsJSON,
+	}
+
+	keyHeaderJSONStr, err := json.Marshal(keyHeaderJSON)
+	if err != nil {
+		return err
+	}
+
+	mac := Sha3(keyHeaderJSONStr, derivedKey[16:32], cipherText)
 
 	cipherStruct := cipherJSON{
 		mac,
 		salt,
 		iv,
+		keyHeaderJSON,
 		cipherText,
 	}
 	keyStruct := encryptedKeyJSON{
@@ -185,15 +207,28 @@ func DecryptKey(ks keyStorePassphrase, keyAddr common.Address, auth string) (key
 	mac := keyProtected.Crypto.MAC
 	salt := keyProtected.Crypto.Salt
 	iv := keyProtected.Crypto.IV
+	keyHeader := keyProtected.Crypto.KeyHeader
 	cipherText := keyProtected.Crypto.CipherText
 
+	// used in MAC
+	keyHeaderJSONStr, err := json.Marshal(keyHeader)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	// TODO: make this more generic when we support different KDF params / key versions
+	n := keyHeader.KdfParams.N
+	r := keyHeader.KdfParams.R
+	p := keyHeader.KdfParams.P
+	dkLen := keyHeader.KdfParams.DkLen
+
 	authArray := []byte(auth)
-	derivedKey, err := scrypt.Key(authArray, salt, scryptN, scryptr, scryptp, scryptdkLen)
+	derivedKey, err := scrypt.Key(authArray, salt, n, r, p, dkLen)
 	if err != nil {
 		return nil, nil, err
 	}
 
-	calculatedMAC := Sha3(derivedKey[16:32], cipherText)
+	calculatedMAC := Sha3(keyHeaderJSONStr, derivedKey[16:32], cipherText)
 	if !bytes.Equal(calculatedMAC, mac) {
 		err = errors.New("Decryption failed: MAC mismatch")
 		return nil, nil, err