mailgoserver/internal/relay/tls_test.go

package relay

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"io"
	"math/big"
	"net"
	"strconv"
	"sync/atomic"
	"testing"
	"time"

	"github.com/emersion/go-smtp"

	"mailgoserver/internal/toolbox"
)

// acceptAllBackend is a minimal go-smtp Backend/Session that accepts every command —
// a stand-in "receiving MTA" for trySend tests, not testing go-smtp itself.
type acceptAllBackend struct{ received chan string }

func (b *acceptAllBackend) NewSession(*smtp.Conn) (smtp.Session, error) {
	return &acceptAllSession{received: b.received}, nil
}

type acceptAllSession struct{ received chan string }

func (s *acceptAllSession) Mail(from string, opts *smtp.MailOptions) error { return nil }
func (s *acceptAllSession) Rcpt(to string, opts *smtp.RcptOptions) error  { return nil }
func (s *acceptAllSession) Data(r io.Reader) error {
	b, err := io.ReadAll(r)
	if err != nil {
		return err
	}
	if s.received != nil {
		s.received <- string(b)
	}
	return nil
}
func (s *acceptAllSession) Reset()        {}
func (s *acceptAllSession) Logout() error { return nil }

// countingListener counts every accepted TCP connection — used to distinguish "one
// connection, verified on the first try" from "two connections, the first abandoned
// after a failed TLS handshake and a second one succeeding via the unverified
// fallback" without needing to inspect trySend's internals or log output directly.
type countingListener struct {
	net.Listener
	accepted *atomic.Int32
}

func (c *countingListener) Accept() (net.Conn, error) {
	conn, err := c.Listener.Accept()
	if err == nil {
		c.accepted.Add(1)
	}
	return conn, err
}

// genTestCert generates a throwaway self-signed CA and a leaf certificate for
// "127.0.0.1" signed by it — enough to exercise real x509 verification (via
// Relay.rootCAs) without needing a system-trusted cert in a test environment.
func genTestCert(t *testing.T) (leafCert tls.Certificate, caPool *x509.CertPool) {
	t.Helper()
	caKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	caTemplate := &x509.Certificate{
		SerialNumber:          big.NewInt(1),
		Subject:               pkix.Name{CommonName: "test CA"},
		NotBefore:             time.Now().Add(-time.Hour),
		NotAfter:              time.Now().Add(time.Hour),
		IsCA:                  true,
		KeyUsage:              x509.KeyUsageCertSign,
		BasicConstraintsValid: true,
	}
	caDER, err := x509.CreateCertificate(rand.Reader, caTemplate, caTemplate, &caKey.PublicKey, caKey)
	if err != nil {
		t.Fatal(err)
	}
	caCert, err := x509.ParseCertificate(caDER)
	if err != nil {
		t.Fatal(err)
	}

	leafKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	leafTemplate := &x509.Certificate{
		SerialNumber: big.NewInt(2),
		Subject:      pkix.Name{CommonName: "127.0.0.1"},
		NotBefore:    time.Now().Add(-time.Hour),
		NotAfter:     time.Now().Add(time.Hour),
		IPAddresses:  []net.IP{net.ParseIP("127.0.0.1")},
		KeyUsage:     x509.KeyUsageDigitalSignature,
		ExtKeyUsage:  []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
	}
	leafDER, err := x509.CreateCertificate(rand.Reader, leafTemplate, caCert, &leafKey.PublicKey, caKey)
	if err != nil {
		t.Fatal(err)
	}

	pool := x509.NewCertPool()
	pool.AddCert(caCert)
	cert := tls.Certificate{Certificate: [][]byte{leafDER}, PrivateKey: leafKey}
	return cert, pool
}

// startTestMTA starts a real go-smtp server (STARTTLS-capable) presenting cert on
// 127.0.0.1, returning its port, a channel receiving each DATA payload, and a counter
// of accepted TCP connections.
func startTestMTA(t *testing.T, cert tls.Certificate) (port int, received chan string, accepted *atomic.Int32) {
	t.Helper()
	received = make(chan string, 2)
	accepted = &atomic.Int32{}
	be := &acceptAllBackend{received: received}
	s := smtp.NewServer(be)
	s.Domain = "mx.example.com"
	s.TLSConfig = &tls.Config{Certificates: []tls.Certificate{cert}}
	s.AllowInsecureAuth = true

	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	cl := &countingListener{Listener: l, accepted: accepted}
	t.Cleanup(func() { s.Close() })
	go s.Serve(cl)

	_, portStr, _ := net.SplitHostPort(l.Addr().String())
	port, err = strconv.Atoi(portStr)
	if err != nil {
		t.Fatal(err)
	}
	return port, received, accepted
}

func TestTrySendVerifiesRealCertOnFirstTry(t *testing.T) {
	cert, caPool := genTestCert(t)
	port, received, accepted := startTestMTA(t, cert)

	r := &Relay{Hostname: "sender.example.com", Timeout: 5 * time.Second, Logger: toolbox.GetLogger("relay_test"), port: port, rootCAs: caPool}

	resp, err := r.trySend("127.0.0.1", "test.invalid", "from@example.com", []string{"to@example.com"}, "Subject: hi\r\n\r\nbody")
	if err != nil {
		t.Fatalf("expected delivery to succeed with a CA-trusted cert, got: %v", err)
	}
	if resp == "" {
		t.Fatal("expected a non-empty server response")
	}
	select {
	case <-received:
	case <-time.After(2 * time.Second):
		t.Fatal("expected the test MTA to receive the message")
	}
	// Exactly one connection: verification succeeded on the very first try, no
	// unverified-fallback retry (which would show up as a second accepted connection).
	if got := accepted.Load(); got != 1 {
		t.Errorf("expected exactly 1 accepted connection (verified first try), got %d", got)
	}
}

func TestTrySendFallsBackToUnverifiedOnUntrustedCert(t *testing.T) {
	cert, _ := genTestCert(t) // leaf cert generated, but its CA is deliberately not trusted below
	port, received, accepted := startTestMTA(t, cert)

	// An empty pool (not nil) so this can't accidentally pass by falling through to a
	// real system-trusted cert.
	r := &Relay{Hostname: "sender.example.com", Timeout: 5 * time.Second, Logger: toolbox.GetLogger("relay_test"), port: port, rootCAs: x509.NewCertPool()}

	resp, err := r.trySend("127.0.0.1", "test.invalid", "from@example.com", []string{"to@example.com"}, "Subject: hi\r\n\r\nbody")
	if err != nil {
		t.Fatalf("expected delivery to still succeed via the unverified fallback, got: %v", err)
	}
	if resp == "" {
		t.Fatal("expected a non-empty server response")
	}
	select {
	case <-received:
	case <-time.After(2 * time.Second):
		t.Fatal("expected the test MTA to receive the message despite the untrusted cert")
	}
	// Two connections: the first is abandoned after a failed verified STARTTLS
	// handshake, the second succeeds via the InsecureSkipVerify fallback.
	if got := accepted.Load(); got != 2 {
		t.Errorf("expected exactly 2 accepted connections (verified attempt + unverified fallback), got %d", got)
	}
}

// TestTrySendSkipsFallbackForMTASTSEnforcedDomain confirms a domain whose policy
// resolves as MTA-STS-enforced never gets the unverified-fallback retry — it should
// fail outright on the untrusted cert, matching what a domain that opted into strict
// enforcement asked for.
func TestTrySendSkipsFallbackForMTASTSEnforcedDomain(t *testing.T) {
	cert, _ := genTestCert(t) // untrusted below, same as the fallback test
	port, _, accepted := startTestMTA(t, cert)

	r := &Relay{
		Hostname: "sender.example.com", Timeout: 5 * time.Second, Logger: toolbox.GetLogger("relay_test"),
		port: port, rootCAs: x509.NewCertPool(),
		mtaSTSCheck: func(domain string) bool { return domain == "strict.example" },
	}

	_, err := r.trySend("127.0.0.1", "strict.example", "from@example.com", []string{"to@example.com"}, "Subject: hi\r\n\r\nbody")
	if err == nil {
		t.Fatal("expected delivery to fail outright for an MTA-STS-enforced domain with an untrusted cert, not fall back")
	}
	// Exactly one connection: the verified attempt, abandoned — no second
	// (unverified-fallback) connection should ever have been made.
	if got := accepted.Load(); got != 1 {
		t.Errorf("expected exactly 1 accepted connection (no fallback attempt), got %d", got)
	}
}