5 files changed, 193 insertions, 10 deletions

diff --git a/worker/imap/extensions/xgmext/client.go b/worker/imap/extensions/xgmext/client.go
index 3d9ce1e9..65f11e74 100644
--- a/worker/imap/extensions/xgmext/client.go
+++ b/worker/imap/extensions/xgmext/client.go
@@ -73,20 +73,22 @@ func (h handler) searchUids(thrid []string) ([]uint32, error) {
 	if len(thrid) == 0 {
 		return nil, errors.New("no thread IDs provided")
 	}
+	return h.runSearch(NewThreadIDSearch(thrid))
+}
+
+func (h handler) RawSearch(rawSearch string) ([]uint32, error) {
+	return h.runSearch(NewRawSearch(rawSearch))
+}
 
+func (h handler) runSearch(cmd imap.Commander) ([]uint32, error) {
 	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
index 49b3448e..42a4f2ca 100644
--- a/worker/imap/extensions/xgmext/search.go
+++ b/worker/imap/extensions/xgmext/search.go
@@ -42,3 +42,33 @@ func (cmd *threadIDSearch) Command() *imap.Command {
 		Arguments: args,
 	}
 }
+
+type rawSearch struct {
+	Charset string
+	Search  string
+}
+
+func NewRawSearch(search string) *rawSearch {
+	return &rawSearch{
+		Charset: "UTF-8",
+		Search:  search,
+	}
+}
+
+func (cmd *rawSearch) Command() *imap.Command {
+	const key = "X-GM-RAW"
+
+	var args []interface{}
+	if cmd.Charset != "" {
+		args = append(args, imap.RawString("CHARSET"))
+		args = append(args, imap.RawString(cmd.Charset))
+	}
+
+	args = append(args, imap.RawString(key))
+	args = append(args, imap.RawString(cmd.Search))
+
+	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
index 8eb90e3c..a2a2791e 100644
--- a/worker/imap/extensions/xgmext/search_test.go
+++ b/worker/imap/extensions/xgmext/search_test.go
@@ -8,7 +8,7 @@ import (
 	"github.com/emersion/go-imap"
 )
 
-func TestXGMEXT_Search(t *testing.T) {
+func TestXGMEXT_ThreadIDSearch(t *testing.T) {
 	tests := []struct {
 		name string
 		ids  []string
@@ -38,3 +38,39 @@ func TestXGMEXT_Search(t *testing.T) {
 		}
 	}
 }
+
+func TestXGMEXT_RawSearch(t *testing.T) {
+	tests := []struct {
+		name   string
+		search string
+		want   string
+	}{
+		{
+			name:   "search messages from mailing list",
+			search: "list:info@example.com",
+			want:   "* SEARCH CHARSET UTF-8 X-GM-RAW list:info@example.com\r\n",
+		},
+		{
+			name:   "search for an exact phrase",
+			search: "\"good morning\"",
+			want:   "* SEARCH CHARSET UTF-8 X-GM-RAW \"good morning\"\r\n",
+		},
+		{
+			name:   "group multiple search terms together",
+			search: "subject:(dinner movie)",
+			want:   "* SEARCH CHARSET UTF-8 X-GM-RAW subject:(dinner movie)\r\n",
+		},
+	}
+	for _, test := range tests {
+		cmd := xgmext.NewRawSearch(test.search).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/extensions/xgmext/terms.go b/worker/imap/extensions/xgmext/terms.go
new file mode 100644
index 00000000..b7dcfd3a
--- /dev/null
+++ b/worker/imap/extensions/xgmext/terms.go
@@ -0,0 +1,46 @@
+package xgmext
+
+var Terms = []string{
+	"from:",
+	"to:",
+	"cc:",
+	"bcc:",
+	"subject:",
+	"label:",
+	"deliveredto:",
+	"category:primary",
+	"category:social",
+	"category:promotions",
+	"category:updates",
+	"category:forums",
+	"category:reservations",
+	"category:purchases",
+	"has:",
+	"has:attachment",
+	"has:drive",
+	"has:document",
+	"has:spreadsheet",
+	"has:presentation",
+	"has:youtube",
+	"list:",
+	"filename:",
+	"in:",
+	"is:",
+	"is:important",
+	"is:read",
+	"is:unread",
+	"is:starred",
+	"after:",
+	"before:",
+	"older:",
+	"newer:",
+	"older_than:",
+	"newer_than:",
+	"size:",
+	"larger:",
+	"smaller:",
+	"rfc822msgid:",
+	"OR",
+	"AND",
+	"AROUND",
+}
diff --git a/worker/middleware/gmailworker.go b/worker/middleware/gmailworker.go
index f9924732..4f5f5456 100644
--- a/worker/middleware/gmailworker.go
+++ b/worker/middleware/gmailworker.go
@@ -1,6 +1,8 @@
 package middleware
 
 import (
+	"strconv"
+	"strings"
 	"sync"
 
 	"git.sr.ht/~rjarry/aerc/worker/imap/extensions/xgmext"
@@ -45,20 +47,87 @@ func (g *gmailWorker) reset(c *client.Client) error {
 }
 
 func (g *gmailWorker) ProcessAction(msg types.WorkerMessage) types.WorkerMessage {
-	if msg, ok := msg.(*types.FetchMessageHeaders); ok && len(msg.Uids) > 0 {
-		g.mu.Lock()
-
+	switch msg := msg.(type) {
+	case *types.FetchMessageHeaders:
 		handler := xgmext.NewHandler(g.client)
+
+		g.mu.Lock()
 		uids, err := handler.FetchEntireThreads(msg.Uids)
+		g.mu.Unlock()
 		if err != nil {
-			g.Errorf("failed to fetch entire threads: %v", err)
+			g.Warnf("failed to fetch entire threads: %v", err)
 		}
 
 		if len(uids) > 0 {
 			msg.Uids = uids
 		}
 
+	case *types.FetchDirectoryContents:
+		if msg.Filter == nil || (msg.Filter != nil &&
+			len(msg.Filter.Terms) == 0) {
+			break
+		}
+		if !msg.Filter.UseExtension {
+			g.Debugf("use regular imap filter instead of X-GM-EXT1: " +
+				"extension flag not set")
+			break
+		}
+
+		search := strings.Join(msg.Filter.Terms, " ")
+		g.Debugf("X-GM-EXT1 filter term: '%s'", search)
+
+		handler := xgmext.NewHandler(g.client)
+
+		g.mu.Lock()
+		uids, err := handler.RawSearch(strconv.Quote(search))
+		g.mu.Unlock()
+		if err != nil {
+			g.Errorf("X-GM-EXT1 filter failed: %v", err)
+			g.Warnf("falling back to imap filtering")
+			break
+		}
+
+		g.PostMessage(&types.DirectoryContents{
+			Message: types.RespondTo(msg),
+			Uids:    uids,
+		}, nil)
+
+		g.PostMessage(&types.Done{Message: types.RespondTo(msg)}, nil)
+
+		return &types.Unsupported{}
+
+	case *types.SearchDirectory:
+		if msg.Criteria == nil || (msg.Criteria != nil &&
+			len(msg.Criteria.Terms) == 0) {
+			break
+		}
+		if !msg.Criteria.UseExtension {
+			g.Debugf("use regular imap search instead of X-GM-EXT1: " +
+				"extension flag not set")
+			break
+		}
+
+		search := strings.Join(msg.Criteria.Terms, " ")
+		g.Debugf("X-GM-EXT1 search term: '%s'", search)
+		handler := xgmext.NewHandler(g.client)
+
+		g.mu.Lock()
+		uids, err := handler.RawSearch(strconv.Quote(search))
 		g.mu.Unlock()
+		if err != nil {
+			g.Errorf("X-GM-EXT1 search failed: %v", err)
+			g.Warnf("falling back to regular imap search.")
+			break
+		}
+
+		g.PostMessage(&types.SearchResults{
+			Message: types.RespondTo(msg),
+			Uids:    uids,
+		}, nil)
+
+		g.PostMessage(&types.Done{Message: types.RespondTo(msg)}, nil)
+
+		return &types.Unsupported{}
 	}
 	return g.WorkerInteractor.ProcessAction(msg)
 }