feat(GODT-2759): CLI debug commands

Add debug commands to CLI to diagnose potential bride problems.
Currently we only have a command which validates whether the state of
all the mailboxes reported by IMAP matches what is currently available
on the proton servers.

internal/bridge/debug.go

@@ -0,0 +1,220 @@
+// 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 <https://www.gnu.org/licenses/>.
+
+package bridge
+
+import (
+	"context"
+	"fmt"
+	"io"
+
+	"github.com/ProtonMail/gluon/imap"
+	"github.com/ProtonMail/gluon/rfc822"
+	"github.com/ProtonMail/proton-bridge/v3/internal/user"
+	"github.com/bradenaw/juniper/iterator"
+	goimap "github.com/emersion/go-imap"
+	goimapclient "github.com/emersion/go-imap/client"
+	"github.com/sirupsen/logrus"
+	"golang.org/x/exp/maps"
+)
+
+// 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 {
+	bridge.usersLock.RLock()
+	defer bridge.usersLock.RUnlock()
+
+	users := maps.Values(bridge.users)
+
+	for _, usr := range users {
+		if progressCB != nil {
+			progressCB(fmt.Sprintf("Checking state for user %v", usr.Name()))
+		}
+		log := logrus.WithField("user", usr.Name()).WithField("diag", "state-check")
+		log.Debug("Retrieving all server metadata")
+		meta, err := usr.GetDiagnosticMetadata(ctx)
+		if err != nil {
+			return err
+		}
+
+		success := true
+
+		if len(meta.Metadata) != len(meta.MessageIDs) {
+			log.Errorf("Metadata (%v) and message(%v) list sizes do not match", len(meta.Metadata), len(meta.MessageIDs))
+		}
+
+		log.Debug("Building state")
+		state, err := meta.BuildMailboxToMessageMap(usr)
+		if err != nil {
+			log.WithError(err).Error("Failed to build state")
+			return err
+		}
+
+		info, err := bridge.GetUserInfo(usr.ID())
+		if err != nil {
+			log.WithError(err).Error("Failed to get user info")
+			return err
+		}
+
+		addr := fmt.Sprintf("127.0.0.1:%v", bridge.GetIMAPPort())
+
+		for account, mboxMap := range state {
+			if progressCB != nil {
+				progressCB(fmt.Sprintf("Checking state for user %v's account '%v'", usr.Name(), account))
+			}
+			if err := func(account string, mboxMap user.AccountMailboxMap) error {
+				client, err := goimapclient.Dial(addr)
+				if err != nil {
+					log.WithError(err).Error("Failed to connect to imap client")
+					return err
+				}
+
+				defer func() {
+					_ = client.Logout()
+				}()
+
+				if err := client.Login(account, string(info.BridgePass)); err != nil {
+					return fmt.Errorf("failed to login for user %v:%w", usr.Name(), err)
+				}
+
+				log := log.WithField("account", account)
+				for mboxName, messageList := range mboxMap {
+					log := log.WithField("mbox", mboxName)
+					status, err := client.Select(mboxName, true)
+					if err != nil {
+						log.WithError(err).Errorf("Failed to select mailbox %v", messageList)
+						return fmt.Errorf("failed to select '%v':%w", mboxName, err)
+					}
+
+					log.Debug("Checking message count")
+
+					if int(status.Messages) != len(messageList) {
+						success = false
+						log.Errorf("Message count doesn't match, got '%v' expected '%v'", status.Messages, len(messageList))
+					}
+
+					ids, err := clientGetMessageIDs(client, mboxName)
+					if err != nil {
+						return fmt.Errorf("failed to get message ids for mbox '%v': %w", mboxName, err)
+					}
+
+					for _, msg := range messageList {
+						imapFlags, ok := ids[msg.ID]
+						if !ok {
+							if meta.FailedMessageIDs.Contains(msg.ID) {
+								log.Warningf("Missing message '%v', but it is part of failed message set", msg.ID)
+							} else {
+								log.Errorf("Missing message '%v'", msg.ID)
+							}
+
+							continue
+						}
+
+						if checkFlags {
+							if !imapFlags.Equals(msg.Flags) {
+								log.Errorf("Message '%v' flags do mot match, got=%v, expected=%v",
+									msg.ID,
+									imapFlags.ToSlice(),
+									msg.Flags.ToSlice(),
+								)
+							}
+						}
+					}
+				}
+
+				if !success {
+					log.Errorf("State does not match")
+				} else {
+					log.Info("State matches")
+				}
+
+				return nil
+			}(account, mboxMap); err != nil {
+				return err
+			}
+		}
+	}
+
+	return nil
+}
+
+func clientGetMessageIDs(client *goimapclient.Client, mailbox string) (map[string]imap.FlagSet, error) {
+	status, err := client.Select(mailbox, false)
+	if err != nil {
+		return nil, err
+	}
+
+	if status.Messages == 0 {
+		return nil, nil
+	}
+
+	resCh := make(chan *goimap.Message)
+
+	section, err := goimap.ParseBodySectionName("BODY[HEADER]")
+	if err != nil {
+		return nil, err
+	}
+
+	fetchItems := []goimap.FetchItem{"BODY[HEADER]", goimap.FetchFlags}
+
+	seq, err := goimap.ParseSeqSet("1:*")
+	if err != nil {
+		return nil, err
+	}
+
+	go func() {
+		if err := client.Fetch(
+			seq,
+			fetchItems,
+			resCh,
+		); err != nil {
+			panic(err)
+		}
+	}()
+
+	messages := iterator.Collect(iterator.Chan(resCh))
+
+	ids := make(map[string]imap.FlagSet, len(messages))
+
+	for i, m := range messages {
+		literal, err := io.ReadAll(m.GetBody(section))
+		if err != nil {
+			return nil, err
+		}
+
+		header, err := rfc822.NewHeader(literal)
+		if err != nil {
+			return nil, fmt.Errorf("failed to parse header for msg %v: %w", i, err)
+		}
+
+		internalID, ok := header.GetChecked("X-Pm-Internal-Id")
+		if !ok {
+			logrus.Errorf("Message %v does not have internal id", internalID)
+			continue
+		}
+
+		messageFlags := imap.NewFlagSet(m.Flags...)
+
+		// Recent and Deleted are not part of the proton flag set.
+		messageFlags.RemoveFromSelf("\\Recent")
+		messageFlags.RemoveFromSelf("\\Deleted")
+
+		ids[internalID] = messageFlags
+	}
+
+	return ids, nil
+}

internal/frontend/cli/debug.go

@@ -0,0 +1,42 @@
+// 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 <https://www.gnu.org/licenses/>.
+
+package cli
+
+import (
+	"context"
+
+	"github.com/abiosoft/ishell"
+)
+
+func (f *frontendCLI) debugMailboxState(c *ishell.Context) {
+	f.ShowPrompt(false)
+	defer f.ShowPrompt(true)
+
+	checkFlags := f.yesNoQuestion("Also check message flags")
+
+	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) {
+		c.Println(s)
+	}); err != nil {
+		c.Printf("State check failed : %v", err)
+		return
+	}
+
+	c.Println("State check finished, see log for more details.")
+}

internal/frontend/cli/frontend.go

@@ -312,6 +312,19 @@ func New(
 	})
 	fe.AddCmd(telemetryCmd)
 
+	dbgCmd := &ishell.Cmd{
+		Name: "debug",
+		Help: "Debug diagnostics ",
+	}
+
+	dbgCmd.AddCmd(&ishell.Cmd{
+		Name: "mailbox-state",
+		Help: "Verify local mailbox state against proton server state",
+		Func: fe.debugMailboxState,
+	})
+
+	fe.AddCmd(dbgCmd)
+
 	go fe.watchEvents(eventCh)
 
 	go func() {

internal/user/debug.go

@@ -0,0 +1,140 @@
+// 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 <https://www.gnu.org/licenses/>.
+
+package user
+
+import (
+	"context"
+	"fmt"
+	"strings"
+
+	"github.com/ProtonMail/gluon/imap"
+	"github.com/ProtonMail/go-proton-api"
+	"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/sirupsen/logrus"
+)
+
+type DiagnosticMetadata struct {
+	MessageIDs       []string
+	Metadata         []proton.MessageMetadata
+	FailedMessageIDs xmaps.Set[string]
+}
+
+type AccountMailboxMap map[string][]MailboxMessage
+
+type MailboxMessage struct {
+	ID    string
+	Flags imap.FlagSet
+}
+
+func (apm DiagnosticMetadata) BuildMailboxToMessageMap(user *User) (map[string]AccountMailboxMap, error) {
+	return safe.RLockRetErr(func() (map[string]AccountMailboxMap, error) {
+		result := make(map[string]AccountMailboxMap)
+
+		mode := user.GetAddressMode()
+		primaryAddrID, err := getPrimaryAddr(user.apiAddrs)
+		if err != nil {
+			return nil, fmt.Errorf("failed to get primary addr for user: %w", err)
+		}
+
+		getAccount := func(addrID string) (AccountMailboxMap, bool) {
+			if mode == vault.CombinedMode {
+				addrID = primaryAddrID.ID
+			}
+
+			addr := user.apiAddrs[addrID]
+			if addr.Status != proton.AddressStatusEnabled {
+				return nil, false
+			}
+
+			v, ok := result[addr.Email]
+			if !ok {
+				result[addr.Email] = make(AccountMailboxMap)
+				v = result[addr.Email]
+			}
+
+			return v, true
+		}
+
+		for _, metadata := range apm.Metadata {
+			for _, label := range metadata.LabelIDs {
+				details, ok := user.apiLabels[label]
+				if !ok {
+					logrus.Warnf("User %v has message with unknown label '%v'", user.Name(), label)
+					continue
+				}
+
+				if !wantLabel(details) {
+					continue
+				}
+
+				account, enabled := getAccount(metadata.AddressID)
+				if !enabled {
+					continue
+				}
+
+				var mboxName string
+				if details.Type == proton.LabelTypeSystem {
+					mboxName = details.Name
+				} else {
+					mboxName = strings.Join(getMailboxName(details), "/")
+				}
+
+				mboxMessage := MailboxMessage{
+					ID:    metadata.ID,
+					Flags: buildFlagSetFromMessageMetadata(metadata),
+				}
+
+				if v, ok := account[mboxName]; ok {
+					account[mboxName] = append(v, mboxMessage)
+				} else {
+					account[mboxName] = []MailboxMessage{mboxMessage}
+				}
+			}
+		}
+		return result, nil
+	}, user.apiAddrsLock, user.apiLabelsLock)
+}
+
+func (user *User) GetDiagnosticMetadata(ctx context.Context) (DiagnosticMetadata, error) {
+	failedMessages := xmaps.SetFromSlice(user.vault.SyncStatus().FailedMessageIDs)
+
+	messageIDs, err := user.client.GetMessageIDs(ctx, "")
+	if err != nil {
+		return DiagnosticMetadata{}, err
+	}
+
+	meta := make([]proton.MessageMetadata, 0, len(messageIDs))
+
+	for _, m := range xslices.Chunk(messageIDs, 100) {
+		metadata, err := user.client.GetMessageMetadataPage(ctx, 0, len(m), proton.MessageFilter{ID: m})
+		if err != nil {
+			return DiagnosticMetadata{}, err
+		}
+
+		meta = append(meta, metadata...)
+	}
+
+	return DiagnosticMetadata{
+		MessageIDs:       messageIDs,
+		Metadata:         meta,
+		FailedMessageIDs: failedMessages,
+	}, nil
+}

internal/user/events.go

@@ -648,23 +648,7 @@ func (user *User) handleUpdateMessageEvent(_ context.Context, message proton.Mes
 			"subject":   logging.Sensitive(message.Subject),
 		}).Info("Handling message updated event")
 
-		flags := imap.NewFlagSet()
+		flags := buildFlagSetFromMessageMetadata(message)
-
-		if message.Seen() {
-			flags.AddToSelf(imap.FlagSeen)
-		}
-
-		if message.Starred() {
-			flags.AddToSelf(imap.FlagFlagged)
-		}
-
-		if message.IsDraft() {
-			flags.AddToSelf(imap.FlagDraft)
-		}
-
-		if message.IsRepliedAll == true || message.IsReplied == true { //nolint: gosimple
-			flags.AddToSelf(imap.FlagAnswered)
-		}
 
 		update := imap.NewMessageMailboxesUpdated(
 			imap.MessageID(message.ID),
@@ -867,3 +851,25 @@ func safePublishMessageUpdate(user *User, addressID string, update imap.Update)
 
 	return true, nil
 }
+
+func buildFlagSetFromMessageMetadata(message proton.MessageMetadata) imap.FlagSet {
+	flags := imap.NewFlagSet()
+
+	if message.Seen() {
+		flags.AddToSelf(imap.FlagSeen)
+	}
+
+	if message.Starred() {
+		flags.AddToSelf(imap.FlagFlagged)
+	}
+
+	if message.IsDraft() {
+		flags.AddToSelf(imap.FlagDraft)
+	}
+
+	if message.IsRepliedAll == true || message.IsReplied == true { //nolint: gosimple
+		flags.AddToSelf(imap.FlagAnswered)
+	}
+
+	return flags
+}