diff --git a/cli/exp_scaletest.go b/cli/exp_scaletest.go index fc21d47fdb..2515012b28 100644 --- a/cli/exp_scaletest.go +++ b/cli/exp_scaletest.go @@ -66,6 +66,7 @@ func (r *RootCmd) scaletestCmd() *serpent.Command { r.scaletestWorkspaceTraffic(), r.scaletestAutostart(), r.scaletestNotifications(), + r.scaletestSMTP(), }, } diff --git a/cli/exp_scaletest_smtp.go b/cli/exp_scaletest_smtp.go new file mode 100644 index 0000000000..3713005de5 --- /dev/null +++ b/cli/exp_scaletest_smtp.go @@ -0,0 +1,112 @@ +//go:build !slim + +package cli + +import ( + "fmt" + "os/signal" + "time" + + "golang.org/x/xerrors" + + "cdr.dev/slog" + "cdr.dev/slog/sloggers/sloghuman" + "github.com/coder/coder/v2/scaletest/smtpmock" + "github.com/coder/serpent" +) + +func (*RootCmd) scaletestSMTP() *serpent.Command { + var ( + hostAddress string + smtpPort int64 + apiPort int64 + purgeAtCount int64 + ) + cmd := &serpent.Command{ + Use: "smtp", + Short: "Start a mock SMTP server for testing", + Long: `Start a mock SMTP server with an HTTP API server that can be used to purge +messages and get messages by email.`, + Handler: func(inv *serpent.Invocation) error { + ctx := inv.Context() + notifyCtx, stop := signal.NotifyContext(ctx, StopSignals...) + defer stop() + ctx = notifyCtx + + logger := slog.Make(sloghuman.Sink(inv.Stderr)).Leveled(slog.LevelInfo) + config := smtpmock.Config{ + HostAddress: hostAddress, + SMTPPort: int(smtpPort), + APIPort: int(apiPort), + Logger: logger, + } + srv := new(smtpmock.Server) + + if err := srv.Start(ctx, config); err != nil { + return xerrors.Errorf("start mock SMTP server: %w", err) + } + defer func() { + _ = srv.Stop() + }() + + _, _ = fmt.Fprintf(inv.Stdout, "Mock SMTP server started on %s\n", srv.SMTPAddress()) + _, _ = fmt.Fprintf(inv.Stdout, "HTTP API server started on %s\n", srv.APIAddress()) + if purgeAtCount > 0 { + _, _ = fmt.Fprintf(inv.Stdout, " Auto-purge when message count reaches %d\n", purgeAtCount) + } + + ticker := time.NewTicker(10 * time.Second) + defer ticker.Stop() + + for { + select { + case <-ctx.Done(): + _, _ = fmt.Fprintf(inv.Stdout, "\nTotal messages received since last purge: %d\n", srv.MessageCount()) + return nil + case <-ticker.C: + count := srv.MessageCount() + if count > 0 { + _, _ = fmt.Fprintf(inv.Stdout, "Messages received: %d\n", count) + } + + if purgeAtCount > 0 && int64(count) >= purgeAtCount { + _, _ = fmt.Fprintf(inv.Stdout, "Message count (%d) reached threshold (%d). Purging...\n", count, purgeAtCount) + srv.Purge() + continue + } + } + } + }, + } + + cmd.Options = []serpent.Option{ + { + Flag: "host-address", + Env: "CODER_SCALETEST_SMTP_HOST_ADDRESS", + Default: "localhost", + Description: "Host address to bind the mock SMTP and API servers.", + Value: serpent.StringOf(&hostAddress), + }, + { + Flag: "smtp-port", + Env: "CODER_SCALETEST_SMTP_PORT", + Description: "Port for the mock SMTP server. Uses a random port if not specified.", + Value: serpent.Int64Of(&smtpPort), + }, + { + Flag: "api-port", + Env: "CODER_SCALETEST_SMTP_API_PORT", + Description: "Port for the HTTP API server. Uses a random port if not specified.", + Value: serpent.Int64Of(&apiPort), + }, + { + Flag: "purge-at-count", + Env: "CODER_SCALETEST_SMTP_PURGE_AT_COUNT", + Default: "100000", + Description: "Maximum number of messages to keep before auto-purging. Set to 0 to disable.", + Value: serpent.Int64Of(&purgeAtCount), + }, + } + + return cmd +} diff --git a/scaletest/smtpmock/server.go b/scaletest/smtpmock/server.go new file mode 100644 index 0000000000..26f5b65ffb --- /dev/null +++ b/scaletest/smtpmock/server.go @@ -0,0 +1,247 @@ +package smtpmock + +import ( + "bufio" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "mime/quotedprintable" + "net" + "net/http" + "net/mail" + "regexp" + "slices" + "strings" + "time" + + "github.com/google/uuid" + smtpmocklib "github.com/mocktools/go-smtp-mock/v2" + "golang.org/x/xerrors" + + "cdr.dev/slog" +) + +// Server wraps the SMTP mock server and provides an HTTP API to retrieve emails. +type Server struct { + smtpServer *smtpmocklib.Server + httpServer *http.Server + httpListener net.Listener + logger slog.Logger + + hostAddress string + smtpPort int + apiPort int +} + +type Config struct { + HostAddress string + SMTPPort int + APIPort int + Logger slog.Logger +} + +type EmailSummary struct { + Subject string `json:"subject"` + Date time.Time `json:"date"` + NotificationTemplateID uuid.UUID `json:"notification_template_id,omitempty"` +} + +var notificationTemplateIDRegex = regexp.MustCompile(`notifications\?disabled=([a-f0-9-]+)`) + +func (s *Server) Start(ctx context.Context, cfg Config) error { + s.hostAddress = cfg.HostAddress + s.smtpPort = cfg.SMTPPort + s.apiPort = cfg.APIPort + s.logger = cfg.Logger + + s.smtpServer = smtpmocklib.New(smtpmocklib.ConfigurationAttr{ + LogToStdout: false, + LogServerActivity: true, + HostAddress: s.hostAddress, + PortNumber: s.smtpPort, + }) + if err := s.smtpServer.Start(); err != nil { + return xerrors.Errorf("start SMTP server: %w", err) + } + s.smtpPort = s.smtpServer.PortNumber() + + if err := s.startAPIServer(ctx); err != nil { + _ = s.smtpServer.Stop() + return xerrors.Errorf("start API server: %w", err) + } + + return nil +} + +func (s *Server) Stop() error { + var httpErr, smtpErr error + + if s.httpServer != nil { + shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + if err := s.httpServer.Shutdown(shutdownCtx); err != nil { + httpErr = xerrors.Errorf("shutdown HTTP server: %w", err) + } + } + + if s.smtpServer != nil { + if err := s.smtpServer.Stop(); err != nil { + smtpErr = xerrors.Errorf("stop SMTP server: %w", err) + } + } + + return errors.Join(httpErr, smtpErr) +} + +func (s *Server) SMTPAddress() string { + return fmt.Sprintf("%s:%d", s.hostAddress, s.smtpPort) +} + +func (s *Server) APIAddress() string { + return fmt.Sprintf("http://%s:%d", s.hostAddress, s.apiPort) +} + +func (s *Server) MessageCount() int { + if s.smtpServer == nil { + return 0 + } + return len(s.smtpServer.Messages()) +} + +func (s *Server) Purge() { + if s.smtpServer != nil { + s.smtpServer.MessagesAndPurge() + } +} + +func (s *Server) startAPIServer(ctx context.Context) error { + mux := http.NewServeMux() + mux.HandleFunc("POST /purge", s.handlePurge) + mux.HandleFunc("GET /messages", s.handleMessages) + + s.httpServer = &http.Server{ + Handler: mux, + ReadHeaderTimeout: 10 * time.Second, + } + + listener, err := net.Listen("tcp", fmt.Sprintf("%s:%d", s.hostAddress, s.apiPort)) + if err != nil { + return xerrors.Errorf("listen on %s:%d: %w", s.hostAddress, s.apiPort, err) + } + s.httpListener = listener + + tcpAddr, valid := listener.Addr().(*net.TCPAddr) + if !valid { + err := listener.Close() + if err != nil { + s.logger.Error(ctx, "failed to close listener", slog.Error(err)) + } + return xerrors.Errorf("listener returned invalid address: %T", listener.Addr()) + } + s.apiPort = tcpAddr.Port + + go func() { + if err := s.httpServer.Serve(listener); err != nil && !errors.Is(err, http.ErrServerClosed) { + s.logger.Error(ctx, "http API server error", slog.Error(err)) + } + }() + + return nil +} + +func (s *Server) handlePurge(w http.ResponseWriter, _ *http.Request) { + s.smtpServer.MessagesAndPurge() + w.WriteHeader(http.StatusOK) +} + +func (s *Server) handleMessages(w http.ResponseWriter, r *http.Request) { + email := r.URL.Query().Get("email") + msgs := s.smtpServer.Messages() + + var summaries []EmailSummary + for _, msg := range msgs { + recipients := msg.RcpttoRequestResponse() + if !matchesRecipient(recipients, email) { + continue + } + + summary, err := parseEmailSummary(msg.MsgRequest()) + if err != nil { + s.logger.Warn(r.Context(), "failed to parse email summary", slog.Error(err)) + continue + } + summaries = append(summaries, summary) + } + + w.Header().Set("Content-Type", "application/json") + if err := json.NewEncoder(w).Encode(summaries); err != nil { + s.logger.Warn(r.Context(), "failed to encode JSON response", slog.Error(err)) + } +} + +func matchesRecipient(recipients [][]string, email string) bool { + if email == "" { + return true + } + return slices.ContainsFunc(recipients, func(rcptPair []string) bool { + if len(rcptPair) == 0 { + return false + } + + addrPart, ok := strings.CutPrefix(rcptPair[0], "RCPT TO:") + if !ok { + return false + } + + addr, err := mail.ParseAddress(addrPart) + if err != nil { + return false + } + + return strings.EqualFold(addr.Address, email) + }) +} + +func parseEmailSummary(message string) (EmailSummary, error) { + var summary EmailSummary + + // Decode quoted-printable message + reader := quotedprintable.NewReader(strings.NewReader(message)) + content, err := io.ReadAll(reader) + if err != nil { + return summary, xerrors.Errorf("decode email content: %w", err) + } + + contentStr := string(content) + scanner := bufio.NewScanner(strings.NewReader(contentStr)) + + // Extract Subject and Date from headers. + // Date is used to measure latency. + for scanner.Scan() { + line := scanner.Text() + if line == "" { + break + } + if prefix, found := strings.CutPrefix(line, "Subject: "); found { + summary.Subject = prefix + } else if prefix, found := strings.CutPrefix(line, "Date: "); found { + if parsedDate, err := time.Parse(time.RFC1123Z, prefix); err == nil { + summary.Date = parsedDate + } + } + } + + // Extract notification ID from decoded email content + // Notification ID is present in the email footer like this + //
+ if matches := notificationTemplateIDRegex.FindStringSubmatch(contentStr); len(matches) > 1 { + summary.NotificationTemplateID, err = uuid.Parse(matches[1]) + if err != nil { + return summary, xerrors.Errorf("parse notification ID: %w", err) + } + } + + return summary, nil +} diff --git a/scaletest/smtpmock/server_test.go b/scaletest/smtpmock/server_test.go new file mode 100644 index 0000000000..7136c5ab9e --- /dev/null +++ b/scaletest/smtpmock/server_test.go @@ -0,0 +1,203 @@ +package smtpmock_test + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/smtp" + "strings" + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + + "cdr.dev/slog/sloggers/slogtest" + "github.com/coder/coder/v2/scaletest/smtpmock" + "github.com/coder/coder/v2/testutil" +) + +func TestServer_StartStop(t *testing.T) { + t.Parallel() + + ctx := context.Background() + srv := new(smtpmock.Server) + err := srv.Start(ctx, smtpmock.Config{ + HostAddress: "127.0.0.1", + SMTPPort: 0, + APIPort: 0, + Logger: slogtest.Make(t, nil), + }) + require.NoError(t, err) + require.NotEmpty(t, srv.SMTPAddress()) + require.NotEmpty(t, srv.APIAddress()) + + err = srv.Stop() + require.NoError(t, err) +} + +func TestServer_SendAndReceiveEmail(t *testing.T) { + t.Parallel() + + ctx := context.Background() + srv := new(smtpmock.Server) + err := srv.Start(ctx, smtpmock.Config{ + HostAddress: "127.0.0.1", + SMTPPort: 0, + APIPort: 0, + Logger: slogtest.Make(t, nil), + }) + require.NoError(t, err) + defer srv.Stop() + + err = sendTestEmail(srv.SMTPAddress(), "test@example.com", "Test Subject", "Test Body") + require.NoError(t, err) + + require.Eventually(t, func() bool { + return srv.MessageCount() == 1 + }, testutil.WaitShort, testutil.IntervalMedium) + + url := fmt.Sprintf("%s/messages", srv.APIAddress()) + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + require.NoError(t, err) + + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + + require.Equal(t, http.StatusOK, resp.StatusCode) + + var summaries []smtpmock.EmailSummary + err = json.NewDecoder(resp.Body).Decode(&summaries) + require.NoError(t, err) + require.Len(t, summaries, 1) + require.Equal(t, "Test Subject", summaries[0].Subject) +} + +func TestServer_FilterByEmail(t *testing.T) { + t.Parallel() + + ctx := context.Background() + srv := new(smtpmock.Server) + err := srv.Start(ctx, smtpmock.Config{ + HostAddress: "127.0.0.1", + SMTPPort: 0, + APIPort: 0, + Logger: slogtest.Make(t, nil), + }) + require.NoError(t, err) + defer srv.Stop() + + err = sendTestEmail(srv.SMTPAddress(), "admin@coder.com", "Email for admin", "Body 1") + require.NoError(t, err) + + err = sendTestEmail(srv.SMTPAddress(), "test-user@coder.com", "Email for test-user", "Body 2") + require.NoError(t, err) + + require.Eventually(t, func() bool { + return srv.MessageCount() == 2 + }, testutil.WaitShort, testutil.IntervalMedium) + + url := fmt.Sprintf("%s/messages?email=admin@coder.com", srv.APIAddress()) + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + require.NoError(t, err) + + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + + var summaries []smtpmock.EmailSummary + err = json.NewDecoder(resp.Body).Decode(&summaries) + require.NoError(t, err) + require.Len(t, summaries, 1) + require.Equal(t, "Email for admin", summaries[0].Subject) +} + +func TestServer_NotificationTemplateID(t *testing.T) { + t.Parallel() + + ctx := context.Background() + srv := new(smtpmock.Server) + err := srv.Start(ctx, smtpmock.Config{ + HostAddress: "127.0.0.1", + SMTPPort: 0, + APIPort: 0, + Logger: slogtest.Make(t, nil), + }) + require.NoError(t, err) + defer srv.Stop() + + notificationID := uuid.New() + body := fmt.Sprintf(``, notificationID.String()) + + err = sendTestEmail(srv.SMTPAddress(), "test-user@coder.com", "Notification", body) + require.NoError(t, err) + + require.Eventually(t, func() bool { + return srv.MessageCount() == 1 + }, testutil.WaitShort, testutil.IntervalMedium) + + url := fmt.Sprintf("%s/messages", srv.APIAddress()) + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + require.NoError(t, err) + + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + + var summaries []smtpmock.EmailSummary + err = json.NewDecoder(resp.Body).Decode(&summaries) + require.NoError(t, err) + require.Len(t, summaries, 1) + require.Equal(t, notificationID, summaries[0].NotificationTemplateID) +} + +func TestServer_Purge(t *testing.T) { + t.Parallel() + + ctx := context.Background() + srv := new(smtpmock.Server) + err := srv.Start(ctx, smtpmock.Config{ + HostAddress: "127.0.0.1", + SMTPPort: 0, + APIPort: 0, + Logger: slogtest.Make(t, nil), + }) + require.NoError(t, err) + defer srv.Stop() + + err = sendTestEmail(srv.SMTPAddress(), "test-user@coder.com", "Test", "Body") + require.NoError(t, err) + + require.Eventually(t, func() bool { + return srv.MessageCount() == 1 + }, testutil.WaitShort, testutil.IntervalMedium) + + url := fmt.Sprintf("%s/purge", srv.APIAddress()) + req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, nil) + require.NoError(t, err) + + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusOK, resp.StatusCode) + + require.Equal(t, 0, srv.MessageCount()) +} + +func sendTestEmail(smtpAddr, to, subject, body string) error { + from := "noreply@coder.com" + now := time.Now().Format(time.RFC1123Z) + + msg := strings.Builder{} + _, _ = msg.WriteString(fmt.Sprintf("From: %s\r\n", from)) + _, _ = msg.WriteString(fmt.Sprintf("To: %s\r\n", to)) + _, _ = msg.WriteString(fmt.Sprintf("Subject: %s\r\n", subject)) + _, _ = msg.WriteString(fmt.Sprintf("Date: %s\r\n", now)) + _, _ = msg.WriteString("Content-Type: text/html; charset=UTF-8\r\n") + _, _ = msg.WriteString("\r\n") + _, _ = msg.WriteString(body) + + return smtp.SendMail(smtpAddr, nil, from, []string{to}, []byte(msg.String())) +}