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>

3 files changed, 170 insertions, 0 deletions

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)
+		}
+	}
+}