From 0c63533aa7a0f7679e3517554191c867b46ba407 Mon Sep 17 00:00:00 2001 From: Atanas Janeshliev Date: Fri, 11 Apr 2025 16:40:32 +0200 Subject: [PATCH] fix(BRIDGE-351): allow draft creation and import to BYOE addresses in combined mode --- internal/services/imapservice/connector.go | 53 +-- internal/services/imapservice/utils.go | 100 ++++++ internal/services/imapservice/utils_test.go | 380 ++++++++++++++++++++ 3 files changed, 481 insertions(+), 52 deletions(-) create mode 100644 internal/services/imapservice/utils.go create mode 100644 internal/services/imapservice/utils_test.go diff --git a/internal/services/imapservice/connector.go b/internal/services/imapservice/connector.go index ca55ef81..e5b29670 100644 --- a/internal/services/imapservice/connector.go +++ b/internal/services/imapservice/connector.go @@ -680,7 +680,7 @@ func (s *Connector) importMessage( } isDraft := slices.Contains(labelIDs, proton.DraftsLabel) - addr, err := s.getImportAddress(p, isDraft) + addr, err := getImportAddress(p, isDraft, s.addrID, s) if err != nil { return imap.Message{}, nil, err } @@ -871,57 +871,6 @@ func equalAddresses(a, b string) bool { return strings.EqualFold(stripPlusAlias(a), stripPlusAlias(b)) } -func (s *Connector) getImportAddress(p *parser.Parser, isDraft bool) (proton.Address, error) { - // addr is primary for combined mode or active for split mode - address, ok := s.identityState.GetAddress(s.addrID) - if !ok { - return proton.Address{}, errors.New("could not find account address") - } - - // If the address is external and not BYOE - with sending enabled, then use the primary address as an import target. - if address.Type == proton.AddressTypeExternal && !address.Send { - var err error - address, err = s.identityState.GetPrimaryAddress() - if err != nil { - return proton.Address{}, errors.New("could not get primary account address") - } - } - - inCombinedMode := s.addressMode == usertypes.AddressModeCombined - if !inCombinedMode { - return address, nil - } - - senderAddr, err := s.getSenderProtonAddress(p) - if err != nil { - if !errors.Is(err, errNoSenderAddressMatch) { - s.log.WithError(err).Warn("Could not get import address") - } - - // We did not find a match, so we use the default address. - return address, nil - } - - if senderAddr.ID == address.ID { - return address, nil - } - - // GODT-3185 / BRIDGE-120 In combined mode, in certain cases we adapt the address used for encryption. - // - draft with non-default address in combined mode: using sender address - // - import with non-default address in combined mode: using sender address - // - import with non-default disabled address in combined mode: using sender address - - isSenderAddressDisabled := (!bool(senderAddr.Send)) || (senderAddr.Status != proton.AddressStatusEnabled) - // BRIDGE-301: forbid imports via sender address if it's external. - isSenderExternal := senderAddr.Type == proton.AddressTypeExternal - - if (isDraft && isSenderAddressDisabled) || isSenderExternal { - return address, nil - } - - return senderAddr, nil -} - func (s *Connector) getSenderProtonAddress(p *parser.Parser) (proton.Address, error) { // Step 1: extract sender email address from message if (p == nil) || (p.Root() == nil) || (p.Root().Header.Len() == 0) { diff --git a/internal/services/imapservice/utils.go b/internal/services/imapservice/utils.go new file mode 100644 index 00000000..4675a675 --- /dev/null +++ b/internal/services/imapservice/utils.go @@ -0,0 +1,100 @@ +// Copyright (c) 2025 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 imapservice + +import ( + "errors" + + "github.com/ProtonMail/go-proton-api" + "github.com/ProtonMail/proton-bridge/v3/internal/usertypes" + "github.com/ProtonMail/proton-bridge/v3/pkg/message/parser" +) + +type connectorInterface interface { + getSenderProtonAddress(p *parser.Parser) (proton.Address, error) + getAddress(id string) (proton.Address, bool) + getPrimaryAddress() (proton.Address, error) + getAddressMode() usertypes.AddressMode + logError(err error, errMsg string) +} + +func (s *Connector) logError(err error, errMsg string) { + s.log.WithError(err).Warn(errMsg) +} + +func (s *Connector) getAddressMode() usertypes.AddressMode { + return s.addressMode +} + +func (s *Connector) getPrimaryAddress() (proton.Address, error) { + return s.identityState.GetPrimaryAddress() +} + +func (s *Connector) getAddress(id string) (proton.Address, bool) { + return s.identityState.GetAddress(id) +} + +func getImportAddress(p *parser.Parser, isDraft bool, id string, conn connectorInterface) (proton.Address, error) { + // addr is primary for combined mode or active for split mode + address, ok := conn.getAddress(id) + if !ok { + return proton.Address{}, errors.New("could not find account address") + } + + // If the address is external and not BYOE - with sending enabled, then use the primary address as an import target. + if address.Type == proton.AddressTypeExternal && !address.Send { + var err error + address, err = conn.getPrimaryAddress() + if err != nil { + return proton.Address{}, errors.New("could not get primary account address") + } + } + + inCombinedMode := conn.getAddressMode() == usertypes.AddressModeCombined + if !inCombinedMode { + return address, nil + } + + senderAddr, err := conn.getSenderProtonAddress(p) + if err != nil { + if !errors.Is(err, errNoSenderAddressMatch) { + conn.logError(err, "Could not get import address") + } + + // We did not find a match, so we use the default address. + return address, nil + } + + if senderAddr.ID == address.ID { + return address, nil + } + + // GODT-3185 / BRIDGE-120 In combined mode, in certain cases we adapt the address used for encryption. + // - draft with non-default address in combined mode: using sender address + // - import with non-default address in combined mode: using sender address + // - import with non-default disabled address in combined mode: using sender address + isSenderAddressDisabled := (!bool(senderAddr.Send)) || (senderAddr.Status != proton.AddressStatusEnabled) + isSenderExternalNonBYOE := senderAddr.Type == proton.AddressTypeExternal && !bool(senderAddr.Send) + + // Forbid drafts/imports for external non-BYOE addresses + if isSenderExternalNonBYOE || (isDraft && isSenderAddressDisabled) { + return address, nil + } + + return senderAddr, nil +} diff --git a/internal/services/imapservice/utils_test.go b/internal/services/imapservice/utils_test.go new file mode 100644 index 00000000..0715746d --- /dev/null +++ b/internal/services/imapservice/utils_test.go @@ -0,0 +1,380 @@ +// Copyright (c) 2025 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 imapservice + +import ( + "errors" + "testing" + + "github.com/ProtonMail/go-proton-api" + "github.com/ProtonMail/proton-bridge/v3/internal/usertypes" + "github.com/ProtonMail/proton-bridge/v3/pkg/message/parser" + "github.com/stretchr/testify/require" +) + +type testConnector struct { + addressMode usertypes.AddressMode + primaryAddress proton.Address + senderAddress proton.Address + imapAddress proton.Address + senderAddressError error +} + +func (t *testConnector) getSenderProtonAddress(_ *parser.Parser) (proton.Address, error) { + return t.senderAddress, t.senderAddressError +} + +func (t *testConnector) getAddress(_ string) (proton.Address, bool) { + return t.imapAddress, true +} + +func (t *testConnector) getPrimaryAddress() (proton.Address, error) { + return t.primaryAddress, nil +} + +func (t *testConnector) getAddressMode() usertypes.AddressMode { + return t.addressMode +} + +func (t *testConnector) logError(_ error, _ string) { +} + +func Test_GetImportAddress_SplitMode(t *testing.T) { + primaryAddress := proton.Address{ + ID: "1", + Email: "primary@proton.me", + Send: true, + Receive: true, + Type: proton.AddressTypeOriginal, + Status: proton.AddressStatusEnabled, + } + + imapAddressProton := proton.Address{ + ID: "2", + Email: "imap@proton.me", + Send: true, + Receive: true, + Type: proton.AddressTypeOriginal, + } + + testConn := &testConnector{ + addressMode: usertypes.AddressModeSplit, + primaryAddress: primaryAddress, + imapAddress: imapAddressProton, + } + + // Import address is internal, we're creating a draft. + // Expected: returned address is internal. + addr, err := getImportAddress(nil, true, imapAddressProton.ID, testConn) + require.NoError(t, err) + require.Equal(t, imapAddressProton.ID, addr.ID) + require.Equal(t, imapAddressProton.Email, addr.Email) + + // Import address is internal, we're attempting to import a message. + // Expected: returned address is internal. + addr, err = getImportAddress(nil, false, imapAddressProton.ID, testConn) + require.NoError(t, err) + require.Equal(t, imapAddressProton.ID, addr.ID) + require.Equal(t, imapAddressProton.Email, addr.Email) + + imapAddressBYOE := proton.Address{ + ID: "3", + Email: "byoe@external.com", + Send: true, + Receive: true, + Type: proton.AddressTypeExternal, + } + + // IMAP address is BYOE, we're creating a draft + // Expected: returned address is BYOE. + testConn.imapAddress = imapAddressBYOE + addr, err = getImportAddress(nil, true, imapAddressBYOE.ID, testConn) + require.NoError(t, err) + require.Equal(t, imapAddressBYOE.ID, addr.ID) + require.Equal(t, imapAddressBYOE.Email, addr.Email) + // IMAP address is BYOE, we're importing a message + // Expected: returned address is BYOE. + addr, err = getImportAddress(nil, false, imapAddressBYOE.ID, testConn) + require.NoError(t, err) + require.Equal(t, imapAddressBYOE.ID, addr.ID) + require.Equal(t, imapAddressBYOE.Email, addr.Email) + + imapAddressExternal := proton.Address{ + ID: "4", + Email: "external@external.com", + Send: false, + Receive: false, + Type: proton.AddressTypeExternal, + } + + // IMAP address is external, we're creating a draft. + // Expected: returned address is primary. + testConn.imapAddress = imapAddressExternal + addr, err = getImportAddress(nil, true, imapAddressExternal.ID, testConn) + require.NoError(t, err) + require.Equal(t, primaryAddress.ID, addr.ID) + require.Equal(t, primaryAddress.Email, addr.Email) + // IMAP address is external, we're trying to import. + // Expected: returned address is primary. + addr, err = getImportAddress(nil, false, imapAddressExternal.ID, testConn) + require.NoError(t, err) + require.Equal(t, primaryAddress.ID, addr.ID) + require.Equal(t, primaryAddress.Email, addr.Email) +} + +func Test_GetImportAddress_CombinedMode_ProtonAddresses(t *testing.T) { + primaryAddress := proton.Address{ + ID: "1", + Email: "primary@proton.me", + Send: true, + Receive: true, + Type: proton.AddressTypeOriginal, + Status: proton.AddressStatusEnabled, + } + + imapAddressProton := proton.Address{ + ID: "2", + Email: "imap@proton.me", + Send: true, + Receive: true, + Type: proton.AddressTypeOriginal, + } + + senderAddress := proton.Address{ + ID: "3", + Email: "sender@proton.me", + Send: true, + Receive: true, + Type: proton.AddressTypeOriginal, + Status: proton.AddressStatusEnabled, + } + + testConn := &testConnector{ + addressMode: usertypes.AddressModeCombined, + primaryAddress: primaryAddress, + imapAddress: imapAddressProton, + senderAddress: senderAddress, + } + + // Both the sender address and the imap address are the same. We're creating a draft. + // Expected: IMAP address is returned. + testConn.senderAddress = imapAddressProton + testConn.imapAddress = imapAddressProton + addr, err := getImportAddress(nil, true, imapAddressProton.ID, testConn) + require.NoError(t, err) + require.Equal(t, imapAddressProton.ID, addr.ID) + require.Equal(t, imapAddressProton.Email, addr.Email) + // Both the sender address and the imap address are the same. We're trying to import + // Expected: IMAP address is returned. + addr, err = getImportAddress(nil, false, imapAddressProton.ID, testConn) + require.NoError(t, err) + require.Equal(t, imapAddressProton.ID, addr.ID) + require.Equal(t, imapAddressProton.Email, addr.Email) + + // Sender address and imap address are different. Sender address is enabled and has sending enabled. + // We're creating a draft. + // Expected: Sender address is returned. + testConn.senderAddress = senderAddress + testConn.imapAddress = imapAddressProton + addr, err = getImportAddress(nil, true, imapAddressProton.ID, testConn) + require.NoError(t, err) + require.Equal(t, senderAddress.ID, addr.ID) + require.Equal(t, senderAddress.Email, addr.Email) + // Sender address and imap address are different. Sender address is enabled and has sending enabled. + // We're importing a message. + // Expected: Sender address is returned. + addr, err = getImportAddress(nil, false, imapAddressProton.ID, testConn) + require.NoError(t, err) + require.Equal(t, senderAddress.ID, addr.ID) + require.Equal(t, senderAddress.Email, addr.Email) + + // Sender address and imap address are different. Sender address is disabled, but has sending enabled. + // We're creating a draft message. + // Expected: IMAP address is returned. + senderAddress.Status = proton.AddressStatusDisabled + testConn.senderAddress = senderAddress + addr, err = getImportAddress(nil, true, imapAddressProton.ID, testConn) + require.NoError(t, err) + require.Equal(t, imapAddressProton.ID, addr.ID) + require.Equal(t, imapAddressProton.Email, addr.Email) + // Sender address and imap address are different. Sender address is disabled, but has sending enabled. + // We're importing a message. + // Expected: IMAP address is returned. + addr, err = getImportAddress(nil, false, imapAddressProton.ID, testConn) + require.NoError(t, err) + require.Equal(t, senderAddress.ID, addr.ID) + require.Equal(t, senderAddress.Email, addr.Email) + + // Sender address and imap address are different. Sender address is enabled, but has sending disabled. + // We're creating a draft. + // Expected: IMAP address is returned. + senderAddress.Status = proton.AddressStatusEnabled + senderAddress.Send = false + testConn.senderAddress = senderAddress + addr, err = getImportAddress(nil, true, imapAddressProton.ID, testConn) + require.NoError(t, err) + require.Equal(t, imapAddressProton.ID, addr.ID) + require.Equal(t, imapAddressProton.Email, addr.Email) + // Sender address and imap address are different. Sender address is enabled, but has sending disabled. + // We're importing a message. + // Expected: IMAP address is returned. + addr, err = getImportAddress(nil, false, imapAddressProton.ID, testConn) + require.NoError(t, err) + require.Equal(t, senderAddress.ID, addr.ID) + require.Equal(t, senderAddress.Email, addr.Email) + + // Sender address and imap address are different. But sender address is not an associated proton address. + // We're creating a draft. + // Expected: Sender address is returned. + testConn.senderAddressError = errors.New("sender address is not associated with the account") + addr, err = getImportAddress(nil, true, imapAddressProton.ID, testConn) + require.NoError(t, err) + require.Equal(t, imapAddressProton.ID, addr.ID) + require.Equal(t, imapAddressProton.Email, addr.Email) + // Sender address and imap address are different. But sender address is not an associated proton address. + // We're importing a message. + // Expected: Sender address is returned. + addr, err = getImportAddress(nil, false, imapAddressProton.ID, testConn) + require.NoError(t, err) + require.Equal(t, imapAddressProton.ID, addr.ID) + require.Equal(t, imapAddressProton.Email, addr.Email) +} + +func Test_GetImportAddress_CombinedMode_ExternalAddresses(t *testing.T) { + primaryAddress := proton.Address{ + ID: "1", + Email: "primary@proton.me", + Send: true, + Receive: true, + Type: proton.AddressTypeOriginal, + Status: proton.AddressStatusEnabled, + } + + imapAddressProton := proton.Address{ + ID: "2", + Email: "imap@proton.me", + Send: true, + Receive: true, + Type: proton.AddressTypeOriginal, + } + + senderAddressExternal := proton.Address{ + ID: "3", + Email: "sender@external.me", + Send: false, + Receive: false, + Type: proton.AddressTypeExternal, + Status: proton.AddressStatusEnabled, + } + + senderAddressExternalSecondary := proton.Address{ + ID: "4", + Email: "sender2@external.me", + Send: false, + Receive: false, + Type: proton.AddressTypeExternal, + Status: proton.AddressStatusEnabled, + } + + testConn := &testConnector{ + addressMode: usertypes.AddressModeCombined, + primaryAddress: primaryAddress, + imapAddress: imapAddressProton, + senderAddress: senderAddressExternal, + } + + // Sender address is external, and we're creating a draft. + // Expected: IMAP address is returned. + testConn.senderAddress = senderAddressExternal + testConn.imapAddress = imapAddressProton + addr, err := getImportAddress(nil, true, imapAddressProton.ID, testConn) + require.NoError(t, err) + require.Equal(t, imapAddressProton.ID, addr.ID) + require.Equal(t, imapAddressProton.Email, addr.Email) + // Sender address is external, and we're importing a message. + // Expected: IMAP address is returned. + addr, err = getImportAddress(nil, false, imapAddressProton.ID, testConn) + require.NoError(t, err) + require.Equal(t, imapAddressProton.ID, addr.ID) + require.Equal(t, imapAddressProton.Email, addr.Email) + + // Sender and IMAP address are external, and we're trying to import. + // Expected: Primary address is returned. + testConn.imapAddress = senderAddressExternalSecondary + addr, err = getImportAddress(nil, false, imapAddressProton.ID, testConn) + require.NoError(t, err) + require.Equal(t, primaryAddress.ID, addr.ID) + require.Equal(t, primaryAddress.Email, addr.Email) + // Sender and IMAP address are external, and we're trying to create a draft. + // Expected: Primary address is returned. + addr, err = getImportAddress(nil, true, imapAddressProton.ID, testConn) + require.NoError(t, err) + require.Equal(t, primaryAddress.ID, addr.ID) + require.Equal(t, primaryAddress.Email, addr.Email) +} + +func Test_GetImportAddress_CombinedMode_BYOEAddresses(t *testing.T) { + primaryAddress := proton.Address{ + ID: "1", + Email: "primary@proton.me", + Send: true, + Receive: true, + Type: proton.AddressTypeOriginal, + Status: proton.AddressStatusEnabled, + } + + imapAddressProton := proton.Address{ + ID: "2", + Email: "imap@proton.me", + Send: true, + Receive: true, + Type: proton.AddressTypeOriginal, + } + + senderAddressBYOE := proton.Address{ + ID: "3", + Email: "sender@external.me", + Send: true, + Receive: true, + Type: proton.AddressTypeExternal, + Status: proton.AddressStatusEnabled, + } + + testConn := &testConnector{ + addressMode: usertypes.AddressModeCombined, + primaryAddress: primaryAddress, + imapAddress: imapAddressProton, + senderAddress: senderAddressBYOE, + } + + // Sender address is BYOE, and we're creating a draft. + // Expected: BYOE address is returned. + testConn.senderAddress = senderAddressBYOE + testConn.imapAddress = imapAddressProton + addr, err := getImportAddress(nil, true, imapAddressProton.ID, testConn) + require.NoError(t, err) + require.Equal(t, senderAddressBYOE.ID, addr.ID) + require.Equal(t, senderAddressBYOE.Email, addr.Email) + + // Sender address is BYOE, and we're importing a message. + // Expected: BYOE address is returned. + addr, err = getImportAddress(nil, false, imapAddressProton.ID, testConn) + require.NoError(t, err) + require.Equal(t, senderAddressBYOE.ID, addr.ID) + require.Equal(t, senderAddressBYOE.Email, addr.Email) +}