Optimize encrypt attachment

src/gitlab.com/ProtonMail/go-pm-crypto/armor/armor.go

@@ -4,13 +4,11 @@ import (
 	"bytes"
 	"errors"
 	"gitlab.com/ProtonMail/go-pm-crypto/internal"
-	"gitlab.com/ProtonMail/go-pm-crypto/models"
-	"golang.org/x/crypto/openpgp/armor"
+		"golang.org/x/crypto/openpgp/armor"
 	"golang.org/x/crypto/openpgp/clearsign"
-	"golang.org/x/crypto/openpgp/packet"
-	"io"
 	"io/ioutil"
-)
+	"gitlab.com/ProtonMail/go-pm-crypto/models"
+			)
 
 // ArmorKey make bytes input key to armor format
 func ArmorKey(input []byte) (string, error) {
@@ -50,80 +48,15 @@ func ReadClearSignedMessage(signedMessage string) (string, error) {
 	return string(modulusBlock.Bytes), nil
 }
 
-//SeparateKeyAndData ...
 func SplitArmor(encrypted string) (*models.EncryptedSplit, error) {
-
 	var err error
-
-	encryptedRaw, err := Unarmor(encrypted)
+	b, err := internal.Unarmor(encrypted)
 	if err != nil {
 		return nil, err
 	}
-
-	encryptedReader := bytes.NewReader(encryptedRaw)
-
-	//kr *KeyRing, r io.Reader) (key *SymmetricKey, symEncryptedData []byte,
-	packets := packet.NewReader(encryptedReader)
-
-	outSplit := &models.EncryptedSplit{}
-
-	// Save encrypted key and signature apart
-	var ek *packet.EncryptedKey
-	// var decryptErr error
-	for {
-		var p packet.Packet
-		if p, err = packets.Next(); err == io.EOF {
-			err = nil
-			break
-		}
-		switch p := p.(type) {
-		case *packet.EncryptedKey:
-			// We got an encrypted key. Try to decrypt it with each available key
-			if ek != nil && ek.Key != nil {
-				break
-			}
-			ek = p
-			break
-		case *packet.SymmetricallyEncrypted:
-			var packetContents []byte
-			if packetContents, err = ioutil.ReadAll(p.Contents); err != nil {
-				return nil, err
-			}
-
-			encodedLength := encodedLength(len(packetContents) + 1)
-			var symEncryptedData []byte
-			symEncryptedData = append(symEncryptedData, byte(210))
-			symEncryptedData = append(symEncryptedData, encodedLength...)
-			symEncryptedData = append(symEncryptedData, byte(1))
-			symEncryptedData = append(symEncryptedData, packetContents...)
-
-			outSplit.DataPacket = symEncryptedData
-			break
-
-		}
+	split, err := internal.SplitPackets(b.Body, len(encrypted))
+	if err != nil {
+		return nil, err
 	}
-
-	var buf bytes.Buffer
-	ek.Serialize(&buf)
-	outSplit.KeyPacket = buf.Bytes()
-
-	return outSplit, err
-}
-
-//encode length based on 4.2.2. in the RFC
-func encodedLength(length int) (b []byte) {
-	if length < 192 {
-		b = append(b, byte(length))
-	} else if length < 8384 {
-		length = length - 192
-		b = append(b, 192+byte(length>>8))
-		b = append(b, byte(length))
-	} else {
-		b = append(b, byte(255))
-		b = append(b, byte(length>>24))
-		b = append(b, byte(length>>16))
-		b = append(b, byte(length>>8))
-		b = append(b, byte(length))
-	}
-	return
-}
+	return split, nil
+}

src/gitlab.com/ProtonMail/go-pm-crypto/crypto/attachment.go

@@ -11,17 +11,15 @@ import (
 	armorUtils "gitlab.com/ProtonMail/go-pm-crypto/armor"
 	"gitlab.com/ProtonMail/go-pm-crypto/internal"
 	"gitlab.com/ProtonMail/go-pm-crypto/models"
-)
+		"sync"
+	)
 
 //EncryptAttachmentBinKey ...
 func (pm *PmCrypto) EncryptAttachmentBinKey(plainData []byte, fileName string, publicKey []byte) (*models.EncryptedSplit, error) {
-
-	var outBuf bytes.Buffer
-	w, err := armor.Encode(&outBuf, armorUtils.MESSAGE_HEADER, internal.ArmorHeaders)
-	if err != nil {
-		return nil, err
-	}
-
+	var wg sync.WaitGroup
+	// you can also add these one at
+	// a time if you need to
+	wg.Add(1)
 	pubKeyReader := bytes.NewReader(publicKey)
 	pubKeyEntries, err := openpgp.ReadKeyRing(pubKeyReader)
 	if err != nil {
@@ -36,13 +34,26 @@ func (pm *PmCrypto) EncryptAttachmentBinKey(plainData []byte, fileName string, p
 		Time:          pm.getTimeGenerator(),
 	}
 
-	ew, err := openpgp.Encrypt(w, pubKeyEntries, nil, hints, config)
+	reader, writer := io.Pipe()
+	var encryptErr error
+	go func() {
+		defer wg.Done()
+		var ew io.WriteCloser
+		ew, encryptErr = openpgp.Encrypt(writer, pubKeyEntries, nil, hints, config)
+		_, _ = ew.Write(plainData)
+		plainData = nil
+		// clear the buffer
+		ew.Close()
+		writer.Close()
+	}()
 
-	_, _ = ew.Write(plainData)
-	ew.Close()
-	w.Close()
+	split, err := internal.SplitPackets(reader, len(plainData))
+
+	wg.Wait()
+	if encryptErr != nil {
+		return nil, encryptErr
+	}
 
-	split, err := armorUtils.SplitArmor(outBuf.String())
 	if err != nil {
 		return nil, err
 	}

src/gitlab.com/ProtonMail/go-pm-crypto/crypto/attachment_test.go

@@ -0,0 +1,67 @@
+package crypto
+
+import (
+	"testing"
+)
+
+const publicKeya = `-----BEGIN PGP PUBLIC KEY BLOCK-----
+
+mQINBFs6C3wBEAD9PXYGS+psd1dPq4sYSmG1Q9K4fjQ2Lks4Xvr1ohvM7V4vPwma
+fllG9DbZ9mCf4+3Okrk4NhvPuVhsOhhecld2UlAEs0Y1WQJgx9Z23Yi5Fvtg428J
+fUCBKaa/GqWXw+/Y1q6ln+FAcxD2BH964V49Mv5dsouLl7FtOLCNuGmgFMQBrbZu
+uSZowSsByN3nsuVutHpNRiarhMda2dfSekti1i4ywxMUVK0xNpT7baZgQqGphQOV
+Xyc8BniSUq9/TCfPJ9Or5BuykokhVM9xufCysKdTlzXGF8Yb3xhQwmQG3d8S6OLz
+H5N12Kiqhmb/BUZWMQ0ESE6izg8IOgk/rJZuHCM61nIpH2jRMygx2j23lwye/A0B
+NH/fZlo21McXEgkWGVS1EiK1QY5IfAdk3v9DhwLsIWftRFp1Tg4kTh0oqUfILH1Z
+6CPHsJR9aOf8a95Y8czBOYji+j5L0I7BRpcMnWwdELjVlJAzvgkdBz1KOskSu8Tb
+3mITzbi05EDhjQnMCUj2pETS9ujZns6Q90kh441wYTQ1Gckl2zVHQg3+A60M32Ys
+IQbBlOY0Di27e8dRhGNuhGdMTtHFq8e0UgBnHEPCVoQJicg/QXZM7F3S2cUTCW8n
+fCBmMJ0Tnaio/iXbhH7qWSaajpJ6AFALl/ZB8/wmJi/dbpoM5OtEHVHOIwARAQAB
+tB5QR1AgVGVzdCA8cG1wZ3B0ZXN0QGdtYWlsLmNvbT6JAlQEEwEIAD4WIQTap+3+
+yz+21dA8kYg3QTCzLuHl6gUCWzoLfAIbIwUJCWYBgAULCQgHAgYVCgkICwIEFgID
+AQIeAQIXgAAKCRA3QTCzLuHl6qJJEACBQh0wcegTU90nOeDLWK6UKDb7vSmo2PBl
+pq/MowBIMLeRBrnCUg+j7F2R6xXJJyJCnkjxmKcA5ZtGASE8l3iHIOJTZQbQMcVE
+eowfnq6bRdjCZcEsbCnWrw2TAlz6Wq0ZblDv7EkBKAl3Uq2SDM5BRTddZSpGdLRF
+4e5TiHf+ddT2rkTNPAdm151fO6rXd4Kh+6gAwYPPv15qUB4KpfJ4SAXNhESN9yes
+zfs0xXVu7ekHF2M551qQWGRpfhXNRcbJLr2mOHEdERsCPxMD6HqTcn4TTt4hmoeG
+kk9RTalfBWHo99nyay3Pc49wMPRP+l9b9ptJ+t/5I1pIvdL0XKHohdXF3KHt6ATX
+2uwASoKcl6FQwpzRmSenYL0vad2Pjziqpy+pC3KOg3r7Iq4hNLP/AnlyOQKozx4d
+OEseWeGqLsDkI23noXcEVib36mXcKCln3xtDednq99e2a8Y673BUwAuta90n1pUO
+K+/aCQ0T7WJQV2fru13IPjGSkFjDh4s+d3IsWysNqx0Shcz26/HMP0XmfAAF0CSW
+JwtzDRvgxocrtGFu3noit1B6ncYqpKrCXDDc8RtvuyJzdzd2V91dP+quBck4R0Pl
+8r99fkJL6ijeahN6QoIapfghyz9qxVqiR8Eii44A0YFvhCLdPbuRlBQcC0+7n/zR
+4F+doiBLTbkCDQRbOgt8ARAAt98HXbVOwtRiXkfC6m6zIFnAgHBVfHDhzBwl2zDo
+83R58W2TlKZetWQApd4+3RKEiilUbrrItO7eLWcF4uFFsjvL5iYOBCO/I+eSpwHO
+Ey11yPZlaR4Nf0VJ4kxRD62Oeriy8WaHG9hok1JNSg6LVdLawZsvApcHNnGbyY8s
+VHA2HA9qiTwpI6EzziebSKpdZJdqnR5F53rvkC7aXMoY4V6WcVQASqBjOuUbMSFG
+a0ZRnaUgHBaqPKup6T2OibRaEHZi/MXKYVKH75Ry4sxIe50uxSstMcpFJm+J0STm
+xNMeWxc7wXusgT3Dbn7goJOISIhUtwYTdGvom1W3DzTCFWyhEnFdPIwpJ/rCSk2Z
+W3qiapt6827mSW4eBl0XmtCBdN/JszcJPRML9+xPcPWSwB8hPmzqZbemnnaJx570
+TC3p0D6BKWvvMNEuBlZJ9Ez+fuVYM6f2wnfssbzuC38D+We4nouqiftuVXdm144H
+vXtLvHsPQA6Er2DF2BclyPh7l2MbifxP7p9X7Nup5Qr7ek/THj8jEBBTbCwLM9by
+o/Gu3ahjmdb3W9cgwEDRAZvWHFtrif1FZ3CF5cdLMrCrnlIpCtDJ/weGxhKv2XW4
+YXAgoYH16kCiGdRyGZ+DsN0aT0fFYjE6LuUMqQKHiLAgaZyC3w53PfeulWdJt1BV
+6nEAEQEAAYkCPAQYAQgAJhYhBNqn7f7LP7bV0DyRiDdBMLMu4eXqBQJbOgt8AhsM
+BQkJZgGAAAoJEDdBMLMu4eXqdUwQAINZSS85w7U/Ghwx6SNL2k8gARxE9ShOq42p
+dcjZzf3ZIfyNVszwZJEpxcnqqyMRZJXx1iOIN2dGOFdYL+bOPxTk0St4k/zpGyLM
+9G6WPuvaqNvqShaSDXi7V+UF/uGcB3KKTA3/4n08t6Yq5Xh93n1roCu/9P3g9xbf
+mls/l/PUkjKCpJHm/1FCejizfw+/QvQpuzy1vU4on1g0U7pJ0R1GiU45vffrV7xa
+bozh2eO0+vo5vd12fIkSLDT8NOwhWQ+BeM5/zze+GZaDvNEcM0eo8jardU4GZqjD
+JWd0Rqr05uXf1THVmjzYXyfRy+/RQaFWoSUo1UWIad58DR3ND3SkeJAqQRx+3hd2
+VRvvcI17qPwo6MZ9eu70ezt0w/IXAc1iLlxPy0tI5oE0bXjc/pGmJLnGT7cYxAQr
+BbzYQNlhrRTo8Ou8JebwiHESCqPRos1FfALckfulsKIVPm0QDRLi7S/qkGNZ0daa
+geeHFhi1vHwLh7L9INcT2npSDO6xDJHuTK+v8Fna8LaqLk/Q2e+77zd2Nen5UoCt
+6rClf3RLPbT1TtfPHkMDcmrKnesdkpSWdZV0S7m/nAqfcjNgRZ/4uoH1SLYLPRCG
+VeiHC6ENvuti8fyWpcfYwjvBoP2xcq2D8n7HbJjPR+LR1z1lgpdDDN3LlD8ggy/y
+OQTmvM6R
+=DvFv
+-----END PGP PUBLIC KEY BLOCK-----`
+
+
+func TestEncrypt(t *testing.T) {
+	o := PmCrypto{}
+	split, err := o.EncryptAttachment([]byte("Hello World!"), "File name", publicKeya)
+	if err != nil || split == nil {
+		panic("Encryption failed")
+	}
+}

src/gitlab.com/ProtonMail/go-pm-crypto/internal/packets.go

@@ -0,0 +1,105 @@
+package internal
+
+import (
+	"bytes"
+	"golang.org/x/crypto/openpgp/packet"
+	"gitlab.com/ProtonMail/go-pm-crypto/models"
+	"io"
+			)
+
+func SplitPackets(encryptedReader io.Reader, estimatedLength int) (*models.EncryptedSplit, error){
+	var err error
+
+	//kr *KeyRing, r io.Reader) (key *SymmetricKey, symEncryptedData []byte,
+	packets := packet.NewReader(encryptedReader)
+
+	outSplit := &models.EncryptedSplit{}
+
+	// Save encrypted key and signature apart
+	var ek *packet.EncryptedKey
+	// var decryptErr error
+	for {
+		var p packet.Packet
+		if p, err = packets.Next(); err == io.EOF {
+			err = nil
+			break
+		}
+		switch p := p.(type) {
+		case *packet.EncryptedKey:
+			// We got an encrypted key. Try to decrypt it with each available key
+			if ek != nil && ek.Key != nil {
+				break
+			}
+			ek = p
+			break
+		case *packet.SymmetricallyEncrypted:
+			// The code below is optimized to not
+			var b bytes.Buffer
+			// 128 is an estimation of the size difference between input and output, the size difference is most probably
+			// 16 bytes at a maximum though.
+			b.Grow(128 + estimatedLength)
+			// empty encoded length + start byte
+			b.Write(make([]byte, 6))
+			b.WriteByte(byte(1))
+			actualLength := 1
+			block := make([]byte, 128)
+			for {
+				n, err := p.Contents.Read(block)
+				if err == io.EOF {
+					break
+				}
+				b.Write(block[:n])
+				actualLength += n
+			}
+
+			// quick encoding
+			symEncryptedData := b.Bytes()
+			if actualLength < 192 {
+				symEncryptedData[4] = byte(210)
+				symEncryptedData[5] = byte(actualLength)
+				symEncryptedData = symEncryptedData[4:]
+			} else if actualLength < 8384 {
+				actualLength = actualLength - 192
+				symEncryptedData[3] = byte(210)
+				symEncryptedData[4] = 192+byte(actualLength>>8)
+				symEncryptedData[5] = byte(actualLength)
+				symEncryptedData = symEncryptedData[3:]
+			} else {
+				symEncryptedData[0] = byte(210)
+				symEncryptedData[1] = byte(255)
+				symEncryptedData[2] = byte(actualLength>>24)
+				symEncryptedData[3] = byte(actualLength>>16)
+				symEncryptedData[4] = byte(actualLength>>8)
+				symEncryptedData[5] = byte(actualLength)
+			}
+
+			outSplit.DataPacket = symEncryptedData
+			break
+
+		}
+	}
+
+	var buf bytes.Buffer
+	ek.Serialize(&buf)
+	outSplit.KeyPacket = buf.Bytes()
+
+	return outSplit, err
+}
+
+//encode length based on 4.2.2. in the RFC
+func encodedLength(length int) (b []byte) {
+	if length < 192 {
+		b = append(b, byte(length))
+	} else if length < 8384 {
+		length = length - 192
+		b = append(b, 192+byte(length>>8))
+		b = append(b, byte(length))
+	} else {
+		b = append(b, byte(255))
+		b = append(b, byte(length>>24))
+		b = append(b, byte(length>>16))
+		b = append(b, byte(length>>8))
+		b = append(b, byte(length))
+	}
+	return
+}