package smtpsrv

import (
	"crypto/tls"
	"errors"
	"flag"
	"fmt"
	"io"
	"io/fs"
	"net"
	"net/smtp"
	"net/textproto"
	"os"
	"strings"
	"syscall"
	"testing"
	"time"

	"blitiri.com.ar/go/chasquid/internal/aliases"
	"blitiri.com.ar/go/chasquid/internal/auth"
	"blitiri.com.ar/go/chasquid/internal/domaininfo"
	"blitiri.com.ar/go/chasquid/internal/maillog"
	"blitiri.com.ar/go/chasquid/internal/testlib"
	"blitiri.com.ar/go/chasquid/internal/userdb"
)

// Flags.
var (
	externalSMTPAddr = flag.String("external_smtp_addr", "",
		"SMTP server address to test (defaults to use internal)")
	externalSubmissionAddr = flag.String("external_submission_addr", "",
		"submission server address to test (defaults to use internal)")
	externalSubmissionTLSAddr = flag.String("external_submission_tls_addr", "",
		"submission+TLS server address to test (defaults to use internal)")
)

var (
	// Server addresses. Will be filled in at init time.
	// We default to internal ones, but may get overridden via flags.
	smtpAddr          = ""
	submissionAddr    = ""
	submissionTLSAddr = ""

	// TLS configuration to use in the clients.
	// Will contain the generated server certificate as root CA.
	tlsConfig *tls.Config

	// Test couriers, so we can validate that emails got sent.
	localC  = testlib.NewTestCourier()
	remoteC = testlib.NewTestCourier()

	// Max data size, in MiB.
	maxDataSizeMiB = 5
)

//
// === Tests ===
//

func mustDial(tb testing.TB, mode SocketMode, startTLS bool) *smtp.Client {
	addr := ""
	switch mode {
	case ModeSMTP:
		addr = smtpAddr
	case ModeSubmission:
		addr = submissionAddr
	case ModeSubmissionTLS:
		addr = submissionTLSAddr
	}

	var err error
	var conn net.Conn
	if mode.TLS {
		conn, err = tls.Dial("tcp", addr, tlsConfig)
	} else {
		conn, err = net.Dial("tcp", addr)
	}
	if err != nil {
		tb.Fatalf("(net||tls).Dial: %v", err)
	}
	c, err := smtp.NewClient(conn, "127.0.0.1")
	if err != nil {
		tb.Fatalf("smtp.Dial: %v", err)
	}

	if err = c.Hello("test"); err != nil {
		tb.Fatalf("c.Hello: %v", err)
	}

	if startTLS {
		if ok, _ := c.Extension("STARTTLS"); !ok {
			tb.Fatalf("STARTTLS not advertised in EHLO")
		}

		if err = c.StartTLS(tlsConfig); err != nil {
			tb.Fatalf("StartTLS: %v", err)
		}
	}

	return c
}

func sendEmail(tb testing.TB, c *smtp.Client) {
	sendEmailWithAuth(tb, c, nil)
}

func sendEmailWithAuth(tb testing.TB, c *smtp.Client, auth smtp.Auth) {
	var err error
	from := "from@from"

	if auth != nil {
		if err = c.Auth(auth); err != nil {
			tb.Errorf("Auth: %v", err)
		}

		// If we authenticated, we must use the user as from, as the server
		// checks otherwise.
		from = "testuser@localhost"
	}

	if err = c.Mail(from); err != nil {
		tb.Errorf("Mail: %v", err)
	}

	if err = c.Rcpt("to@localhost"); err != nil {
		tb.Errorf("Rcpt: %v", err)
	}

	w, err := c.Data()
	if err != nil {
		tb.Fatalf("Data: %v", err)
	}

	msg := []byte("Subject: Hi!\n\n This is an email\n")
	if _, err = w.Write(msg); err != nil {
		tb.Errorf("Data write: %v", err)
	}

	localC.Expect(1)

	if err = w.Close(); err != nil {
		tb.Errorf("Data close: %v", err)
	}

	localC.Wait()
}

// Checks err is a textproto.Error with the given code and message.
func expectTPErr(tb testing.TB, err error, code int, msg string) {
	tb.Helper()
	if err == nil {
		tb.Fatalf("Expected error %d %q, got nil", code, msg)
	}
	netErr, ok := err.(*textproto.Error)
	if !ok {
		tb.Fatalf("Expected textproto.Error(%d %q), got %T: %v",
			code, msg, err, err)
	}
	if netErr.Code != code || netErr.Msg != msg {
		tb.Fatalf("Expected textproto.Error(%d %q), got %d %q",
			code, msg, netErr.Code, netErr.Msg)
	}
}

func TestSimple(t *testing.T) {
	c := mustDial(t, ModeSMTP, false)
	defer c.Close()
	sendEmail(t, c)
}

func TestSimpleTLS(t *testing.T) {
	c := mustDial(t, ModeSMTP, true)
	defer c.Close()
	sendEmail(t, c)
}

func TestManyEmails(t *testing.T) {
	c := mustDial(t, ModeSMTP, true)
	defer c.Close()
	sendEmail(t, c)
	sendEmail(t, c)
	sendEmail(t, c)
}

func TestAuth(t *testing.T) {
	c := mustDial(t, ModeSubmission, true)
	defer c.Close()

	auth := smtp.PlainAuth("", "testuser@localhost", "testpasswd", "127.0.0.1")
	sendEmailWithAuth(t, c, auth)
}

func TestSubmissionWithoutAuth(t *testing.T) {
	c := mustDial(t, ModeSubmission, true)
	defer c.Close()

	if err := c.Mail("from@from"); err == nil {
		t.Errorf("Mail not failed as expected")
	}
}

func TestAuthOnTLS(t *testing.T) {
	c := mustDial(t, ModeSubmissionTLS, false)
	defer c.Close()

	auth := smtp.PlainAuth("", "testuser@localhost", "testpasswd", "127.0.0.1")
	sendEmailWithAuth(t, c, auth)
}

func TestAuthOnSMTP(t *testing.T) {
	c := mustDial(t, ModeSMTP, true)
	defer c.Close()

	auth := smtp.PlainAuth("", "testuser@localhost", "testpasswd", "127.0.0.1")

	// At least for now, we allow AUTH over the SMTP port to avoid unnecessary
	// complexity, so we expect it to work.
	sendEmailWithAuth(t, c, auth)
}

func TestBrokenAuth(t *testing.T) {
	c := mustDial(t, ModeSubmission, true)
	defer c.Close()

	auth := smtp.PlainAuth("", "user@broken", "passwd", "127.0.0.1")
	err := c.Auth(auth)
	expectTPErr(t, err, 454, "4.7.0 Temporary authentication failure")
}

func TestWrongMailParsing(t *testing.T) {
	addrs := []string{"from", "a b c", "a @ b", "<x>", "<x y>", "><",
		"a@b/c", "a@[1.2.3.4]"}
	for _, addr := range addrs {
		c := mustDial(t, ModeSMTP, false)

		if err := c.Mail(addr); err == nil {
			t.Errorf("Mail not failed as expected with %q", addr)
		}

		if err := c.Mail("from@plain"); err != nil {
			t.Errorf("Mail: %v", err)
		}

		for _, addr := range addrs {
			if err := c.Rcpt(addr); err == nil {
				t.Errorf("Rcpt not failed as expected with %q", addr)
			}
		}

		c.Close()
	}
}

func TestMalformedDomain(t *testing.T) {
	// Unauthenticated MAIL FROM.
	c := mustDial(t, ModeSMTP, false)
	defer c.Close()

	err := c.Mail("a@b/c")
	expectTPErr(t, err, 501, "5.1.8 Malformed sender domain")

	// Authenticated RCPT TO. We need to authenticate, otherwise it may get
	// rejected as relaying.
	c = mustDial(t, ModeSubmission, true)
	defer c.Close()

	auth := smtp.PlainAuth("", "testuser@localhost", "testpasswd", "127.0.0.1")
	if err := c.Auth(auth); err != nil {
		t.Fatalf("Auth: %v", err)
	}

	if err := c.Mail("testuser@localhost"); err != nil {
		t.Fatalf("Mail: %v", err)
	}

	err = c.Rcpt("a@b/c")
	expectTPErr(t, err, 501, "5.1.2 Malformed destination domain")
}

func TestUnsafeAddress(t *testing.T) {
	// Unauthenticated MAIL FROM.
	c := mustDial(t, ModeSMTP, false)
	defer c.Close()

	err := c.Mail(`"a@b"@from`)
	expectTPErr(t, err, 501, "5.1.7 Sender address malformed")

	err = c.Mail("a\u009bb@from")
	expectTPErr(t, err, 501, "5.1.7 Sender address malformed")

	// Characters that need quoting are fine.
	if err := c.Mail(`"a;b"@from`); err != nil {
		t.Errorf("Mail: %v", err)
	}

	// Authenticated RCPT TO. We need to authenticate, otherwise it may get
	// rejected as relaying.
	c = mustDial(t, ModeSubmission, true)
	defer c.Close()

	auth := smtp.PlainAuth("", "testuser@localhost", "testpasswd", "127.0.0.1")
	if err := c.Auth(auth); err != nil {
		t.Fatalf("Auth: %v", err)
	}

	if err := c.Mail("testuser@localhost"); err != nil {
		t.Fatalf("Mail: %v", err)
	}

	err = c.Rcpt(`"a@b"@remote`)
	expectTPErr(t, err, 501, "5.1.3 Malformed destination address")

	err = c.Rcpt("a\u009bb@remote")
	expectTPErr(t, err, 501, "5.1.3 Malformed destination address")
}

func TestNullMailFrom(t *testing.T) {
	c := mustDial(t, ModeSMTP, false)
	defer c.Close()

	addrs := []string{"<>", "  <>", "<> OPTION"}
	for _, addr := range addrs {
		simpleCmd(t, c, fmt.Sprintf("MAIL FROM:%s", addr), 250)
	}
}

func TestRcptBeforeMail(t *testing.T) {
	c := mustDial(t, ModeSMTP, false)
	defer c.Close()

	if err := c.Rcpt("to@to"); err == nil {
		t.Errorf("Rcpt not failed as expected")
	}
}

func TestRcptOption(t *testing.T) {
	c := mustDial(t, ModeSMTP, true)
	defer c.Close()

	if err := c.Mail("from@localhost"); err != nil {
		t.Fatalf("Mail: %v", err)
	}

	params := []string{
		"<to@localhost>", "  <to@localhost>", "<to@localhost> OPTION"}
	for _, p := range params {
		simpleCmd(t, c, fmt.Sprintf("RCPT TO:%s", p), 250)
	}
}

func TestRelayForbidden(t *testing.T) {
	c := mustDial(t, ModeSMTP, false)
	defer c.Close()

	if err := c.Mail("from@somewhere"); err != nil {
		t.Errorf("Mail: %v", err)
	}

	if err := c.Rcpt("to@somewhere"); err == nil {
		t.Errorf("Accepted relay email")
	}
}

func TestTooManyRecipients(t *testing.T) {
	c := mustDial(t, ModeSubmission, true)
	defer c.Close()

	auth := smtp.PlainAuth("", "testuser@localhost", "testpasswd", "127.0.0.1")
	if err := c.Auth(auth); err != nil {
		t.Fatalf("Auth: %v", err)
	}

	if err := c.Mail("testuser@localhost"); err != nil {
		t.Fatalf("Mail: %v", err)
	}

	for i := range 101 {
		if err := c.Rcpt(fmt.Sprintf("to%d@somewhere", i)); err != nil {
			t.Fatalf("Rcpt: %v", err)
		}
	}

	err := c.Rcpt("to102@somewhere")
	expectTPErr(t, err, 452, "4.5.3 Too many recipients")
}

func TestRcptBrokenExists(t *testing.T) {
	c := mustDial(t, ModeSMTP, true)
	defer c.Close()

	if err := c.Mail("from@localhost"); err != nil {
		t.Fatalf("Mail: %v", err)
	}

	err := c.Rcpt("to@broken")
	expectTPErr(t, err, 451, "4.4.3 Temporary error checking address")
}

func TestRcptUserDoesNotExist(t *testing.T) {
	c := mustDial(t, ModeSMTP, true)
	defer c.Close()

	if err := c.Mail("from@localhost"); err != nil {
		t.Fatalf("Mail: %v", err)
	}

	err := c.Rcpt("doesnotexist@localhost")
	expectTPErr(t, err, 550,
		"5.1.1 Destination address is unknown (user does not exist)")
}

var str1MiB string

func sendLargeEmail(tb testing.TB, c *smtp.Client, sizeMiB int) error {
	tb.Helper()
	if err := c.Mail("from@from"); err != nil {
		tb.Fatalf("Mail: %v", err)
	}
	if err := c.Rcpt("to@localhost"); err != nil {
		tb.Fatalf("Rcpt: %v", err)
	}

	w, err := c.Data()
	if err != nil {
		tb.Fatalf("Data: %v", err)
	}

	if _, err := w.Write([]byte("Subject: I ate too much\n\n")); err != nil {
		tb.Fatalf("Data write: %v", err)
	}

	// Write the 1 MiB string sizeMiB times.
	for range sizeMiB {
		if _, err := w.Write([]byte(str1MiB)); err != nil {
			tb.Fatalf("Data write: %v", err)
		}
	}

	return w.Close()
}

func TestTooMuchData(t *testing.T) {
	c := mustDial(t, ModeSMTP, true)
	defer c.Close()

	localC.Expect(1)
	err := sendLargeEmail(t, c, maxDataSizeMiB-1)
	if err != nil {
		t.Errorf("Error sending large but ok email: %v", err)
	}
	localC.Wait()

	// Repeat the test - we want to check that the limit applies to each
	// message, not the entire connection.
	localC.Expect(1)
	err = sendLargeEmail(t, c, maxDataSizeMiB-1)
	if err != nil {
		t.Errorf("Error sending large but ok email: %v", err)
	}
	localC.Wait()

	err = sendLargeEmail(t, c, maxDataSizeMiB+1)
	expectTPErr(t, err, 552, "5.3.4 Message too big")

	// Repeat the test once again, the limit should not prevent connection
	// from continuing.
	localC.Expect(1)
	err = sendLargeEmail(t, c, maxDataSizeMiB-1)
	if err != nil {
		t.Errorf("Error sending large but ok email: %v", err)
	}
	localC.Wait()
}

func simpleCmd(t *testing.T, c *smtp.Client, cmd string, expected int) string {
	t.Helper()
	if err := c.Text.PrintfLine("%s", cmd); err != nil {
		t.Fatalf("Failed to write %s: %v", cmd, err)
	}

	_, msg, err := c.Text.ReadResponse(expected)
	if err != nil {
		t.Errorf("Incorrect %s response: %v", cmd, err)
	}
	return msg
}

func TestSimpleCommands(t *testing.T) {
	c := mustDial(t, ModeSMTP, false)
	defer c.Close()
	simpleCmd(t, c, "HELP", 214)
	simpleCmd(t, c, "NOOP", 250)
	simpleCmd(t, c, "VRFY", 502)
	simpleCmd(t, c, "EXPN", 502)
}

func TestLongLines(t *testing.T) {
	c := mustDial(t, ModeSMTP, false)
	defer c.Close()

	// Send a not-too-long line.
	simpleCmd(t, c, fmt.Sprintf("%1000s", "x"), 500)

	// Send a very long line, expect an error.
	msg := simpleCmd(t, c, fmt.Sprintf("%1001s", "x"), 554)
	if msg != "error reading command: line too long" {
		t.Errorf("Expected 'line too long', got %v", msg)
	}
}

func TestReset(t *testing.T) {
	c := mustDial(t, ModeSMTP, false)
	defer c.Close()

	if err := c.Mail("from@plain"); err != nil {
		t.Fatalf("MAIL FROM: %v", err)
	}

	if err := c.Reset(); err != nil {
		t.Errorf("RSET: %v", err)
	}

	if err := c.Mail("from@plain"); err != nil {
		t.Errorf("MAIL after RSET: %v", err)
	}
}

func TestRepeatedStartTLS(t *testing.T) {
	c, err := smtp.Dial(smtpAddr)
	if err != nil {
		t.Fatalf("smtp.Dial: %v", err)
	}

	if err = c.StartTLS(tlsConfig); err != nil {
		t.Fatalf("StartTLS: %v", err)
	}

	if err = c.StartTLS(tlsConfig); err == nil {
		t.Errorf("Second STARTTLS did not fail as expected")
	}
}

// Test that STARTTLS fails on a TLS connection.
func TestStartTLSOnTLS(t *testing.T) {
	c := mustDial(t, ModeSubmissionTLS, false)
	defer c.Close()

	if err := c.StartTLS(tlsConfig); err == nil {
		t.Errorf("STARTTLS did not fail as expected")
	}
}

// The server name requested via SNI must only be used as our hostname if one
// of our certificates is valid for it.
func TestServerNameFromSNI(t *testing.T) {
	cases := []struct {
		sni  string
		want string
	}{
		{"localhost", "localhost"},
		{"other.example.com", "localhost"},
		{"localhost; dkim=pass header.d=example.com", "localhost"},
	}

	for _, c := range cases {
		cfg := tlsConfig.Clone()
		cfg.ServerName = c.sni
		cfg.InsecureSkipVerify = true

		// STARTTLS, and then EHLO, whose reply starts with our hostname.
		sc, err := smtp.Dial(smtpAddr)
		if err != nil {
			t.Fatalf("smtp.Dial: %v", err)
		}
		if err := sc.StartTLS(cfg); err != nil {
			t.Fatalf("%q: StartTLS: %v", c.sni, err)
		}
		msg := simpleCmd(t, sc, "EHLO test", 250)
		if !strings.HasPrefix(msg, c.want+" ") {
			t.Errorf("%q: STARTTLS: EHLO reply %q, expected hostname %q",
				c.sni, msg, c.want)
		}
		sc.Close()

		// TLS-wrapped, where the greeting contains our hostname.
		conn, err := tls.Dial("tcp", submissionTLSAddr, cfg)
		if err != nil {
			t.Fatalf("%q: tls.Dial: %v", c.sni, err)
		}
		_, msg, err = textproto.NewConn(conn).ReadResponse(220)
		if err != nil {
			t.Fatalf("%q: reading greeting: %v", c.sni, err)
		}
		if !strings.HasPrefix(msg, c.want+" ") {
			t.Errorf("%q: TLS: greeting %q, expected hostname %q",
				c.sni, msg, c.want)
		}
		conn.Close()
	}
}

func TestAddDKIMSigner(t *testing.T) {
	s := NewServer()
	err := s.AddDKIMSigner("example.com", "selector", "keyfile-does-not-exist")
	if !errors.Is(err, fs.ErrNotExist) {
		t.Errorf("AddDKIMSigner: expected not exist, got %v", err)
	}

	tmpDir := t.ArtifactDir()

	// Invalid PEM file.
	kf1 := tmpDir + "/key1-bad_pem.pem"
	testlib.Rewrite(t, kf1, "not a valid PEM file")
	err = s.AddDKIMSigner("example.com", "selector", kf1)
	if !errors.Is(err, errDecodingPEMBlock) {
		t.Errorf("AddDKIMSigner: expected %v, got %v",
			errDecodingPEMBlock, err)
	}

	// Unsupported block type.
	kf2 := tmpDir + "/key2.pem"
	testlib.Rewrite(t, kf2,
		"-----BEGIN TEST KEY-----\n-----END TEST KEY-----")
	err = s.AddDKIMSigner("example.com", "selector", kf2)
	if !errors.Is(err, errUnsupportedBlockType) {
		t.Errorf("AddDKIMSigner: expected %v, got %v",
			errUnsupportedBlockType, err)
	}

	// x509 error: this is an ed448 key, which is not supported.
	kf3 := tmpDir + "/key3.pem"
	testlib.Rewrite(t, kf3, `-----BEGIN PRIVATE KEY-----
MEcCAQAwBQYDK2VxBDsEOSBHT9DNG6/FNBnRGrLay+jIrK8WrViiVMz9AoXqYSb6
ghwTZSd3E0X8oIFTgs9ch3pxJM1KDrs4NA==
-----END PRIVATE KEY-----`)
	err = s.AddDKIMSigner("example.com", "selector", kf3)
	if !strings.Contains(err.Error(),
		"x509: PKCS#8 wrapping contained private key with unknown algorithm") {
		t.Errorf("AddDKIMSigner: expected x509 error, got %q", err.Error())
	}

	// Unsupported key type: X25519.
	kf4 := tmpDir + "/key4.pem"
	testlib.Rewrite(t, kf4, `-----BEGIN PRIVATE KEY-----
MC4CAQAwBQYDK2VuBCIEIKBUDwEDc5cCv/yEvnA93yk0gXyiTZe7Qip8QU3rJuZC
-----END PRIVATE KEY-----`)
	err = s.AddDKIMSigner("example.com", "selector", kf4)
	if !errors.Is(err, errUnsupportedKeyType) {
		t.Errorf("AddDKIMSigner: expected %v, got %v",
			errUnsupportedKeyType, err)
	}

	// Successful.
	kf5 := tmpDir + "/key5.pem"
	testlib.Rewrite(t, kf5, `-----BEGIN PRIVATE KEY-----
MC4CAQAwBQYDK2VwBCIEID6bjSoiW6g6NJA67RNl0SZ7zpylVOq9w/VGAXF5whnS
-----END PRIVATE KEY-----`)
	err = s.AddDKIMSigner("example.com", "selector", kf5)
	if err != nil {
		t.Errorf("AddDKIMSigner: %v", err)
	}
}

// chanListener is a net.Listener that returns the results sent to it from
// Accept, and net.ErrClosed once the channel is closed.
type chanListener chan acceptResult

type acceptResult struct {
	conn net.Conn
	err  error
}

func (l chanListener) Accept() (net.Conn, error) {
	r, ok := <-l
	if !ok {
		return nil, net.ErrClosed
	}
	return r.conn, r.err
}

func (l chanListener) Close() error   { return nil }
func (l chanListener) Addr() net.Addr { return &net.TCPAddr{} }

// serveInBackground runs s.serve(l) in a goroutine, and returns a channel
// that is closed once it returns.
func serveInBackground(s *Server, l net.Listener) chan struct{} {
	done := make(chan struct{})
	go func() {
		s.serve(l, ModeSMTP)
		close(done)
	}()
	return done
}

func waitForServe(t *testing.T, done chan struct{}) {
	t.Helper()
	select {
	case <-done:
	case <-time.After(5 * time.Second):
		t.Fatalf("serve did not return after net.ErrClosed")
	}
}

// newPipeConn sends a new connection to the listener, and returns the client
// side of it.
func newPipeConn(l chanListener) *textproto.Conn {
	srvConn, cliConn := net.Pipe()
	l <- acceptResult{conn: srvConn}
	return textproto.NewConn(cliConn)
}

func expectGreeting(t *testing.T, c *textproto.Conn, hostname string) {
	t.Helper()
	greeting, err := c.ReadLine()
	if err != nil {
		t.Fatalf("error reading greeting: %v", err)
	}
	if want := "220 " + hostname + " ESMTP chasquid"; greeting != want {
		t.Errorf("got greeting %q, expected %q", greeting, want)
	}
}

func TestAcceptErrors(t *testing.T) {
	s := NewServer()
	s.Hostname = "acceptsrv"
	prevErrors := acceptErrors.Value()

	emfile := &net.OpError{Op: "accept", Net: "tcp",
		Err: os.NewSyscallError("accept4", syscall.EMFILE)}
	srvConn, cliConn := net.Pipe()
	defer cliConn.Close()

	l := make(chanListener, 4)
	l <- acceptResult{err: emfile}
	l <- acceptResult{err: emfile}
	l <- acceptResult{conn: srvConn}
	l <- acceptResult{err: emfile}
	close(l)

	// serve should retry on errors, and return once the listener reports
	// net.ErrClosed.
	done := serveInBackground(s, l)

	// The connection we got in between the errors must have been handled.
	expectGreeting(t, textproto.NewConn(cliConn), "acceptsrv")

	waitForServe(t, done)

	if n := acceptErrors.Value() - prevErrors; n != 3 {
		t.Errorf("acceptErrors increased by %d, expected 3", n)
	}
}

func TestConnectionLimit(t *testing.T) {
	s := NewServer()
	s.Hostname = "limitsrv"
	s.maxConns = 2
	prevOverLimit := connectionsOverLimit.Value()

	l := make(chanListener)
	done := serveInBackground(s, l)

	// The first two connections are within the limit.
	c1 := newPipeConn(l)
	expectGreeting(t, c1, "limitsrv")
	c2 := newPipeConn(l)
	defer c2.Close()
	expectGreeting(t, c2, "limitsrv")

	// The third one is over the limit, and gets closed without a greeting.
	c3 := newPipeConn(l)
	if line, err := c3.ReadLine(); err != io.EOF {
		t.Errorf("over the limit: got %q, %v; expected EOF", line, err)
	}
	if n := connectionsOverLimit.Value() - prevOverLimit; n != 1 {
		t.Errorf("connectionsOverLimit increased by %d, expected 1", n)
	}

	// Once a connection finishes, we can accept new ones again.
	c1.Close()
	deadline := time.Now().Add(5 * time.Second)
	for s.activeConns.Load() != 1 {
		if time.Now().After(deadline) {
			t.Fatalf("activeConns = %d, expected 1", s.activeConns.Load())
		}
		time.Sleep(time.Millisecond)
	}

	c4 := newPipeConn(l)
	defer c4.Close()
	expectGreeting(t, c4, "limitsrv")

	close(l)
	waitForServe(t, done)
}

//
// === Benchmarks ===
//

func BenchmarkManyEmails(b *testing.B) {
	c := mustDial(b, ModeSMTP, false)
	defer c.Close()

	for b.Loop() {
		sendEmail(b, c)
	}
}

func BenchmarkManyEmailsParallel(b *testing.B) {
	b.RunParallel(func(pb *testing.PB) {
		c := mustDial(b, ModeSMTP, false)
		defer c.Close()

		for pb.Next() {
			sendEmail(b, c)
		}
	})
}

//
// === Test environment ===
//

type brokenAuthBE struct{}

func (b brokenAuthBE) Authenticate(user, password string) (bool, error) {
	return false, fmt.Errorf("failed to auth")
}

func (b brokenAuthBE) Exists(user string) (bool, error) {
	return false, fmt.Errorf("failed to check if user exists")
}

func (b brokenAuthBE) Reload() error {
	return fmt.Errorf("failed to reload")
}

// realMain is the real main function, which returns the value to pass to
// os.Exit(). We have to do this so we can use defer.
func realMain(m *testing.M) int {
	flag.Parse()

	// Create a 1MiB string, which the large message tests use.
	str1MiB = strings.Repeat("a", 1024*1024)

	// Set up the mail log to stdout, which is captured by the test runner,
	// so we have better debugging information on failures.
	maillog.Default = maillog.New(os.Stdout)

	if *externalSMTPAddr != "" {
		smtpAddr = *externalSMTPAddr
		submissionAddr = *externalSubmissionAddr
		submissionTLSAddr = *externalSubmissionTLSAddr
		tlsConfig = &tls.Config{
			InsecureSkipVerify: true,
		}
	} else {
		// Generate certificates in a temporary directory.
		tmpDir, err := os.MkdirTemp("", "chasquid_test:")
		if err != nil {
			fmt.Printf("Failed to create temp dir: %v\n", tmpDir)
			return 1
		}
		defer os.RemoveAll(tmpDir)

		tlsConfig, err = testlib.GenerateCert(tmpDir)
		if err != nil {
			fmt.Printf("Failed to generate cert for testing: %v\n", err)
			return 1
		}

		s := NewServer()
		s.Hostname = "localhost"
		s.MaxDataSize = int64(maxDataSizeMiB) * 1024 * 1024
		s.AddCerts(tmpDir+"/cert.pem", tmpDir+"/key.pem")

		// Open the listeners ourselves, so the ports are already bound
		// (and we know their addresses) by the time the tests run.
		for _, l := range []struct {
			addr *string
			mode SocketMode
		}{
			{&smtpAddr, ModeSMTP},
			{&submissionAddr, ModeSubmission},
			{&submissionTLSAddr, ModeSubmissionTLS},
		} {
			ln, err := net.Listen("tcp", "localhost:0")
			if err != nil {
				fmt.Printf("Failed to listen: %v\n", err)
				return 1
			}
			*l.addr = ln.Addr().String()
			s.AddListeners([]net.Listener{ln}, l.mode)
		}

		s.InitQueue(tmpDir+"/queue", localC, remoteC)

		dinfo, err := domaininfo.New(tmpDir + "/domaininfo")
		if err != nil {
			fmt.Printf("Error initializing domaininfo: %v", err)
			return 1
		}
		s.SetDomainInfo(dinfo)

		udb := userdb.New("/dev/null")
		udb.AddUser("testuser", "testpasswd")
		s.aliasesR.AddAliasForTesting(
			"to@localhost", "testuser@localhost", nil, aliases.EMAIL)
		s.authr.Register("localhost", auth.WrapNoErrorBackend(udb))
		s.AddDomain("localhost")

		s.AddDomain("broken")
		s.authr.Register("broken", &brokenAuthBE{})

		// Disable SPF lookups, to avoid leaking DNS queries.
		disableSPFForTesting = true

		// Disable reloading.
		reloadEvery = nil

		go s.ListenAndServe()
	}

	return m.Run()
}

func TestMain(m *testing.M) {
	os.Exit(realMain(m))
}