From 8c04ff64a52f34203944860b15baefcd6ff481ec Mon Sep 17 00:00:00 2001
From: wussler <aron@wussler.it>
Date: Mon, 20 Jul 2020 11:43:36 +0200
Subject: [PATCH] Add session key size check (#62)

* Add session key size check

Co-authored-by: Daniel Huigens <d.huigens@protonmail.com>
---
 crypto/password.go        | 14 ++++++++++++--
 crypto/sessionkey.go      | 21 +++++++++++++++++++--
 crypto/sessionkey_test.go | 19 +++++++++++++++++++
 3 files changed, 50 insertions(+), 4 deletions(-)

diff --git a/crypto/password.go b/crypto/password.go
index 91f2dd6..051e031 100644
--- a/crypto/password.go
+++ b/crypto/password.go
@@ -62,10 +62,16 @@ func DecryptSessionKeyWithPassword(keyPacket, password []byte) (*SessionKey, err
 		for _, s := range symKeys {
 			key, cipherFunc, err := s.Decrypt(password)
 			if err == nil {
-				return &SessionKey{
+				sk := &SessionKey{
 					Key:  key,
 					Algo: getAlgo(cipherFunc),
-				}, nil
+				}
+
+				if err = sk.checkSize(); err != nil {
+					return nil, errors.Wrap(err, "gopenpgp: unable to decrypt session key with password")
+				}
+
+				return sk, nil
 			}
 		}
 	}
@@ -87,6 +93,10 @@ func EncryptSessionKeyWithPassword(sk *SessionKey, password []byte) ([]byte, err
 		return nil, errors.New("gopenpgp: password can't be empty")
 	}
 
+	if err = sk.checkSize(); err != nil {
+		return nil, errors.Wrap(err, "gopenpgp: unable to encrypt session key with password")
+	}
+
 	config := &packet.Config{
 		DefaultCipher: cf,
 	}
diff --git a/crypto/sessionkey.go b/crypto/sessionkey.go
index 6751119..ce31dfb 100644
--- a/crypto/sessionkey.go
+++ b/crypto/sessionkey.go
@@ -98,12 +98,16 @@ func newSessionKeyFromEncrypted(ek *packet.EncryptedKey) (*SessionKey, error) {
 		return nil, fmt.Errorf("gopenpgp: unsupported cipher function: %v", ek.CipherFunc)
 	}
 
-	symmetricKey := &SessionKey{
+	sk := &SessionKey{
 		Key:  ek.Key,
 		Algo: algo,
 	}
 
-	return symmetricKey, nil
+	if err := sk.checkSize(); err != nil {
+		return nil, errors.Wrap(err, "gopenpgp: unable to decrypt session key")
+	}
+
+	return sk, nil
 }
 
 // Encrypt encrypts a PlainMessage to PGPMessage with a SessionKey.
@@ -209,6 +213,19 @@ func (sk *SessionKey) Decrypt(dataPacket []byte) (*PlainMessage, error) {
 	return NewPlainMessage(messageBuf.Bytes()), nil
 }
 
+func (sk *SessionKey) checkSize() error {
+	cf, ok := symKeyAlgos[sk.Algo]
+	if !ok {
+		return errors.New("unknown symmetric key algorithm")
+	}
+
+	if cf.KeySize() != len(sk.Key) {
+		return errors.New("wrong session key size")
+	}
+
+	return nil
+}
+
 func getAlgo(cipher packet.CipherFunction) string {
 	algo := constants.AES256
 	for k, v := range symKeyAlgos {
diff --git a/crypto/sessionkey_test.go b/crypto/sessionkey_test.go
index 42ed26e..b090085 100644
--- a/crypto/sessionkey_test.go
+++ b/crypto/sessionkey_test.go
@@ -81,6 +81,25 @@ func TestSymmetricKeyPacket(t *testing.T) {
 	assert.Exactly(t, testSessionKey, outputSymmetricKey)
 }
 
+func TestSymmetricKeyPacketWrongSize(t *testing.T) {
+	r, err := RandomToken(symKeyAlgos[constants.AES256].KeySize())
+	if err != nil {
+		t.Fatal("Expected no error while generating session key, got:", err)
+	}
+
+	sk := &SessionKey{
+		Key:  r,
+		Algo: constants.AES128,
+	}
+
+	password := []byte("I like encryption")
+
+	_, err = EncryptSessionKeyWithPassword(sk, password)
+	if err == nil {
+		t.Fatal("Expected error while generating key packet with wrong sized key")
+	}
+}
+
 func TestDataPacketEncryption(t *testing.T) {
 	var message = NewPlainMessageFromString("The secret code is... 1, 2, 3, 4, 5")