diff --git a/apps/api/internal/app/app.go b/apps/api/internal/app/app.go index 80bfc72..41fbbf8 100644 --- a/apps/api/internal/app/app.go +++ b/apps/api/internal/app/app.go @@ -15,6 +15,7 @@ import ( "os" "path/filepath" "strings" + "sync" "time" "golang.org/x/crypto/bcrypt" @@ -28,6 +29,7 @@ type App struct { now func() time.Time policy *HTMLPolicy workerCancel context.CancelFunc + workerWG sync.WaitGroup maildirHealth *maildirSyncHealthTracker externalIMAP externalIMAPClientFactory } @@ -73,16 +75,24 @@ func New(cfg Config, logger *slog.Logger) (*App, error) { } workerCtx, cancel := context.WithCancel(context.Background()) a.workerCancel = cancel - go a.scheduledSendWorker(workerCtx) + a.startWorker(func() { a.scheduledSendWorker(workerCtx) }) if strings.TrimSpace(a.cfg.MaildirRoot) != "" { - go a.maildirWorker(workerCtx) + a.startWorker(func() { a.maildirWorker(workerCtx) }) } - go a.sendQueueWorker(workerCtx) - go a.externalIMAPWorker(workerCtx) - go a.smtpEventsCleanupWorker(workerCtx) + a.startWorker(func() { a.sendQueueWorker(workerCtx) }) + a.startWorker(func() { a.externalIMAPWorker(workerCtx) }) + a.startWorker(func() { a.smtpEventsCleanupWorker(workerCtx) }) return a, nil } +func (a *App) startWorker(fn func()) { + a.workerWG.Add(1) + go func() { + defer a.workerWG.Done() + fn() + }() +} + func (a *App) Close() error { if a == nil || a.db == nil { return nil @@ -90,6 +100,7 @@ func (a *App) Close() error { if a.workerCancel != nil { a.workerCancel() } + a.workerWG.Wait() return a.db.Close() } diff --git a/apps/api/internal/app/app_test.go b/apps/api/internal/app/app_test.go index dc84dc1..38a3997 100644 --- a/apps/api/internal/app/app_test.go +++ b/apps/api/internal/app/app_test.go @@ -62,6 +62,13 @@ func newTestAppWithConfig(t *testing.T, cfg Config) *App { return a } +func stopTestWorkers(a *App) { + if a != nil && a.workerCancel != nil { + a.workerCancel() + a.workerWG.Wait() + } +} + func defaultAdminUserAndMailbox(t *testing.T, a *App) (*User, *Mailbox) { t.Helper() ctx := context.Background() @@ -581,6 +588,21 @@ func TestExternalIMAPOAuthEmailFromIDToken(t *testing.T) { } } +func TestExternalIMAPXOAUTH2ClientFormat(t *testing.T) { + client := newExternalIMAPXOAUTH2Client("user@example.com", "access-token") + mech, initialResponse, err := client.Start() + if err != nil { + t.Fatal(err) + } + if mech != "XOAUTH2" { + t.Fatalf("mechanism=%q, want XOAUTH2", mech) + } + want := "user=user@example.com\x01auth=Bearer access-token\x01\x01" + if string(initialResponse) != want { + t.Fatalf("initial response=%q, want %q", string(initialResponse), want) + } +} + func mustOAuthStateFromURL(t *testing.T, rawURL string) string { t.Helper() u, err := url.Parse(rawURL) @@ -1588,6 +1610,7 @@ func TestMailSendRollsBackSentCopyWhenQueueInsertFails(t *testing.T) { func TestSendQueueRecoversStaleSendingItems(t *testing.T) { a := newTestApp(t) + stopTestWorkers(a) host, port, received := startCapturingSMTP(t, 1) a.cfg.SMTPHost = host a.cfg.SMTPPort = port @@ -1637,6 +1660,7 @@ func TestSendQueueRecoversStaleSendingItems(t *testing.T) { func TestSendQueueStaleDeliveredMarkerDoesNotRedeliver(t *testing.T) { a := newTestApp(t) + stopTestWorkers(a) host, port, received := startCapturingSMTP(t, 1) a.cfg.SMTPHost = host a.cfg.SMTPPort = port diff --git a/apps/api/internal/app/external_imap.go b/apps/api/internal/app/external_imap.go index 77dcea8..10abc39 100644 --- a/apps/api/internal/app/external_imap.go +++ b/apps/api/internal/app/external_imap.go @@ -24,7 +24,6 @@ import ( "github.com/emersion/go-imap/v2" "github.com/emersion/go-imap/v2/imapclient" - "github.com/emersion/go-sasl" "github.com/go-chi/chi/v5" "golang.org/x/oauth2" ) @@ -1498,7 +1497,7 @@ func (a *App) openExternalIMAPClient(ctx context.Context, account externalIMAPAc c.Close() return nil, err } - if err := c.Authenticate(sasl.NewOAuthBearerClient(&sasl.OAuthBearerOptions{Username: account.Username, Token: token, Host: account.Host, Port: account.Port})); err != nil { + if err := c.Authenticate(newExternalIMAPXOAUTH2Client(account.Username, token)); err != nil { c.Close() return nil, err } @@ -1516,6 +1515,23 @@ func (a *App) openExternalIMAPClient(ctx context.Context, account externalIMAPAc return &goExternalIMAPClient{client: c}, nil } +type externalIMAPXOAUTH2Client struct { + username string + token string +} + +func newExternalIMAPXOAUTH2Client(username, token string) externalIMAPXOAUTH2Client { + return externalIMAPXOAUTH2Client{username: username, token: token} +} + +func (c externalIMAPXOAUTH2Client) Start() (string, []byte, error) { + return "XOAUTH2", []byte("user=" + c.username + "\x01auth=Bearer " + c.token + "\x01\x01"), nil +} + +func (c externalIMAPXOAUTH2Client) Next(challenge []byte) ([]byte, error) { + return []byte{}, nil +} + func (a *App) externalIMAPOAuthAccessToken(ctx context.Context, account externalIMAPAccountRecord) (string, error) { access, err := a.decryptExternalIMAPPassword(account.OAuthAccessTokenCiphertext) if err != nil { diff --git a/apps/api/internal/app/send_queue.go b/apps/api/internal/app/send_queue.go index a4e32ae..a8038ca 100644 --- a/apps/api/internal/app/send_queue.go +++ b/apps/api/internal/app/send_queue.go @@ -126,6 +126,12 @@ func (a *App) sendQueueWorker(ctx context.Context) { ticker := time.NewTicker(10 * time.Second) defer ticker.Stop() for { + select { + case <-ctx.Done(): + a.log.Info("send queue worker stopped") + return + default: + } if err := a.processDueSendQueue(ctx); err != nil { a.log.Warn("send queue worker failed", "error", err) } @@ -232,12 +238,12 @@ func (a *App) recoverStaleSendQueueItems(ctx context.Context) error { } continue } - res, err := a.db.ExecContext(ctx, `UPDATE send_queue SET status=?,next_attempt_at=?,last_error=?,updated_at=? WHERE id=? AND status=?`, sendQueueStatusFailed, now, "send attempt interrupted", now, item.ID, sendQueueStatusSending) + res, err := a.db.ExecContext(ctx, `UPDATE send_queue SET status=?,next_attempt_at=?,last_error=?,updated_at=? WHERE id=? AND status=?`, sendQueueStatusQueued, now, "send attempt interrupted", now, item.ID, sendQueueStatusSending) if err != nil { return err } if n, _ := res.RowsAffected(); n > 0 { - a.recordSendAudit(ctx, sendAuditRetry, sendQueueStatusFailed, sendAuditInputFromQueue(item, "send attempt interrupted")) + a.recordSendAudit(ctx, sendAuditRetry, sendQueueStatusQueued, sendAuditInputFromQueue(item, "send attempt interrupted")) } } return nil