feat(GODT-2759): Add prompt to download missing messages for analysis

This will download the missing messages into a temporary directory and
decrypt them along with the metadata so we can attempt analyze them once
submitted to see what is going wrong.

internal/bridge/debug.go

@@ -21,6 +21,8 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"os"
+	"path/filepath"
 
 	"github.com/ProtonMail/gluon/imap"
 	"github.com/ProtonMail/gluon/rfc822"
@@ -32,14 +34,31 @@ import (
 	"golang.org/x/exp/maps"
 )
 
+type CheckClientStateResult struct {
+	MissingMessages map[string]map[string]user.DiagMailboxMessage
+}
+
+func (c *CheckClientStateResult) AddMissingMessage(userID string, message user.DiagMailboxMessage) {
+	v, ok := c.MissingMessages[userID]
+	if !ok {
+		c.MissingMessages[userID] = map[string]user.DiagMailboxMessage{message.ID: message}
+	} else {
+		v[message.ID] = message
+	}
+}
+
 // CheckClientState checks the current IMAP client reported state against the proton server state and reports
 // anything that is out of place.
-func (bridge *Bridge) CheckClientState(ctx context.Context, checkFlags bool, progressCB func(string)) error {
+func (bridge *Bridge) CheckClientState(ctx context.Context, checkFlags bool, progressCB func(string)) (CheckClientStateResult, error) {
 	bridge.usersLock.RLock()
 	defer bridge.usersLock.RUnlock()
 
 	users := maps.Values(bridge.users)
 
+	result := CheckClientStateResult{
+		MissingMessages: make(map[string]map[string]user.DiagMailboxMessage),
+	}
+
 	for _, usr := range users {
 		if progressCB != nil {
 			progressCB(fmt.Sprintf("Checking state for user %v", usr.Name()))
@@ -48,7 +67,7 @@ func (bridge *Bridge) CheckClientState(ctx context.Context, checkFlags bool, pro
 		log.Debug("Retrieving all server metadata")
 		meta, err := usr.GetDiagnosticMetadata(ctx)
 		if err != nil {
-			return err
+			return result, err
 		}
 
 		success := true
@@ -61,13 +80,13 @@ func (bridge *Bridge) CheckClientState(ctx context.Context, checkFlags bool, pro
 		state, err := meta.BuildMailboxToMessageMap(usr)
 		if err != nil {
 			log.WithError(err).Error("Failed to build state")
-			return err
+			return result, err
 		}
 
 		info, err := bridge.GetUserInfo(usr.ID())
 		if err != nil {
 			log.WithError(err).Error("Failed to get user info")
-			return err
+			return result, err
 		}
 
 		addr := fmt.Sprintf("127.0.0.1:%v", bridge.GetIMAPPort())
@@ -121,6 +140,7 @@ func (bridge *Bridge) CheckClientState(ctx context.Context, checkFlags bool, pro
 								log.Errorf("Missing message '%v'", msg.ID)
 							}
 
+							result.AddMissingMessage(msg.UserID, msg)
 							continue
 						}
 
@@ -144,11 +164,39 @@ func (bridge *Bridge) CheckClientState(ctx context.Context, checkFlags bool, pro
 
 				return nil
 			}(account, mboxMap); err != nil {
-				return err
+				return result, err
 			}
 		}
 	}
 
+	return result, nil
+}
+
+func (bridge *Bridge) DebugDownloadFailedMessages(
+	ctx context.Context,
+	result CheckClientStateResult,
+	exportPath string,
+	progressCB func(string, int, int),
+) error {
+	bridge.usersLock.RLock()
+	defer bridge.usersLock.RUnlock()
+
+	for userID, messages := range result.MissingMessages {
+		usr, ok := bridge.users[userID]
+		if !ok {
+			return fmt.Errorf("failed to find user with id %v", userID)
+		}
+
+		userDir := filepath.Join(exportPath, userID)
+		if err := os.MkdirAll(userDir, 0o700); err != nil {
+			return fmt.Errorf("failed to create directory '%v': %w", userDir, err)
+		}
+
+		if err := usr.DebugDownloadMessages(ctx, userDir, messages, progressCB); err != nil {
+			return err
+		}
+	}
+
 	return nil
 }
 

internal/frontend/cli/debug.go

@@ -19,6 +19,7 @@ package cli
 
 import (
 	"context"
+	"os"
 
 	"github.com/abiosoft/ishell"
 )
@@ -31,12 +32,43 @@ func (f *frontendCLI) debugMailboxState(c *ishell.Context) {
 
 	c.Println("Starting state check. Note that depending on your message count this may take a while.")
 
-	if err := f.bridge.CheckClientState(context.Background(), checkFlags, func(s string) {
+	result, err := f.bridge.CheckClientState(context.Background(), checkFlags, func(s string) {
 		c.Println(s)
-	}); err != nil {
+	})
+	if err != nil {
 		c.Printf("State check failed : %v", err)
 		return
 	}
 
 	c.Println("State check finished, see log for more details.")
+
+	if len(result.MissingMessages) == 0 {
+		return
+	}
+
+	f.Println("\n\nSome missing messages were detected. Bridge can download these messages for you")
+	f.Println("in a directory which you can later send to the developers for analysis.\n")
+	f.Println(bold("Note that the Messages will be stored unencrypted on disk.") + " If you do not wish")
+	f.Println("to continue, input no in the prompt below.\n")
+
+	if !f.yesNoQuestion("Would you like to proceed") {
+		return
+	}
+
+	location, err := os.MkdirTemp("", "debug-state-check-*")
+	if err != nil {
+		f.Printf("Failed to create temporary directory: %v\n", err)
+		return
+	}
+
+	c.Printf("Messages will be downloaded to: %v\n\n", bold(location))
+
+	if err := f.bridge.DebugDownloadFailedMessages(context.Background(), result, location, func(s string, i int, i2 int) {
+		f.Printf("[%v] Retrieving message %v of %v\n", s, i, i2)
+	}); err != nil {
+		f.Println(err)
+		return
+	}
+
+	c.Printf("\nMessage download finished. Data is available at %v\n", bold(location))
 }

internal/user/debug.go

@@ -18,16 +18,27 @@
 package user
 
 import (
+	"bytes"
 	"context"
+	"encoding/base64"
+	"encoding/json"
 	"fmt"
+	"io"
+	"mime"
+	"os"
+	"path/filepath"
 	"strings"
 
 	"github.com/ProtonMail/gluon/imap"
+	"github.com/ProtonMail/gluon/rfc822"
 	"github.com/ProtonMail/go-proton-api"
+	"github.com/ProtonMail/gopenpgp/v2/constants"
+	"github.com/ProtonMail/gopenpgp/v2/crypto"
 	"github.com/ProtonMail/proton-bridge/v3/internal/safe"
 	"github.com/ProtonMail/proton-bridge/v3/internal/vault"
 	"github.com/bradenaw/juniper/xmaps"
 	"github.com/bradenaw/juniper/xslices"
+	"github.com/emersion/go-message"
 	"github.com/sirupsen/logrus"
 )
 
@@ -37,11 +48,13 @@ type DiagnosticMetadata struct {
 	FailedMessageIDs xmaps.Set[string]
 }
 
-type AccountMailboxMap map[string][]MailboxMessage
+type AccountMailboxMap map[string][]DiagMailboxMessage
 
-type MailboxMessage struct {
+type DiagMailboxMessage struct {
-	ID    string
+	AddressID string
-	Flags imap.FlagSet
+	UserID    string
+	ID        string
+	Flags     imap.FlagSet
 }
 
 func (apm DiagnosticMetadata) BuildMailboxToMessageMap(user *User) (map[string]AccountMailboxMap, error) {
@@ -97,15 +110,17 @@ func (apm DiagnosticMetadata) BuildMailboxToMessageMap(user *User) (map[string]A
 					mboxName = strings.Join(getMailboxName(details), "/")
 				}
 
-				mboxMessage := MailboxMessage{
+				mboxMessage := DiagMailboxMessage{
-					ID:    metadata.ID,
+					UserID:    user.ID(),
-					Flags: buildFlagSetFromMessageMetadata(metadata),
+					ID:        metadata.ID,
+					AddressID: metadata.AddressID,
+					Flags:     buildFlagSetFromMessageMetadata(metadata),
 				}
 
 				if v, ok := account[mboxName]; ok {
 					account[mboxName] = append(v, mboxMessage)
 				} else {
-					account[mboxName] = []MailboxMessage{mboxMessage}
+					account[mboxName] = []DiagMailboxMessage{mboxMessage}
 				}
 			}
 		}
@@ -138,3 +153,267 @@ func (user *User) GetDiagnosticMetadata(ctx context.Context) (DiagnosticMetadata
 		FailedMessageIDs: failedMessages,
 	}, nil
 }
+
+func (user *User) DebugDownloadMessages(
+	ctx context.Context,
+	path string,
+	msgs map[string]DiagMailboxMessage,
+	progressCB func(string, int, int),
+) error {
+	var err error
+	safe.RLock(func() {
+		err = func() error {
+			total := len(msgs)
+			userID := user.ID()
+
+			counter := 1
+			for _, msg := range msgs {
+				if progressCB != nil {
+					progressCB(userID, counter, total)
+					counter++
+				}
+
+				msgDir := filepath.Join(path, msg.ID)
+				if err := os.MkdirAll(msgDir, 0o700); err != nil {
+					return fmt.Errorf("failed to create directory '%v':%w", msgDir, err)
+				}
+
+				message, err := user.client.GetFullMessage(ctx, msg.ID, newProtonAPIScheduler(user.panicHandler), proton.NewDefaultAttachmentAllocator())
+				if err != nil {
+					return fmt.Errorf("failed to download message '%v':%w", msg.ID, err)
+				}
+
+				if err := writeMetadata(msgDir, message.Message); err != nil {
+					return err
+				}
+
+				if err := withAddrKR(user.apiUser, user.apiAddrs[msg.AddressID], user.vault.KeyPass(), func(_, addrKR *crypto.KeyRing) error {
+					switch {
+					case len(message.Attachments) > 0:
+						return decodeMultipartMessage(msgDir, addrKR, message.Message, message.AttData)
+
+					case message.MIMEType == "multipart/mixed":
+						return decodePGPMessage(msgDir, addrKR, message.Message)
+
+					default:
+						return decodeSimpleMessage(msgDir, addrKR, message.Message)
+					}
+				}); err != nil {
+					return err
+				}
+			}
+			return nil
+		}()
+	}, user.apiAddrsLock, user.apiUserLock)
+	return err
+}
+
+func getBodyName(path string) string {
+	return filepath.Join(path, "body.txt")
+}
+
+func getBodyNameFailed(path string) string {
+	return filepath.Join(path, "body_failed.txt")
+}
+
+func getBodyNamePGP(path string) string {
+	return filepath.Join(path, "body.pgp")
+}
+
+func getMetadataPath(path string) string {
+	return filepath.Join(path, "metadata.json")
+}
+
+func getAttachmentPathSuccess(path, id, name string) string {
+	return filepath.Join(path, fmt.Sprintf("attachment_%v_%v", id, name))
+}
+
+func getAttachmentPathFailure(path, id string) string {
+	return filepath.Join(path, fmt.Sprintf("attachment_%v_failed.pgp", id))
+}
+
+func decodeMultipartMessage(outPath string, kr *crypto.KeyRing, msg proton.Message, attData [][]byte) error {
+	for idx, attachment := range msg.Attachments {
+		if err := decodeAttachment(outPath, kr, attachment, attData[idx]); err != nil {
+			return fmt.Errorf("failed to decode attachment %v of message %v: %w", attachment.ID, msg.ID, err)
+		}
+	}
+
+	return decodeSimpleMessage(outPath, kr, msg)
+}
+
+func decodePGPMessage(outPath string, kr *crypto.KeyRing, msg proton.Message) error {
+	var decrypted bytes.Buffer
+	decrypted.Grow(len(msg.Body))
+
+	if err := msg.DecryptInto(kr, &decrypted); err != nil {
+		logrus.Warnf("Failed to decrypt pgp message %v, storing as is: %v", msg.ID, err)
+		bodyPath := getBodyNamePGP(outPath)
+		if err := os.WriteFile(bodyPath, []byte(msg.Body), 0o600); err != nil {
+			return fmt.Errorf("failed to write pgp body to '%v': %w", bodyPath, err)
+		}
+	}
+
+	bodyPath := getBodyName(outPath)
+
+	if err := os.WriteFile(bodyPath, decrypted.Bytes(), 0o600); err != nil {
+		return fmt.Errorf("failed to write pgp body to '%v': %w", bodyPath, err)
+	}
+
+	return nil
+}
+
+func decodeSimpleMessage(outPath string, kr *crypto.KeyRing, msg proton.Message) error {
+	var decrypted bytes.Buffer
+	decrypted.Grow(len(msg.Body))
+
+	if err := msg.DecryptInto(kr, &decrypted); err != nil {
+		logrus.Warnf("Failed to decrypt simple message %v, will try again as attachment : %v", msg.ID, err)
+		return writeCustomTextPart(getBodyNameFailed(outPath), msg, err)
+	}
+
+	bodyPath := getBodyName(outPath)
+
+	if err := os.WriteFile(bodyPath, decrypted.Bytes(), 0o600); err != nil {
+		return fmt.Errorf("failed to write simple body to '%v': %w", bodyPath, err)
+	}
+
+	return nil
+}
+
+func writeMetadata(outPath string, msg proton.Message) error {
+	type CustomMetadata struct {
+		proton.MessageMetadata
+		Header        string
+		ParsedHeaders proton.Headers
+		MIMEType      rfc822.MIMEType
+		Attachments   []proton.Attachment
+	}
+
+	metadata := CustomMetadata{
+		MessageMetadata: msg.MessageMetadata,
+		Header:          msg.Header,
+		ParsedHeaders:   msg.ParsedHeaders,
+		MIMEType:        msg.MIMEType,
+		Attachments:     msg.Attachments,
+	}
+
+	j, err := json.MarshalIndent(metadata, "", "  ")
+	if err != nil {
+		return fmt.Errorf("failed to encode json for message %v: %w", msg.ID, err)
+	}
+
+	metaPath := getMetadataPath(outPath)
+
+	if err := os.WriteFile(metaPath, j, 0o600); err != nil {
+		return fmt.Errorf("failed to write metadata to '%v': %w", metaPath, err)
+	}
+
+	return nil
+}
+
+func decodeAttachment(outPath string, kr *crypto.KeyRing,
+	att proton.Attachment,
+	attData []byte) error {
+	kps, err := base64.StdEncoding.DecodeString(att.KeyPackets)
+	if err != nil {
+		return err
+	}
+
+	// Use io.Multi
+	attachmentReader := io.MultiReader(bytes.NewReader(kps), bytes.NewReader(attData))
+
+	stream, err := kr.DecryptStream(attachmentReader, nil, crypto.GetUnixTime())
+	if err != nil {
+		logrus.
+			WithField("attID", att.ID).
+			WithError(err).
+			Warn("Attachment decryption failed - construct")
+
+		var pgpMessageBuffer bytes.Buffer
+		pgpMessageBuffer.Grow(len(kps) + len(attData))
+		pgpMessageBuffer.Write(kps)
+		pgpMessageBuffer.Write(attData)
+
+		return writeCustomAttachmentPart(getAttachmentPathFailure(outPath, att.ID), att, &crypto.PGPMessage{Data: pgpMessageBuffer.Bytes()}, err)
+	}
+
+	var decryptBuffer bytes.Buffer
+	decryptBuffer.Grow(len(kps) + len(attData))
+
+	if _, err := decryptBuffer.ReadFrom(stream); err != nil {
+		logrus.
+			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(getAttachmentPathFailure(outPath, att.ID), att, &crypto.PGPMessage{Data: pgpMessageBuffer.Bytes()}, err)
+	}
+
+	attachmentPath := getAttachmentPathSuccess(outPath, att.ID, att.Name)
+
+	if err := os.WriteFile(attachmentPath, decryptBuffer.Bytes(), 0o600); err != nil {
+		return fmt.Errorf("failed to write attachment %v to '%v': %w", att.ID, attachmentPath, err)
+	}
+
+	return nil
+}
+
+func writeCustomTextPart(
+	outPath string,
+	msg proton.Message,
+	decError error,
+) error {
+	enc, err := crypto.NewPGPMessageFromArmored(msg.Body)
+	if err != nil {
+		return err
+	}
+
+	arm, err := enc.GetArmoredWithCustomHeaders(
+		fmt.Sprintf("This message could not be decrypted: %v", decError),
+		constants.ArmorHeaderVersion,
+	)
+	if err != nil {
+		return err
+	}
+
+	if err := os.WriteFile(outPath, []byte(arm), 0o600); err != nil {
+		return fmt.Errorf("failed to write custom message %v data to '%v': %w", msg.ID, outPath, err)
+	}
+
+	return nil
+}
+
+// writeCustomAttachmentPart writes an armored-PGP data part for an attachment that couldn't be decrypted.
+func writeCustomAttachmentPart(
+	outPath string,
+	att proton.Attachment,
+	msg *crypto.PGPMessage,
+	decError error,
+) error {
+	arm, err := msg.GetArmoredWithCustomHeaders(
+		fmt.Sprintf("This attachment could not be decrypted: %v", decError),
+		constants.ArmorHeaderVersion,
+	)
+	if err != nil {
+		return err
+	}
+
+	filename := mime.QEncoding.Encode("utf-8", att.Name+".pgp")
+
+	var hdr message.Header
+
+	hdr.SetContentType("application/octet-stream", map[string]string{"name": filename})
+	hdr.SetContentDisposition(string(att.Disposition), map[string]string{"filename": filename})
+
+	if err := os.WriteFile(outPath, []byte(arm), 0o600); err != nil {
+		return fmt.Errorf("failed to write custom attachment %v part  to '%v': %w", att.ID, outPath, err)
+	}
+
+	return nil
+}