1 files changed, 73 insertions, 4 deletions

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