mailgoserver/internal/smtpserver/domain_ratelimit_test.go

package smtpserver

import (
	"net/smtp"
	"strings"
	"testing"
	"time"

	"mailgoserver/internal/db"
)

// TestDomainSendRateLimited is a white-box unit test of the core rate-limit decision
// (no real SMTP transaction, no network) — see TestSendRateLimitRejectsOverLimitMail
// below for the live end-to-end wiring check.
func TestDomainSendRateLimited(t *testing.T) {
	backend := newTestBackend(t)
	domainID, err := backend.DB.CreateDomain("ratelimited.example")
	if err != nil {
		t.Fatal(err)
	}
	limit := 2
	if err := backend.DB.SetDomainSendRateLimit(domainID, &limit); err != nil {
		t.Fatal(err)
	}

	if limited, err := backend.domainSendRateLimited("ratelimited.example"); err != nil || limited {
		t.Fatalf("expected not limited with zero sends so far, limited=%v err=%v", limited, err)
	}

	for i := 0; i < 2; i++ {
		if _, err := backend.DB.InsertEmailLog(db.EmailLog{
			MessageID: "m" + string(rune('a'+i)) + "@ratelimited.example", Timestamp: time.Now(),
			MailFrom: "sender@ratelimited.example", EmailHeaders: "h", Status: "relayed",
		}); err != nil {
			t.Fatal(err)
		}
	}
	if limited, err := backend.domainSendRateLimited("ratelimited.example"); err != nil || !limited {
		t.Fatalf("expected limited after hitting the cap of 2, limited=%v err=%v", limited, err)
	}

	// An unconfigured domain (no limit set) must never limit, regardless of volume.
	if limited, err := backend.domainSendRateLimited("example.com"); err != nil || limited {
		t.Fatalf("expected no limit for a domain with send_rate_limit_per_hour unset, limited=%v err=%v", limited, err)
	}

	// An unrecognized domain must never limit either (not this server's problem to
	// cap, and GetDomainByName returning nil must fail open, not error).
	if limited, err := backend.domainSendRateLimited("nowhere.invalid"); err != nil || limited {
		t.Fatalf("expected no limit for an unrecognized domain, limited=%v err=%v", limited, err)
	}
}

// TestSendRateLimitRejectsOverLimitMail confirms the real wiring: an authorized
// sender's relay recipient gets a 450 with a rate-limit reason recorded once the
// domain's cap is hit, via a real SMTP transaction end to end.
func TestSendRateLimitRejectsOverLimitMail(t *testing.T) {
	backend, _ := newTestBackendWithMailbox(t)
	domains, err := backend.DB.ListDomains()
	if err != nil || len(domains) == 0 {
		t.Fatal(err)
	}
	limit := 0
	if err := backend.DB.SetDomainSendRateLimit(domains[0].ID, &limit); err != nil {
		t.Fatal(err)
	}
	addr := startTestServer(t, backend)

	c, err := smtp.Dial(addr)
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()
	if err := c.Auth(smtp.PlainAuth("", "test@example.com", "testpass123", "127.0.0.1")); err != nil {
		t.Fatalf("auth: %v", err)
	}
	if err := c.Mail("test@example.com"); err != nil {
		t.Fatalf("MAIL FROM: %v", err)
	}
	if err := c.Rcpt("someone@elsewhere.example"); err != nil {
		t.Fatalf("RCPT: %v", err)
	}
	w, err := c.Data()
	if err != nil {
		t.Fatal(err)
	}
	w.Write([]byte("Subject: hi\r\n\r\nhi"))
	// The final response is the same generic 550 either way (see Data()'s fixed
	// message for a total failure) — what actually distinguishes "rate limited" from
	// "a real relay failure" is the recorded per-recipient reason, checked below.
	w.Close()

	var errMsg string
	if err := backend.DB.QueryRow(`SELECT error_message FROM esrv_email_recipient_logs ORDER BY id DESC LIMIT 1`).Scan(&errMsg); err != nil {
		t.Fatal(err)
	}
	if errMsg == "" || !strings.Contains(errMsg, "rate limit") {
		t.Fatalf("expected the recipient log's error to mention the rate limit, got %q", errMsg)
	}
}