feat(BRIDGE-120): use appropriate address key when importing / saving draft.

internal/services/imapservice/connector.go

@@ -31,6 +31,7 @@ import (
 	"github.com/ProtonMail/gluon/connector"
 	"github.com/ProtonMail/gluon/imap"
 	"github.com/ProtonMail/gluon/reporter"
+	"github.com/ProtonMail/gluon/rfc5322"
 	"github.com/ProtonMail/gluon/rfc822"
 	"github.com/ProtonMail/go-proton-api"
 	"github.com/ProtonMail/gopenpgp/v2/crypto"
@@ -69,6 +70,8 @@ type Connector struct {
 	syncState   *SyncState
 }
 
+var errNoSenderAddressMatch = errors.New("no matching sender found in address list")
+
 func NewConnector(
 	addrID string,
 	apiClient APIClient,
@@ -107,6 +110,7 @@ func NewConnector(
 		labels:      labels,
 		addressMode: addressMode,
 		log: logrus.WithFields(logrus.Fields{
+			"pkg":             "imapservice",
 			"gluon-connector": addressMode,
 			"addr-id":         addrID,
 			"user-id":         userID,
@@ -676,20 +680,18 @@ func (s *Connector) importMessage(
 ) (imap.Message, []byte, error) {
 	var full proton.FullMessage
 
-	// addr is primary for combined mode or active for split mode
-	addr, ok := s.identityState.GetAddress(s.addrID)
-	if !ok {
-		return imap.Message{}, nil, fmt.Errorf("could not find address")
-	}
-
 	p, err2 := parser.New(bytes.NewReader(literal))
 	if err2 != nil {
 		return imap.Message{}, nil, fmt.Errorf("failed to parse literal: %w", err2)
 	}
 
 	isDraft := slices.Contains(labelIDs, proton.DraftsLabel)
+	addr, err := s.getImportAddress(p, isDraft)
+	if err != nil {
+		return imap.Message{}, nil, err
+	}
 
-	if err := s.identityState.WithAddrKR(s.addrID, func(_, addrKR *crypto.KeyRing) error {
+	if err := s.identityState.WithAddrKR(addr.ID, func(_, addrKR *crypto.KeyRing) error {
 		primaryKey, errKey := addrKR.FirstKey()
 		if errKey != nil {
 			return fmt.Errorf("failed to get primary key for import: %w", errKey)
@@ -716,7 +718,7 @@ func (s *Connector) importMessage(
 			}
 			str, err := s.client.ImportMessages(ctx, primaryKey, 1, 1, []proton.ImportReq{{
 				Metadata: proton.ImportMetadata{
-					AddressID: s.addrID,
+					AddressID: addr.ID,
 					LabelIDs:  labelIDs,
 					Unread:    proton.Bool(unread),
 					Flags:     flags,
@@ -874,3 +876,75 @@ func stripPlusAlias(a string) string {
 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")
+	}
+
+	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)
+	if isDraft && isSenderAddressDisabled {
+		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) {
+		return proton.Address{}, errors.New("invalid message encountered while trying to extract sender address")
+	}
+
+	addrField := p.Root().Header.Get("From")
+	if len(addrField) == 0 {
+		addrField = p.Root().Header.Get("Sender")
+	}
+	if len(addrField) == 0 {
+		return proton.Address{}, errors.New("no sender found in message headers")
+	}
+
+	sender, err := rfc5322.ParseAddressList(addrField)
+	if (err != nil) || (len(sender) == 0) {
+		return proton.Address{}, fmt.Errorf("invalid sender address in message: %w", err)
+	}
+
+	addrStr := sender[0].Address
+
+	// Step 2: match email with the user address list.
+	addressList := s.identityState.GetAddresses()
+	index := slices.IndexFunc(addressList, func(a proton.Address) bool {
+		return equalAddresses(a.Email, addrStr)
+	})
+	if index < 0 {
+		return proton.Address{}, errNoSenderAddressMatch
+	}
+
+	return addressList[index], nil
+}

tests/environment_test.go

@@ -20,11 +20,15 @@ package tests
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"net/http"
 	"strings"
 
+	"github.com/ProtonMail/go-proton-api"
+	"github.com/ProtonMail/gopenpgp/v2/crypto"
 	"github.com/cucumber/godog"
 )
 
@@ -116,7 +120,7 @@ func (s *scenario) checkParsedMultipartFormForFile(method, path, file string, ha
 	}
 
 	if _, ok := req.MultipartForm.File[file]; hasFile != ok {
-		return fmt.Errorf("Multipart file in bug report is %t, want it to be %t", ok, hasFile)
+		return fmt.Errorf("multipart file in bug report is %t, want it to be %t", ok, hasFile)
 	}
 
 	return nil
@@ -240,3 +244,57 @@ func (s *scenario) theMessageUsedKeyForSending(address string) error {
 
 	return nil
 }
+
+func (s *scenario) theKeyForAddressWasUsedToImport(address string) error {
+	// Response does not include the address ID, only the messageID, so we extract the messageID from the response body
+	call, err := s.t.getLastCallExcludingHTTPOverride("POST", "/mail/v4/messages/import")
+	if err != nil {
+		return err
+	}
+
+	var resp struct {
+		Responses []struct{ Response struct{ MessageID string } }
+	}
+
+	if err := json.Unmarshal(call.ResponseBody, &resp); err != nil {
+		return err
+	}
+
+	return s.checkMessageIsEncryptedForAddress(resp.Responses[0].Response.MessageID, address)
+}
+
+func (s *scenario) theKeyForAddressWasUsedToCreateDraft(address string) error {
+	call, err := s.t.getLastCallExcludingHTTPOverride("POST", "/mail/v4/messages")
+	if err != nil {
+		return err
+	}
+
+	var resp struct{ Message struct{ ID string } }
+	if err := json.Unmarshal(call.ResponseBody, &resp); err != nil {
+		return err
+	}
+
+	return s.checkMessageIsEncryptedForAddress(resp.Message.ID, address)
+}
+
+func (s *scenario) checkMessageIsEncryptedForAddress(messageID string, address string) error {
+	user := s.t.getUserByAddress(address)
+	addrID := user.getAddrID(address)
+	return s.t.withClient(context.Background(), user.getName(), func(ctx context.Context, client *proton.Client) error {
+		message, err := client.GetMessage(ctx, messageID)
+		if err != nil {
+			return err
+		}
+
+		// Check 1: the address associated with the message is the one we expect.
+		if message.AddressID != addrID {
+			return errors.New("the message is not encrypted with the specified address")
+		}
+
+		// Check 2: we indeed encrypted the message with this address' key ring.
+		return s.t.withAddrKR(ctx, client, user.name, addrID, func(_ context.Context, kr *crypto.KeyRing) error {
+			_, err := message.Decrypt(kr)
+			return err
+		})
+	})
+}

tests/features/imap/message/import_key.feature

@@ -0,0 +1,136 @@
+Feature: IMAP import messages
+
+  Background:
+    Given there exists an account with username "[user:user]" and password "password"
+    And the account "[user:user]" has additional address "[alias:secondary]@[domain]"
+    And the account "[user:user]" has additional disabled address "[alias:disabled]@[domain]"
+    Then it succeeds
+    When bridge starts
+    And the user logs in with username "[user:user]" and password "password"
+    And user "[user:user]" finishes syncing
+    And user "[user:user]" connects and authenticates IMAP client "1"
+    Then it succeeds
+
+  @skip-black
+  Scenario: Messages imported with default address as sender are encrypted with the default address key
+    When IMAP client "1" appends the following message to "INBOX":
+      """
+      From: Bridge Test <[user:user]@[domain]>
+      Date: 01 Jan 1980 00:00:00 +0000
+      To: Internal Bridge <bridgetest@example.com>
+      Received: by 2002:0:0:0:0:0:0:0 with SMTP id 0123456789abcdef; Wed, 30 Dec 2020 01:23:45 0000
+      Subject: Basic text/plain message
+      Content-Type: text/plain
+
+      Hello
+      """
+    Then it succeeds
+    And the key for address "[user:user]@[domain]" was used to import
+
+  @skip-black
+  Scenario: Messages imported with alias as sender are encrypted with secondary address key
+    When IMAP client "1" appends the following message to "INBOX":
+      """
+      From: Bridge Test <[alias:secondary]@[domain]>
+      Date: 01 Jan 1980 00:00:00 +0000
+      To: Internal Bridge <bridgetest@example.com>
+      Received: by 2002:0:0:0:0:0:0:0 with SMTP id 0123456789abcdef; Wed, 30 Dec 2020 01:23:45 0000
+      Subject: Basic text/plain message
+      Content-Type: text/plain
+
+      Hello
+      """
+    Then it succeeds
+    And the key for address "[alias:secondary]@[domain]" was used to import
+
+  @skip-black
+  Scenario: Messages imported with a disabled alias as sender are encrypted with the disabled address key
+    When IMAP client "1" appends the following message to "INBOX":
+      """
+      From: Bridge Test <[alias:disabled]@[domain]>
+      Date: 01 Jan 1980 00:00:00 +0000
+      To: Internal Bridge <bridgetest@example.com>
+      Received: by 2002:0:0:0:0:0:0:0 with SMTP id 0123456789abcdef; Wed, 30 Dec 2020 01:23:45 0000
+      Subject: Basic text/plain message
+      Content-Type: text/plain
+
+      Hello
+      """
+    Then it succeeds
+    And the key for address "[alias:disabled]@[domain]" was used to import
+
+  @skip-black
+  Scenario: Messages imported with an unknown address as sender are encrypted with primary address key
+    When IMAP client "1" appends the following message to "INBOX":
+      """
+      From: Bridge Test <bridgeqa@example.com>
+      Date: 01 Jan 1980 00:00:00 +0000
+      To: Internal Bridge <bridgetest@example.com>
+      Received: by 2002:0:0:0:0:0:0:0 with SMTP id 0123456789abcdef; Wed, 30 Dec 2020 01:23:45 0000
+      Subject: Basic text/plain message
+      Content-Type: text/plain
+
+      Hello
+      """
+    Then it succeeds
+    And the key for address "[user:user]@[domain]" was used to import
+
+  @skip-black
+  Scenario: Drafts imported with default address as sender are encrypted with the default address key
+    When IMAP client "1" appends the following message to "Drafts":
+      """
+      From: Bridge Test <[user:user]@[domain]>
+      Date: 01 Jan 1980 00:00:00 +0000
+      To: Internal Bridge <bridgetest@example.com>
+      Subject: Basic text/plain message
+      Content-Type: text/plain
+
+      Hello
+      """
+    Then it succeeds
+    And the key for address "[user:user]@[domain]" was used to create draft
+
+  @skip-black
+  Scenario: Drafts imported with alias as sender are encrypted with secondary key
+    When IMAP client "1" appends the following message to "Drafts":
+      """
+      From: Bridge Test <[alias:secondary]@[domain]>
+      Date: 01 Jan 1980 00:00:00 +0000
+      To: Internal Bridge <bridgetest@example.com>
+      Subject: Basic text/plain message
+      Content-Type: text/plain
+
+      Hello
+      """
+    Then it succeeds
+    And the key for address "[alias:secondary]@[domain]" was used to create draft
+
+  @skip-black
+  Scenario: Drafts imported with a disabled alias as sender are encrypted with the disabled address key
+    When IMAP client "1" appends the following message to "Drafts":
+      """
+      From: Bridge Test <[alias:disabled]@[domain]>
+      Date: 01 Jan 1980 00:00:00 +0000
+      To: Internal Bridge <bridgetest@example.com>
+      Subject: Basic text/plain message
+      Content-Type: text/plain
+
+      Hello
+      """
+    Then it succeeds
+    And the key for address "[user:user]@[domain]" was used to create drafts
+
+  @skip-black
+  Scenario: Drafts imported with an unknown address as sender are encrypted with primary address key
+    When IMAP client "1" appends the following message to "Drafts":
+      """
+      From: Bridge Test <bridgeqa@example.com>
+      Date: 01 Jan 1980 00:00:00 +0000
+      To: Internal Bridge <bridgetest@example.com>
+      Subject: Basic text/plain message
+      Content-Type: text/plain
+
+      Hello
+      """
+    Then it succeeds
+    And the key for address "[user:user]@[domain]" was used to create draft

tests/steps_test.go

@@ -39,6 +39,8 @@ func (s *scenario) steps(ctx *godog.ScenarioContext) {
 	ctx.Step(`^bridge IMAP port is (\d+)`, s.bridgeIMAPPortIs)
 	ctx.Step(`^bridge SMTP port is (\d+)`, s.bridgeSMTPPortIs)
 	ctx.Step(`^the message used "([^"]*)" key for sending$`, s.theMessageUsedKeyForSending)
+	ctx.Step(`^the key for address "([^"]*)" was used to import`, s.theKeyForAddressWasUsedToImport)
+	ctx.Step(`^the key for address "([^"]*)" was used to create draft`, s.theKeyForAddressWasUsedToCreateDraft)
 
 	// ==== SETUP ====
 	ctx.Step(`^there exists an account with username "([^"]*)" and password "([^"]*)"$`, s.thereExistsAnAccountWithUsernameAndPassword)