imap: support the Gmail extension (X-GM-EXT-1)

Support the IMAP Gmail extension (X-GM-EXT-1) to fetch all messages for
a given thread. This allows client-side threading to display a full
message thread. Obviously, it requires a Gmail account to work.

The extension is only used when requested in accounts.conf with:
"use-gmail-ext = true" (default: false)

Signed-off-by: Koni Marti <koni.marti@gmail.com>
Acked-by: Robin Jarry <robin@jarry.cc>
Tested-by: Tristan Partin <tristan@partin.io>

8 files changed, 273 insertions, 0 deletions

diff --git a/worker/imap/configure.go b/worker/imap/configure.go
index 94b5ac60..c325de23 100644
--- a/worker/imap/configure.go
+++ b/worker/imap/configure.go
@@ -155,6 +155,12 @@ func (w *IMAPWorker) handleConfigure(msg *types.Configure) error {
 				return fmt.Errorf("invalid cache-max-age value %v: %w", value, err)
 			}
 			w.config.cacheMaxAge = val
+		case "use-gmail-ext":
+			val, err := strconv.ParseBool(value)
+			if err != nil {
+				return fmt.Errorf("invalid use-gmail-ext value %v: %w", value, err)
+			}
+			w.config.useXGMEXT = val
 		}
 	}
 	if w.config.cacheEnabled {
diff --git a/worker/imap/extensions/xgmext/client.go b/worker/imap/extensions/xgmext/client.go
new file mode 100644
index 00000000..3107e642
--- /dev/null
+++ b/worker/imap/extensions/xgmext/client.go
@@ -0,0 +1,86 @@
+package xgmext
+
+import (
+	"errors"
+	"fmt"
+
+	"git.sr.ht/~rjarry/aerc/log"
+	"github.com/emersion/go-imap"
+	"github.com/emersion/go-imap/client"
+	"github.com/emersion/go-imap/commands"
+	"github.com/emersion/go-imap/responses"
+)
+
+type handler struct {
+	client *client.Client
+}
+
+func NewHandler(c *client.Client) *handler {
+	return &handler{client: c}
+}
+
+func (h handler) FetchEntireThreads(requested []uint32) ([]uint32, error) {
+	threadIds, err := h.fetchThreadIds(requested)
+	if err != nil {
+		return nil,
+			fmt.Errorf("faild to fetch thread IDs: %w", err)
+	}
+	uids, err := h.searchUids(threadIds)
+	if err != nil {
+		return nil,
+			fmt.Errorf("faild to search for thread IDs: %w", err)
+	}
+	return uids, nil
+}
+
+func (h handler) fetchThreadIds(uids []uint32) ([]string, error) {
+	messages := make(chan *imap.Message)
+	done := make(chan error)
+
+	thriditem := imap.FetchItem("X-GM-THRID")
+	items := []imap.FetchItem{
+		thriditem,
+	}
+
+	m := make(map[string]struct{}, len(uids))
+	go func() {
+		defer log.PanicHandler()
+		for msg := range messages {
+			m[msg.Items[thriditem].(string)] = struct{}{}
+		}
+		done <- nil
+	}()
+
+	var set imap.SeqSet
+	set.AddNum(uids...)
+	err := h.client.UidFetch(&set, items, messages)
+	<-done
+
+	thrid := make([]string, 0, len(m))
+	for id := range m {
+		thrid = append(thrid, id)
+	}
+	return thrid, err
+}
+
+func (h handler) searchUids(thrid []string) ([]uint32, error) {
+	if len(thrid) == 0 {
+		return nil, errors.New("no thread IDs provided")
+	}
+
+	if h.client.State() != imap.SelectedState {
+		return nil, errors.New("no mailbox selected")
+	}
+
+	var cmd imap.Commander = NewThreadIDSearch(thrid)
+	cmd = &commands.Uid{Cmd: cmd}
+
+	res := new(responses.Search)
+
+	status, err := h.client.Execute(cmd, res)
+	if err != nil {
+		return nil, fmt.Errorf("imap execute failed: %w", err)
+	}
+
+	return res.Ids, status.Err()
+}
diff --git a/worker/imap/extensions/xgmext/search.go b/worker/imap/extensions/xgmext/search.go
new file mode 100644
index 00000000..49b3448e
--- /dev/null
+++ b/worker/imap/extensions/xgmext/search.go
@@ -0,0 +1,44 @@
+package xgmext
+
+import "github.com/emersion/go-imap"
+
+type threadIDSearch struct {
+	Charset   string
+	ThreadIDs []string
+}
+
+// NewThreadIDSearch return an imap.Command to search UIDs for the provided
+// thread IDs using the X-GM-EXT-1 (Gmail extension)
+func NewThreadIDSearch(threadIDs []string) *threadIDSearch {
+	return &threadIDSearch{
+		Charset:   "UTF-8",
+		ThreadIDs: threadIDs,
+	}
+}
+
+func (cmd *threadIDSearch) Command() *imap.Command {
+	const threadSearchKey = "X-GM-THRID"
+
+	var args []interface{}
+	if cmd.Charset != "" {
+		args = append(args, imap.RawString("CHARSET"))
+		args = append(args, imap.RawString(cmd.Charset))
+	}
+
+	// we want to produce a search query that looks like this:
+	// SEARCH CHARSET UTF-8 OR OR X-GM-THRID 1771431779961568536 \
+	// X-GM-THRID 1765355745646219617 X-GM-THRID 1771500774375286796
+	for i := 0; i < len(cmd.ThreadIDs)-1; i++ {
+		args = append(args, imap.RawString("OR"))
+	}
+
+	for _, thrid := range cmd.ThreadIDs {
+		args = append(args, imap.RawString(threadSearchKey))
+		args = append(args, imap.RawString(thrid))
+	}
+
+	return &imap.Command{
+		Name:      "SEARCH",
+		Arguments: args,
+	}
+}
diff --git a/worker/imap/extensions/xgmext/search_test.go b/worker/imap/extensions/xgmext/search_test.go
new file mode 100644
index 00000000..8eb90e3c
--- /dev/null
+++ b/worker/imap/extensions/xgmext/search_test.go
@@ -0,0 +1,40 @@
+package xgmext_test
+
+import (
+	"bytes"
+	"testing"
+
+	"git.sr.ht/~rjarry/aerc/worker/imap/extensions/xgmext"
+	"github.com/emersion/go-imap"
+)
+
+func TestXGMEXT_Search(t *testing.T) {
+	tests := []struct {
+		name string
+		ids  []string
+		want string
+	}{
+		{
+			name: "search for single id",
+			ids:  []string{"1234"},
+			want: "* SEARCH CHARSET UTF-8 X-GM-THRID 1234\r\n",
+		},
+		{
+			name: "search for multiple id",
+			ids:  []string{"1234", "5678", "2345"},
+			want: "* SEARCH CHARSET UTF-8 OR OR X-GM-THRID 1234 X-GM-THRID 5678 X-GM-THRID 2345\r\n",
+		},
+	}
+	for _, test := range tests {
+		cmd := xgmext.NewThreadIDSearch(test.ids).Command()
+		var buf bytes.Buffer
+		err := cmd.WriteTo(imap.NewWriter(&buf))
+		if err != nil {
+			t.Errorf("failed to write command: %v", err)
+		}
+		if got := buf.String(); got != test.want {
+			t.Errorf("test '%s' failed: got: '%s', but wanted: '%s'",
+				test.name, got, test.want)
+		}
+	}
+}
diff --git a/worker/imap/worker.go b/worker/imap/worker.go
index f08c0ec9..7ef759d4 100644
--- a/worker/imap/worker.go
+++ b/worker/imap/worker.go
@@ -15,6 +15,7 @@ import (
 	"git.sr.ht/~rjarry/aerc/models"
 	"git.sr.ht/~rjarry/aerc/worker/handlers"
 	"git.sr.ht/~rjarry/aerc/worker/imap/extensions"
+	"git.sr.ht/~rjarry/aerc/worker/middleware"
 	"git.sr.ht/~rjarry/aerc/worker/types"
 )
 
@@ -58,6 +59,7 @@ type imapConfig struct {
 	keepalive_interval int
 	cacheEnabled       bool
 	cacheMaxAge        time.Duration
+	useXGMEXT          bool
 }
 
 type IMAPWorker struct {
@@ -120,6 +122,14 @@ func (w *IMAPWorker) newClient(c *client.Client) {
 		w.liststatus = true
 		w.worker.Debugf("Server Capability found: LIST-STATUS")
 	}
+	xgmext, err := w.client.Support("X-GM-EXT-1")
+	if err == nil && xgmext && w.config.useXGMEXT {
+		w.worker.Debugf("Server Capability found: X-GM-EXT-1")
+		w.worker = middleware.NewGmailWorker(w.worker, w.client.Client, w.idler)
+	}
+	if err == nil && !xgmext && w.config.useXGMEXT {
+		w.worker.Infof("X-GM-EXT-1 requested, but it is not supported")
+	}
 }
 
 func (w *IMAPWorker) handleMessage(msg types.WorkerMessage) error {
diff --git a/worker/middleware/foldermapper.go b/worker/middleware/foldermapper.go
index ee57b9c8..ba098191 100644
--- a/worker/middleware/foldermapper.go
+++ b/worker/middleware/foldermapper.go
@@ -26,6 +26,10 @@ func NewFolderMapper(base types.WorkerInteractor, mapping map[string]string,
 	}
 }
 
+func (f *folderMapper) Unwrap() types.WorkerInteractor {
+	return f.WorkerInteractor
+}
+
 func (f *folderMapper) incoming(msg types.WorkerMessage, dir string) string {
 	f.Lock()
 	defer f.Unlock()
diff --git a/worker/middleware/gmailworker.go b/worker/middleware/gmailworker.go
new file mode 100644
index 00000000..807f7bff
--- /dev/null
+++ b/worker/middleware/gmailworker.go
@@ -0,0 +1,78 @@
+package middleware
+
+import (
+	"sync"
+
+	"git.sr.ht/~rjarry/aerc/worker/imap/extensions/xgmext"
+	"git.sr.ht/~rjarry/aerc/worker/types"
+	"github.com/emersion/go-imap/client"
+)
+
+type idler interface {
+	Start()
+	Stop() error
+}
+
+type gmailWorker struct {
+	types.WorkerInteractor
+	mu     sync.Mutex
+	client *client.Client
+	idler  idler
+}
+
+// NewGmailWorker returns an IMAP middleware for the X-GM-EXT-1 extension
+func NewGmailWorker(base types.WorkerInteractor, c *client.Client, i idler,
+) types.WorkerInteractor {
+	base.Infof("loading worker middleware: X-GM-EXT-1")
+
+	// avoid double wrapping; unwrap and check for another gmail handler
+	for iter := base; iter != nil; iter = iter.Unwrap() {
+		if g, ok := iter.(*gmailWorker); ok {
+			base.Infof("already loaded; resetting")
+			err := g.reset(c, i)
+			if err != nil {
+				base.Errorf("reset failed: %v", err)
+			}
+			return base
+		}
+	}
+	return &gmailWorker{WorkerInteractor: base, client: c, idler: i}
+}
+
+func (g *gmailWorker) Unwrap() types.WorkerInteractor {
+	return g.WorkerInteractor
+}
+
+func (g *gmailWorker) reset(c *client.Client, i idler) error {
+	g.mu.Lock()
+	defer g.mu.Unlock()
+	g.client = c
+	g.idler = i
+	return nil
+}
+
+func (g *gmailWorker) ProcessAction(msg types.WorkerMessage) types.WorkerMessage {
+	switch msg := msg.(type) {
+	case *types.FetchMessageHeaders:
+		g.mu.Lock()
+		err := g.idler.Stop()
+		if err != nil {
+			g.Errorf("idler reported an error: %v", err)
+			break
+		}
+
+		handler := xgmext.NewHandler(g.client)
+		uids, err := handler.FetchEntireThreads(msg.Uids)
+		if err != nil {
+			g.Errorf("failed to fetch entire threads: %v", err)
+		}
+
+		if len(uids) > 0 {
+			msg.Uids = uids
+		}
+
+		g.idler.Start()
+		g.mu.Unlock()
+	}
+	return g.WorkerInteractor.ProcessAction(msg)
+}
diff --git a/worker/types/worker.go b/worker/types/worker.go
index f57f9fc3..37e0e2ad 100644
--- a/worker/types/worker.go
+++ b/worker/types/worker.go
@@ -15,6 +15,7 @@ type WorkerInteractor interface {
 	ProcessAction(WorkerMessage) WorkerMessage
 	PostAction(WorkerMessage, func(msg WorkerMessage))
 	PostMessage(WorkerMessage, func(msg WorkerMessage))
+	Unwrap() WorkerInteractor
 }
 
 var lastId int64 = 1 // access via atomic
@@ -50,6 +51,10 @@ func NewWorker(name string) *Worker {
 	}
 }
 
+func (worker *Worker) Unwrap() WorkerInteractor {
+	return nil
+}
+
 func (worker *Worker) Actions() chan WorkerMessage {
 	return worker.actions
 }