Skip to content

Commit e593322

Browse files
salmonumbrellasalmonumbrella
authored
feat(gmail): add --exclude-labels to watch serve (#194)
1 parent 3463597 commit e593322

File tree

6 files changed

+274
-48
lines changed

6 files changed

+274
-48
lines changed

‎internal/cmd/gmail_watch_cmds.go‎

Lines changed: 40 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -191,20 +191,21 @@ func (c *GmailWatchStopCmd) Run(ctx context.Context, flags *RootFlags) error {
191191
}
192192

193193
type GmailWatchServeCmd struct {
194-
Bind string `name:"bind" help:"Bind address" default:"127.0.0.1"`
195-
Port int `name:"port" help:"Listen port" default:"8788"`
196-
Path string `name:"path" help:"Push handler path" default:"/gmail-pubsub"`
197-
Timezone string `name:"timezone" short:"z" help:"Output timezone (IANA name, e.g. America/New_York, UTC). Default: local"`
198-
Local bool `name:"local" help:"Use local timezone (default behavior, useful to override --timezone)"`
199-
VerifyOIDC bool `name:"verify-oidc" help:"Verify Pub/Sub OIDC tokens"`
200-
OIDCEmail string `name:"oidc-email" help:"Expected service account email"`
201-
OIDCAudience string `name:"oidc-audience" help:"Expected OIDC audience"`
202-
SharedToken string `name:"token" help:"Shared token for x-gog-token or ?token="`
203-
HookURL string `name:"hook-url" help:"Webhook URL to forward messages"`
204-
HookToken string `name:"hook-token" help:"Webhook bearer token"`
205-
IncludeBody bool `name:"include-body" help:"Include text/plain body in hook payload"`
206-
MaxBytes int `name:"max-bytes" help:"Max bytes of body to include" default:"20000"`
207-
SaveHook bool `name:"save-hook" help:"Persist hook settings to watch state"`
194+
Bind string `name:"bind" help:"Bind address" default:"127.0.0.1"`
195+
Port int `name:"port" help:"Listen port" default:"8788"`
196+
Path string `name:"path" help:"Push handler path" default:"/gmail-pubsub"`
197+
Timezone string `name:"timezone" short:"z" help:"Output timezone (IANA name, e.g. America/New_York, UTC). Default: local"`
198+
Local bool `name:"local" help:"Use local timezone (default behavior, useful to override --timezone)"`
199+
VerifyOIDC bool `name:"verify-oidc" help:"Verify Pub/Sub OIDC tokens"`
200+
OIDCEmail string `name:"oidc-email" help:"Expected service account email"`
201+
OIDCAudience string `name:"oidc-audience" help:"Expected OIDC audience"`
202+
SharedToken string `name:"token" help:"Shared token for x-gog-token or ?token="`
203+
HookURL string `name:"hook-url" help:"Webhook URL to forward messages"`
204+
HookToken string `name:"hook-token" help:"Webhook bearer token"`
205+
IncludeBody bool `name:"include-body" help:"Include text/plain body in hook payload"`
206+
MaxBytes int `name:"max-bytes" help:"Max bytes of body to include" default:"20000"`
207+
ExcludeLabels string `name:"exclude-labels" help:"Comma-separated list of Gmail label IDs to exclude from hook payload (e.g. SPAM,TRASH). Set to empty string to disable." default:"SPAM,TRASH"`
208+
SaveHook bool `name:"save-hook" help:"Persist hook settings to watch state"`
208209
}
209210

210211
func (c *GmailWatchServeCmd) Run(ctx context.Context, kctx *kong.Context, flags *RootFlags) error {
@@ -286,21 +287,23 @@ func (c *GmailWatchServeCmd) Run(ctx context.Context, kctx *kong.Context, flags
286287
}
287288

288289
cfg := gmailWatchServeConfig{
289-
Account: account,
290-
Bind: c.Bind,
291-
Port: c.Port,
292-
Path: c.Path,
293-
VerifyOIDC: c.VerifyOIDC,
294-
OIDCEmail: c.OIDCEmail,
295-
OIDCAudience: c.OIDCAudience,
296-
SharedToken: c.SharedToken,
297-
HookTimeout: defaultHookRequestTimeoutSec * time.Second,
298-
HistoryMax: defaultHistoryMaxResults,
299-
ResyncMax: defaultHistoryResyncMax,
300-
AllowNoHook: hook == nil,
301-
IncludeBody: includeBody,
302-
MaxBodyBytes: maxBytes,
303-
DateLocation: loc,
290+
Account: account,
291+
Bind: c.Bind,
292+
Port: c.Port,
293+
Path: c.Path,
294+
VerifyOIDC: c.VerifyOIDC,
295+
OIDCEmail: c.OIDCEmail,
296+
OIDCAudience: c.OIDCAudience,
297+
SharedToken: c.SharedToken,
298+
HookTimeout: defaultHookRequestTimeoutSec * time.Second,
299+
HistoryMax: defaultHistoryMaxResults,
300+
ResyncMax: defaultHistoryResyncMax,
301+
AllowNoHook: hook == nil,
302+
IncludeBody: includeBody,
303+
MaxBodyBytes: maxBytes,
304+
DateLocation: loc,
305+
ExcludeLabels: splitCommaList(c.ExcludeLabels),
306+
VerboseOutput: flags.Verbose,
304307
}
305308
if hook != nil {
306309
cfg.HookURL = hook.URL
@@ -315,13 +318,14 @@ func (c *GmailWatchServeCmd) Run(ctx context.Context, kctx *kong.Context, flags
315318

316319
hookClient := &http.Client{Timeout: cfg.HookTimeout}
317320
server := &gmailWatchServer{
318-
cfg: cfg,
319-
store: store,
320-
validator: validator,
321-
newService: newGmailService,
322-
hookClient: hookClient,
323-
logf: u.Err().Printf,
324-
warnf: u.Err().Printf,
321+
cfg: cfg,
322+
store: store,
323+
validator: validator,
324+
newService: newGmailService,
325+
hookClient: hookClient,
326+
excludeLabelIDs: lowerStringSet(cfg.ExcludeLabels),
327+
logf: u.Err().Printf,
328+
warnf: u.Err().Printf,
325329
}
326330

327331
addr := net.JoinHostPort(c.Bind, strconv.Itoa(c.Port))

‎internal/cmd/gmail_watch_serve_test.go‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,52 @@ func TestGmailWatchServeCmd_DefaultMaxBytes(t *testing.T) {
112112
if got.cfg.MaxBodyBytes != defaultHookMaxBytes {
113113
t.Fatalf("expected default max bytes, got %d", got.cfg.MaxBodyBytes)
114114
}
115+
if len(got.cfg.ExcludeLabels) != 2 || got.cfg.ExcludeLabels[0] != "SPAM" || got.cfg.ExcludeLabels[1] != "TRASH" {
116+
t.Fatalf("unexpected exclude labels: %#v", got.cfg.ExcludeLabels)
117+
}
118+
}
119+
120+
func TestGmailWatchServeCmd_ExcludeLabels_Disable(t *testing.T) {
121+
origListen := listenAndServe
122+
t.Cleanup(func() { listenAndServe = origListen })
123+
124+
home := t.TempDir()
125+
t.Setenv("HOME", home)
126+
127+
store, err := newGmailWatchStore("a@b.com")
128+
if err != nil {
129+
t.Fatalf("store: %v", err)
130+
}
131+
updateErr := store.Update(func(s *gmailWatchState) error {
132+
s.Account = "a@b.com"
133+
return nil
134+
})
135+
if updateErr != nil {
136+
t.Fatalf("seed: %v", updateErr)
137+
}
138+
139+
flags := &RootFlags{Account: "a@b.com"}
140+
var got *gmailWatchServer
141+
listenAndServe = func(srv *http.Server) error {
142+
if gs, ok := srv.Handler.(*gmailWatchServer); ok {
143+
got = gs
144+
}
145+
return nil
146+
}
147+
148+
u, err := ui.New(ui.Options{Stdout: io.Discard, Stderr: io.Discard, Color: "never"})
149+
if err != nil {
150+
t.Fatalf("ui.New: %v", err)
151+
}
152+
if execErr := runKong(t, &GmailWatchServeCmd{}, []string{"--port", "9999", "--path", "/hook", "--exclude-labels", ""}, ui.WithUI(context.Background(), u), flags); execErr != nil {
153+
t.Fatalf("execute: %v", execErr)
154+
}
155+
if got == nil {
156+
t.Fatalf("expected server")
157+
}
158+
if len(got.cfg.ExcludeLabels) != 0 {
159+
t.Fatalf("expected exclude labels disabled, got: %#v", got.cfg.ExcludeLabels)
160+
}
115161
}
116162

117163
func TestGmailWatchServeCmd_SaveHookAndOIDC(t *testing.T) {

‎internal/cmd/gmail_watch_server.go‎

Lines changed: 51 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -26,13 +26,14 @@ const (
2626
)
2727

2828
type gmailWatchServer struct {
29-
cfg gmailWatchServeConfig
30-
store *gmailWatchStore
31-
validator *idtoken.Validator
32-
newService func(context.Context, string) (*gmail.Service, error)
33-
hookClient *http.Client
34-
logf func(string, ...any)
35-
warnf func(string, ...any)
29+
cfg gmailWatchServeConfig
30+
store *gmailWatchStore
31+
validator *idtoken.Validator
32+
newService func(context.Context, string) (*gmail.Service, error)
33+
hookClient *http.Client
34+
excludeLabelIDs map[string]struct{}
35+
logf func(string, ...any)
36+
warnf func(string, ...any)
3637
}
3738

3839
func (s *gmailWatchServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
@@ -188,7 +189,7 @@ func (s *gmailWatchServer) handlePush(ctx context.Context, payload gmailPushPayl
188189
}
189190

190191
messageIDs := collectHistoryMessageIDs(historyResp)
191-
msgs, err := s.fetchMessages(ctx, svc, messageIDs)
192+
msgs, excluded, err := s.fetchMessages(ctx, svc, messageIDs)
192193
if err != nil {
193194
return nil, err
194195
}
@@ -214,6 +215,13 @@ func (s *gmailWatchServer) handlePush(ctx context.Context, payload gmailPushPayl
214215
s.warnf("watch: failed to update state: %v", err)
215216
}
216217

218+
if excluded > 0 && len(msgs) == 0 {
219+
if s.cfg.VerboseOutput {
220+
s.logf("watch: skipping hook; all messages excluded")
221+
}
222+
return nil, errNoNewMessages
223+
}
224+
217225
return &gmailHookPayload{
218226
Source: "gmail",
219227
Account: s.cfg.Account,
@@ -233,7 +241,7 @@ func (s *gmailWatchServer) resyncHistory(ctx context.Context, svc *gmail.Service
233241
ids = append(ids, m.Id)
234242
}
235243
}
236-
msgs, err := s.fetchMessages(ctx, svc, ids)
244+
msgs, excluded, err := s.fetchMessages(ctx, svc, ids)
237245
if err != nil {
238246
return nil, err
239247
}
@@ -255,6 +263,13 @@ func (s *gmailWatchServer) resyncHistory(ctx context.Context, svc *gmail.Service
255263
s.warnf("watch: failed to update state after resync: %v", err)
256264
}
257265

266+
if excluded > 0 && len(msgs) == 0 {
267+
if s.cfg.VerboseOutput {
268+
s.logf("watch: skipping hook; all messages excluded")
269+
}
270+
return nil, errNoNewMessages
271+
}
272+
258273
return &gmailHookPayload{
259274
Source: "gmail",
260275
Account: s.cfg.Account,
@@ -263,8 +278,9 @@ func (s *gmailWatchServer) resyncHistory(ctx context.Context, svc *gmail.Service
263278
}, nil
264279
}
265280

266-
func (s *gmailWatchServer) fetchMessages(ctx context.Context, svc *gmail.Service, ids []string) ([]gmailHookMessage, error) {
281+
func (s *gmailWatchServer) fetchMessages(ctx context.Context, svc *gmail.Service, ids []string) ([]gmailHookMessage, int, error) {
267282
messages := make([]gmailHookMessage, 0, len(ids))
283+
excluded := 0
268284
format := gmailWatchFormatMetadata
269285
if s.cfg.IncludeBody {
270286
format = "full"
@@ -282,11 +298,18 @@ func (s *gmailWatchServer) fetchMessages(ctx context.Context, svc *gmail.Service
282298
if isNotFoundAPIError(err) {
283299
continue
284300
}
285-
return nil, err
301+
return nil, excluded, err
286302
}
287303
if msg == nil {
288304
continue
289305
}
306+
if s.isExcludedLabel(msg.LabelIds) {
307+
excluded++
308+
if s.cfg.VerboseOutput {
309+
s.logf("watch: excluded message %s labels=%v", msg.Id, msg.LabelIds)
310+
}
311+
continue
312+
}
290313
item := gmailHookMessage{
291314
ID: msg.Id,
292315
ThreadID: msg.ThreadId,
@@ -303,7 +326,23 @@ func (s *gmailWatchServer) fetchMessages(ctx context.Context, svc *gmail.Service
303326
}
304327
messages = append(messages, item)
305328
}
306-
return messages, nil
329+
return messages, excluded, nil
330+
}
331+
332+
func (s *gmailWatchServer) isExcludedLabel(labelIDs []string) bool {
333+
if len(labelIDs) == 0 || len(s.excludeLabelIDs) == 0 {
334+
return false
335+
}
336+
for _, id := range labelIDs {
337+
trimmed := strings.TrimSpace(id)
338+
if trimmed == "" {
339+
continue
340+
}
341+
if _, ok := s.excludeLabelIDs[strings.ToLower(trimmed)]; ok {
342+
return true
343+
}
344+
}
345+
return false
307346
}
308347

309348
func (s *gmailWatchServer) sendHook(ctx context.Context, payload *gmailHookPayload) error {

‎internal/cmd/gmail_watch_server_exclude_labels_test.go‎

Lines changed: 118 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,118 @@
1+
package cmd
2+
3+
import (
4+
"bytes"
5+
"context"
6+
"encoding/base64"
7+
"encoding/json"
8+
"net/http"
9+
"net/http/httptest"
10+
"strings"
11+
"testing"
12+
13+
"google.golang.org/api/gmail/v1"
14+
"google.golang.org/api/option"
15+
)
16+
17+
func TestGmailWatchServer_ServeHTTP_ExcludeLabels_SkipsHook(t *testing.T) {
18+
home := t.TempDir()
19+
t.Setenv("HOME", home)
20+
21+
store, err := newGmailWatchStore("a@b.com")
22+
if err != nil {
23+
t.Fatalf("store: %v", err)
24+
}
25+
// Seed state so StartHistoryID returns non-zero.
26+
if updateErr := store.Update(func(s *gmailWatchState) error {
27+
s.Account = "a@b.com"
28+
s.HistoryID = "100"
29+
return nil
30+
}); updateErr != nil {
31+
t.Fatalf("seed: %v", updateErr)
32+
}
33+
34+
var hookCalls int
35+
hookSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
36+
hookCalls++
37+
w.WriteHeader(http.StatusOK)
38+
}))
39+
defer hookSrv.Close()
40+
41+
gmailSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
42+
switch {
43+
case strings.Contains(r.URL.Path, "/gmail/v1/users/me/history"):
44+
w.Header().Set("Content-Type", "application/json")
45+
_ = json.NewEncoder(w).Encode(map[string]any{
46+
"historyId": "200",
47+
"history": []map[string]any{
48+
{"messagesAdded": []map[string]any{{"message": map[string]any{"id": "m1"}}}},
49+
},
50+
})
51+
return
52+
case strings.Contains(r.URL.Path, "/gmail/v1/users/me/messages/m1"):
53+
w.Header().Set("Content-Type", "application/json")
54+
_ = json.NewEncoder(w).Encode(map[string]any{
55+
"id": "m1",
56+
"threadId": "t1",
57+
"snippet": "spam",
58+
"labelIds": []string{"SPAM"},
59+
"payload": map[string]any{
60+
"headers": []map[string]any{
61+
{"name": "Subject", "value": "S"},
62+
},
63+
},
64+
})
65+
return
66+
default:
67+
http.NotFound(w, r)
68+
return
69+
}
70+
}))
71+
defer gmailSrv.Close()
72+
73+
gsvc, err := gmail.NewService(context.Background(),
74+
option.WithoutAuthentication(),
75+
option.WithHTTPClient(gmailSrv.Client()),
76+
option.WithEndpoint(gmailSrv.URL+"/"),
77+
)
78+
if err != nil {
79+
t.Fatalf("NewService: %v", err)
80+
}
81+
82+
s := &gmailWatchServer{
83+
cfg: gmailWatchServeConfig{
84+
Account: "a@b.com",
85+
Path: "/gmail-pubsub",
86+
SharedToken: "tok",
87+
HookURL: hookSrv.URL,
88+
HistoryMax: 100,
89+
ResyncMax: 10,
90+
},
91+
store: store,
92+
newService: func(context.Context, string) (*gmail.Service, error) { return gsvc, nil },
93+
hookClient: hookSrv.Client(),
94+
excludeLabelIDs: map[string]struct{}{"spam": {}},
95+
logf: func(string, ...any) {},
96+
warnf: func(string, ...any) {},
97+
}
98+
99+
push := pubsubPushEnvelope{}
100+
push.Message.Data = base64.StdEncoding.EncodeToString([]byte(`{"emailAddress":"a@b.com","historyId":"200"}`))
101+
body, _ := json.Marshal(push)
102+
103+
req := httptest.NewRequest(http.MethodPost, "/gmail-pubsub?token=tok", bytes.NewReader(body))
104+
rr := httptest.NewRecorder()
105+
s.ServeHTTP(rr, req)
106+
107+
if rr.Code != http.StatusAccepted {
108+
t.Fatalf("status: %d body=%q", rr.Code, rr.Body.String())
109+
}
110+
if hookCalls != 0 {
111+
t.Fatalf("expected no hook calls, got %d", hookCalls)
112+
}
113+
114+
st := store.Get()
115+
if st.HistoryID != "200" {
116+
t.Fatalf("expected history updated, got %q", st.HistoryID)
117+
}
118+
}

0 commit comments

Comments
 (0)