From e464e11ab924e8009dda37ebea60082e7671cd59 Mon Sep 17 00:00:00 2001 From: Leander Beernaert Date: Tue, 24 Jan 2023 13:41:03 +0100 Subject: [PATCH] GODT-2224: Refactor bridge sync to use less memory Updates go-proton-api and Gluon to includes memory reduction changes and modify the sync process to take into account how much memory is used during the sync stage. The sync process now has an extra stage which first download the message metada to ensure that we only download up to `syncMaxDownloadRequesMem` messages or 250 messages total. This allows for scaling the download request automatically to accommodate many small or few very large messages. The IDs are then sent to a download go-routine which downloads the message and its attachments. The result is then forwarded to another go-routine which builds the actual message. This stage tries to ensure that we don't use more than `syncMaxMessageBuildingMem` to build these messages. Finally the result is sent to a last go-routine which applies the changes to Gluon and waits for them to be completed. The new process is currently limited to 2GB. Dynamic scaling will be implemented in a follow up. For systems with less than 2GB of memory we limit the values to a set of values that is known to work. --- COPYING_NOTES.md | 2 + go.mod | 6 +- go.sum | 12 +- internal/app/app.go | 26 ++- internal/app/migration.go | 8 - internal/app/migration_test.go | 2 - internal/bridge/api.go | 2 - internal/bridge/api_default.go | 3 +- internal/bridge/api_qa.go | 3 +- internal/bridge/bridge.go | 2 +- internal/bridge/imap.go | 1 - internal/bridge/user.go | 1 - internal/user/events.go | 9 +- internal/user/imap.go | 4 +- internal/user/sync.go | 360 ++++++++++++++++++++++++++----- internal/user/sync_build.go | 8 +- internal/user/sync_flusher.go | 63 ------ internal/user/types.go | 5 + internal/user/user.go | 3 - internal/user/user_test.go | 2 +- internal/vault/settings.go | 24 --- internal/vault/settings_test.go | 9 - internal/vault/types_settings.go | 21 -- pkg/message/build.go | 208 ++++++++++-------- 24 files changed, 481 insertions(+), 303 deletions(-) delete mode 100644 internal/user/sync_flusher.go diff --git a/COPYING_NOTES.md b/COPYING_NOTES.md index 2269bf30..ce628759 100644 --- a/COPYING_NOTES.md +++ b/COPYING_NOTES.md @@ -53,6 +53,7 @@ Proton Mail Bridge includes the following 3rd party software: * [html2text](https://github.com/jaytaylor/html2text) available under [license](https://github.com/jaytaylor/html2text/blob/master/LICENSE) * [go-keychain](https://github.com/keybase/go-keychain) available under [license](https://github.com/keybase/go-keychain/blob/master/LICENSE) * [dns](https://github.com/miekg/dns) available under [license](https://github.com/miekg/dns/blob/master/LICENSE) +* [memory](https://github.com/pbnjay/memory) available under [license](https://github.com/pbnjay/memory/blob/master/LICENSE) * [errors](https://github.com/pkg/errors) available under [license](https://github.com/pkg/errors/blob/master/LICENSE) * [profile](https://github.com/pkg/profile) available under [license](https://github.com/pkg/profile/blob/master/LICENSE) * [logrus](https://github.com/sirupsen/logrus) available under [license](https://github.com/sirupsen/logrus/blob/master/LICENSE) @@ -114,6 +115,7 @@ Proton Mail Bridge includes the following 3rd party software: * [reflect2](https://github.com/modern-go/reflect2) available under [license](https://github.com/modern-go/reflect2/blob/master/LICENSE) * [tablewriter](https://github.com/olekukonko/tablewriter) available under [license](https://github.com/olekukonko/tablewriter/blob/master/LICENSE) * [go-toml](https://github.com/pelletier/go-toml/v2) available under [license](https://github.com/pelletier/go-toml/v2/blob/master/LICENSE) +* [lz4](https://github.com/pierrec/lz4/v4) available under [license](https://github.com/pierrec/lz4/v4/blob/master/LICENSE) * [go-difflib](https://github.com/pmezard/go-difflib) available under [license](https://github.com/pmezard/go-difflib/blob/master/LICENSE) * [procfs](https://github.com/prometheus/procfs) available under [license](https://github.com/prometheus/procfs/blob/master/LICENSE) * [uniseg](https://github.com/rivo/uniseg) available under [license](https://github.com/rivo/uniseg/blob/master/LICENSE) diff --git a/go.mod b/go.mod index 00cc0a86..06393d87 100644 --- a/go.mod +++ b/go.mod @@ -5,9 +5,9 @@ go 1.18 require ( github.com/0xAX/notificator v0.0.0-20220220101646-ee9b8921e557 github.com/Masterminds/semver/v3 v3.1.1 - github.com/ProtonMail/gluon v0.14.2-0.20230127085305-bc2d818d9d13 + github.com/ProtonMail/gluon v0.14.2-0.20230130104154-2c64e59b8f54 github.com/ProtonMail/go-autostart v0.0.0-20210130080809-00ed301c8e9a - github.com/ProtonMail/go-proton-api v0.3.1-0.20230126112849-3c1ac277855e + github.com/ProtonMail/go-proton-api v0.3.1-0.20230130093944-dd1190680368 github.com/ProtonMail/go-rfc5322 v0.11.0 github.com/ProtonMail/gopenpgp/v2 v2.4.10 github.com/PuerkitoBio/goquery v1.8.0 @@ -35,6 +35,7 @@ require ( github.com/jaytaylor/html2text v0.0.0-20211105163654-bc68cce691ba github.com/keybase/go-keychain v0.0.0 github.com/miekg/dns v1.1.50 + github.com/pbnjay/memory v0.0.0-20210728143218-7b4eea64cf58 github.com/pkg/errors v0.9.1 github.com/pkg/profile v1.6.0 github.com/sirupsen/logrus v1.9.0 @@ -99,6 +100,7 @@ require ( github.com/modern-go/reflect2 v1.0.2 // indirect github.com/olekukonko/tablewriter v0.0.5 // indirect github.com/pelletier/go-toml/v2 v2.0.5 // indirect + github.com/pierrec/lz4/v4 v4.1.17 // indirect github.com/pmezard/go-difflib v1.0.0 // indirect github.com/prometheus/procfs v0.8.0 // indirect github.com/rivo/uniseg v0.4.2 // indirect diff --git a/go.sum b/go.sum index 6456748f..7ed408bd 100644 --- a/go.sum +++ b/go.sum @@ -28,8 +28,8 @@ github.com/ProtonMail/bcrypt v0.0.0-20211005172633-e235017c1baf h1:yc9daCCYUefEs github.com/ProtonMail/bcrypt v0.0.0-20211005172633-e235017c1baf/go.mod h1:o0ESU9p83twszAU8LBeJKFAAMX14tISa0yk4Oo5TOqo= github.com/ProtonMail/docker-credential-helpers v1.1.0 h1:+kvUIpwWcbtP3WFv5sSvkFn/XLzSqPOB5AAthuk9xPk= github.com/ProtonMail/docker-credential-helpers v1.1.0/go.mod h1:mK0aBveCxhnQ756AmaTfXMZDeULvheYVhF/MWMErN5g= -github.com/ProtonMail/gluon v0.14.2-0.20230127085305-bc2d818d9d13 h1:rljNZVgfq/F1LLyJ4NmCfEzWayC/rk+l9QgJjtQTLKI= -github.com/ProtonMail/gluon v0.14.2-0.20230127085305-bc2d818d9d13/go.mod h1:z2AxLIiBCT1K+0OBHyaDI7AEaO5qI6/BEC2TE42vs4Q= +github.com/ProtonMail/gluon v0.14.2-0.20230130104154-2c64e59b8f54 h1:uUg8CDiYTMlbvGijzoN0fb72vwDJD7hMjgNTbmAHxRc= +github.com/ProtonMail/gluon v0.14.2-0.20230130104154-2c64e59b8f54/go.mod h1:HYHr7hG7LPWI1S50M8NfHRb1kYi5B+Yu4/N/H+y+JUY= github.com/ProtonMail/go-autostart v0.0.0-20210130080809-00ed301c8e9a h1:D+aZah+k14Gn6kmL7eKxoo/4Dr/lK3ChBcwce2+SQP4= github.com/ProtonMail/go-autostart v0.0.0-20210130080809-00ed301c8e9a/go.mod h1:oTGdE7/DlWIr23G0IKW3OXK9wZ5Hw1GGiaJFccTvZi4= github.com/ProtonMail/go-crypto v0.0.0-20210428141323-04723f9f07d7/go.mod h1:z4/9nQmJSSwwds7ejkxaJwO37dru3geImFUdJlaLzQo= @@ -41,8 +41,8 @@ github.com/ProtonMail/go-message v0.0.0-20210611055058-fabeff2ec753/go.mod h1:NB github.com/ProtonMail/go-mime v0.0.0-20220302105931-303f85f7fe0f/go.mod h1:NYt+V3/4rEeDuaev/zw1zCq8uqVEuPHzDPo3OZrlGJ4= github.com/ProtonMail/go-mime v0.0.0-20220429130430-2192574d760f h1:4IWzKjHzZxdrW9k4zl/qCwenOVHDbVDADPPHFLjs0Oc= github.com/ProtonMail/go-mime v0.0.0-20220429130430-2192574d760f/go.mod h1:qRZgbeASl2a9OwmsV85aWwRqic0NHPh+9ewGAzb4cgM= -github.com/ProtonMail/go-proton-api v0.3.1-0.20230126112849-3c1ac277855e h1:UkfLQc44UvknNCLoBEZb1qg7zfVWVLMvCE/LtdVEcAw= -github.com/ProtonMail/go-proton-api v0.3.1-0.20230126112849-3c1ac277855e/go.mod h1:JUo5IQG0hNuPRuDpOUsCOvtee6UjTEHHF1QN2i8RSos= +github.com/ProtonMail/go-proton-api v0.3.1-0.20230130093944-dd1190680368 h1:XWPaCK8ctgpvI+ZAYcsFiYWbuvkv3T4+WZq0o6Gnt2s= +github.com/ProtonMail/go-proton-api v0.3.1-0.20230130093944-dd1190680368/go.mod h1:JUo5IQG0hNuPRuDpOUsCOvtee6UjTEHHF1QN2i8RSos= github.com/ProtonMail/go-rfc5322 v0.11.0 h1:o5Obrm4DpmQEffvgsVqG6S4BKwC1Wat+hYwjIp2YcCY= github.com/ProtonMail/go-rfc5322 v0.11.0/go.mod h1:6oOKr0jXvpoE6pwTx/HukigQpX2J9WUf6h0auplrFTw= github.com/ProtonMail/go-srp v0.0.5 h1:xhUioxZgDbCnpo9JehyFhwwsn9JLWkUGfB0oiKXgiGg= @@ -314,9 +314,13 @@ github.com/oklog/ulid v1.3.1/go.mod h1:CirwcVhetQ6Lv90oh/F+FBtV6XMibvdAFo93nm5qn github.com/olekukonko/tablewriter v0.0.5 h1:P2Ga83D34wi1o9J6Wh1mRuqd4mF/x/lgBS7N7AbDhec= github.com/olekukonko/tablewriter v0.0.5/go.mod h1:hPp6KlRPjbx+hW8ykQs1w3UBbZlj6HuIJcUGPhkA7kY= github.com/pascaldekloe/goe v0.0.0-20180627143212-57f6aae5913c/go.mod h1:lzWF7FIEvWOWxwDKqyGYQf6ZUaNfKdP144TG7ZOy1lc= +github.com/pbnjay/memory v0.0.0-20210728143218-7b4eea64cf58 h1:onHthvaw9LFnH4t2DcNVpwGmV9E1BkGknEliJkfwQj0= +github.com/pbnjay/memory v0.0.0-20210728143218-7b4eea64cf58/go.mod h1:DXv8WO4yhMYhSNPKjeNKa5WY9YCIEBRbNzFFPJbWO6Y= github.com/pelletier/go-toml v1.2.0/go.mod h1:5z9KED0ma1S8pY6P1sdut58dfprrGBbd/94hg7ilaic= github.com/pelletier/go-toml/v2 v2.0.5 h1:ipoSadvV8oGUjnUbMub59IDPPwfxF694nG/jwbMiyQg= github.com/pelletier/go-toml/v2 v2.0.5/go.mod h1:OMHamSCAODeSsVrwwvcJOaoN0LIUIaFVNZzmWyNfXas= +github.com/pierrec/lz4/v4 v4.1.17 h1:kV4Ip+/hUBC+8T6+2EgburRtkE9ef4nbY3f4dFhGjMc= +github.com/pierrec/lz4/v4 v4.1.17/go.mod h1:gZWDp/Ze/IJXGXf23ltt2EXimqmTUXEy0GFuRQyBid4= github.com/pingcap/errors v0.11.4 h1:lFuQV/oaUMGcD2tqt+01ROSmJs75VG1ToEOkZIZ4nE4= github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA= github.com/pkg/errors v0.8.0/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= diff --git a/internal/app/app.go b/internal/app/app.go index 92289afa..db93442f 100644 --- a/internal/app/app.go +++ b/internal/app/app.go @@ -205,7 +205,7 @@ func run(c *cli.Context) error { //nolint:funlen // Ensure we are the only instance running. return withSingleInstance(locations, version, func() error { // Unlock the encrypted vault. - return WithVault(locations, func(vault *vault.Vault, insecure, corrupt bool) error { + return WithVault(locations, func(v *vault.Vault, insecure, corrupt bool) error { // Report insecure vault. if insecure { _ = reporter.ReportMessageWithContext("Vault is insecure", map[string]interface{}{}) @@ -216,27 +216,39 @@ func run(c *cli.Context) error { //nolint:funlen _ = reporter.ReportMessageWithContext("Vault is corrupt", map[string]interface{}{}) } - if !vault.Migrated() { + // Force re-sync if last version <= 3.0.12 due to chances in the gluon cache format. + if lastVersion := v.GetLastVersion(); lastVersion != nil { + versionWithLZ4Cache := semver.MustParse("3.0.13") + if lastVersion.LessThan(versionWithLZ4Cache) { + if err := v.ForUser(1, func(user *vault.User) error { + return user.ClearSyncStatus() + }); err != nil { + logrus.WithError(err).Error("Failed to force resync on user") + } + } + } + + if !v.Migrated() { // Migrate old settings into the vault. - if err := migrateOldSettings(vault); err != nil { + if err := migrateOldSettings(v); err != nil { logrus.WithError(err).Error("Failed to migrate old settings") } // Migrate old accounts into the vault. - if err := migrateOldAccounts(locations, vault); err != nil { + if err := migrateOldAccounts(locations, v); err != nil { logrus.WithError(err).Error("Failed to migrate old accounts") } // The vault has been migrated. - if err := vault.SetMigrated(); err != nil { + if err := v.SetMigrated(); err != nil { logrus.WithError(err).Error("Failed to mark vault as migrated") } } // Load the cookies from the vault. - return withCookieJar(vault, func(cookieJar http.CookieJar) error { + return withCookieJar(v, func(cookieJar http.CookieJar) error { // Create a new bridge instance. - return withBridge(c, exe, locations, version, identifier, crashHandler, reporter, vault, cookieJar, func(b *bridge.Bridge, eventCh <-chan events.Event) error { + return withBridge(c, exe, locations, version, identifier, crashHandler, reporter, v, cookieJar, func(b *bridge.Bridge, eventCh <-chan events.Event) error { if insecure { logrus.Warn("The vault key could not be retrieved; the vault will not be encrypted") b.PushError(bridge.ErrVaultInsecure) diff --git a/internal/app/migration.go b/internal/app/migration.go index 1124de9b..7f187b1b 100644 --- a/internal/app/migration.go +++ b/internal/app/migration.go @@ -265,14 +265,6 @@ func migratePrefsToVault(vault *vault.Vault, b []byte) error { errs = multierror.Append(errs, fmt.Errorf("failed to migrate show all mail: %w", err)) } - if err := vault.SetSyncWorkers(prefs.FetchWorkers); err != nil { - errs = multierror.Append(errs, fmt.Errorf("failed to migrate sync workers: %w", err)) - } - - if err := vault.SetSyncAttPool(prefs.AttachmentWorkers); err != nil { - errs = multierror.Append(errs, fmt.Errorf("failed to migrate sync attachment pool: %w", err)) - } - if err := vault.SetCookies([]byte(prefs.Cookies)); err != nil { errs = multierror.Append(errs, fmt.Errorf("failed to migrate cookies: %w", err)) } diff --git a/internal/app/migration_test.go b/internal/app/migration_test.go index a9e128e9..fd0b41e1 100644 --- a/internal/app/migration_test.go +++ b/internal/app/migration_test.go @@ -68,8 +68,6 @@ func TestMigratePrefsToVault(t *testing.T) { require.True(t, vault.GetAutostart()) // Check that the other app settings have been migrated. - require.Equal(t, 16, vault.SyncWorkers()) - require.Equal(t, 16, vault.SyncAttPool()) require.False(t, vault.GetProxyAllowed()) require.False(t, vault.GetShowAllMail()) diff --git a/internal/bridge/api.go b/internal/bridge/api.go index 1f0edaca..cbe2f2be 100644 --- a/internal/bridge/api.go +++ b/internal/bridge/api.go @@ -32,14 +32,12 @@ func defaultAPIOptions( version *semver.Version, cookieJar http.CookieJar, transport http.RoundTripper, - poolSize int, ) []proton.Option { return []proton.Option{ proton.WithHostURL(apiURL), proton.WithAppVersion(constants.AppVersion(version.Original())), proton.WithCookieJar(cookieJar), proton.WithTransport(transport), - proton.WithAttPoolSize(poolSize), proton.WithLogger(logrus.StandardLogger()), } } diff --git a/internal/bridge/api_default.go b/internal/bridge/api_default.go index a71433cc..6f800d77 100644 --- a/internal/bridge/api_default.go +++ b/internal/bridge/api_default.go @@ -32,7 +32,6 @@ func newAPIOptions( version *semver.Version, cookieJar http.CookieJar, transport http.RoundTripper, - poolSize int, ) []proton.Option { - return defaultAPIOptions(apiURL, version, cookieJar, transport, poolSize) + return defaultAPIOptions(apiURL, version, cookieJar, transport) } diff --git a/internal/bridge/api_qa.go b/internal/bridge/api_qa.go index cf182330..618ce42e 100644 --- a/internal/bridge/api_qa.go +++ b/internal/bridge/api_qa.go @@ -33,9 +33,8 @@ func newAPIOptions( version *semver.Version, cookieJar http.CookieJar, transport http.RoundTripper, - poolSize int, ) []proton.Option { - opt := defaultAPIOptions(apiURL, version, cookieJar, transport, poolSize) + opt := defaultAPIOptions(apiURL, version, cookieJar, transport) if host := os.Getenv("BRIDGE_API_HOST"); host != "" { opt = append(opt, proton.WithHostURL(host)) diff --git a/internal/bridge/bridge.go b/internal/bridge/bridge.go index 886c740b..1a6de766 100644 --- a/internal/bridge/bridge.go +++ b/internal/bridge/bridge.go @@ -145,7 +145,7 @@ func New( //nolint:funlen logSMTP bool, // whether to log SMTP activity ) (*Bridge, <-chan events.Event, error) { // api is the user's API manager. - api := proton.New(newAPIOptions(apiURL, curVersion, cookieJar, roundTripper, vault.SyncAttPool())...) + api := proton.New(newAPIOptions(apiURL, curVersion, cookieJar, roundTripper)...) // tasks holds all the bridge's background tasks. tasks := async.NewGroup(context.Background(), crashHandler) diff --git a/internal/bridge/imap.go b/internal/bridge/imap.go index 928f852f..b9fd540f 100644 --- a/internal/bridge/imap.go +++ b/internal/bridge/imap.go @@ -310,7 +310,6 @@ func (*storeBuilder) New(path, userID string, passphrase []byte) (store.Store, e return store.NewOnDiskStore( filepath.Join(path, userID), passphrase, - store.WithCompressor(new(store.GZipCompressor)), ) } diff --git a/internal/bridge/user.go b/internal/bridge/user.go index 25a697b7..fa86a313 100644 --- a/internal/bridge/user.go +++ b/internal/bridge/user.go @@ -461,7 +461,6 @@ func (bridge *Bridge) addUserWithVault( bridge.reporter, apiUser, bridge.crashHandler, - bridge.vault.SyncWorkers(), bridge.vault.GetShowAllMail(), ) if err != nil { diff --git a/internal/user/events.go b/internal/user/events.go index 1f75c043..5311c8bd 100644 --- a/internal/user/events.go +++ b/internal/user/events.go @@ -18,6 +18,7 @@ package user import ( + "bytes" "context" "errors" "fmt" @@ -502,7 +503,7 @@ func (user *User) handleMessageEvents(ctx context.Context, messageEvents []proto } func (user *User) handleCreateMessageEvent(ctx context.Context, event proton.MessageEvent) ([]imap.Update, error) { - full, err := user.client.GetFullMessage(ctx, event.Message.ID) + full, err := user.client.GetFullMessage(ctx, event.Message.ID, newProtonAPIScheduler(), proton.NewDefaultAttachmentAllocator()) if err != nil { // If the message is not found, it means that it has been deleted before we could fetch it. if apiErr := new(proton.APIError); errors.As(err, &apiErr) && apiErr.Status == http.StatusUnprocessableEntity { @@ -521,7 +522,7 @@ func (user *User) handleCreateMessageEvent(ctx context.Context, event proton.Mes var update imap.Update if err := withAddrKR(user.apiUser, user.apiAddrs[event.Message.AddressID], user.vault.KeyPass(), func(_, addrKR *crypto.KeyRing) error { - res := buildRFC822(user.apiLabels, full, addrKR) + res := buildRFC822(user.apiLabels, full, addrKR, new(bytes.Buffer)) if res.err != nil { user.log.WithError(err).Error("Failed to build RFC822 message") @@ -599,7 +600,7 @@ func (user *User) handleUpdateDraftEvent(ctx context.Context, event proton.Messa "subject": logging.Sensitive(event.Message.Subject), }).Info("Handling draft updated event") - full, err := user.client.GetFullMessage(ctx, event.Message.ID) + full, err := user.client.GetFullMessage(ctx, event.Message.ID, newProtonAPIScheduler(), proton.NewDefaultAttachmentAllocator()) if err != nil { // If the message is not found, it means that it has been deleted before we could fetch it. if apiErr := new(proton.APIError); errors.As(err, &apiErr) && apiErr.Status == http.StatusUnprocessableEntity { @@ -613,7 +614,7 @@ func (user *User) handleUpdateDraftEvent(ctx context.Context, event proton.Messa var update imap.Update if err := withAddrKR(user.apiUser, user.apiAddrs[event.Message.AddressID], user.vault.KeyPass(), func(_, addrKR *crypto.KeyRing) error { - res := buildRFC822(user.apiLabels, full, addrKR) + res := buildRFC822(user.apiLabels, full, addrKR, new(bytes.Buffer)) if res.err != nil { logrus.WithError(err).Error("Failed to build RFC822 message") diff --git a/internal/user/imap.go b/internal/user/imap.go index ab17f2f8..9de879e2 100644 --- a/internal/user/imap.go +++ b/internal/user/imap.go @@ -336,7 +336,7 @@ func (conn *imapConnector) CreateMessage( } func (conn *imapConnector) GetMessageLiteral(ctx context.Context, id imap.MessageID) ([]byte, error) { - msg, err := conn.client.GetFullMessage(ctx, string(id)) + msg, err := conn.client.GetFullMessage(ctx, string(id), newProtonAPIScheduler(), proton.NewDefaultAttachmentAllocator()) if err != nil { return nil, err } @@ -538,7 +538,7 @@ func (conn *imapConnector) importMessage( var err error - if full, err = conn.client.GetFullMessage(ctx, messageID); err != nil { + if full, err = conn.client.GetFullMessage(ctx, messageID, newProtonAPIScheduler(), proton.NewDefaultAttachmentAllocator()); err != nil { return fmt.Errorf("failed to fetch message: %w", err) } diff --git a/internal/user/sync.go b/internal/user/sync.go index 6e4faeea..0735ed83 100644 --- a/internal/user/sync.go +++ b/internal/user/sync.go @@ -18,6 +18,7 @@ package user import ( + "bytes" "context" "fmt" "runtime" @@ -25,6 +26,7 @@ import ( "time" "github.com/ProtonMail/gluon/imap" + "github.com/ProtonMail/gluon/logging" "github.com/ProtonMail/gluon/queue" "github.com/ProtonMail/gluon/reporter" "github.com/ProtonMail/go-proton-api" @@ -35,16 +37,12 @@ import ( "github.com/bradenaw/juniper/parallel" "github.com/bradenaw/juniper/xslices" "github.com/google/uuid" + "github.com/pbnjay/memory" "github.com/sirupsen/logrus" "golang.org/x/exp/maps" "golang.org/x/exp/slices" ) -const ( - maxUpdateSize = 1 << 27 // 128 MiB - maxBatchSize = 1 << 8 // 256 -) - // doSync begins syncing the users data. // It first ensures the latest event ID is known; if not, it fetches it. // It sends a SyncStarted event and then either SyncFinished or SyncFailed @@ -143,7 +141,6 @@ func (user *User) sync(ctx context.Context) error { addrKRs, user.updateCh, user.eventCh, - user.syncWorkers, ); err != nil { return fmt.Errorf("failed to sync messages: %w", err) } @@ -212,7 +209,15 @@ func syncLabels(ctx context.Context, apiLabels map[string]proton.Label, updateCh return nil } -// nolint:funlen +const Kilobyte = uint64(1024) +const Megabyte = 1024 * Kilobyte +const Gigabyte = 1024 * Megabyte + +func toMB(v uint64) float64 { + return float64(v) / float64(Megabyte) +} + +// nolint:funlen,gocyclo func syncMessages( ctx context.Context, userID string, @@ -224,7 +229,6 @@ func syncMessages( addrKRs map[string]*crypto.KeyRing, updateCh map[string]*queue.QueuedChannel[imap.Update], eventCh *queue.QueuedChannel[events.Event], - syncWorkers int, ) error { ctx, cancel := context.WithCancel(ctx) defer cancel() @@ -235,78 +239,319 @@ func syncMessages( logrus.WithFields(logrus.Fields{ "messages": len(messageIDs), - "workers": syncWorkers, "numCPU": runtime.NumCPU(), }).Info("Starting message sync") // Create the flushers, one per update channel. - flushers := make(map[string]*flusher, len(updateCh)) - - for addrID, updateCh := range updateCh { - flushers[addrID] = newFlusher(updateCh, maxUpdateSize) - } // Create a reporter to report sync progress updates. syncReporter := newSyncReporter(userID, eventCh, len(messageIDs), time.Second) defer syncReporter.done() - type flushUpdate struct { - messageID string - pushedUpdates []imap.Update - batchLen int + // Expected mem usage for this whole process should be the sum of MaxMessageBuildingMem and MaxDownloadRequestMem + // times x due to pipeline and all additional memory used by network requests and compression+io. + + // There's no point in using more than 128MB of download data per stage, after that we reach a point of diminishing + // returns as we can't keep the pipeline fed fast enough. + const MaxDownloadRequestMem = 128 * Megabyte + + // Any lower than this and we may fail to download messages. + const MinDownloadRequestMem = 40 * Megabyte + + // This value can be increased to your hearts content. The more system memory the user has, the more messages + // we can build in parallel. + const MaxMessageBuildingMem = 128 * Megabyte + const MinMessageBuildingMem = 64 * Megabyte + + // Maximum recommend value for parallel downloads by the API team. + const maxParallelDownloads = 20 + + totalMemory := memory.TotalMemory() + logrus.Debugf("Total System Memory: %v", toMB(totalMemory)) + + syncMaxDownloadRequestMem := MaxDownloadRequestMem + syncMaxMessageBuildingMem := MaxMessageBuildingMem + + // If less than 2GB available try and limit max memory to 512 MB + if totalMemory < 2*Gigabyte { + if totalMemory < 800*Megabyte { + logrus.Warnf("System has less than 800MB of memory, you may experience issues sycing large mailboxes") + } + syncMaxDownloadRequestMem = MinDownloadRequestMem + syncMaxMessageBuildingMem = MinMessageBuildingMem + } else { + // Increasing the max download capacity has very little effect on sync speed. We could increase the download + // memory but the user would see less sync notifications. A smaller value here leads to more frequent + // updates. Additionally, most of ot sync time is spent in the message building. + syncMaxDownloadRequestMem = MaxDownloadRequestMem + // Currently limited so that if a user has multiple accounts active it also doesn't cause excessive memory usage. + syncMaxMessageBuildingMem = MaxMessageBuildingMem } + logrus.Debugf("Max memory usage for sync Download=%vMB Building=%vMB Predicted Max Total=%vMB", + toMB(syncMaxDownloadRequestMem), + toMB(syncMaxMessageBuildingMem), + toMB((syncMaxMessageBuildingMem*4)+(syncMaxDownloadRequestMem*4)), + ) + + type flushUpdate struct { + messageID string + err error + batchLen int + } + + type downloadRequest struct { + ids []string + expectedSize uint64 + err error + } + + type downloadedMessageBatch struct { + batch []proton.FullMessage + } + + type builtMessageBatch struct { + batch []*buildRes + } + + downloadCh := make(chan downloadRequest) + + buildCh := make(chan downloadedMessageBatch) + // The higher this value, the longer we can continue our download iteration before being blocked on channel writes // to the update flushing goroutine. - flushCh := make(chan []*buildRes, 2) + flushCh := make(chan builtMessageBatch) - // Allow up to 4 batched wait requests. - flushUpdateCh := make(chan flushUpdate, 4) + flushUpdateCh := make(chan flushUpdate) - errorCh := make(chan error, syncWorkers) + errorCh := make(chan error, maxParallelDownloads+2) + + // Go routine in charge of downloading message metadata + logging.GoAnnotated(ctx, func(ctx context.Context) { + defer close(downloadCh) + const MetadataDataPageSize = 150 + + var downloadReq downloadRequest + downloadReq.ids = make([]string, 0, MetadataDataPageSize) + + metadataChunks := xslices.Chunk(messageIDs, MetadataDataPageSize) + for i, metadataChunk := range metadataChunks { + logrus.Debugf("Metadata Request (%v of %v), previous: %v", i, len(metadataChunks), len(downloadReq.ids)) + metadata, err := client.GetMessageMetadataPage(ctx, 0, len(metadataChunk), proton.MessageFilter{ID: metadataChunk}) + if err != nil { + downloadReq.err = err + downloadCh <- downloadReq + return + } + + if ctx.Err() != nil { + downloadReq.err = err + downloadCh <- downloadReq + return + } + + // Build look up table so that messages are processed in the same order. + metadataMap := make(map[string]int, len(metadata)) + for i, v := range metadata { + metadataMap[v.ID] = i + } + + for i, id := range metadataChunk { + m := &metadata[metadataMap[id]] + nextSize := downloadReq.expectedSize + uint64(m.Size) + if nextSize >= syncMaxDownloadRequestMem || len(downloadReq.ids) >= 256 { + logrus.Debugf("Download Request Sent at %v of %v", i, len(metadata)) + select { + case downloadCh <- downloadReq: + + case <-ctx.Done(): + return + } + downloadReq.expectedSize = 0 + downloadReq.ids = make([]string, 0, MetadataDataPageSize) + nextSize = uint64(m.Size) + } + downloadReq.ids = append(downloadReq.ids, id) + downloadReq.expectedSize = nextSize + } + } + + if len(downloadReq.ids) != 0 { + logrus.Debugf("Sending remaining download request") + select { + case downloadCh <- downloadReq: + + case <-ctx.Done(): + return + } + } + }, logging.Labels{"sync-stage": "meta-data"}) // Goroutine in charge of downloading and building messages in maxBatchSize batches. - go func() { - defer close(flushCh) + logging.GoAnnotated(ctx, func(ctx context.Context) { + defer close(buildCh) defer close(errorCh) + defer func() { + logrus.Debugf("sync downloader exit") + }() + + for request := range downloadCh { + logrus.Debugf("Download request: %v MB:%v", len(request.ids), toMB(request.expectedSize)) + if request.err != nil { + errorCh <- request.err + return + } - for _, batch := range xslices.Chunk(messageIDs, maxBatchSize) { if ctx.Err() != nil { errorCh <- ctx.Err() return } - result, err := parallel.MapContext(ctx, syncWorkers, batch, func(ctx context.Context, id string) (*buildRes, error) { - msg, err := client.GetFullMessage(ctx, id) + result, err := parallel.MapContext(ctx, maxParallelDownloads, request.ids, func(ctx context.Context, id string) (proton.FullMessage, error) { + var result proton.FullMessage + + msg, err := client.GetMessage(ctx, id) if err != nil { - return nil, err + return proton.FullMessage{}, err } - if ctx.Err() != nil { - return nil, ctx.Err() + var attachmentSize int64 + for _, a := range msg.Attachments { + attachmentSize += a.Size } - return buildRFC822(apiLabels, msg, addrKRs[msg.AddressID]), nil + // allocate attachment data. + result.AttData = make([][]byte, len(msg.Attachments)) + + for i, a := range msg.Attachments { + var buffer bytes.Buffer + buffer.Grow(int(a.Size)) + if err := client.GetAttachmentInto(ctx, a.ID, &buffer); err != nil { + return proton.FullMessage{}, err + } + + result.AttData[i] = buffer.Bytes() + } + + result.Message = msg + + return result, nil }) if err != nil { errorCh <- err return } + select { + case buildCh <- downloadedMessageBatch{ + batch: result, + }: + + case <-ctx.Done(): + return + } + } + }, logging.Labels{"sync-stage": "download"}) + + // Goroutine which builds messages after they have been downloaded + logging.GoAnnotated(ctx, func(ctx context.Context) { + defer close(flushCh) + defer func() { + logrus.Debugf("sync builder exit") + }() + + maxMessagesInParallel := runtime.NumCPU() + + for buildBatch := range buildCh { if ctx.Err() != nil { - errorCh <- ctx.Err() return } - flushCh <- result + var expectedMemUsage uint64 + var chunks [][]proton.FullMessage + + { + var lastIndex int + var index int + for _, v := range buildBatch.batch { + var dataSize uint64 + for _, a := range v.Attachments { + dataSize += uint64(a.Size) + } + + // 2x increase for attachment due to extra memory needed for decrypting and writing + // in memory buffer. + dataSize *= 2 + dataSize += uint64(len(v.Body)) + + nextMemSize := expectedMemUsage + dataSize + if nextMemSize >= syncMaxMessageBuildingMem { + chunks = append(chunks, buildBatch.batch[lastIndex:index]) + lastIndex = index + expectedMemUsage = dataSize + } else { + expectedMemUsage = nextMemSize + } + + index++ + } + + if index < len(buildBatch.batch) { + chunks = append(chunks, buildBatch.batch[index:]) + } else if index == len(buildBatch.batch) && len(chunks) == 0 { + chunks = [][]proton.FullMessage{buildBatch.batch} + } + } + + for index, chunk := range chunks { + result, err := parallel.MapContext(ctx, maxMessagesInParallel, chunk, func(ctx context.Context, msg proton.FullMessage) (*buildRes, error) { + return buildRFC822(apiLabels, msg, addrKRs[msg.AddressID], new(bytes.Buffer)), nil + }) + if err != nil { + return + } + + logrus.Debugf("Build request: %v of %v", index, len(chunks)) + + select { + case flushCh <- builtMessageBatch{result}: + + case <-ctx.Done(): + return + } + } } - }() + }, logging.Labels{"sync-stage": "builder"}) // Goroutine which converts the messages into updates and builds a waitable structure for progress tracking. - go func() { + logging.GoAnnotated(ctx, func(ctx context.Context) { defer close(flushUpdateCh) - for batch := range flushCh { - for _, res := range batch { + defer func() { + logrus.Debugf("sync flush exit") + }() + + type updateTargetInfo struct { + queueIndex int + ch *queue.QueuedChannel[imap.Update] + } + + pendingUpdates := make([][]*imap.MessageCreated, len(updateCh)) + addressToIndex := make(map[string]updateTargetInfo) + + { + i := 0 + for addrID, updateCh := range updateCh { + addressToIndex[addrID] = updateTargetInfo{ + ch: updateCh, + queueIndex: i, + } + i++ + } + } + + for downloadBatch := range flushCh { + logrus.Debugf("Flush batch: %v", len(downloadBatch.batch)) + for _, res := range downloadBatch.batch { if res.err != nil { if err := vault.AddFailedMessageID(res.messageID); err != nil { logrus.WithError(err).Error("Failed to add failed message ID") @@ -327,31 +572,38 @@ func syncMessages( } } - flushers[res.addressID].push(res.update) + targetInfo := addressToIndex[res.addressID] + pendingUpdates[targetInfo.queueIndex] = append(pendingUpdates[targetInfo.queueIndex], res.update) } - var pushedUpdates []imap.Update - for _, flusher := range flushers { - flusher.flush() - pushedUpdates = append(pushedUpdates, flusher.collectPushedUpdates()...) + for _, info := range addressToIndex { + up := imap.NewMessagesCreated(true, pendingUpdates[info.queueIndex]...) + info.ch.Enqueue(up) + + err, ok := up.WaitContext(ctx) + if ok && err != nil { + flushUpdateCh <- flushUpdate{ + err: fmt.Errorf("failed to apply sync update to gluon %v: %w", up.String(), err), + } + return + } + + pendingUpdates[info.queueIndex] = pendingUpdates[info.queueIndex][:0] } - flushUpdateCh <- flushUpdate{ - messageID: batch[0].messageID, - pushedUpdates: pushedUpdates, - batchLen: len(batch), + select { + case flushUpdateCh <- flushUpdate{ + messageID: downloadBatch.batch[0].messageID, + err: nil, + batchLen: len(downloadBatch.batch), + }: + case <-ctx.Done(): + return } } - }() + }, logging.Labels{"sync-stage": "flush"}) for flushUpdate := range flushUpdateCh { - for _, up := range flushUpdate.pushedUpdates { - err, ok := up.WaitContext(ctx) - if ok && err != nil { - return fmt.Errorf("failed to apply sync update to gluon %v: %w", up.String(), err) - } - } - if err := vault.SetLastMessageID(flushUpdate.messageID); err != nil { return fmt.Errorf("failed to set last synced message ID: %w", err) } diff --git a/internal/user/sync_build.go b/internal/user/sync_build.go index e3d87a5e..93309c35 100644 --- a/internal/user/sync_build.go +++ b/internal/user/sync_build.go @@ -48,16 +48,18 @@ func defaultJobOpts() message.JobOptions { } } -func buildRFC822(apiLabels map[string]proton.Label, full proton.FullMessage, addrKR *crypto.KeyRing) *buildRes { +func buildRFC822(apiLabels map[string]proton.Label, full proton.FullMessage, addrKR *crypto.KeyRing, buffer *bytes.Buffer) *buildRes { var ( update *imap.MessageCreated err error ) - if literal, buildErr := message.BuildRFC822(addrKR, full.Message, full.AttData, defaultJobOpts()); buildErr != nil { + buffer.Grow(full.Size) + + if buildErr := message.BuildRFC822Into(addrKR, full.Message, full.AttData, defaultJobOpts(), buffer); buildErr != nil { update = newMessageCreatedFailedUpdate(apiLabels, full.MessageMetadata, buildErr) err = buildErr - } else if created, parseErr := newMessageCreatedUpdate(apiLabels, full.MessageMetadata, literal); parseErr != nil { + } else if created, parseErr := newMessageCreatedUpdate(apiLabels, full.MessageMetadata, buffer.Bytes()); parseErr != nil { update = newMessageCreatedFailedUpdate(apiLabels, full.MessageMetadata, parseErr) err = parseErr } else { diff --git a/internal/user/sync_flusher.go b/internal/user/sync_flusher.go deleted file mode 100644 index 0c0dd2dd..00000000 --- a/internal/user/sync_flusher.go +++ /dev/null @@ -1,63 +0,0 @@ -// Copyright (c) 2023 Proton AG -// -// This file is part of Proton Mail Bridge. -// -// Proton Mail Bridge is free software: you can redistribute it and/or modify -// it under the terms of the GNU General Public License as published by -// the Free Software Foundation, either version 3 of the License, or -// (at your option) any later version. -// -// Proton Mail Bridge is distributed in the hope that it will be useful, -// but WITHOUT ANY WARRANTY; without even the implied warranty of -// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -// GNU General Public License for more details. -// -// You should have received a copy of the GNU General Public License -// along with Proton Mail Bridge. If not, see . - -package user - -import ( - "github.com/ProtonMail/gluon/imap" - "github.com/ProtonMail/gluon/queue" -) - -type flusher struct { - updateCh *queue.QueuedChannel[imap.Update] - updates []*imap.MessageCreated - pushedUpdates []imap.Update - - maxUpdateSize int - curChunkSize int -} - -func newFlusher(updateCh *queue.QueuedChannel[imap.Update], maxUpdateSize int) *flusher { - return &flusher{ - updateCh: updateCh, - maxUpdateSize: maxUpdateSize, - } -} - -func (f *flusher) push(update *imap.MessageCreated) { - f.updates = append(f.updates, update) - - if f.curChunkSize += len(update.Literal); f.curChunkSize >= f.maxUpdateSize { - f.flush() - } -} - -func (f *flusher) flush() { - if len(f.updates) > 0 { - update := imap.NewMessagesCreated(true, f.updates...) - f.updateCh.Enqueue(update) - f.updates = nil - f.curChunkSize = 0 - f.pushedUpdates = append(f.pushedUpdates, update) - } -} - -func (f *flusher) collectPushedUpdates() []imap.Update { - updates := f.pushedUpdates - f.pushedUpdates = nil - return updates -} diff --git a/internal/user/types.go b/internal/user/types.go index 254bf8b9..09478f5e 100644 --- a/internal/user/types.go +++ b/internal/user/types.go @@ -20,6 +20,7 @@ package user import ( "fmt" "reflect" + "runtime" "strings" "github.com/ProtonMail/go-proton-api" @@ -91,3 +92,7 @@ func sortSlice[Item any](items []Item, less func(Item, Item) bool) []Item { return sorted } + +func newProtonAPIScheduler() proton.Scheduler { + return proton.NewParallelScheduler(runtime.NumCPU() / 2) +} diff --git a/internal/user/user.go b/internal/user/user.go index 2c48c97d..abdef220 100644 --- a/internal/user/user.go +++ b/internal/user/user.go @@ -84,7 +84,6 @@ type User struct { pollAPIEventsCh chan chan struct{} goPollAPIEvents func(wait bool) - syncWorkers int showAllMail uint32 } @@ -98,7 +97,6 @@ func New( reporter reporter.Reporter, apiUser proton.User, crashHandler async.PanicHandler, - syncWorkers int, showAllMail bool, ) (*User, error) { //nolint:funlen logrus.WithField("userID", apiUser.ID).Info("Creating new user") @@ -142,7 +140,6 @@ func New( tasks: async.NewGroup(context.Background(), crashHandler), pollAPIEventsCh: make(chan chan struct{}), - syncWorkers: syncWorkers, showAllMail: b32(showAllMail), } diff --git a/internal/user/user_test.go b/internal/user/user_test.go index a1dcd786..13efde1c 100644 --- a/internal/user/user_test.go +++ b/internal/user/user_test.go @@ -216,7 +216,7 @@ func withUser(tb testing.TB, ctx context.Context, _ *server.Server, m *proton.Ma vaultUser, err := vault.AddUser(apiUser.ID, username, username+"@pm.me", apiAuth.UID, apiAuth.RefreshToken, saltedKeyPass) require.NoError(tb, err) - user, err := New(ctx, vaultUser, client, nil, apiUser, nil, vault.SyncWorkers(), true) + user, err := New(ctx, vaultUser, client, nil, apiUser, nil, true) require.NoError(tb, err) defer user.Close() diff --git a/internal/vault/settings.go b/internal/vault/settings.go index c067b33f..904b0374 100644 --- a/internal/vault/settings.go +++ b/internal/vault/settings.go @@ -189,27 +189,3 @@ func (vault *Vault) SetFirstStart(firstStart bool) error { data.Settings.FirstStart = firstStart }) } - -// SyncWorkers returns the number of workers to use for syncing. -func (vault *Vault) SyncWorkers() int { - return vault.get().Settings.SyncWorkers -} - -// SetSyncWorkers sets the number of workers to use for syncing. -func (vault *Vault) SetSyncWorkers(workers int) error { - return vault.mod(func(data *Data) { - data.Settings.SyncWorkers = workers - }) -} - -// SyncAttPool returns the size of the attachment pool. -func (vault *Vault) SyncAttPool() int { - return vault.get().Settings.SyncAttPool -} - -// SetSyncAttPool sets the size of the attachment pool. -func (vault *Vault) SetSyncAttPool(pool int) error { - return vault.mod(func(data *Data) { - data.Settings.SyncAttPool = pool - }) -} diff --git a/internal/vault/settings_test.go b/internal/vault/settings_test.go index 7a661309..4b90b3ba 100644 --- a/internal/vault/settings_test.go +++ b/internal/vault/settings_test.go @@ -202,12 +202,3 @@ func TestVault_Settings_FirstStart(t *testing.T) { // Check the new first start value. require.Equal(t, false, s.GetFirstStart()) } - -func TestVault_Settings_SyncWorkers(t *testing.T) { - // create a new test vault. - s := newVault(t) - - syncWorkers := vault.GetDefaultSyncWorkerCount() - require.Equal(t, syncWorkers, s.SyncWorkers()) - require.Equal(t, syncWorkers, s.SyncAttPool()) -} diff --git a/internal/vault/types_settings.go b/internal/vault/types_settings.go index 47c70654..4ba00f7e 100644 --- a/internal/vault/types_settings.go +++ b/internal/vault/types_settings.go @@ -19,7 +19,6 @@ package vault import ( "math/rand" - "runtime" "github.com/ProtonMail/proton-bridge/v3/internal/updater" ) @@ -43,26 +42,9 @@ type Settings struct { LastVersion string FirstStart bool - - SyncWorkers int - SyncAttPool int -} - -func GetDefaultSyncWorkerCount() int { - const minSyncWorkers = 16 - - syncWorkers := runtime.NumCPU() * 4 - - if syncWorkers < minSyncWorkers { - syncWorkers = minSyncWorkers - } - - return syncWorkers } func newDefaultSettings(gluonDir string) Settings { - syncWorkers := GetDefaultSyncWorkerCount() - return Settings{ GluonDir: gluonDir, @@ -82,8 +64,5 @@ func newDefaultSettings(gluonDir string) Settings { LastVersion: "0.0.0", FirstStart: true, - - SyncWorkers: syncWorkers, - SyncAttPool: syncWorkers, } } diff --git a/pkg/message/build.go b/pkg/message/build.go index f31b20f7..4811b48e 100644 --- a/pkg/message/build.go +++ b/pkg/message/build.go @@ -20,6 +20,7 @@ package message import ( "bytes" "encoding/base64" + "io" "mime" "net/mail" "strings" @@ -46,65 +47,73 @@ var ( const InternalIDDomain = `protonmail.internalid` func BuildRFC822(kr *crypto.KeyRing, msg proton.Message, attData [][]byte, opts JobOptions) ([]byte, error) { - switch { - case len(msg.Attachments) > 0: - return buildMultipartRFC822(kr, msg, attData, opts) - - case msg.MIMEType == "multipart/mixed": - return buildPGPRFC822(kr, msg, opts) - - default: - return buildSimpleRFC822(kr, msg, opts) - } -} - -func buildSimpleRFC822(kr *crypto.KeyRing, msg proton.Message, opts JobOptions) ([]byte, error) { - dec, err := msg.Decrypt(kr) - if err != nil { - if !opts.IgnoreDecryptionErrors { - return nil, errors.Wrap(ErrDecryptionFailed, err.Error()) - } - - return buildMultipartRFC822(kr, msg, nil, opts) - } - - hdr := getTextPartHeader(getMessageHeader(msg, opts), dec, msg.MIMEType) - buf := new(bytes.Buffer) - - w, err := message.CreateWriter(buf, hdr) - if err != nil { - return nil, err - } - - if _, err := w.Write(dec); err != nil { - return nil, err - } - - if err := w.Close(); err != nil { + if err := BuildRFC822Into(kr, msg, attData, opts, buf); err != nil { return nil, err } return buf.Bytes(), nil } +func BuildRFC822Into(kr *crypto.KeyRing, msg proton.Message, attData [][]byte, opts JobOptions, buf *bytes.Buffer) error { + switch { + case len(msg.Attachments) > 0: + return buildMultipartRFC822(kr, msg, attData, opts, buf) + + case msg.MIMEType == "multipart/mixed": + return buildPGPRFC822(kr, msg, opts, buf) + + default: + return buildSimpleRFC822(kr, msg, opts, buf) + } +} + +func buildSimpleRFC822(kr *crypto.KeyRing, msg proton.Message, opts JobOptions, buf *bytes.Buffer) error { + var decrypted bytes.Buffer + decrypted.Grow(len(msg.Body)) + + if err := msg.DecryptInto(kr, &decrypted); err != nil { + if !opts.IgnoreDecryptionErrors { + return errors.Wrap(ErrDecryptionFailed, err.Error()) + } + + return buildMultipartRFC822(kr, msg, nil, opts, buf) + } + + hdr := getTextPartHeader(getMessageHeader(msg, opts), decrypted.Bytes(), msg.MIMEType) + + w, err := message.CreateWriter(buf, hdr) + if err != nil { + return err + } + + if _, err := w.Write(decrypted.Bytes()); err != nil { + return err + } + + if err := w.Close(); err != nil { + return err + } + + return nil +} + func buildMultipartRFC822( kr *crypto.KeyRing, msg proton.Message, attData [][]byte, opts JobOptions, -) ([]byte, error) { + buf *bytes.Buffer, +) error { boundary := newBoundary(msg.ID) hdr := getMessageHeader(msg, opts) hdr.SetContentType("multipart/mixed", map[string]string{"boundary": boundary.gen()}) - buf := new(bytes.Buffer) - w, err := message.CreateWriter(buf, hdr) if err != nil { - return nil, err + return err } var ( @@ -126,23 +135,23 @@ func buildMultipartRFC822( if len(inlineAtts) > 0 { if err := writeRelatedParts(w, kr, boundary, msg, inlineAtts, inlineData, opts); err != nil { - return nil, err + return err } } else if err := writeTextPart(w, kr, msg, opts); err != nil { - return nil, err + return err } for i, att := range attachAtts { if err := writeAttachmentPart(w, kr, att, attachData[i], opts); err != nil { - return nil, err + return err } } if err := w.Close(); err != nil { - return nil, err + return err } - return buf.Bytes(), nil + return nil } func writeTextPart( @@ -151,8 +160,10 @@ func writeTextPart( msg proton.Message, opts JobOptions, ) error { - dec, err := msg.Decrypt(kr) - if err != nil { + var decrypted bytes.Buffer + decrypted.Grow(len(msg.Body)) + + if err := msg.DecryptInto(kr, &decrypted); err != nil { if !opts.IgnoreDecryptionErrors { return errors.Wrap(ErrDecryptionFailed, err.Error()) } @@ -160,7 +171,7 @@ func writeTextPart( return writeCustomTextPart(w, msg, err) } - return writePart(w, getTextPartHeader(message.Header{}, dec, msg.MIMEType), dec) + return writePart(w, getTextPartHeader(message.Header{}, decrypted.Bytes(), msg.MIMEType), decrypted.Bytes()) } func writeAttachmentPart( @@ -175,9 +186,10 @@ func writeAttachmentPart( return err } - msg := crypto.NewPGPSplitMessage(kps, attData).GetPGPMessage() + // Use io.Multi + attachmentReader := io.MultiReader(bytes.NewReader(kps), bytes.NewReader(attData)) - dec, err := kr.Decrypt(msg, nil, crypto.GetUnixTime()) + stream, err := kr.DecryptStream(attachmentReader, nil, crypto.GetUnixTime()) if err != nil { if !opts.IgnoreDecryptionErrors { return errors.Wrap(ErrDecryptionFailed, err.Error()) @@ -186,12 +198,38 @@ func writeAttachmentPart( log. WithField("attID", att.ID). WithError(err). - Warn("Attachment decryption failed") + Warn("Attachment decryption failed - construct") - return writeCustomAttachmentPart(w, att, msg, err) + var pgpMessageBuffer bytes.Buffer + pgpMessageBuffer.Grow(len(kps) + len(attData)) + pgpMessageBuffer.Write(kps) + pgpMessageBuffer.Write(attData) + + return writeCustomAttachmentPart(w, att, &crypto.PGPMessage{Data: pgpMessageBuffer.Bytes()}, err) } - return writePart(w, getAttachmentPartHeader(att), dec.GetBinary()) + var decryptBuffer bytes.Buffer + decryptBuffer.Grow(len(kps) + len(attData)) + + if _, err := decryptBuffer.ReadFrom(stream); err != nil { + if !opts.IgnoreDecryptionErrors { + return errors.Wrap(ErrDecryptionFailed, err.Error()) + } + + log. + WithField("attID", att.ID). + WithError(err). + Warn("Attachment decryption failed - stream") + + var pgpMessageBuffer bytes.Buffer + pgpMessageBuffer.Grow(len(kps) + len(attData)) + pgpMessageBuffer.Write(kps) + pgpMessageBuffer.Write(attData) + + return writeCustomAttachmentPart(w, att, &crypto.PGPMessage{Data: pgpMessageBuffer.Bytes()}, err) + } + + return writePart(w, getAttachmentPartHeader(att), decryptBuffer.Bytes()) } func writeRelatedParts( @@ -222,14 +260,16 @@ func writeRelatedParts( }) } -func buildPGPRFC822(kr *crypto.KeyRing, msg proton.Message, opts JobOptions) ([]byte, error) { - dec, err := msg.Decrypt(kr) - if err != nil { +func buildPGPRFC822(kr *crypto.KeyRing, msg proton.Message, opts JobOptions, buf *bytes.Buffer) error { + var decrypted bytes.Buffer + decrypted.Grow(len(msg.Body)) + + if err := msg.DecryptInto(kr, &decrypted); err != nil { if !opts.IgnoreDecryptionErrors { - return nil, errors.Wrap(ErrDecryptionFailed, err.Error()) + return errors.Wrap(ErrDecryptionFailed, err.Error()) } - return buildPGPMIMEFallbackRFC822(msg, opts) + return buildPGPMIMEFallbackRFC822(msg, opts, buf) } hdr := getMessageHeader(msg, opts) @@ -240,13 +280,13 @@ func buildPGPRFC822(kr *crypto.KeyRing, msg proton.Message, opts JobOptions) ([] } if len(sigs) > 0 { - return writeMultipartSignedRFC822(hdr, dec, sigs[0]) + return writeMultipartSignedRFC822(hdr, decrypted.Bytes(), sigs[0], buf) } - return writeMultipartEncryptedRFC822(hdr, dec) + return writeMultipartEncryptedRFC822(hdr, decrypted.Bytes(), buf) } -func buildPGPMIMEFallbackRFC822(msg proton.Message, opts JobOptions) ([]byte, error) { +func buildPGPMIMEFallbackRFC822(msg proton.Message, opts JobOptions, buf *bytes.Buffer) error { hdr := getMessageHeader(msg, opts) hdr.SetContentType("multipart/encrypted", map[string]string{ @@ -254,11 +294,9 @@ func buildPGPMIMEFallbackRFC822(msg proton.Message, opts JobOptions) ([]byte, er "protocol": "application/pgp-encrypted", }) - buf := new(bytes.Buffer) - w, err := message.CreateWriter(buf, hdr) if err != nil { - return nil, err + return err } var encHdr message.Header @@ -267,7 +305,7 @@ func buildPGPMIMEFallbackRFC822(msg proton.Message, opts JobOptions) ([]byte, er encHdr.Set("Content-Description", "PGP/MIME version identification") if err := writePart(w, encHdr, []byte("Version: 1")); err != nil { - return nil, err + return err } var dataHdr message.Header @@ -277,19 +315,17 @@ func buildPGPMIMEFallbackRFC822(msg proton.Message, opts JobOptions) ([]byte, er dataHdr.Set("Content-Description", "OpenPGP encrypted message") if err := writePart(w, dataHdr, []byte(msg.Body)); err != nil { - return nil, err + return err } if err := w.Close(); err != nil { - return nil, err + return err } - return buf.Bytes(), nil + return nil } -func writeMultipartSignedRFC822(header message.Header, body []byte, sig proton.Signature) ([]byte, error) { //nolint:funlen - buf := new(bytes.Buffer) - +func writeMultipartSignedRFC822(header message.Header, body []byte, sig proton.Signature, buf *bytes.Buffer) error { //nolint:funlen boundary := newBoundary("").gen() header.SetContentType("multipart/signed", map[string]string{ @@ -299,27 +335,27 @@ func writeMultipartSignedRFC822(header message.Header, body []byte, sig proton.S }) if err := textproto.WriteHeader(buf, header.Header); err != nil { - return nil, err + return err } mw := textproto.NewMultipartWriter(buf) if err := mw.SetBoundary(boundary); err != nil { - return nil, err + return err } bodyHeader, bodyData, err := readHeaderBody(body) if err != nil { - return nil, err + return err } bodyPart, err := mw.CreatePart(*bodyHeader) if err != nil { - return nil, err + return err } if _, err := bodyPart.Write(bodyData); err != nil { - return nil, err + return err } var sigHeader message.Header @@ -330,31 +366,29 @@ func writeMultipartSignedRFC822(header message.Header, body []byte, sig proton.S sigPart, err := mw.CreatePart(sigHeader.Header) if err != nil { - return nil, err + return err } sigData, err := sig.Data.GetArmored() if err != nil { - return nil, err + return err } if _, err := sigPart.Write([]byte(sigData)); err != nil { - return nil, err + return err } if err := mw.Close(); err != nil { - return nil, err + return err } - return buf.Bytes(), nil + return nil } -func writeMultipartEncryptedRFC822(header message.Header, body []byte) ([]byte, error) { - buf := new(bytes.Buffer) - +func writeMultipartEncryptedRFC822(header message.Header, body []byte, buf *bytes.Buffer) error { bodyHeader, bodyData, err := readHeaderBody(body) if err != nil { - return nil, err + return err } // Remove old content type header as it is non-standard. Ensure that messages @@ -371,14 +405,14 @@ func writeMultipartEncryptedRFC822(header message.Header, body []byte) ([]byte, } if err := textproto.WriteHeader(buf, header.Header); err != nil { - return nil, err + return err } if _, err := buf.Write(bodyData); err != nil { - return nil, err + return err } - return buf.Bytes(), nil + return nil } func addressEmpty(address *mail.Address) bool {