Allow passing message size hint

Also, set the AEAD chunk size to twice the message size hint (or the
actual message size, when known), so that we don't generate an overly
large chunk size when the message could fit in a smaller chunk size.

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 int) *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 int)
 	// 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 int
 }
 
 // --- 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 int) {
+	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 = 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 int) *packet.Config {
 	config := &packet.Config{
 		DefaultHash:                            p.Hash,
 		DefaultCipher:                          p.CipherEncryption,
 		AEADConfig:                             p.AeadEncryption,
 		S2KConfig:                              p.S2kEncryption,
 		InsecureAllowDecryptionWithSigningKeys: p.InsecureAllowDecryptionWithSigningKeys,
 	}
+	if config.AEADConfig != nil {
+		chunkSize := config.AEADConfig.ChunkSize
+		if messageSizeHint*2 < 1<<(config.AEADConfig.ChunkSizeByte()+6) {
+			chunkSize = uint64(messageSizeHint * 2)
+		}
+		config.AEADConfig = &packet.AEADConfig{
+			DefaultMode: config.AEADConfig.DefaultMode,
+			ChunkSize:   chunkSize,
+		}
+	}
 	if p.DisableIntendedRecipients {
 		intendedRecipients := false
 		config.CheckIntendedRecipients = &intendedRecipients