Shamelessly dumb SMTP server stress-tester

smtp_spin.go

package main

import (
	"bytes"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/smtp"
	"os"
	"os/signal"
	"strings"
	"sync/atomic"
	"syscall"
	"time"
)

func main() {
	addr := flag.String("addr", "127.0.0.1:1125", "Address of server to stress-test")
	msgFile := flag.String("msg", "msg.eml", "Message sample to use")
	conns := flag.Int("conns", 100, "Number of concurrent connections to open")
	limit := flag.Int("limit", 0, "Stop after sending N messages, default is no limit")
	sender := flag.String("sender", "foxcpp@foxcpp.dev", "Address to use in MAIL FROM")
	rcptsStr := flag.String("rcpts", "foxcpp@foxcpp.dev", "Comma-separated list of addresses to use RCPT TO")
	flag.Parse()
	
	rcpts := strings.Split(*rcptsStr, ",")

	f, err := ioutil.ReadFile(*msgFile)
	if err != nil {
		log.Fatalln(err)
	}

	ch := make(chan os.Signal)
	signal.Notify(ch, syscall.SIGINT)
	stop := false
	
	fmt.Println("Bass-boosted 'Dead Or Alive - You Spin Me Round' starts playing...")

	N := uint64(0)
	start := time.Now()
	for i := 0; i < *conns; i++ {
		go func() {
			cl, err := smtp.Dial(*addr)
			if err != nil {
				log.Fatalln(err)
			}
			defer cl.Close()
			if err := cl.Hello("localhost"); err != nil {
				log.Fatalln(err)
			}

			for !stop || (atomic.LoadUint64(&N) > uint64(*limit) && *limit != 0) {
				if err := cl.Mail(*sender); err != nil {
					log.Fatalln(err)
				}
				for _, rcpt := range rcpts {
					if err := cl.Rcpt(rcpt); err != nil {
						log.Fatalln(err)
					}
				}
				wr, err := cl.Data()
				if err != nil {
					log.Fatalln(err)
				}
				if _, err := io.Copy(wr, bytes.NewReader(f)); err != nil {
					log.Fatalln(err)
				}
				wr.Close()
				atomic.AddUint64(&N, 1)
			}
		}()
	}

	<-ch
	stop = true

	end := time.Now()

	fmt.Println("time elapsed:", end.Sub(start))
	fmt.Println("messages sent:", N)
}