From 026172acb5af476a33b8f4d5dae3cfe61dc32327 Mon Sep 17 00:00:00 2001
From: Felix Lange <fjl@twurst.com>
Date: Wed, 1 Apr 2020 12:45:48 +0200
Subject: [PATCH] crypto/ecies: extract KDF call into shared function

This fixes the hash.Reset issue in a different way and removes
some duplication.
---
 crypto/ecies/ecies.go | 31 +++++++++++++++----------------
 1 file changed, 15 insertions(+), 16 deletions(-)

diff --git a/crypto/ecies/ecies.go b/crypto/ecies/ecies.go
index c7169a770e..9f94047864 100644
--- a/crypto/ecies/ecies.go
+++ b/crypto/ecies/ecies.go
@@ -149,11 +149,11 @@ func concatKDF(hash hash.Hash, z, s1 []byte, kdLen int) []byte {
 	k := make([]byte, 0, roundup(kdLen, hash.Size()))
 	for counter := uint32(1); len(k) < kdLen; counter++ {
 		binary.BigEndian.PutUint32(counterBytes, counter)
+		hash.Reset()
 		hash.Write(counterBytes)
 		hash.Write(z)
 		hash.Write(s1)
 		k = hash.Sum(k)
-		hash.Reset()
 	}
 	return k[:kdLen]
 }
@@ -163,6 +163,17 @@ func roundup(size, blocksize int) int {
 	return size + blocksize - (size % blocksize)
 }
 
+// deriveKeys creates the encryption and MAC keys using concatKDF.
+func deriveKeys(hash hash.Hash, z, s1 []byte, keyLen int) (Ke, Km []byte) {
+	K := concatKDF(hash, z, s1, 2*keyLen)
+	Ke = K[:keyLen]
+	Km = K[keyLen:]
+	hash.Reset()
+	hash.Write(Km)
+	Km = hash.Sum(Km[:0])
+	return Ke, Km
+}
+
 // messageTag computes the MAC of a message (called the tag) as per
 // SEC 1, 3.5.
 func messageTag(hash func() hash.Hash, km, msg, shared []byte) []byte {
@@ -238,12 +249,7 @@ func Encrypt(rand io.Reader, pub *PublicKey, m, s1, s2 []byte) (ct []byte, err e
 	if err != nil {
 		return
 	}
-	K := concatKDF(hash, z, s1, params.KeyLen+params.KeyLen)
-	Ke := K[:params.KeyLen]
-	Km := K[params.KeyLen:]
-	hash.Write(Km)
-	Km = hash.Sum(nil)
-	hash.Reset()
+	Ke, Km := deriveKeys(hash, z, s1, params.KeyLen)
 
 	em, err := symEncrypt(rand, params, Ke, m)
 	if err != nil || len(em) <= params.BlockSize {
@@ -310,22 +316,15 @@ func (prv *PrivateKey) Decrypt(c, s1, s2 []byte) (m []byte, err error) {
 
 	z, err := prv.GenerateShared(R, params.KeyLen, params.KeyLen)
 	if err != nil {
-		return
+		return nil, err
 	}
-
-	K := concatKDF(hash, z, s1, params.KeyLen+params.KeyLen)
-	Ke := K[:params.KeyLen]
-	Km := K[params.KeyLen:]
-	hash.Write(Km)
-	Km = hash.Sum(nil)
-	hash.Reset()
+	Ke, Km := deriveKeys(hash, z, s1, params.KeyLen)
 
 	d := messageTag(params.Hash, Km, c[mStart:mEnd], s2)
 	if subtle.ConstantTimeCompare(c[mEnd:], d) != 1 {
 		err = ErrInvalidMessage
 		return
 	}
-
 	m, err = symDecrypt(params, Ke, c[mStart:mEnd])
 	return
 }