From d7f0550a4b1d1d6b809b3b1137619b1de8088eca Mon Sep 17 00:00:00 2001 From: Jakub Lehotsky Date: Thu, 22 Nov 2018 10:53:14 +0100 Subject: [PATCH] Low-memory garbage collector --- crypto/attachment.go | 88 ++++++++++++++++++++++++++++++++++++-------- crypto/key.go | 61 ++++++++++++++++++++++++------ crypto/keyring.go | 2 +- crypto/message.go | 1 + 4 files changed, 125 insertions(+), 27 deletions(-) diff --git a/crypto/attachment.go b/crypto/attachment.go index 21e0dda..ee7ceb1 100644 --- a/crypto/attachment.go +++ b/crypto/attachment.go @@ -4,22 +4,49 @@ 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) } diff --git a/crypto/key.go b/crypto/key.go index de17d35..e885dd8 100644 --- a/crypto/key.go +++ b/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 { diff --git a/crypto/keyring.go b/crypto/keyring.go index 687e9c1..2d6e172 100644 --- a/crypto/keyring.go +++ b/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 } diff --git a/crypto/message.go b/crypto/message.go index 8f140ad..a8fd28d 100644 --- a/crypto/message.go +++ b/crypto/message.go @@ -3,6 +3,7 @@ package crypto import ( "bytes" "errors" + "io" "io/ioutil" "time"