Allow passing a message size hint

crypto/crypto.go

@@ -74,6 +74,6 @@ func (p *PGPHandle) LockKey(key *Key, passphrase []byte) (*Key, error) {
 
 // GenerateSessionKey generates a random session key for the profile.
 func (p *PGPHandle) GenerateSessionKey() (*SessionKey, error) {
-	config := p.profile.EncryptionConfig()
+	config := p.profile.EncryptionConfig(0)
 	return generateSessionKey(config)
 }

crypto/decryption_core.go

@@ -330,7 +330,7 @@ func createPasswordPrompt(password []byte) func(keys []openpgp.Key, symmetric bo
 }
 
 func (dh *decryptionHandle) decryptionConfig(configTime int64) *packet.Config {
-	config := dh.profile.EncryptionConfig()
+	config := dh.profile.EncryptionConfig(0)
 
 	// Check intended recipients in signatures.
 	checkIntendedRecipients := !dh.DisableIntendedRecipients

crypto/encryption.go

@@ -3,13 +3,18 @@ package crypto
 import "github.com/ProtonMail/go-crypto/openpgp/packet"
 
 type EncryptionProfile interface {
-	EncryptionConfig() *packet.Config
+	EncryptionConfig(messageSizeHint uint64) *packet.Config
 	CompressionConfig() *packet.Config
 }
 
 // PGPEncryption is an interface for encrypting messages with GopenPGP.
 // Use an EncryptionHandleBuilder to create a PGPEncryption handle.
 type PGPEncryption interface {
+	// SetMessageSizeHint gives the encryption handle a hint about the
+	// expected size of the message, in order to set an appropriate chunk
+	// size when using AEAD. Nothing will break when the message size hint
+	// turns out to be wrong.
+	SetMessageSizeHint(messageSizeHint uint64)
 	// EncryptingWriter returns a wrapper around underlying output Writer,
 	// such that any write-operation via the wrapper results in a write to an encrypted pgp message.
 	// If the output Writer is of type PGPSplitWriter, the output can be split to multiple writers

crypto/encryption_core.go

@@ -91,7 +91,7 @@ func (eh *encryptionHandle) prepareEncryptAndSign(
 		ModTime:  time.Unix(plainMessageMetadata.Time(), 0),
 	}
 
-	config = eh.profile.EncryptionConfig()
+	config = eh.profile.EncryptionConfig(eh.messageSizeHint)
 	config.Time = eh.clock
 
 	compressionConfig := eh.selectCompression()
@@ -324,7 +324,7 @@ func (eh *encryptionHandle) encryptSignDetachedStreamWithSessionKey(
 		eh.IsUTF8,
 		eh.SigningContext,
 		eh.clock,
-		eh.profile.EncryptionConfig(),
+		eh.profile.EncryptionConfig(eh.messageSizeHint),
 	)
 	if err != nil {
 		return nil, err
@@ -345,7 +345,7 @@ func (eh *encryptionHandle) encryptSignDetachedStreamToRecipients(
 	keyPacketWriter io.Writer,
 	encryptSignature bool,
 ) (plaintextWriter io.WriteCloser, err error) {
-	configInput := eh.profile.EncryptionConfig()
+	configInput := eh.profile.EncryptionConfig(eh.messageSizeHint)
 	configInput.Time = NewConstantClock(eh.clock().Unix())
 	// Generate a session key for encryption.
 	if eh.SessionKey == nil {

crypto/encryption_handle.go

@@ -62,6 +62,8 @@ type encryptionHandle struct {
 
 	encryptionTimeOverride Clock
 	clock                  Clock
+
+	messageSizeHint uint64
 }
 
 // --- Default decryption handle to build from
@@ -74,6 +76,13 @@ func defaultEncryptionHandle(profile EncryptionProfile, clock Clock) *encryption
 }
 
 // --- Implements PGPEncryption interface
+// SetMessageSizeHint gives the encryption handle a hint about the
+// expected size of the message, in order to set an appropriate chunk
+// size when using AEAD. Nothing will break when the message size hint
+// turns out to be wrong.
+func (eh *encryptionHandle) SetMessageSizeHint(messageSizeHint uint64) {
+	eh.messageSizeHint = messageSizeHint
+}
 
 // EncryptingWriter returns a wrapper around underlying output Writer,
 // such that any write-operation via the wrapper results in a write to an encrypted pgp message.
@@ -95,6 +104,7 @@ func (eh *encryptionHandle) EncryptingWriter(outputWriter Writer, encoding int8)
 
 // Encrypt encrypts a plaintext message.
 func (eh *encryptionHandle) Encrypt(message []byte) (*PGPMessage, error) {
+	eh.messageSizeHint = uint64(len(message))
 	pgpMessageBuffer := NewPGPMessageBuffer()
 	// Enforce that for a PGPMessage struct the output should not be armored.
 	encryptingWriter, err := eh.EncryptingWriter(pgpMessageBuffer, Bytes)
@@ -116,7 +126,7 @@ func (eh *encryptionHandle) Encrypt(message []byte) (*PGPMessage, error) {
 // EncryptSessionKey encrypts a session key with the encryption handle.
 // To encrypt a session key, the handle must contain either recipients or a password.
 func (eh *encryptionHandle) EncryptSessionKey(sessionKey *SessionKey) ([]byte, error) {
-	config := eh.profile.EncryptionConfig()
+	config := eh.profile.EncryptionConfig(0)
 	config.Time = NewConstantClock(eh.clock().Unix())
 	switch {
 	case eh.Password != nil:
@@ -159,7 +169,7 @@ func (eh *encryptionHandle) armorChecksumRequired() bool {
 		// the logic for the RFC9580 check.
 		return false
 	}
-	encryptionConfig := eh.profile.EncryptionConfig()
+	encryptionConfig := eh.profile.EncryptionConfig(0)
 	if encryptionConfig.AEADConfig == nil {
 		return true
 	}

crypto/sessionkey_test.go

@@ -104,7 +104,7 @@ func TestSymmetricKeyPacketWrongSize(t *testing.T) {
 
 	password := []byte("I like encryption")
 
-	_, err = encryptSessionKeyWithPassword(sk, password, testPGP.profile.EncryptionConfig())
+	_, err = encryptSessionKeyWithPassword(sk, password, testPGP.profile.EncryptionConfig(0))
 	if err == nil {
 		t.Fatal("Expected error while generating key packet with wrong sized key")
 	}

profile/profile.go

@@ -68,14 +68,24 @@ func (p *Custom) KeyGenerationConfig(securityLevel int8) *packet.Config {
 	return cfg
 }
 
-func (p *Custom) EncryptionConfig() *packet.Config {
+func (p *Custom) EncryptionConfig(messageSizeHint uint64) *packet.Config {
 	config := &packet.Config{
 		DefaultHash:                            p.Hash,
 		DefaultCipher:                          p.CipherEncryption,
 		AEADConfig:                             p.AeadEncryption,
 		S2KConfig:                              p.S2kEncryption,
 		InsecureAllowDecryptionWithSigningKeys: p.InsecureAllowDecryptionWithSigningKeys,
 	}
+	if config.AEADConfig != nil && messageSizeHint != 0 {
+		chunkSize := config.AEADConfig.ChunkSize
+		if messageSizeHint*2 < 1<<(config.AEADConfig.ChunkSizeByte()+6) {
+			chunkSize = messageSizeHint * 2
+		}
+		config.AEADConfig = &packet.AEADConfig{
+			DefaultMode: config.AEADConfig.DefaultMode,
+			ChunkSize:   chunkSize,
+		}
+	}
 	if p.DisableIntendedRecipients {
 		intendedRecipients := false
 		config.CheckIntendedRecipients = &intendedRecipients