GODT-1056 Check encrypted size of the message before upload

internal/smtp/user.go

@@ -24,7 +24,6 @@ import (
 	"encoding/base64"
 	"fmt"
 	"io"
-	"io/ioutil"
 	"mime"
 	"net/mail"
 	"strings"
@@ -325,12 +324,6 @@ func (su *smtpUser) Send(returnPath string, to []string, messageReader io.Reader
 		return nil
 	}
 
-	if ok, err := su.isTotalSizeOkay(message, attReaders); err != nil {
-		return err
-	} else if !ok {
-		return errors.New("message is too large")
-	}
-
 	su.backend.sendRecorder.addMessage(sendRecorderMessageHash)
 	message, atts, err := su.storeUser.CreateDraft(kr, message, attReaders, attachedPublicKey, attachedPublicKeyName, parentID)
 	if err != nil {
@@ -541,24 +534,3 @@ func (su *smtpUser) Logout() error {
 	log.Debug("SMTP client logged out user ", su.addressID)
 	return nil
 }
-
-func (su *smtpUser) isTotalSizeOkay(message *pmapi.Message, attReaders []io.Reader) (bool, error) {
-	maxUpload, err := su.storeUser.GetMaxUpload()
-	if err != nil {
-		return false, err
-	}
-
-	var attSize int64
-
-	for i := range attReaders {
-		b, err := ioutil.ReadAll(attReaders[i])
-		if err != nil {
-			return false, err
-		}
-
-		attSize += int64(len(b))
-		attReaders[i] = bytes.NewBuffer(b)
-	}
-
-	return message.Size+attSize <= maxUpload, nil
-}

internal/store/event_loop.go

@@ -99,6 +99,11 @@ func (loop *eventLoop) setFirstEventID() (err error) {
 // pollNow starts polling events right away and waits till the events are
 // processed so we are sure updates are propagated to the database.
 func (loop *eventLoop) pollNow() {
+	// When event loop is not running, it would cause infinite wait.
+	if !loop.isRunning {
+		return
+	}
+
 	eventProcessedCh := make(chan struct{})
 	loop.pollCh <- eventProcessedCh
 	<-eventProcessedCh

internal/store/store_test.go

@@ -25,6 +25,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/ProtonMail/gopenpgp/v2/crypto"
 	storemocks "github.com/ProtonMail/proton-bridge/internal/store/mocks"
 	"github.com/ProtonMail/proton-bridge/pkg/pmapi"
 	pmapimocks "github.com/ProtonMail/proton-bridge/pkg/pmapi/mocks"
@@ -39,8 +40,92 @@ const (
 
 	addr2   = "jamesandmichalarecool@pm.me"
 	addrID2 = "jamesandmichalarecool"
+
+	testPrivateKeyPassword = "apple"
+	testPrivateKey         = `-----BEGIN PGP PRIVATE KEY BLOCK-----
+Version: OpenPGP.js v0.7.1
+Comment: http://openpgpjs.org
+
+xcMGBFRJbc0BCAC0mMLZPDBbtSCWvxwmOfXfJkE2+ssM3ux21LhD/bPiWefE
+WSHlCjJ8PqPHy7snSiUuxuj3f9AvXPvg+mjGLBwu1/QsnSP24sl3qD2onl39
+vPiLJXUqZs20ZRgnvX70gjkgEzMFBxINiy2MTIG+4RU8QA7y8KzWev0btqKi
+MeVa+GLEHhgZ2KPOn4Jv1q4bI9hV0C9NUe2tTXS6/Vv3vbCY7lRR0kbJ65T5
+c8CmpqJuASIJNrSXM/Q3NnnsY4kBYH0s5d2FgbASQvzrjuC2rngUg0EoPsrb
+DEVRA2/BCJonw7aASiNCrSP92lkZdtYlax/pcoE/mQ4WSwySFmcFT7yFABEB
+AAH+CQMIvzcDReuJkc9gnxAkfgmnkBFwRQrqT/4UAPOF8WGVo0uNvDo7Snlk
+qWsJS+54+/Xx6Jur/PdBWeEu+6+6GnppYuvsaT0D0nFdFhF6pjng+02IOxfG
+qlYXYcW4hRru3BfvJlSvU2LL/Z/ooBnw3T5vqd0eFHKrvabUuwf0x3+K/sru
+Fp24rl2PU+bzQlUgKpWzKDmO+0RdKQ6KVCyCDMIXaAkALwNffAvYxI0wnb2y
+WAV/bGn1ODnszOYPk3pEMR6kKSxLLaO69kYx4eTERFyJ+1puAxEPCk3Cfeif
+yDWi4rU03YB16XH7hQLSFl61SKeIYlkKmkO5Hk1ybi/BhvOGBPVeGGbxWnwI
+46G8DfBHW0+uvD5cAQtk2d/q3Ge1I+DIyvuRCcSu0XSBNv/Bkpp4IbAUPBaW
+TIvf5p9oxw+AjrMtTtcdSiee1S6CvMMaHhVD7SI6qGA8GqwaXueeLuEXa0Ok
+BWlehx8wibMi4a9fLcQZtzJkmGhR1WzXcJfiEg32srILwIzPQYxuFdZZ2elb
+gYp/bMEIp4LKhi43IyM6peCDHDzEba8NuOSd0heEqFIm0vlXujMhkyMUvDBv
+H0V5On4aMuw/aSEKcAdbazppOru/W1ndyFa5ZHQIC19g72ZaDVyYjPyvNgOV
+AFqO4o3IbC5z31zMlTtMbAq2RG9svwUVejn0tmF6UPluTe0U1NuXFpLK6TCH
+wqocLz4ecptfJQulpYjClVLgzaYGDuKwQpIwPWg5G/DtKSCGNtEkfqB3aemH
+V5xmoYm1v5CQZAEvvsrLA6jxCk9lzqYV8QMivWNXUG+mneIEM35G0HOPzXca
+LLyB+N8Zxioc9DPGfdbcxXuVgOKRepbkq4xv1pUpMQ4BUmlkejDRSP+5SIR3
+iEthg+FU6GRSQbORE6nhrKjGBk8fpNpozQZVc2VySUTCwHIEEAEIACYFAlRJ
+bc8GCwkIBwMCCRA+tiWe3yHfJAQVCAIKAxYCAQIbAwIeAQAA9J0H/RLR/Uwt
+CakrPKtfeGaNuOI45SRTNxM8TklC6tM28sJSzkX8qKPzvI1PxyLhs/i0/fCQ
+7Z5bU6n41oLuqUt2S9vy+ABlChKAeziOqCHUcMzHOtbKiPkKW88aO687nx+A
+ol2XOnMTkVIC+edMUgnKp6tKtZnbO4ea6Cg88TFuli4hLHNXTfCECswuxHOc
+AO1OKDRrCd08iPI5CLNCIV60QnduitE1vF6ehgrH25Vl6LEdd8vPVlTYAvsa
+6ySk2RIrHNLUZZ3iII3MBFL8HyINp/XA1BQP+QbH801uSLq8agxM4iFT9C+O
+D147SawUGhjD5RG7T+YtqItzgA1V9l277EXHwwYEVEltzwEIAJD57uX6bOc4
+Tgf3utfL/4hdyoqIMVHkYQOvE27wPsZxX08QsdlaNeGji9Ap2ifIDuckUqn6
+Ji9jtZDKtOzdTBm6rnG5nPmkn6BJXPhnecQRP8N0XBISnAGmE4t+bxtts5Wb
+qeMdxJYqMiGqzrLBRJEIDTcg3+QF2Y3RywOqlcXqgG/xX++PsvR1Jiz0rEVP
+TcBc7ytyb/Av7mx1S802HRYGJHOFtVLoPTrtPCvv+DRDK8JzxQW2XSQLlI0M
+9s1tmYhCogYIIqKx9qOTd5mFJ1hJlL6i9xDkvE21qPFASFtww5tiYmUfFaxI
+LwbXPZlQ1I/8fuaUdOxctQ+g40ZgHPcAEQEAAf4JAwgdUg8ubE2BT2DITBD+
+XFgjrnUlQBilbN8/do/36KHuImSPO/GGLzKh4+oXxrvLc5fQLjeO+bzeen4u
+COCBRO0hG7KpJPhQ6+T02uEF6LegE1sEz5hp6BpKUdPZ1+8799Rylb5kubC5
+IKnLqqpGDbH3hIsmSV3CG/ESkaGMLc/K0ZPt1JRWtUQ9GesXT0v6fdM5GB/L
+cZWFdDoYgZAw5BtymE44knIodfDAYJ4DHnPCh/oilWe1qVTQcNMdtkpBgkuo
+THecqEmiODQz5EX8pVmS596XsnPO299Lo3TbaHUQo7EC6Au1Au9+b5hC1pDa
+FVCLcproi/Cgch0B/NOCFkVLYmp6BEljRj2dSZRWbO0vgl9kFmJEeiiH41+k
+EAI6PASSKZs3BYLFc2I8mBkcvt90kg4MTBjreuk0uWf1hdH2Rv8zprH4h5Uh
+gjx5nUDX8WXyeLxTU5EBKry+A2DIe0Gm0/waxp6lBlUl+7ra28KYEoHm8Nq/
+N9FCuEhFkFgw6EwUp7jsrFcqBKvmni6jyplm+mJXi3CK+IiNcqub4XPnBI97
+lR19fupB/Y6M7yEaxIM8fTQXmP+x/fe8zRphdo+7o+pJQ3hk5LrrNPK8GEZ6
+DLDOHjZzROhOgBvWtbxRktHk+f5YpuQL+xWd33IV1xYSSHuoAm0Zwt0QJxBs
+oFBwJEq1NWM4FxXJBogvzV7KFhl/hXgtvx+GaMv3y8gucj+gE89xVv0XBXjl
+5dy5/PgCI0Id+KAFHyKpJA0N0h8O4xdJoNyIBAwDZ8LHt0vlnLGwcJFR9X7/
+PfWe0PFtC3d7cYY3RopDhnRP7MZs1Wo9nZ4IvlXoEsE2nPkWcns+Wv5Yaewr
+s2ra9ZIK7IIJhqKKgmQtCeiXyFwTq+kfunDnxeCavuWL3HuLKIOZf7P9vXXt
+XgEir9rCwF8EGAEIABMFAlRJbdIJED62JZ7fId8kAhsMAAD+LAf+KT1EpkwH
+0ivTHmYako+6qG6DCtzd3TibWw51cmbY20Ph13NIS/MfBo828S9SXm/sVUzN
+/r7qZgZYfI0/j57tG3BguVGm53qya4bINKyi1RjK6aKo/rrzRkh5ZVD5rVNO
+E2zzvyYAnLUWG9AV1OYDxcgLrXqEMWlqZAo+Wmg7VrTBmdCGs/BPvscNgQRr
+6Gpjgmv9ru6LjRL7vFhEcov/tkBLj+CtaWWFTd1s2vBLOs4rCsD9TT/23vfw
+CnokvvVjKYN5oviy61yhpqF1rWlOsxZ4+2sKW3Pq7JLBtmzsZegTONfcQAf7
+qqGRQm3MxoTdgQUShAwbNwNNQR9cInfMnA==
+=2wIY
+-----END PGP PRIVATE KEY BLOCK-----
+`
 )
 
+var testPrivateKeyRing *crypto.KeyRing
+
+func init() {
+	privKey, err := crypto.NewKeyFromArmored(testPrivateKey)
+	if err != nil {
+		panic(err)
+	}
+
+	privKeyUnlocked, err := privKey.Unlock([]byte(testPrivateKeyPassword))
+	if err != nil {
+		panic(err)
+	}
+
+	if testPrivateKeyRing, err = crypto.NewKeyRing(privKeyUnlocked); err != nil {
+		panic(err)
+	}
+}
+
 type mocksForStore struct {
 	tb testing.TB
 

internal/store/user_message.go

@@ -44,15 +44,17 @@ func (store *Store) CreateDraft(
 	attachedPublicKey,
 	attachedPublicKeyName string,
 	parentID string) (*pmapi.Message, []*pmapi.Attachment, error) {
-	defer store.eventLoop.pollNow()
+	attachments := store.prepareDraftAttachments(message, attachmentReaders, attachedPublicKey, attachedPublicKeyName)
 
-	// Since this is a draft, we don't need to sign it.
-	if err := message.Encrypt(kr, nil); err != nil {
+	if err := encryptDraft(kr, message, attachments); err != nil {
 		return nil, nil, errors.Wrap(err, "failed to encrypt draft")
 	}
 
-	attachments := message.Attachments
-	message.Attachments = nil
+	if ok, err := store.checkDraftTotalSize(message, attachments); err != nil {
+		return nil, nil, err
+	} else if !ok {
+		return nil, nil, errors.New("message is too large")
+	}
 
 	draftAction := store.getDraftAction(message)
 	draft, err := store.client().CreateDraft(message, parentID, draftAction)
@@ -60,29 +62,114 @@ func (store *Store) CreateDraft(
 		return nil, nil, errors.Wrap(err, "failed to create draft")
 	}
 
+	// Do poll only when call to API succeeded.
+	defer store.eventLoop.pollNow()
+
+	createdAttachments := []*pmapi.Attachment{}
+	for _, att := range attachments {
+		att.attachment.MessageID = draft.ID
+
+		createdAttachment, err := store.client().CreateAttachment(att.attachment, att.encReader, att.sigReader)
+		if err != nil {
+			return nil, nil, errors.Wrap(err, "failed to create attachment")
+		}
+		createdAttachments = append(createdAttachments, createdAttachment)
+	}
+
+	return draft, createdAttachments, nil
+}
+
+type draftAttachment struct {
+	attachment *pmapi.Attachment
+	reader     io.Reader
+	sigReader  io.Reader
+	encReader  io.Reader
+}
+
+func (store *Store) prepareDraftAttachments(
+	message *pmapi.Message,
+	attachmentReaders []io.Reader,
+	attachedPublicKey,
+	attachedPublicKeyName string) []*draftAttachment {
+	attachments := []*draftAttachment{}
+	for idx, attachment := range message.Attachments {
+		attachments = append(attachments, &draftAttachment{
+			attachment: attachment,
+			reader:     attachmentReaders[idx],
+		})
+	}
+
+	message.Attachments = nil
+
 	if attachedPublicKey != "" {
-		attachmentReaders = append(attachmentReaders, strings.NewReader(attachedPublicKey))
 		publicKeyAttachment := &pmapi.Attachment{
 			Name:     attachedPublicKeyName + ".asc",
 			MIMEType: "application/pgp-keys",
 			Header:   textproto.MIMEHeader{},
 		}
-		attachments = append(attachments, publicKeyAttachment)
+		attachments = append(attachments, &draftAttachment{
+			attachment: publicKeyAttachment,
+			reader:     strings.NewReader(attachedPublicKey),
+		})
 	}
 
-	for idx, attachment := range attachments {
-		attachment.MessageID = draft.ID
-		attachmentBody, _ := ioutil.ReadAll(attachmentReaders[idx])
+	return attachments
+}
 
-		createdAttachment, err := store.createAttachment(kr, attachment, attachmentBody)
+func encryptDraft(kr *crypto.KeyRing, message *pmapi.Message, attachments []*draftAttachment) error {
+	// Since this is a draft, we don't need to sign it.
+	if err := message.Encrypt(kr, nil); err != nil {
+		return errors.Wrap(err, "failed to encrypt message")
+	}
+
+	for _, att := range attachments {
+		attachment := att.attachment
+		attachmentBody, err := ioutil.ReadAll(att.reader)
 		if err != nil {
-			return nil, nil, errors.Wrap(err, "failed to create attachment for draft")
+			return errors.Wrap(err, "failed to read attachment")
 		}
 
-		attachments[idx] = createdAttachment
+		r := bytes.NewReader(attachmentBody)
+		sigReader, err := attachment.DetachedSign(kr, r)
+		if err != nil {
+			return errors.Wrap(err, "failed to sign attachment")
+		}
+		att.sigReader = sigReader
+
+		r = bytes.NewReader(attachmentBody)
+		encReader, err := attachment.Encrypt(kr, r)
+		if err != nil {
+			return errors.Wrap(err, "failed to encrypt attachment")
+		}
+		att.encReader = encReader
+
+		att.reader = nil
+	}
+	return nil
 }
 
-	return draft, attachments, nil
+func (store *Store) checkDraftTotalSize(message *pmapi.Message, attachments []*draftAttachment) (bool, error) {
+	maxUpload, err := store.GetMaxUpload()
+	if err != nil {
+		return false, err
+	}
+
+	msgSize := message.Size
+	if msgSize == 0 {
+		msgSize = int64(len(message.Body))
+	}
+
+	var attSize int64
+	for _, att := range attachments {
+		b, err := ioutil.ReadAll(att.encReader)
+		if err != nil {
+			return false, err
+		}
+		attSize += int64(len(b))
+		att.encReader = bytes.NewBuffer(b)
+	}
+
+	return msgSize+attSize <= maxUpload, nil
 }
 
 func (store *Store) getDraftAction(message *pmapi.Message) int {
@@ -93,27 +180,6 @@ func (store *Store) getDraftAction(message *pmapi.Message) int {
 	return pmapi.DraftActionReply
 }
 
-func (store *Store) createAttachment(kr *crypto.KeyRing, attachment *pmapi.Attachment, attachmentBody []byte) (*pmapi.Attachment, error) {
-	r := bytes.NewReader(attachmentBody)
-	sigReader, err := attachment.DetachedSign(kr, r)
-	if err != nil {
-		return nil, errors.Wrap(err, "failed to sign attachment")
-	}
-
-	r = bytes.NewReader(attachmentBody)
-	encReader, err := attachment.Encrypt(kr, r)
-	if err != nil {
-		return nil, errors.Wrap(err, "failed to encrypt attachment")
-	}
-
-	createdAttachment, err := store.client().CreateAttachment(attachment, encReader, sigReader)
-	if err != nil {
-		return nil, errors.Wrap(err, "failed to create attachment")
-	}
-
-	return createdAttachment, nil
-}
-
 // SendMessage sends the message.
 func (store *Store) SendMessage(messageID string, req *pmapi.SendMessageReq) error {
 	defer store.eventLoop.pollNow()

internal/store/user_message_test.go

@@ -18,7 +18,9 @@
 package store
 
 import (
+	"io"
 	"net/mail"
+	"strings"
 	"testing"
 
 	"github.com/ProtonMail/proton-bridge/pkg/pmapi"
@@ -154,3 +156,47 @@ func checkAllMessageIDs(t *testing.T, m *mocksForStore, wantIDs []string) {
 	require.Nil(t, allErr)
 	require.Equal(t, wantIDs, allIds)
 }
+
+func TestCreateDraftCheckMessageSize(t *testing.T) {
+	m, clear := initMocks(t)
+	defer clear()
+
+	m.newStoreNoEvents(false)
+	m.client.EXPECT().CurrentUser().Return(&pmapi.User{
+		MaxUpload: 100, // Decrypted message 5 chars, encrypted 500+.
+	}, nil)
+
+	// Even small body is bloated to at least about 500 chars of basic pgp message.
+	message := &pmapi.Message{
+		Body: strings.Repeat("a", 5),
+	}
+	attachmentReaders := []io.Reader{}
+	_, _, err := m.store.CreateDraft(testPrivateKeyRing, message, attachmentReaders, "", "", "")
+
+	require.EqualError(t, err, "message is too large")
+}
+
+func TestCreateDraftCheckMessageWithAttachmentSize(t *testing.T) {
+	m, clear := initMocks(t)
+	defer clear()
+
+	m.newStoreNoEvents(false)
+	m.client.EXPECT().CurrentUser().Return(&pmapi.User{
+		MaxUpload: 800, // Decrypted message 5 chars + 5 chars of attachment, encrypted 500+ + 300+.
+	}, nil)
+
+	// Even small body is bloated to at least about 500 chars of basic pgp message.
+	message := &pmapi.Message{
+		Body: strings.Repeat("a", 5),
+		Attachments: []*pmapi.Attachment{
+			{Name: "name"},
+		},
+	}
+	// Even small attachment is bloated to about 300 chars of encrypted text.
+	attachmentReaders := []io.Reader{
+		strings.NewReader(strings.Repeat("b", 5)),
+	}
+	_, _, err := m.store.CreateDraft(testPrivateKeyRing, message, attachmentReaders, "", "", "")
+
+	require.EqualError(t, err, "message is too large")
+}