diff --git a/internal/app/worker/event_add_email.go b/internal/app/worker/event_add_email.go index 55a91d0ac..4051bb4b1 100644 --- a/internal/app/worker/event_add_email.go +++ b/internal/app/worker/event_add_email.go @@ -18,6 +18,7 @@ import ( type mailboxLoad struct { removed atomic.Bool done chan struct{} + prev *mailboxLoad } // HandleAddEmail loads a mailbox off the bus loop. Loading dials the mail @@ -28,13 +29,13 @@ func (w *WorkerService) HandleAddEmail(ctx context.Context, e *models.AddWorkerE return nil } - if w.mailManager.Has(e.ID) { + if _, loading := w.loads.Load(e.ID); !loading { // Already loaded: keep the handler idempotent, but take a changed sync // budget so an operator's settings change reaches this mailbox. - if mail := w.mailManager.Get(e.ID); mail != nil { + if mail := w.mailManager.Get(e.ID); mail != nil && mail.SameExecution(e) { mail.ApplySyncPolicy(e.Sync) + return nil } - return nil } load := &mailboxLoad{done: make(chan struct{})} @@ -45,10 +46,8 @@ func (w *WorkerService) HandleAddEmail(ctx context.Context, e *models.AddWorkerE break } old := v.(*mailboxLoad) - if !old.removed.Load() { - return nil - } - // A removal is pending on the load in flight: this add runs after it rather than being dropped. + // Queue after in-flight loads; AddWMail keeps unchanged adds idempotent. + load.prev = old if w.loads.CompareAndSwap(e.ID, old, load) { prev = old break @@ -73,10 +72,18 @@ func (w *WorkerService) HandleAddEmail(ctx context.Context, e *models.AddWorkerE // loadedMailbox returns a mailbox, waiting on a load of it still in flight so a // command queued behind its ADD_EMAIL finds it as it did when loads were inline. func (w *WorkerService) loadedMailbox(ctx context.Context, id uuid.UUID) (*wmail.WMail, bool) { - if v, ok := w.loads.Load(id); ok { + for { + v, ok := w.loads.Load(id) + if !ok { + break + } select { case <-v.(*mailboxLoad).done: case <-ctx.Done(): + return nil, false + } + if current, exists := w.loads.Load(id); !exists || current == v { + break } } mail := w.mailManager.Get(id) @@ -85,6 +92,10 @@ func (w *WorkerService) loadedMailbox(ctx context.Context, id uuid.UUID) (*wmail // loadMailbox connects one mailbox and starts its sync, or reports why it could not. func (w *WorkerService) loadMailbox(ctx context.Context, e *models.AddWorkerEmail, load *mailboxLoad) { + if load.removed.Load() { + return + } + previous := w.mailManager.Get(e.ID) if err := w.mailManager.AddWMail(ctx, e); err != nil { log.Error().Err(err).Str("email_id", e.ID.String()).Msg("failed to add email account to worker") w.reportLoadFailure(e, err) @@ -95,8 +106,10 @@ func (w *WorkerService) loadMailbox(ctx context.Context, e *models.AddWorkerEmai return } if load.removed.Load() { - mail.Discard() - w.mailManager.Terminate(e.ID) + w.mailManager.TerminateMailbox(e.ID, mail) + return + } + if previous == mail { return } @@ -146,8 +159,5 @@ func (w *WorkerService) reportLoadFailure(e *models.AddWorkerEmail, err error) { // dropMailbox removes a loaded mailbox and stops its work. func (w *WorkerService) dropMailbox(id uuid.UUID, mail *wmail.WMail) { - if mail.Cancel != nil { - mail.Cancel() - } - w.mailManager.Terminate(id) + w.mailManager.TerminateMailbox(id, mail) } diff --git a/internal/app/worker/event_mailbox_identity.go b/internal/app/worker/event_mailbox_identity.go index 80562aacb..68c3013df 100644 --- a/internal/app/worker/event_mailbox_identity.go +++ b/internal/app/worker/event_mailbox_identity.go @@ -58,6 +58,13 @@ func (w *WorkerService) HandleMailboxIdentity(ctx context.Context, data models.E if mail.GoogleData == nil || mail.GoogleData.Client == nil { return reply(models.MailboxIdentityResult{Error: "provider has no send-as list"}) } + done, ok := mail.BeginExecution() + if !ok { + return reply(models.MailboxIdentityResult{Error: "mailbox is not loaded on this worker"}) + } + defer done() + ctx, cancel = mail.ExecutionContext(ctx) + defer cancel() rows, err := mail.GoogleData.Client.ListSendAs(ctx) if err != nil { diff --git a/internal/app/worker/event_message_folder.go b/internal/app/worker/event_message_folder.go index ed882950e..6a0624640 100644 --- a/internal/app/worker/event_message_folder.go +++ b/internal/app/worker/event_message_folder.go @@ -29,6 +29,13 @@ func (w *WorkerService) HandleMessageFolder(ctx context.Context, action models.M log.Warn().Str("email_id", action.EmailID.String()).Msg("Mailbox not loaded; filing not relayed to the provider") return nil } + done, ok := mail.BeginExecution() + if !ok { + return nil + } + defer done() + ctx, cancel := mail.ExecutionContext(ctx) + defer cancel() log.Info(). Str("email_id", action.EmailID.String()). diff --git a/internal/app/worker/event_message_seen.go b/internal/app/worker/event_message_seen.go index fa14e8efd..5eb7073e4 100644 --- a/internal/app/worker/event_message_seen.go +++ b/internal/app/worker/event_message_seen.go @@ -28,6 +28,13 @@ func (w *WorkerService) HandleMessageSeen(ctx context.Context, action models.Mes log.Warn().Str("email_id", action.EmailID.String()).Msg("Mailbox not loaded; read state not relayed to the provider") return nil } + done, ok := mail.BeginExecution() + if !ok { + return nil + } + defer done() + ctx, cancel := mail.ExecutionContext(ctx) + defer cancel() log.Info(). Str("email_id", action.EmailID.String()). diff --git a/internal/app/worker/event_remove_email.go b/internal/app/worker/event_remove_email.go index 57a1bd08f..fab2f7f0d 100644 --- a/internal/app/worker/event_remove_email.go +++ b/internal/app/worker/event_remove_email.go @@ -22,7 +22,9 @@ func (w *WorkerService) HandleRemoveEmail(ctx context.Context, e *models.RemoveW // A load still dialing drops the mailbox when it finishes. if v, ok := w.loads.Load(id); ok { - v.(*mailboxLoad).removed.Store(true) + for load := v.(*mailboxLoad); load != nil; load = load.prev { + load.removed.Store(true) + } } mail := w.mailManager.Get(id) diff --git a/internal/app/worker/event_warmup_action.go b/internal/app/worker/event_warmup_action.go index 99c1c5feb..8cc7ba62f 100644 --- a/internal/app/worker/event_warmup_action.go +++ b/internal/app/worker/event_warmup_action.go @@ -52,6 +52,16 @@ func (w *WorkerService) HandleWarmupAction(ctx context.Context, action models.Wa } mail, exists := w.loadedMailbox(ctx, action.EmailID) + if exists { + if done, ok := mail.BeginExecution(); ok { + defer done() + var cancel func() + ctx, cancel = mail.ExecutionContext(ctx) + defer cancel() + } else { + exists = false + } + } var err error switch { case !exists: diff --git a/internal/app/worker/mailbox_lifecycle_test.go b/internal/app/worker/mailbox_lifecycle_test.go new file mode 100644 index 000000000..20dfcea65 --- /dev/null +++ b/internal/app/worker/mailbox_lifecycle_test.go @@ -0,0 +1,87 @@ +package worker + +import ( + "context" + "sync/atomic" + "testing" + "time" + + "github.com/google/uuid" + "github.com/warmbly/warmbly/internal/app/worker/wmail" + "github.com/warmbly/warmbly/internal/models" +) + +type closedWorkerIMAP struct { + wmail.ImapConn + closes atomic.Int32 +} + +func (c *closedWorkerIMAP) Close() error { c.closes.Add(1); return nil } + +func TestAddEmailReplacesChangedCredentialsAndRemoveClosesInstalledSession(t *testing.T) { + w := newLoadingWorker(&capturedEvents{}) + id := uuid.New() + d := &models.AddWorkerEmail{ID: id, UserID: uuid.New(), Type: models.InboxProviderSMTPIMAP, Email: "a@example.test", SmtpImap: &models.AddWorkerEmailSmtpImapData{Credentials: &models.SmtpImap{SMTP: &models.Service{Host: "smtp.example.test", Port: 587, Password: "old"}, IMAP: &models.Service{Host: "imap.example.test", Port: 993, Password: "old"}}}, Sync: &models.AddWorkerEmailSyncData{State: &models.SyncState{BackfillStatus: models.SyncBackfillComplete, BackfillSynced: 12}}} + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + if err := w.HandleAddEmail(ctx, d); err != nil { + t.Fatal(err) + } + old, ok := w.loadedMailbox(ctx, id) + if !ok { + t.Fatal("initial load failed") + } + conn := &closedWorkerIMAP{} + old.SmtpImapData.ImapClient = conn + for i := 0; i < 3; i++ { + if err := w.HandleAddEmail(ctx, d); err != nil { + t.Fatal(err) + } + } + copyData := *d + copyData.Sync = &models.AddWorkerEmailSyncData{Policy: models.SyncPolicy{DailyMessages: 88}} + if err := w.HandleAddEmail(ctx, ©Data); err != nil { + t.Fatal(err) + } + if got, _ := w.loadedMailbox(ctx, id); got != old || conn.closes.Load() != 0 { + t.Fatal("policy/duplicate ADD_EMAIL retired healthy client") + } + s := *d.SmtpImap + creds := *s.Credentials + smtp := *creds.SMTP + smtp.Password = "new" + creds.SMTP = &smtp + s.Credentials = &creds + copyData.SmtpImap = &s + if err := w.HandleAddEmail(ctx, ©Data); err != nil { + t.Fatal(err) + } + current, ok := w.loadedMailbox(ctx, id) + if !ok || current == old || current.SmtpImapData.SmtpClient.Credentials.Password != "new" || conn.closes.Load() != 1 { + t.Fatal("still-loaded credentials were not replaced and retired") + } + currentConn := &closedWorkerIMAP{} + current.SmtpImapData.ImapClient = currentConn + old.Terminate() + if w.mailManager.Get(id) != current || currentConn.closes.Load() != 0 { + t.Fatal("stale old termination closed replacement") + } + rem := &models.RemoveWorkerEmail{EmailID: id.String()} + for i := 0; i < 2; i++ { + if err := w.HandleRemoveEmail(ctx, rem); err != nil { + t.Fatal(err) + } + } + if w.mailManager.Get(id) != nil || currentConn.closes.Load() != 1 { + t.Fatal("installed removal did not close exactly once") + } + if err := w.HandleAddEmail(ctx, ©Data); err != nil { + t.Fatal(err) + } + if mail, ok := w.loadedMailbox(ctx, id); !ok || mail == current { + t.Fatal("mailbox could not reload after removal") + } + if err := w.HandleRemoveEmail(ctx, rem); err != nil { + t.Fatal(err) + } +} diff --git a/internal/app/worker/mailmanager/add_email.go b/internal/app/worker/mailmanager/add_email.go index cafcd680c..ecdb0b323 100644 --- a/internal/app/worker/mailmanager/add_email.go +++ b/internal/app/worker/mailmanager/add_email.go @@ -15,6 +15,18 @@ func (m *MailManager) AddWMail( ctx context.Context, data *models.AddWorkerEmail, ) error { + dataCopy := *data + data = &dataCopy + old := m.Get(data.ID) + if old != nil { + if old.SameExecution(data) { + old.ApplySyncPolicy(data.Sync) + return nil + } + old.Stop() + old.Drain() + old.ResumeInto(data) + } // Cfg is avro-excluded from the payload, so rebuild it from the worker's // local oauth config for token refresh (no-op for smtp_imap). data.Cfg = m.cfgFor(data.Type) @@ -25,11 +37,12 @@ func (m *MailManager) AddWMail( data.TokenSource = m.tokenBroker.Source(data.ID) } + var newMail *wmail.WMail newMail, err := wmail.NewWMail( data, m.OnEvent, func() { - m.Terminate(data.ID) + m.TerminateMailbox(data.ID, newMail) }, m.cache, m.storage, @@ -47,13 +60,15 @@ func (m *MailManager) AddWMail( // Built before the lock: NewWMail dials the server, and holding the lock // through that stalled every send and lookup on this worker. m.Lock() - defer m.Unlock() - if _, loaded := m.Emails[data.ID]; loaded { + if current := m.Emails[data.ID]; current != old { + m.Unlock() newMail.Discard() return nil } + newMail.PreservePending(old) m.Emails[data.ID] = newMail newMail.ExecutorID = m.ExecutorID + m.Unlock() return nil } diff --git a/internal/app/worker/mailmanager/lifecycle_test.go b/internal/app/worker/mailmanager/lifecycle_test.go new file mode 100644 index 000000000..509dba95b --- /dev/null +++ b/internal/app/worker/mailmanager/lifecycle_test.go @@ -0,0 +1,172 @@ +package mailmanager + +import ( + "context" + "sync/atomic" + "testing" + "time" + + "github.com/google/uuid" + "github.com/warmbly/warmbly/internal/app/worker/wmail" + "github.com/warmbly/warmbly/internal/models" + "golang.org/x/oauth2" +) + +type countedClient struct { + wmail.ImapConn + closes atomic.Int32 +} + +func (c *countedClient) Close() error { c.closes.Add(1); return nil } + +func smtpPayload() *models.AddWorkerEmail { + return &models.AddWorkerEmail{ID: uuid.New(), UserID: uuid.New(), Type: models.InboxProviderSMTPIMAP, Email: "a@example.test", SmtpImap: &models.AddWorkerEmailSmtpImapData{Credentials: &models.SmtpImap{SMTP: &models.Service{Host: "smtp.example.test", Port: 587, Password: "old"}, IMAP: &models.Service{Host: "imap.example.test", Port: 993, Password: "old"}}, Mailboxes: []models.Mailbox{{Name: "Inbox", UIDValidity: 5, UIDNext: 25}}}, Sync: &models.AddWorkerEmailSyncData{State: &models.SyncState{BackfillStatus: models.SyncBackfillComplete, BackfillSynced: 7}}} +} + +func TestMailboxReplacementDrainsClosesAndPreservesCursors(t *testing.T) { + m := NewMailManager(func(models.JobEventType, string, any) error { return nil }, nil, nil, nil, nil, nil, nil) + d := smtpPayload() + if err := m.AddWMail(t.Context(), d); err != nil { + t.Fatal(err) + } + old := m.Get(d.ID) + c := &countedClient{} + old.SmtpImapData.ImapClient = c + for i := 0; i < 3; i++ { + if err := m.AddWMail(t.Context(), d); err != nil { + t.Fatal(err) + } + } + if m.Get(d.ID) != old || c.closes.Load() != 0 { + t.Fatal("unchanged add restarted client") + } + copyData := *d + data := ©Data + s := *d.SmtpImap + creds := *s.Credentials + smtp := *creds.SMTP + smtp.Password = "replacement" + creds.SMTP = &smtp + s.Credentials = &creds + data.SmtpImap = &s + done, ok := old.BeginExecution() + if !ok { + t.Fatal("live executor not admitted") + } + result := make(chan error, 1) + go func() { result <- m.AddWMail(context.Background(), data) }() + select { + case <-old.Ctx.Done(): + case <-time.After(time.Second): + t.Fatal("replacement did not cancel old executor") + } + if m.Get(d.ID) != old { + t.Fatal("new executor installed before old operation drained") + } + done() + select { + case err := <-result: + if err != nil { + t.Fatal(err) + } + case <-time.After(time.Second): + t.Fatal("replacement stuck") + } + next := m.Get(d.ID) + if next == nil || next == old || next.SmtpImapData.SmtpClient.Credentials.Password != "replacement" || c.closes.Load() != 1 { + t.Fatal("replacement credentials/cleanup not installed") + } + if len(next.SmtpImapData.Mailboxes) != 1 || next.SmtpImapData.Mailboxes[0].UIDNext != 25 { + t.Fatal("disabled sync lost saved cursor") + } + old.Terminate() + if m.Get(d.ID) != next { + t.Fatal("old auth callback retired replacement") + } + m.Terminate(d.ID) + m.Terminate(d.ID) + if next.Ctx.Err() == nil || c.closes.Load() != 1 { + t.Fatal("removal did not stop session idempotently") + } +} + +func TestMailboxReplacementLocalOAuthAndPolicy(t *testing.T) { + for _, provider := range []models.InboxProvider{models.InboxProviderGoogle, models.InboxProviderOutlook} { + t.Run(string(provider), func(t *testing.T) { + m := NewMailManager(func(models.JobEventType, string, any) error { return nil }, nil, nil, nil, nil, nil, nil) + d := &models.AddWorkerEmail{ID: uuid.New(), UserID: uuid.New(), Type: provider, Google: &models.AddWorkerEmailGoogleData{Token: &oauth2.Token{AccessToken: "old", RefreshToken: "old-refresh"}, LastHistoryID: 19}, Graph: &models.AddWorkerEmailGraphData{Token: &oauth2.Token{AccessToken: "old", RefreshToken: "old-refresh"}, DeltaLinks: map[string]string{"Inbox": "live"}}} + if err := m.AddWMail(t.Context(), d); err != nil { + t.Fatal(err) + } + old := m.Get(d.ID) + copyData := *d + copyData.Sync = &models.AddWorkerEmailSyncData{Policy: models.SyncPolicy{DailyMessages: 90}} + if err := m.AddWMail(t.Context(), ©Data); err != nil { + t.Fatal(err) + } + if m.Get(d.ID) != old { + t.Fatal("budget change rebuilt provider") + } + g := *d.Google + g.Token = &oauth2.Token{AccessToken: "new", RefreshToken: "new-refresh"} + copyData.Google = &g + graph := *d.Graph + graph.Token = g.Token + copyData.Graph = &graph + if err := m.AddWMail(t.Context(), ©Data); err != nil { + t.Fatal(err) + } + next := m.Get(d.ID) + if next == old || old.Ctx.Err() == nil || !next.SameExecution(©Data) { + t.Fatal("reauthorization kept local OAuth client") + } + if next.GoogleData != nil && next.GoogleData.LastHistoryID != 19 { + t.Fatal("history cursor lost") + } + if next.GraphData != nil && next.GraphData.Client.DeltaLinks["Inbox"] != "live" { + t.Fatal("delta cursor lost") + } + m.Terminate(d.ID) + if err := m.AddWMail(t.Context(), ©Data); err != nil { + t.Fatal(err) + } + if m.Get(d.ID) == nil || m.Get(d.ID) == next { + t.Fatal("reload after removal failed") + } + m.Terminate(d.ID) + }) + } +} + +func TestMailboxReplacementEndpointSkipAndSentSettings(t *testing.T) { + for _, field := range []string{"imap-password", "smtp-host", "smtp-security", "skip-folders", "sent-copy"} { + t.Run(field, func(t *testing.T) { + m := NewMailManager(func(models.JobEventType, string, any) error { return nil }, nil, nil, nil, nil, nil, nil) + d := smtpPayload() + if err := m.AddWMail(t.Context(), d); err != nil { + t.Fatal(err) + } + old := m.Get(d.ID) + switch field { + case "imap-password": + d.SmtpImap.Credentials.IMAP.Password = "new" + case "smtp-host": + d.SmtpImap.Credentials.SMTP.Host = "new.example.test" + case "smtp-security": + d.SmtpImap.Credentials.SMTP.Security = models.MailSecurityTLS + case "skip-folders": + d.Sync.Policy.SkipFolders = []string{"Custom"} + case "sent-copy": + v := false + d.SaveToSent = &v + } + if err := m.AddWMail(t.Context(), d); err != nil { + t.Fatal(err) + } + if m.Get(d.ID) == old { + t.Fatalf("changed %s ignored", field) + } + m.Terminate(d.ID) + }) + } +} diff --git a/internal/app/worker/mailmanager/terminate.go b/internal/app/worker/mailmanager/terminate.go index dbce9d67d..aaa67d726 100644 --- a/internal/app/worker/mailmanager/terminate.go +++ b/internal/app/worker/mailmanager/terminate.go @@ -1,10 +1,28 @@ package mailmanager -import "github.com/google/uuid" +import ( + "github.com/google/uuid" + "github.com/warmbly/warmbly/internal/app/worker/wmail" +) func (m *MailManager) Terminate(id uuid.UUID) { m.Lock() - defer m.Unlock() - + mail := m.Emails[id] delete(m.Emails, id) + m.Unlock() + if mail != nil { + mail.Stop() + } +} + +// TerminateMailbox cannot retire a newer generation through a stale callback. +func (m *MailManager) TerminateMailbox(id uuid.UUID, mail *wmail.WMail) { + m.Lock() + if m.Emails[id] == mail { + delete(m.Emails, id) + } + m.Unlock() + if mail != nil { + mail.Stop() + } } diff --git a/internal/app/worker/wmail/err.go b/internal/app/worker/wmail/err.go index 2b01514e6..428e89ab7 100644 --- a/internal/app/worker/wmail/err.go +++ b/internal/app/worker/wmail/err.go @@ -37,6 +37,12 @@ func reportable(err error) bool { } func (w *WMail) CaptureError(err error) { + w.lifecycleMu.Lock() + stopped := w.stopped + w.lifecycleMu.Unlock() + if stopped { + return + } if reportable(err) { errs.CaptureException(err, errs.Tag("user_id", w.UserID.String()), @@ -71,18 +77,19 @@ func (w *WMail) CaptureError(err error) { Timestamp: time.Now().Unix(), } + w.lifecycleMu.Lock() + if w.stopped { + w.lifecycleMu.Unlock() + return + } _ = w.onEvent(eventType, errorEvent) + w.lifecycleMu.Unlock() // Critical errors should stop the sync loop and remove the account from the // worker's local state until the user re-authenticates. if eventType == models.JobEventTypeEmailAuthError || eventType == models.JobEventTypeEmailDisabled { - if w.Cancel != nil { - w.Cancel() - } - if w.TerminateFunc != nil { - w.TerminateFunc() - } + w.Terminate() } } diff --git a/internal/app/worker/wmail/err_mail.go b/internal/app/worker/wmail/err_mail.go index b304d3442..5dfd4e895 100644 --- a/internal/app/worker/wmail/err_mail.go +++ b/internal/app/worker/wmail/err_mail.go @@ -35,6 +35,8 @@ func (w *WMail) RunMailError(f func() *errx.MailError) { } func (w *WMail) Terminate() { - w.Cancel() - w.TerminateFunc() + w.Stop() + if w.TerminateFunc != nil { + w.TerminateFunc() + } } diff --git a/internal/app/worker/wmail/event_token.go b/internal/app/worker/wmail/event_token.go index 642240e9f..8e699d08f 100644 --- a/internal/app/worker/wmail/event_token.go +++ b/internal/app/worker/wmail/event_token.go @@ -6,6 +6,21 @@ import ( ) func (w *WMail) onTokenUpdate(token *oauth2.Token) error { + w.lifecycleMu.Lock() + defer w.lifecycleMu.Unlock() + if w.stopped { + return nil + } + if data := w.executionData; data != nil && !data.Brokered { + previous := ExecutionKey(data) + if data.Google != nil { + data.Google.Token = token + } + if data.Graph != nil { + data.Graph.Token = token + } + w.executionKeys = map[[32]byte]struct{}{w.initialExecutionKey: {}, previous: {}, ExecutionKey(data): {}} + } return w.onEvent(models.JobEventTypeTokenUpdate, &models.JobEventTokenUpdate{ UserID: w.UserID, EmailID: w.ID, diff --git a/internal/app/worker/wmail/governor.go b/internal/app/worker/wmail/governor.go index 11d754fe1..817703078 100644 --- a/internal/app/worker/wmail/governor.go +++ b/internal/app/worker/wmail/governor.go @@ -304,10 +304,5 @@ func (w *WMail) escalate(mailErr *errx.MailError) { ActionRequired: userInfo.ActionRequired, Timestamp: time.Now().Unix(), }) - if w.Cancel != nil { - w.Cancel() - } - if w.TerminateFunc != nil { - w.TerminateFunc() - } + w.Terminate() } diff --git a/internal/app/worker/wmail/lifecycle.go b/internal/app/worker/wmail/lifecycle.go new file mode 100644 index 000000000..acfa8eaf6 --- /dev/null +++ b/internal/app/worker/wmail/lifecycle.go @@ -0,0 +1,189 @@ +package wmail + +import ( + "context" + "crypto/sha256" + "encoding/json" + "slices" + "strings" + + "github.com/google/uuid" + "github.com/warmbly/warmbly/internal/models" + "golang.org/x/oauth2" +) + +// ExecutionKey excludes checkpoints and policy budgets from client identity. +func ExecutionKey(data *models.AddWorkerEmail) [32]byte { + type credential struct{ Access, Refresh string } + tokenKey := func(t *oauth2.Token) credential { + if t == nil { + return credential{} + } + return credential{t.AccessToken, t.RefreshToken} + } + key := struct { + ID, User uuid.UUID + Org *uuid.UUID + Provider models.InboxProvider + Email, First, Last, GraphUser string + Brokered, ImapSync, SaveToSent bool + Token credential + SMTP, IMAP *models.Service + Skip []string + }{ID: data.ID, User: data.UserID, Org: data.OrganizationID, Provider: data.Type, + Email: data.Email, First: data.FirstName, Last: data.LastName, Brokered: data.Brokered, + ImapSync: data.Type == models.InboxProviderSMTPIMAP && data.ImapSync, SaveToSent: data.Type == models.InboxProviderSMTPIMAP && data.SavesSentCopy()} + if data.Type == models.InboxProviderGoogle && data.Google != nil && !data.Brokered { + key.Token = tokenKey(data.Google.Token) + } + if data.Type == models.InboxProviderOutlook && data.Graph != nil { + key.GraphUser = data.Graph.User + if !data.Brokered { + key.Token = tokenKey(data.Graph.Token) + } + } + if data.Type == models.InboxProviderSMTPIMAP && data.SmtpImap != nil && data.SmtpImap.Credentials != nil { + key.SMTP, key.IMAP = data.SmtpImap.Credentials.SMTP, data.SmtpImap.Credentials.IMAP + } + if data.Sync != nil && data.Type == models.InboxProviderSMTPIMAP { + for _, folder := range data.Sync.Policy.SkipFolders { + if folder = strings.ToLower(strings.TrimSpace(folder)); folder != "" { + key.Skip = append(key.Skip, folder) + } + } + slices.Sort(key.Skip) + key.Skip = slices.Compact(key.Skip) + } + b, _ := json.Marshal(key) + return sha256.Sum256(b) +} + +func (w *WMail) rememberExecution(data *models.AddWorkerEmail) { + // Own the token pointers; ADD_EMAIL payloads may be reused by a publisher. + copyData := *data + if data.Google != nil { + g := *data.Google + copyData.Google = &g + } + if data.Graph != nil { + g := *data.Graph + copyData.Graph = &g + } + w.executionData = ©Data + w.initialExecutionKey = ExecutionKey(data) + w.executionKeys = map[[32]byte]struct{}{w.initialExecutionKey: {}} +} + +func (w *WMail) SameExecution(data *models.AddWorkerEmail) bool { + w.lifecycleMu.Lock() + defer w.lifecycleMu.Unlock() + _, ok := w.executionKeys[ExecutionKey(data)] + return ok && !w.stopped +} + +// BeginExecution rejects stale references and lets replacements drain admitted work. +func (w *WMail) BeginExecution() (func(), bool) { + w.lifecycleMu.Lock() + defer w.lifecycleMu.Unlock() + if w.stopped { + return nil, false + } + w.active.Add(1) + return w.active.Done, true +} + +// Stop cancels scheduling and closes the provider without waiting on its caller. +func (w *WMail) Stop() { + w.lifecycleMu.Lock() + if w.stopped { + w.lifecycleMu.Unlock() + return + } + w.stopped = true + w.stopDone = make(chan struct{}) + done := w.stopDone + w.lifecycleMu.Unlock() + defer close(done) + if w.Cancel != nil { + w.Cancel() + } + if w.SmtpImapData != nil && w.SmtpImapData.ImapClient != nil { + if c, ok := w.SmtpImapData.ImapClient.(interface{ Close() error }); ok { + _ = c.Close() + } + } +} + +// Drain is called outside mailbox execution, after Stop prevents new work. +func (w *WMail) Drain() { + w.lifecycleMu.Lock() + done := w.stopDone + w.lifecycleMu.Unlock() + if done != nil { + <-done + } + w.active.Wait() +} + +// ResumeInto uses the drained live state, which can be newer than ADD_EMAIL. +func (w *WMail) ResumeInto(data *models.AddWorkerEmail) { + if w.UserID != data.UserID || w.EmailType != data.Type { + return + } + if w.tracker != nil { + if data.Sync == nil { + data.Sync = &models.AddWorkerEmailSyncData{} + } else { + s := *data.Sync + data.Sync = &s + } + state := w.tracker.state + state.BackfillCursor.Folders = make(map[string]models.SyncFolderCursor, len(w.tracker.state.BackfillCursor.Folders)) + for k, v := range w.tracker.state.BackfillCursor.Folders { + state.BackfillCursor.Folders[k] = v + } + data.Sync.State = &state + } + if w.GoogleData != nil && data.Google != nil { + g := *data.Google + g.LastHistoryID = w.GoogleData.LastHistoryID + data.Google = &g + } + if w.GraphData != nil && data.Graph != nil { + g := *data.Graph + g.DeltaLinks = cloneStringMap(w.GraphData.Client.DeltaLinks) + data.Graph = &g + } + if w.SmtpImapData != nil && data.SmtpImap != nil { + s := *data.SmtpImap + s.Mailboxes = nil + for _, box := range w.SmtpImapData.Mailboxes { + if box != nil { + s.Mailboxes = append(s.Mailboxes, *box) + } + } + data.SmtpImap = &s + } +} + +// PreservePending retains session-local retry work during same-worker replacement. +func (w *WMail) PreservePending(old *WMail) { + if old == nil || old.UserID != w.UserID || old.EmailType != w.EmailType { + return + } + w.unmapPending = old.unmapPending + w.diagnosticRetries = old.diagnosticRetries + if old.tracker != nil && w.tracker != nil { + w.tracker.dirty = old.tracker.dirty + w.tracker.lastSent = old.tracker.lastSent + } +} + +func (w *WMail) ExecutionContext(ctx context.Context) (context.Context, func()) { + ctx, cancel := context.WithCancel(ctx) + if w.Ctx == nil { + return ctx, cancel + } + stop := context.AfterFunc(w.Ctx, cancel) + return ctx, func() { stop(); cancel() } +} diff --git a/internal/app/worker/wmail/lifecycle_test.go b/internal/app/worker/wmail/lifecycle_test.go new file mode 100644 index 000000000..798740844 --- /dev/null +++ b/internal/app/worker/wmail/lifecycle_test.go @@ -0,0 +1,221 @@ +package wmail + +import ( + "context" + "net" + "sync/atomic" + "testing" + "time" + + goimap "github.com/emersion/go-imap/v2" + "github.com/emersion/go-imap/v2/imapserver" + "github.com/emersion/go-imap/v2/imapserver/imapmemserver" + "github.com/google/uuid" + "github.com/warmbly/warmbly/internal/client/msgraph" + "github.com/warmbly/warmbly/internal/errx" + "github.com/warmbly/warmbly/internal/models" + "golang.org/x/oauth2" +) + +type closingIMAP struct { + ImapConn + closes atomic.Int32 +} + +type observedIMAPSession struct { + imapserver.Session + closed chan struct{} +} + +func (s *observedIMAPSession) Close() error { defer close(s.closed); return s.Session.Close() } + +func TestMailboxLifecycleStopClosesAuthenticatedIdleSocket(t *testing.T) { + t.Setenv("DEPLOYMENT_MODE", "self_hosted") + mem := imapmemserver.New() + user := imapmemserver.NewUser("mailbox@example.test", "local-test") + if err := user.Create("INBOX", nil); err != nil { + t.Fatal(err) + } + mem.AddUser(user) + closed := make(chan struct{}) + srv := imapserver.New(&imapserver.Options{Caps: goimap.CapSet{goimap.CapIMAP4rev1: {}}, InsecureAuth: true, NewSession: func(*imapserver.Conn) (imapserver.Session, *imapserver.GreetingData, error) { + return &observedIMAPSession{Session: mem.NewSession(), closed: closed}, nil, nil + }}) + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + go func() { _ = srv.Serve(ln) }() + defer srv.Close() + d := &models.AddWorkerEmail{ID: uuid.New(), UserID: uuid.New(), Type: models.InboxProviderSMTPIMAP, Email: "mailbox@example.test", ImapSync: true, SmtpImap: &models.AddWorkerEmailSmtpImapData{Credentials: &models.SmtpImap{IMAP: &models.Service{Host: "127.0.0.1", Port: ln.Addr().(*net.TCPAddr).Port, Username: "mailbox@example.test", Password: "local-test", Security: models.MailSecurityNone}}}} + w, merr := NewWMail(d, func(models.JobEventType, string, any) error { return nil }, nil, nil, nil, nil, nil, nil) + if merr != nil { + t.Fatal(merr) + } + w.Stop() + w.Stop() + select { + case <-closed: + case <-time.After(time.Second): + t.Fatal("authenticated idle server session retained after stop") + } +} + +func (c *closingIMAP) Close() error { c.closes.Add(1); return nil } + +func TestMailboxLifecycleStopDrainsAndRejectsStaleWork(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + c := &closingIMAP{} + w := &WMail{Ctx: ctx, Cancel: cancel, SmtpImapData: &SmtpImapData{ImapClient: c}} + done, ok := w.BeginExecution() + if !ok { + t.Fatal("live execution rejected") + } + w.Stop() + w.Stop() + if c.closes.Load() != 1 || ctx.Err() == nil { + t.Fatal("stop must cancel and close once") + } + if _, ok := w.BeginExecution(); ok { + t.Fatal("retired execution admitted") + } + drained := make(chan struct{}) + go func() { w.Drain(); close(drained) }() + select { + case <-drained: + t.Fatal("drained before admitted work returned") + default: + } + done() + select { + case <-drained: + case <-time.After(time.Second): + t.Fatal("drain did not finish") + } +} + +func TestMailboxLifecycleExecutionKeyAndTokenRefresh(t *testing.T) { + token := &oauth2.Token{AccessToken: "old", RefreshToken: "refresh", Expiry: time.Now().Add(time.Hour)} + d := &models.AddWorkerEmail{ID: uuid.New(), UserID: uuid.New(), Type: models.InboxProviderGoogle, Google: &models.AddWorkerEmailGoogleData{Token: token}} + w := &WMail{onEvent: func(models.JobEventType, any) error { return nil }} + w.rememberExecution(d) + changed := *d + changed.Sync = &models.AddWorkerEmailSyncData{Policy: models.SyncPolicy{DailyMessages: 99}, State: &models.SyncState{BackfillStatus: models.SyncBackfillComplete}} + if !w.SameExecution(&changed) { + t.Fatal("policy/checkpoint publication restarted provider") + } + g := *d.Google + changed.Google = &g + g.Token = &oauth2.Token{AccessToken: "fresh", RefreshToken: "rotated", Expiry: time.Now().Add(time.Hour)} + if w.SameExecution(&changed) { + t.Fatal("local reauthorization ignored") + } + if err := w.onTokenUpdate(g.Token); err != nil { + t.Fatal(err) + } + if !w.SameExecution(&changed) || !w.SameExecution(d) { + t.Fatal("routine token refresh/repeated original publication restarted provider") + } + for i := 0; i < 10; i++ { + if err := w.onTokenUpdate(&oauth2.Token{AccessToken: string(rune('a' + i)), RefreshToken: "rotated"}); err != nil { + t.Fatal(err) + } + } + if len(w.executionKeys) > 3 { + t.Fatal("refresh aliases grew without bound") + } + brokered := *d + brokered.Brokered = true + w.rememberExecution(&brokered) + changed.Brokered = true + if !w.SameExecution(&changed) { + t.Fatal("broker-owned token rotation restarted provider") + } + w.Stop() + w.onEvent = func(models.JobEventType, any) error { t.Fatal("stale token callback published"); return nil } + if err := w.onTokenUpdate(token); err != nil { + t.Fatal(err) + } +} + +func TestMailboxLifecycleLegacyDefaultsAndNormalizedSkipSettings(t *testing.T) { + d := &models.AddWorkerEmail{ID: uuid.New(), UserID: uuid.New(), Type: models.InboxProviderSMTPIMAP} + w := &WMail{} + w.rememberExecution(d) + copyData := *d + v := true + copyData.SaveToSent = &v + copyData.Sync = &models.AddWorkerEmailSyncData{Policy: models.SyncPolicy{DailyMessages: 99}} + if !w.SameExecution(©Data) { + t.Fatal("old backend missing sync/sent fields differs from effective defaults") + } + copyData.Sync.Policy.SkipFolders = []string{" Custom ", "CUSTOM", ""} + w.rememberExecution(©Data) + copyData.Sync = &models.AddWorkerEmailSyncData{Policy: models.SyncPolicy{SkipFolders: []string{"custom"}, DailyMessages: 120}} + if !w.SameExecution(©Data) { + t.Fatal("same effective folder skips restarted client") + } +} + +func TestMailboxLifecycleResumePreservesCompletedStateAndProviderCursors(t *testing.T) { + id, user := uuid.New(), uuid.New() + stamp := time.Now().Add(-time.Hour) + state := models.SyncState{BackfillStatus: models.SyncBackfillComplete, BackfillSynced: 12, BackfillCompletedAt: &stamp, LastSyncedAt: &stamp, Deferred: 5, BackfillCursor: models.SyncCursor{PageToken: "page", Folders: map[string]models.SyncFolderCursor{"Inbox": {UID: 8, Done: true}}}} + for _, provider := range []models.InboxProvider{models.InboxProviderGoogle, models.InboxProviderOutlook, models.InboxProviderSMTPIMAP} { + t.Run(string(provider), func(t *testing.T) { + old := &WMail{ID: id, UserID: user, EmailType: provider, tracker: newSyncTracker(&state, func(models.SyncState) error { return nil }), unmapPending: map[string]uuid.UUID{"pending": uuid.New()}} + d := &models.AddWorkerEmail{ID: id, UserID: user, Type: provider} + switch provider { + case models.InboxProviderGoogle: + old.GoogleData = &GoogleData{LastHistoryID: 91} + d.Google = &models.AddWorkerEmailGoogleData{LastHistoryID: 1} + case models.InboxProviderOutlook: + old.GraphData = &GraphData{Client: &msgraph.Client{DeltaLinks: map[string]string{"Inbox": "live"}}} + d.Graph = &models.AddWorkerEmailGraphData{DeltaLinks: map[string]string{"Inbox": "stale"}} + case models.InboxProviderSMTPIMAP: + old.SmtpImapData = &SmtpImapData{Mailboxes: []*models.Mailbox{{Name: "Inbox", UIDValidity: 4, UIDNext: 17}}} + d.SmtpImap = &models.AddWorkerEmailSmtpImapData{} + } + old.Stop() + old.Drain() + old.ResumeInto(d) + if d.Sync.State.BackfillStatus != models.SyncBackfillComplete || d.Sync.State.BackfillSynced != 12 || d.Sync.State.Deferred != 5 || d.Sync.State.LastSyncedAt != &stamp || d.Sync.State.BackfillCursor.PageToken != "page" { + t.Fatal("replacement reset completed/live state") + } + d.Sync.State.BackfillCursor.Folders["Inbox"] = models.SyncFolderCursor{} + if old.tracker.state.BackfillCursor.Folders["Inbox"].UID != 8 { + t.Fatal("replacement aliases retired mutable state") + } + if d.Google != nil && d.Google.LastHistoryID != 91 { + t.Fatal("Gmail cursor lost") + } + if d.Graph != nil && d.Graph.DeltaLinks["Inbox"] != "live" { + t.Fatal("Graph cursor lost") + } + if d.SmtpImap != nil && (len(d.SmtpImap.Mailboxes) != 1 || d.SmtpImap.Mailboxes[0].UIDNext != 17) { + t.Fatal("IMAP cursor lost") + } + next := &WMail{UserID: user, EmailType: provider, tracker: newSyncTracker(d.Sync.State, func(models.SyncState) error { return nil })} + old.tracker.dirty = true + next.PreservePending(old) + if next.unmapPending["pending"] != old.unmapPending["pending"] || !next.tracker.dirty { + t.Fatal("pending retries/state flush lost") + } + }) + } +} + +func TestMailboxLifecycleAuthTerminationClosesOnceAndIgnoresRetiredErrors(t *testing.T) { + c := &closingIMAP{} + events := 0 + w := &WMail{SmtpImapData: &SmtpImapData{ImapClient: c}, onEvent: func(models.JobEventType, any) error { events++; return nil }} + w.CaptureError(errx.ErrMailServerUnreachable) + if c.closes.Load() != 0 { + t.Fatal("transport timeout permanently retired mailbox") + } + w.CaptureError(&errx.MailError{Type: errx.MailErrorCritical, Code: errx.MailErrorCodeInvalidCredentials}) + w.CaptureError(&errx.MailError{Type: errx.MailErrorCritical, Code: errx.MailErrorCodeInvalidCredentials}) + if c.closes.Load() != 1 || events != 2 { + t.Fatalf("closes=%d events=%d", c.closes.Load(), events) + } +} diff --git a/internal/app/worker/wmail/provider_deadline_test.go b/internal/app/worker/wmail/provider_deadline_test.go new file mode 100644 index 000000000..50ec75bdc --- /dev/null +++ b/internal/app/worker/wmail/provider_deadline_test.go @@ -0,0 +1,91 @@ +package wmail + +import ( + "context" + "fmt" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + "time" + + "github.com/warmbly/warmbly/internal/errx" + "github.com/warmbly/warmbly/internal/models" + "golang.org/x/oauth2" +) + +func TestMailboxProviderDeadlineReleasesPassAndAllowsRetry(t *testing.T) { + for _, provider := range []models.InboxProvider{models.InboxProviderGoogle, models.InboxProviderOutlook} { + for _, body := range []bool{false, true} { + t.Run(fmt.Sprintf("%s/body=%v", provider, body), func(t *testing.T) { + var stalled atomic.Bool + stalled.Store(true) + release := make(chan struct{}) + srv := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) { + rw.Header().Set("Content-Type", "application/json") + if stalled.Load() { + if body { + rw.WriteHeader(http.StatusOK) + _, _ = rw.Write([]byte(`{"`)) + rw.(http.Flusher).Flush() + } + select { + case <-release: + case <-r.Context().Done(): + } + return + } + if provider == models.InboxProviderGoogle { + _, _ = rw.Write([]byte(`{"historyId":"91","history":[]}`)) + } else { + _, _ = fmt.Fprintf(rw, `{"value":[],"@odata.deltaLink":%q}`, "https://graph.microsoft.com"+r.URL.Path+"?next=live") + } + })) + defer srv.Close() + defer close(release) + var events []captured + var w *WMail + ctx := context.WithValue(context.Background(), oauth2.HTTPClient, &http.Client{Timeout: 40 * time.Millisecond, Transport: rewriteToTestServer(srv.URL)}) + token := &oauth2.Token{AccessToken: "test", Expiry: time.Now().Add(time.Hour)} + if provider == models.InboxProviderGoogle { + w = newGoogleTestMail(t, srv, &events) + w.GoogleData.LastHistoryID = 91 + if err := w.GoogleData.Client.Init(ctx, token, oauth2.Config{}); err != nil { + t.Fatal(err) + } + w.googleReconciledAt = time.Now() + } else { + var states []models.SyncState + w = newGraphTestMail(t, srv, &events, &states) + if err := w.GraphData.Client.Init(ctx, token, oauth2.Config{}); err != nil { + t.Fatal(err) + } + } + w.tracker.state.BackfillStatus = models.SyncBackfillComplete + mailCtx, cancel := context.WithCancel(context.Background()) + defer cancel() + w.Ctx, w.Cancel = mailCtx, cancel + result := make(chan *errx.MailError, 1) + go func() { result <- w.SyncMail(context.Background()) }() + select { + case err := <-result: + if err == nil { + t.Fatal("stalled provider pass reported success") + } + case <-time.After(time.Second): + stalled.Store(false) + srv.CloseClientConnections() + <-result + t.Fatal("provider request deadline did not release sync pass") + } + if ctx.Err() != nil || w.Ctx.Err() != nil { + t.Fatal("request timeout cancelled mailbox lifetime") + } + stalled.Store(false) + if err := w.SyncMail(context.Background()); err != nil { + t.Fatalf("later retry without reconnect failed: %v", err) + } + }) + } + } +} diff --git a/internal/app/worker/wmail/send.go b/internal/app/worker/wmail/send.go index 3a1eed281..c34b6a0a7 100644 --- a/internal/app/worker/wmail/send.go +++ b/internal/app/worker/wmail/send.go @@ -182,6 +182,13 @@ func providerThrottle(err *errx.MailError) bool { // Send attempts to send an email with retry for transient failures func (w *WMail) Send(ctx context.Context, req *SendRequest) *SendResult { + done, ok := w.BeginExecution() + if !ok { + return &SendResult{Error: errx.ErrMailServerUnreachable} + } + defer done() + ctx, cancel := w.ExecutionContext(ctx) + defer cancel() // For warmup emails, ensure HTML is empty bodyHTML := req.BodyHTML if req.IsWarmup { diff --git a/internal/app/worker/wmail/sync_mail.go b/internal/app/worker/wmail/sync_mail.go index 7badcb353..e2b9b9ee3 100644 --- a/internal/app/worker/wmail/sync_mail.go +++ b/internal/app/worker/wmail/sync_mail.go @@ -8,6 +8,18 @@ import ( ) func (w *WMail) SyncMail(ctx context.Context) *errx.MailError { + done, ok := w.BeginExecution() + if !ok { + return errx.ErrMailServerUnreachable + } + defer done() + w.syncMu.Lock() + defer w.syncMu.Unlock() + ctx, cancel := w.ExecutionContext(ctx) + defer cancel() + if ctx.Err() != nil { + return errx.ErrMailServerUnreachable + } w.retryDiagnostics(ctx) defer w.retryDiagnostics(ctx) switch w.EmailType { diff --git a/internal/app/worker/wmail/wmail.go b/internal/app/worker/wmail/wmail.go index a161045f7..260bf4584 100644 --- a/internal/app/worker/wmail/wmail.go +++ b/internal/app/worker/wmail/wmail.go @@ -2,6 +2,7 @@ package wmail import ( "context" + "sync" "time" "github.com/google/uuid" @@ -58,10 +59,18 @@ type SmtpImapData struct { } type WMail struct { - diagnosticRetries map[uuid.UUID]diagnosticRetry - ExecutorID uuid.UUID - UserID uuid.UUID - ID uuid.UUID + lifecycleMu sync.Mutex + stopped bool + stopDone chan struct{} + active sync.WaitGroup + syncMu sync.Mutex + executionKeys map[[32]byte]struct{} + initialExecutionKey [32]byte + executionData *models.AddWorkerEmail + diagnosticRetries map[uuid.UUID]diagnosticRetry + ExecutorID uuid.UUID + UserID uuid.UUID + ID uuid.UUID // OrgID scopes the organization-wide sync budget; nil for a legacy // personal mailbox. OrgID *uuid.UUID @@ -169,6 +178,13 @@ func NewWMail( SyncContext: syncContext, CipherService: cipherService, } + mail.rememberExecution(data) + ready := false + defer func() { + if !ready { + mail.Discard() + } + }() // A publisher older than the sync policy sends no Sync block; compiled // defaults then apply and the mailbox starts its backfill from scratch. @@ -287,6 +303,7 @@ func NewWMail( if data.ImapSync { conn := &imap.Client{ + Context: mailCtx, Email: data.Email, AuthType: models.AuthPlain, Credentials: data.SmtpImap.Credentials.IMAP, @@ -295,13 +312,11 @@ func NewWMail( return nil, err } mail.SmtpImapData.ImapClient = conn - // Saved folder cursors: live sync resumes from each folder's stored - // HIGHESTMODSEQ instead of re-baselining (and, before this, instead - // of re-walking every folder on every worker restart). - for i := range data.SmtpImap.Mailboxes { - box := data.SmtpImap.Mailboxes[i] - mail.SmtpImapData.Mailboxes = append(mail.SmtpImapData.Mailboxes, &box) - } + } + // Keep saved cursors even while IMAP sync is temporarily disabled. + for i := range data.SmtpImap.Mailboxes { + box := data.SmtpImap.Mailboxes[i] + mail.SmtpImapData.Mailboxes = append(mail.SmtpImapData.Mailboxes, &box) } mail.SmtpImapData.SmtpClient = &smtp.Client{ @@ -320,6 +335,7 @@ func NewWMail( ) } + ready = true return mail, nil } @@ -336,13 +352,5 @@ func (w *WMail) ApplySyncPolicy(data *models.AddWorkerEmailSyncData) { // Discard releases a WMail that was built but never put to work. func (w *WMail) Discard() { - if w.Cancel != nil { - w.Cancel() - } - if w.SmtpImapData == nil || w.SmtpImapData.ImapClient == nil { - return - } - if c, ok := w.SmtpImapData.ImapClient.(interface{ Close() error }); ok { - _ = c.Close() - } + w.Stop() } diff --git a/internal/client/goog/deadline_test.go b/internal/client/goog/deadline_test.go new file mode 100644 index 000000000..8b4152fec --- /dev/null +++ b/internal/client/goog/deadline_test.go @@ -0,0 +1,32 @@ +package goog + +import ( + "context" + "testing" + + "github.com/warmbly/warmbly/internal/pkg/stoken" + "golang.org/x/oauth2" +) + +func TestProviderDefaultRequestDeadline(t *testing.T) { + for _, brokered := range []bool{false, true} { + c := &Client{} + token := &oauth2.Token{AccessToken: "local-test"} + var err error + if brokered { + if merr := c.InitWithSource(context.Background(), oauth2.StaticTokenSource(token)); merr != nil { + err = merr + } + } else { + if merr := c.Init(context.Background(), token, oauth2.Config{}); merr != nil { + err = merr + } + } + if err != nil { + t.Fatal(err) + } + if c.rawClient.Timeout != stoken.ProviderRequestTimeout { + t.Fatalf("brokered=%v default timeout=%v", brokered, c.rawClient.Timeout) + } + } +} diff --git a/internal/client/goog/goog.go b/internal/client/goog/goog.go index 618b14d51..a37436894 100644 --- a/internal/client/goog/goog.go +++ b/internal/client/goog/goog.go @@ -47,6 +47,7 @@ type Client struct { } func (c *Client) Init(ctx context.Context, token *oauth2.Token, cfg oauth2.Config) *errx.MailError { + ctx = stoken.BoundedContext(ctx) ts := cfg.TokenSource(ctx, token) ts = oauth2.ReuseTokenSource(token, ts) // Only wrap when there is somewhere to persist to. stoken calls onUpdate on @@ -63,7 +64,7 @@ func (c *Client) Init(ctx context.Context, token *oauth2.Token, cfg oauth2.Confi // InitWithSource builds the client on a caller-owned token source (brokered // tokens from Warmbly Cloud); nothing is persisted from it. func (c *Client) InitWithSource(ctx context.Context, ts oauth2.TokenSource) *errx.MailError { - httpClient := oauth2.NewClient(ctx, ts) + httpClient := stoken.HTTPClient(ctx, ts) c.rawClient = httpClient var err error c.srv, err = gmail.NewService(ctx, option.WithHTTPClient(httpClient)) diff --git a/internal/client/msgraph/deadline_test.go b/internal/client/msgraph/deadline_test.go new file mode 100644 index 000000000..3a679d400 --- /dev/null +++ b/internal/client/msgraph/deadline_test.go @@ -0,0 +1,32 @@ +package msgraph + +import ( + "context" + "testing" + + "github.com/warmbly/warmbly/internal/pkg/stoken" + "golang.org/x/oauth2" +) + +func TestProviderDefaultRequestDeadline(t *testing.T) { + for _, brokered := range []bool{false, true} { + c := &Client{} + token := &oauth2.Token{AccessToken: "local-test"} + var err error + if brokered { + if merr := c.InitWithSource(context.Background(), oauth2.StaticTokenSource(token)); merr != nil { + err = merr + } + } else { + if merr := c.Init(context.Background(), token, oauth2.Config{}); merr != nil { + err = merr + } + } + if err != nil { + t.Fatal(err) + } + if c.hc.Timeout != stoken.ProviderRequestTimeout { + t.Fatalf("brokered=%v default timeout=%v", brokered, c.hc.Timeout) + } + } +} diff --git a/internal/client/msgraph/msgraph.go b/internal/client/msgraph/msgraph.go index 4cc958e44..0a3876687 100644 --- a/internal/client/msgraph/msgraph.go +++ b/internal/client/msgraph/msgraph.go @@ -11,7 +11,9 @@ import ( "bytes" "context" "encoding/json" + "errors" "io" + "net" "net/http" "net/url" "sync" @@ -85,6 +87,7 @@ type Client struct { // cfg and persists the new token via OnTokenRefresh (the worker relays it back to // the control plane, which owns the encrypted credential store). func (c *Client) Init(ctx context.Context, token *oauth2.Token, cfg oauth2.Config) *errx.MailError { + ctx = stoken.BoundedContext(ctx) ts := cfg.TokenSource(ctx, token) ts = oauth2.ReuseTokenSource(token, ts) // Same guard as goog.Init: a nil OnTokenRefresh would panic per request. @@ -99,7 +102,7 @@ func (c *Client) Init(ctx context.Context, token *oauth2.Token, cfg oauth2.Confi // InitWithSource builds the client on a caller-owned token source (brokered // tokens from Warmbly Cloud); nothing is persisted from it. func (c *Client) InitWithSource(ctx context.Context, ts oauth2.TokenSource) *errx.MailError { - c.hc = oauth2.NewClient(ctx, ts) + c.hc = stoken.HTTPClient(ctx, ts) if c.DeltaLinks == nil { c.DeltaLinks = map[string]string{} } @@ -158,9 +161,16 @@ func (c *Client) doJSON(ctx context.Context, method, url string, in, out any) er return HandleError(resp) } if out != nil { - return json.NewDecoder(resp.Body).Decode(out) + err := json.NewDecoder(resp.Body).Decode(out) + var networkError net.Error + if errors.As(err, &networkError) || errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) { + return transportError(err) + } + return err + } + if _, err := io.Copy(io.Discard, resp.Body); err != nil { + return transportError(err) } - _, _ = io.Copy(io.Discard, resp.Body) return nil } diff --git a/internal/pkg/stoken/http.go b/internal/pkg/stoken/http.go new file mode 100644 index 000000000..6753364d5 --- /dev/null +++ b/internal/pkg/stoken/http.go @@ -0,0 +1,31 @@ +package stoken + +import ( + "context" + "net/http" + "time" + + "golang.org/x/oauth2" +) + +const ProviderRequestTimeout = 45 * time.Second + +// BoundedContext also bounds refresh I/O, which uses the token source's context. +func BoundedContext(ctx context.Context) context.Context { + base := http.DefaultClient + if client, ok := ctx.Value(oauth2.HTTPClient).(*http.Client); ok && client != nil { + base = client + } + client := *base + if client.Timeout <= 0 || client.Timeout > ProviderRequestTimeout { + client.Timeout = ProviderRequestTimeout + } + return context.WithValue(ctx, oauth2.HTTPClient, &client) +} + +func HTTPClient(ctx context.Context, source oauth2.TokenSource) *http.Client { + ctx = BoundedContext(ctx) + client := oauth2.NewClient(ctx, source) + client.Timeout = ctx.Value(oauth2.HTTPClient).(*http.Client).Timeout + return client +} diff --git a/internal/pkg/stoken/http_test.go b/internal/pkg/stoken/http_test.go new file mode 100644 index 000000000..1d87cb68e --- /dev/null +++ b/internal/pkg/stoken/http_test.go @@ -0,0 +1,47 @@ +package stoken + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" + "time" + + "golang.org/x/oauth2" +) + +func TestBoundedProviderAndRefreshHTTPClient(t *testing.T) { + for _, timeout := range []time.Duration{0, time.Hour, 20 * time.Millisecond} { + base := &http.Client{Timeout: timeout} + ctx := BoundedContext(context.WithValue(context.Background(), oauth2.HTTPClient, base)) + bounded := ctx.Value(oauth2.HTTPClient).(*http.Client) + want := ProviderRequestTimeout + if timeout > 0 && timeout < want { + want = timeout + } + if base.Timeout != timeout || bounded.Timeout != want { + t.Fatalf("base=%v bounded=%v want=%v", base.Timeout, bounded.Timeout, want) + } + if c := HTTPClient(ctx, oauth2.StaticTokenSource(&oauth2.Token{AccessToken: "test"})); c.Timeout != want { + t.Fatalf("provider timeout=%v want=%v", c.Timeout, want) + } + } + release := make(chan struct{}) + srv := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) { + select { + case <-r.Context().Done(): + case <-release: + } + })) + defer srv.Close() + defer close(release) + ctx := BoundedContext(context.WithValue(t.Context(), oauth2.HTTPClient, &http.Client{Timeout: 30 * time.Millisecond})) + cfg := oauth2.Config{Endpoint: oauth2.Endpoint{TokenURL: srv.URL, AuthStyle: oauth2.AuthStyleInParams}} + start := time.Now() + if _, err := cfg.TokenSource(ctx, &oauth2.Token{RefreshToken: "test"}).Token(); err == nil { + t.Fatal("stalled refresh succeeded") + } + if time.Since(start) > time.Second || ctx.Err() != nil { + t.Fatal("refresh was unbounded or cancelled mailbox context") + } +}