Use write buffer for KDF serialization

ProtonMail · Jul 31, 2020 · 61bc8c2 · 61bc8c2
1 parent 2ce2266
commit 61bc8c2
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 19 deletions.
diff --git a/openpgp/ecdh/ecdh.go b/openpgp/ecdh/ecdh.go
@@ -19,28 +19,38 @@ import (
 )
 
 type KDF struct {
+	Version                int // Defaults to v1; non-standard v2 allows forwarding
 	Hash                   algorithm.Hash
 	Cipher                 algorithm.Cipher
-	Version                int    // Defaults to v1; non-standard v2 allows forwarding
 	Flags                  byte   // (v2 only)
 	ReplacementFingerprint []byte // (v2 only) fingerprint to use instead of recipient's (for v5 keys, the 20 leftmost bytes only)
 	ReplacementKDFParams   []byte // (v2 only) serialized KDF params to use in KDF digest computation
 }
 
-func (kdf KDF) write() []byte {
+func (kdf *KDF) serialize(w io.Writer) (err error) {
 	if kdf.Version != 2 {
 		// Default version is 1
 		// Length || Version || Hash || Cipher
-		return []byte{3, 1, kdf.Hash.Id(), kdf.Cipher.Id()}
+		if _, err := w.Write([]byte{3, 1, kdf.Hash.Id(), kdf.Cipher.Id()}); err != nil {
+			return err
+		}
+
+		return nil
 	}
 
 	// Length || Version || Hash || Cipher || Flags || (Optional) v2 Fields...
-	v2Fields := []byte{4, 2, kdf.Hash.Id(), kdf.Cipher.Id(), kdf.Flags}
-	v2Fields = append(v2Fields, kdf.ReplacementFingerprint...)
-	v2Fields = append(v2Fields, kdf.ReplacementKDFParams...)
-	// Update length field
-	v2Fields[0] = byte(len(v2Fields) - 1)
-	return v2Fields
+	v2Length := byte(4 + len(kdf.ReplacementFingerprint) + len(kdf.ReplacementKDFParams))
+	if _, err := w.Write([]byte{v2Length, 2, kdf.Hash.Id(), kdf.Cipher.Id(), kdf.Flags}); err != nil {
+		return err
+	}
+	if _, err := w.Write(kdf.ReplacementFingerprint); err != nil {
+		return err
+	}
+	if _, err := w.Write(kdf.ReplacementKDFParams); err != nil {
+		return err
+	}
+
+	return nil
 }
 
 type PublicKey struct {
@@ -145,7 +155,9 @@ func buildKey(pub *PublicKey, zb []byte, curveOID, fingerprint []byte, stripLead
 	}
 	kdf := pub.KDF.ReplacementKDFParams
 	if kdf == nil {
-		kdf = pub.KDF.write()
+		if err := pub.KDF.serialize(param); err != nil {
+			return nil, err
+		}
 	}
 	if _, err := param.Write(kdf); err != nil {
 		return nil, err

diff --git a/openpgp/ecdh/ecdh_test.go b/openpgp/ecdh/ecdh_test.go
@@ -50,46 +50,54 @@ func TestKDFParamsWrite(t *testing.T) {
 		Hash:   algorithm.SHA512,
 		Cipher: algorithm.AES256,
 	}
+	byteBuffer := new(bytes.Buffer)
+
 	expectBytesV1 := []byte{3, 1, kdf.Hash.Id(), kdf.Cipher.Id()}
-	gotBytes := kdf.write()
+	kdf.serialize(byteBuffer)
+	gotBytes := byteBuffer.Bytes()
 	if !bytes.Equal(gotBytes, expectBytesV1) {
 		t.Errorf("error serializing KDF params, got %x, want: %x", gotBytes, expectBytesV1)
 	}
+	byteBuffer.Reset()
 
 	kdfV2Flags0x01 := KDF{
+		Version:                2,
 		Hash:                   algorithm.SHA512,
 		Cipher:                 algorithm.AES256,
-		Version:                2,
 		Flags:                  0x01,
 		ReplacementFingerprint: testFingerprint,
 	}
 	expectBytesV2Flags0x01 := []byte{24, 2, kdfV2Flags0x01.Hash.Id(), kdfV2Flags0x01.Cipher.Id(), 0x01}
 	expectBytesV2Flags0x01 = append(expectBytesV2Flags0x01, testFingerprint...)
 
-	gotBytes = kdfV2Flags0x01.write()
+	kdfV2Flags0x01.serialize(byteBuffer)
+	gotBytes = byteBuffer.Bytes()
 	if !bytes.Equal(gotBytes, expectBytesV2Flags0x01) {
 		t.Errorf("error serializing KDF params v2 (flags 0x01), got %x, want: %x", gotBytes, expectBytesV2Flags0x01)
 	}
+	byteBuffer.Reset()
 
 	kdfV2Flags0x02 := KDF{
+		Version:              2,
 		Hash:                 algorithm.SHA512,
 		Cipher:               algorithm.AES256,
-		Version:              2,
 		Flags:                0x02,
 		ReplacementKDFParams: expectBytesV1,
 	}
 	expectBytesV2Flags0x02 := []byte{8, 2, kdfV2Flags0x02.Hash.Id(), kdfV2Flags0x01.Cipher.Id(), 0x02}
 	expectBytesV2Flags0x02 = append(expectBytesV2Flags0x02, expectBytesV1...)
 
-	gotBytes = kdfV2Flags0x02.write()
+	kdfV2Flags0x02.serialize(byteBuffer)
+	gotBytes = byteBuffer.Bytes()
 	if !bytes.Equal(gotBytes, expectBytesV2Flags0x02) {
 		t.Errorf("error serializing KDF params v2 (flags 0x02), got %x, want: %x", gotBytes, expectBytesV2Flags0x02)
 	}
+	byteBuffer.Reset()
 
 	kdfV2Flags0x03 := KDF{
+		Version:                2,
 		Hash:                   algorithm.SHA512,
 		Cipher:                 algorithm.AES256,
-		Version:                2,
 		Flags:                  0x03,
 		ReplacementFingerprint: testFingerprint,
 		ReplacementKDFParams:   expectBytesV1,
@@ -98,8 +106,10 @@ func TestKDFParamsWrite(t *testing.T) {
 	expectBytesV2Flags0x03 = append(expectBytesV2Flags0x03, testFingerprint...)
 	expectBytesV2Flags0x03 = append(expectBytesV2Flags0x03, expectBytesV1...)
 
-	gotBytes = kdfV2Flags0x03.write()
+	kdfV2Flags0x03.serialize(byteBuffer)
+	gotBytes = byteBuffer.Bytes()
 	if !bytes.Equal(gotBytes, expectBytesV2Flags0x03) {
 		t.Errorf("error serializing KDF params v2 (flags 0x03), got %x, want: %x", gotBytes, expectBytesV2Flags0x03)
 	}
+	byteBuffer.Reset()
 }
diff --git a/openpgp/packet/public_key.go b/openpgp/packet/public_key.go
@@ -356,11 +356,14 @@ func (pk *PublicKey) parseECDH(r io.Reader) (err error) {
 		return errors.UnsupportedError("failed to parse EC point")
 	}
 
-	var kdfLen int
-	if kdfLen = len(pk.kdf.Bytes()); kdfLen < 3 {
+	kdfLen := len(pk.kdf.Bytes())
+	if kdfLen < 3 {
 		return errors.UnsupportedError("unsupported ECDH KDF length: " + strconv.Itoa(kdfLen))
 	}
 	kdfVersion := int(pk.kdf.Bytes()[0])
+	if kdfVersion != 1 && kdfVersion != 2 {
+		return errors.UnsupportedError("unsupported ECDH KDF version: " + strconv.Itoa(int(kdfVersion)))
+	}
 	kdfHash, ok := algorithm.HashById[pk.kdf.Bytes()[1]]
 	if !ok {
 		return errors.UnsupportedError("unsupported ECDH KDF hash: " + strconv.Itoa(int(pk.kdf.Bytes()[1])))