felhom-controller/controller/internal/mailrelay/mailrelay_test.go

package mailrelay

import (
	"bytes"
	"context"
	"crypto/tls"
	"io"
	"log"
	netsmtp "net/smtp"
	"sync"
	"testing"

	"github.com/emersion/go-smtp"
)

// fakeForwarder records every Forward call so tests can assert call count + byte
// fidelity, and returns a programmable status/err.
type fakeForwarder struct {
	mu       sync.Mutex
	calls    int
	lastRaw  []byte
	lastFrom string
	lastTo   []string
	status   int
	body     string
	err      error
}

func (f *fakeForwarder) Forward(_ context.Context, raw []byte, from string, to []string) (int, string, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.calls++
	f.lastRaw = append([]byte(nil), raw...)
	f.lastFrom = from
	f.lastTo = append([]string(nil), to...)
	return f.status, f.body, f.err
}

func (f *fakeForwarder) callCount() int {
	f.mu.Lock()
	defer f.mu.Unlock()
	return f.calls
}

func quietLogger() *log.Logger { return log.New(io.Discard, "", 0) }

func rawMsg(from string) []byte {
	return []byte("From: " + from + "\r\nTo: recipient@example.com\r\nSubject: test\r\n\r\nhello world\r\n")
}

func newTestServer(t *testing.T, allowed []string, fwd Forwarder) *Server {
	t.Helper()
	s, err := New(Options{
		Policy:    NewPolicy(allowed),
		Forwarder: fwd,
		Logger:    quietLogger(),
	})
	if err != nil {
		t.Fatalf("New: %v", err)
	}
	return s
}

// dataResult drives a session.Data with the given raw message and returns the reply.
func dataResult(srv *Server, from string, raw []byte) error {
	sess := &session{srv: srv, from: from, to: []string{"recipient@example.com"}}
	return sess.Data(bytes.NewReader(raw))
}

// --- §7 A: happy path — passthrough, NOT parse (§10 byte-equality companion) ---

func TestData_HappyPath_ForwardsRawBytesUnchanged(t *testing.T) {
	fwd := &fakeForwarder{status: 200}
	srv := newTestServer(t, []string{"felhom.eu"}, fwd)

	raw := rawMsg("vaultwarden@felhom.eu")
	if err := dataResult(srv, "vaultwarden@felhom.eu", raw); err != nil {
		t.Fatalf("expected success (250), got %v", err)
	}
	if fwd.callCount() != 1 {
		t.Fatalf("expected exactly 1 forward, got %d", fwd.callCount())
	}
	// Passthrough: the bytes handed to the hub must be byte-for-byte the input — no
	// MIME parse/re-encode (the parse path silently drops inline CID images; §10).
	if !bytes.Equal(fwd.lastRaw, raw) {
		t.Fatalf("forwarded bytes differ from input:\n got=%q\nwant=%q", fwd.lastRaw, raw)
	}
	if fwd.lastFrom != "vaultwarden@felhom.eu" {
		t.Fatalf("envelope from = %q", fwd.lastFrom)
	}
}

// --- §7 B: From outside the allowlist is refused BEFORE any hub call ---

func TestData_FromOutsideAllowlist_Rejected_NoForward(t *testing.T) {
	fwd := &fakeForwarder{status: 200}
	srv := newTestServer(t, []string{"felhom.eu"}, fwd)

	err := dataResult(srv, "vaultwarden@felhom.eu", rawMsg("evil@notfelhom.example"))
	smtpErr, ok := err.(*smtp.SMTPError)
	if !ok {
		t.Fatalf("expected *smtp.SMTPError, got %T (%v)", err, err)
	}
	if smtpErr.Code != 550 {
		t.Fatalf("expected 550, got %d", smtpErr.Code)
	}
	// The hub/Resend must NEVER be dialed for a doomed message.
	if fwd.callCount() != 0 {
		t.Fatalf("forwarder must not be called on From-reject, got %d calls", fwd.callCount())
	}
}

// Companion red-proof for §7 B: REMOVE the gate (allow the bad domain) and the same
// message reaches the forwarder. Proves the allowlist check is load-bearing.
func TestData_FromReject_CompanionProof_GateRemovedLetsItThrough(t *testing.T) {
	fwd := &fakeForwarder{status: 200}
	srv := newTestServer(t, []string{"notfelhom.example"}, fwd) // gate effectively removed

	if err := dataResult(srv, "x@notfelhom.example", rawMsg("evil@notfelhom.example")); err != nil {
		t.Fatalf("with the domain allowed the message should pass, got %v", err)
	}
	if fwd.callCount() != 1 {
		t.Fatalf("companion: expected the wrong-From to reach the forwarder once, got %d", fwd.callCount())
	}
}

func TestData_UnparseableFrom_Rejected(t *testing.T) {
	fwd := &fakeForwarder{status: 200}
	srv := newTestServer(t, []string{"felhom.eu"}, fwd)

	// No From header at all → fail-closed reject, no forward.
	noFrom := []byte("To: recipient@example.com\r\nSubject: test\r\n\r\nbody\r\n")
	err := dataResult(srv, "x@felhom.eu", noFrom)
	if smtpErr, ok := err.(*smtp.SMTPError); !ok || smtpErr.Code != 550 {
		t.Fatalf("expected 550 for missing From, got %v", err)
	}
	if fwd.callCount() != 0 {
		t.Fatalf("forwarder must not be called when From is unparseable, got %d", fwd.callCount())
	}
}

// --- §7 D: transient hub outage — single-shot, surfaces transient error ---

func TestData_HubUnreachable_SingleShot_Transient(t *testing.T) {
	fwd := &fakeForwarder{err: io.ErrUnexpectedEOF} // simulate hub unreachable
	srv := newTestServer(t, []string{"felhom.eu"}, fwd)

	err := dataResult(srv, "vaultwarden@felhom.eu", rawMsg("vaultwarden@felhom.eu"))
	smtpErr, ok := err.(*smtp.SMTPError)
	if !ok {
		t.Fatalf("expected *smtp.SMTPError, got %T", err)
	}
	if smtpErr.Code/100 != 4 {
		t.Fatalf("expected transient 4xx, got %d", smtpErr.Code)
	}
	// §10 single-shot companion: EXACTLY one attempt — no retry loop.
	if fwd.callCount() != 1 {
		t.Fatalf("v1 forward must be single-shot; got %d attempts", fwd.callCount())
	}
}

// --- status mapping (spike §6) ---

func TestMapStatusToSMTP(t *testing.T) {
	cases := []struct {
		status   int
		wantNil  bool
		wantCode int
	}{
		{200, true, 0},
		{202, true, 0},
		{403, false, 451}, // hub From-reject backstop → transient per the prompt mapping
		{429, false, 451}, // rate limit → retry later
		{451, false, 451},
		{500, false, 554},
		{503, false, 554},
	}
	for _, c := range cases {
		got := mapStatusToSMTP(c.status, "")
		if c.wantNil {
			if got != nil {
				t.Errorf("status %d: expected nil, got %v", c.status, got)
			}
			continue
		}
		smtpErr, ok := got.(*smtp.SMTPError)
		if !ok {
			t.Errorf("status %d: expected *smtp.SMTPError, got %T", c.status, got)
			continue
		}
		if smtpErr.Code != c.wantCode {
			t.Errorf("status %d: expected SMTP %d, got %d", c.status, c.wantCode, smtpErr.Code)
		}
	}
}

func TestSanitizeReason(t *testing.T) {
	if got := sanitizeReason("line one\r\nleaked@addr.example"); got != "line one" {
		t.Fatalf("sanitizeReason should keep only the first line, got %q", got)
	}
	long := make([]byte, 500)
	for i := range long {
		long[i] = 'a'
	}
	if got := sanitizeReason(string(long)); len(got) != 200 {
		t.Fatalf("sanitizeReason should cap at 200, got len %d", len(got))
	}
}

// --- end-to-end over a real socket (Q1/Q2: real SMTP, STARTTLS, no auth) ---

func TestLifecycle_StartStopIdempotent(t *testing.T) {
	fwd := &fakeForwarder{status: 200}
	lc := NewLifecycle(Options{
		PlainAddr:      "127.0.0.1:0",
		TLSAddr:        "127.0.0.1:0",
		PlainNoTLSAddr: "127.0.0.1:0",
		Policy:         NewPolicy([]string{"felhom.eu"}),
		Forwarder:      fwd,
		Logger:         quietLogger(),
	})
	if lc.Running() {
		t.Fatal("should not be running before Apply")
	}
	if err := lc.Apply(true); err != nil {
		t.Fatalf("Apply(true): %v", err)
	}
	if !lc.Running() {
		t.Fatal("should be running after Apply(true)")
	}
	// idempotent enable
	if err := lc.Apply(true); err != nil {
		t.Fatalf("second Apply(true): %v", err)
	}
	if err := lc.Apply(false); err != nil {
		t.Fatalf("Apply(false): %v", err)
	}
	if lc.Running() {
		t.Fatal("should be stopped after Apply(false)")
	}
	// re-enable after disable works (uses a fresh server)
	if err := lc.Apply(true); err != nil {
		t.Fatalf("re-enable: %v", err)
	}
	if !lc.Running() {
		t.Fatal("should run after re-enable")
	}
	lc.Close()
	if lc.Running() {
		t.Fatal("Close should stop the shim")
	}
	// Apply(true) after Close is refused (no resurrection on shutdown)
	if err := lc.Apply(true); err != nil {
		t.Fatalf("Apply after Close should be a no-op nil, got %v", err)
	}
	if lc.Running() {
		t.Fatal("must stay stopped after Close")
	}
}

// Gap-1: the :2526 listener must NOT advertise STARTTLS (TLSConfig nil), while :2525 must.
// Proven both at the config level and over a real EHLO.
func TestServer_PlainNoTLSListener_NoSTARTTLS(t *testing.T) {
	fwd := &fakeForwarder{status: 200}
	s, err := New(Options{
		PlainAddr:      "127.0.0.1:0",
		TLSAddr:        "127.0.0.1:0",
		PlainNoTLSAddr: "127.0.0.1:0",
		Policy:         NewPolicy([]string{"felhom.eu"}),
		Forwarder:      fwd,
		Logger:         quietLogger(),
	})
	if err != nil {
		t.Fatalf("New: %v", err)
	}
	// Config-level: STARTTLS is advertised iff TLSConfig != nil.
	if s.plainNoTLSSrv.TLSConfig != nil {
		t.Fatal(":2526 server must have TLSConfig==nil (so STARTTLS is not advertised)")
	}
	if s.plainSrv.TLSConfig == nil {
		t.Fatal(":2525 server must keep TLSConfig (STARTTLS advertised)")
	}
	if err := s.Start(); err != nil {
		t.Fatalf("Start: %v", err)
	}
	defer s.Close()

	// Real EHLO: :2526 must NOT offer STARTTLS; :2525 must.
	advertises := func(addr string) bool {
		c, err := netsmtp.Dial(addr)
		if err != nil {
			t.Fatalf("dial %s: %v", addr, err)
		}
		defer c.Close()
		if err := c.Hello("test.local"); err != nil {
			t.Fatalf("EHLO %s: %v", addr, err)
		}
		ok, _ := c.Extension("STARTTLS")
		return ok
	}
	if advertises(s.PlainNoTLSAddr()) {
		t.Error(":2526 must NOT advertise STARTTLS over EHLO")
	}
	if !advertises(s.PlainAddr()) {
		t.Error(":2525 must advertise STARTTLS over EHLO")
	}
}

func TestServer_EndToEnd_STARTTLS(t *testing.T) {
	fwd := &fakeForwarder{status: 200}
	s, err := New(Options{
		PlainAddr:      "127.0.0.1:0",
		TLSAddr:        "127.0.0.1:0",
		PlainNoTLSAddr: "127.0.0.1:0",
		Policy:         NewPolicy([]string{"felhom.eu"}),
		Forwarder:      fwd,
		Logger:         quietLogger(),
	})
	if err != nil {
		t.Fatalf("New: %v", err)
	}
	if err := s.Start(); err != nil {
		t.Fatalf("Start: %v", err)
	}
	defer s.Close()

	c, err := netsmtp.Dial(s.PlainAddr())
	if err != nil {
		t.Fatalf("dial: %v", err)
	}
	defer c.Close()
	if err := c.StartTLS(&tls.Config{InsecureSkipVerify: true}); err != nil {
		t.Fatalf("starttls: %v", err)
	}
	if err := c.Mail("vaultwarden@felhom.eu"); err != nil {
		t.Fatalf("mail: %v", err)
	}
	if err := c.Rcpt("recipient@example.com"); err != nil {
		t.Fatalf("rcpt: %v", err)
	}
	wc, err := c.Data()
	if err != nil {
		t.Fatalf("data: %v", err)
	}
	msg := rawMsg("vaultwarden@felhom.eu")
	if _, err := wc.Write(msg); err != nil {
		t.Fatalf("write: %v", err)
	}
	if err := wc.Close(); err != nil {
		t.Fatalf("close data (expected 250): %v", err)
	}
	_ = c.Quit()

	if fwd.callCount() != 1 {
		t.Fatalf("expected 1 forward through the real listener, got %d", fwd.callCount())
	}
}