Low-memory garbage collector

crypto/attachment.go

@@ -4,23 +4,50 @@ import (
 	"bytes"
 	"io"
 	"io/ioutil"
+	"runtime"
+	"sync"
 
 	armorUtils "github.com/ProtonMail/go-pm-crypto/armor"
-	"github.com/ProtonMail/go-pm-crypto/internal"
 	"github.com/ProtonMail/go-pm-crypto/models"
 	"golang.org/x/crypto/openpgp"
-	"golang.org/x/crypto/openpgp/armor"
 	"golang.org/x/crypto/openpgp/packet"
 )
 
-// Encrypt attachment. Takes input data and key data in binary form
-func (pm *PmCrypto) EncryptAttachment(plainData []byte, fileName string, publicKey *KeyRing) (*models.EncryptedSplit, error) {
-	var outBuf bytes.Buffer
-	w, err := armor.Encode(&outBuf, armorUtils.PGP_MESSAGE_HEADER, internal.ArmorHeaders)
-	if err != nil {
-		return nil, err
+//EncryptedSplit when encrypt attachment
+type AttachmentProcessor struct {
+	w                *io.WriteCloser
+	pipe             *io.PipeWriter
+	done             sync.WaitGroup
+	split            *models.EncryptedSplit
+	garbageCollector int
+	err              error
 }
 
+func (ap *AttachmentProcessor) Process(plainData []byte) {
+	(*ap.w).Write(plainData)
+}
+
+func (ap *AttachmentProcessor) Finish() (*models.EncryptedSplit, error) {
+	if ap.err != nil {
+		return nil, ap.err
+	}
+	(*ap.w).Close()
+	(*ap.pipe).Close()
+	ap.done.Wait()
+	if ap.garbageCollector > 0 {
+		runtime.GC()
+	}
+	return ap.split, nil
+}
+
+// Encrypt attachment. Takes input data and key data in binary form
+func (pm *PmCrypto) encryptAttachment(estimatedSize int, fileName string, publicKey *KeyRing, garbageCollector int) (*AttachmentProcessor, error) {
+	attachmentProc := &AttachmentProcessor{}
+	// you can also add these one at
+	// a time if you need to
+	attachmentProc.done.Add(1)
+	attachmentProc.garbageCollector = garbageCollector
+
 	hints := &openpgp.FileHints{
 		FileName: fileName,
 	}
@@ -30,18 +57,49 @@ func (pm *PmCrypto) EncryptAttachment(plainData []byte, fileName string, publicK
 		Time:          pm.getTimeGenerator(),
 	}
 
-	ew, err := openpgp.Encrypt(w, publicKey.entities, nil, hints, config)
+	reader, writer := io.Pipe()
 
-	_, _ = ew.Write(plainData)
-	ew.Close()
-	w.Close()
+	go func() {
+		defer attachmentProc.done.Done()
+		split, splitError := SeparateKeyAndData(nil, reader, estimatedSize, garbageCollector)
+		if attachmentProc.err != nil {
+			attachmentProc.err = splitError
+		}
+		split.Algo = "aes256"
+		attachmentProc.split = split
+	}()
 
-	split, err := SplitArmor(outBuf.String())
+	var ew io.WriteCloser
+	var encryptErr error
+	ew, encryptErr = openpgp.Encrypt(writer, publicKey.entities, nil, hints, config)
+	attachmentProc.w = &ew
+	attachmentProc.pipe = writer
+	if attachmentProc.err != nil {
+		attachmentProc.err = encryptErr
+	}
+
+	return attachmentProc, nil
+
+}
+
+func (pm *PmCrypto) EncryptAttachment(plainData []byte, fileName string, publicKey *KeyRing) (*models.EncryptedSplit, error) {
+	ap, err := pm.encryptAttachment(len(plainData), fileName, publicKey, -1)
+	if err != nil {
+		return nil, err
+	}
+	ap.Process(plainData)
+	split, err := ap.Finish()
 	if err != nil {
 		return nil, err
 	}
-	split.Algo = "aes256"
 	return split, nil
+
+}
+
+//EncryptAttachment ...
+func (pm *PmCrypto) EncryptAttachmentLowMemory(estimatedSize int, fileName string, publicKey *KeyRing) (*AttachmentProcessor, error) {
+	// Garbage collect every megabyte
+	return pm.encryptAttachment(estimatedSize, fileName, publicKey, 1<<20)
 }
 
 // Helper method. Splits armored pgp session into key and packet data
@@ -56,7 +114,7 @@ func SplitArmor(encrypted string) (*models.EncryptedSplit, error) {
 
 	encryptedReader := bytes.NewReader(encryptedRaw)
 
-	return SeparateKeyAndData(nil, encryptedReader)
+	return SeparateKeyAndData(nil, encryptedReader, len(encrypted), -1)
 
 }
 

crypto/key.go

@@ -9,11 +9,11 @@ import (
 	"fmt"
 	"github.com/ProtonMail/go-pm-crypto/armor"
 	"io"
-	"io/ioutil"
 	"math/big"
 	"time"
 	//	"net/http"
 	//	"net/url"
+	"runtime"
 	"strings"
 
 	//"github.com/ProtonMail/go-pm-crypto/armor"
@@ -110,9 +110,12 @@ func DecryptAttKey(kr *KeyRing, keyPacket string) (key *SymmetricKey, err error)
 }
 
 // Separate key and data packets in a pgp session
-func SeparateKeyAndData(kr *KeyRing, r io.Reader) (outSplit *models.EncryptedSplit, err error) {
+func SeparateKeyAndData(kr *KeyRing, r io.Reader, estimatedLength int, garbageCollector int) (outSplit *models.EncryptedSplit, err error) {
+
+	// For info on each, see: https://golang.org/pkg/runtime/#MemStats
 	packets := packet.NewReader(r)
 	outSplit = &models.EncryptedSplit{}
+	gcCounter := 0
 
 	// Save encrypted key and signature apart
 	var ek *packet.EncryptedKey
@@ -144,20 +147,56 @@ func SeparateKeyAndData(kr *KeyRing, r io.Reader) (outSplit *models.EncryptedSpl
 				}
 			}
 		case *packet.SymmetricallyEncrypted:
-			var packetContents []byte
-			if packetContents, err = ioutil.ReadAll(p.Contents); err != nil {
-				return
+			// The code below is optimized to not
+			var b bytes.Buffer
+			// 2^16 is an estimation of the size difference between input and output, the size difference is most probably
+			// 16 bytes at a maximum though.
+			// We need to avoid triggering a grow from the system as this will allocate too much memory causing problems
+			// in low-memory environments
+			b.Grow(1<<16 + 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
+				gcCounter += n
+				if gcCounter > garbageCollector && garbageCollector > 0 {
+					runtime.GC()
+					gcCounter = 0
+				}
 			}
 
-			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...)
+			// 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
+
 		}
 	}
 	if decryptErr != nil {

crypto/keyring.go

@@ -241,7 +241,7 @@ func (kr *KeyRing) EncryptSymmetric(textToEncrypt string, canonicalizeText bool)
 	}
 	encryptedWriter.Close()
 
-	if outSplit, err = SeparateKeyAndData(kr, buffer); err != nil {
+	if outSplit, err = SeparateKeyAndData(kr, buffer, len(textToEncrypt), -1); err != nil {
 		return
 	}
 

crypto/message.go

@@ -3,6 +3,7 @@ package crypto
 import (
 	"bytes"
 	"errors"
+	"io"
 	"io/ioutil"
 	"time"