lunar-mini/server/internal/handler/email_auth_test.go

package handler

import (
	"bufio"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"fmt"
	"math/big"
	"net"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/gouki/lunar-server/internal/config"
	"github.com/joho/godotenv"
)

type smtpTestResult struct {
	message       string
	authenticated bool
	err           error
}

type smtpTestServer struct {
	port    string
	caFile  string
	tlsMode string
	results chan smtpTestResult
}

func startSMTPTestServer(t *testing.T, tlsMode string, rejectAuthentication bool) smtpTestServer {
	t.Helper()
	certificate, caFile := createSMTPTestCertificate(t)
	listener, err := net.Listen("tcp", "localhost:0")
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { listener.Close() })
	results := make(chan smtpTestResult, 1)
	port := fmt.Sprint(listener.Addr().(*net.TCPAddr).Port)
	go serveSMTPTestConnection(listener, certificate, tlsMode, rejectAuthentication, results)
	return smtpTestServer{port: port, caFile: caFile, tlsMode: tlsMode, results: results}
}

func createSMTPTestCertificate(t *testing.T) (tls.Certificate, string) {
	t.Helper()
	now := time.Now()
	rootKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	rootTemplate := &x509.Certificate{
		SerialNumber:          big.NewInt(1),
		NotBefore:             now.Add(-time.Minute),
		NotAfter:              now.Add(time.Hour),
		IsCA:                  true,
		BasicConstraintsValid: true,
		KeyUsage:              x509.KeyUsageCertSign | x509.KeyUsageDigitalSignature,
	}
	rootDER, err := x509.CreateCertificate(rand.Reader, rootTemplate, rootTemplate, &rootKey.PublicKey, rootKey)
	if err != nil {
		t.Fatal(err)
	}
	rootCertificate, err := x509.ParseCertificate(rootDER)
	if err != nil {
		t.Fatal(err)
	}
	serverKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	serverTemplate := &x509.Certificate{
		SerialNumber: big.NewInt(2),
		NotBefore:    now.Add(-time.Minute),
		NotAfter:     now.Add(time.Hour),
		DNSNames:     []string{"localhost"},
		KeyUsage:     x509.KeyUsageDigitalSignature,
		ExtKeyUsage:  []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
	}
	serverDER, err := x509.CreateCertificate(rand.Reader, serverTemplate, rootCertificate, &serverKey.PublicKey, rootKey)
	if err != nil {
		t.Fatal(err)
	}
	privateKeyDER, err := x509.MarshalECPrivateKey(serverKey)
	if err != nil {
		t.Fatal(err)
	}
	certificate, err := tls.X509KeyPair(
		pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: serverDER}),
		pem.EncodeToMemory(&pem.Block{Type: "EC PRIVATE KEY", Bytes: privateKeyDER}),
	)
	if err != nil {
		t.Fatal(err)
	}
	caFile := filepath.Join(t.TempDir(), "smtp-test-ca.pem")
	if err := os.WriteFile(caFile, pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: rootDER}), 0600); err != nil {
		t.Fatal(err)
	}
	return certificate, caFile
}

func serveSMTPTestConnection(listener net.Listener, certificate tls.Certificate, tlsMode string, rejectAuthentication bool, results chan<- smtpTestResult) {
	connection, err := listener.Accept()
	if err != nil {
		results <- smtpTestResult{err: err}
		return
	}
	defer connection.Close()
	serverTLSConfig := &tls.Config{Certificates: []tls.Certificate{certificate}, MinVersion: tls.VersionTLS12}
	tlsActive := false
	if tlsMode == "implicit" {
		tlsConnection := tls.Server(connection, serverTLSConfig)
		if err := tlsConnection.Handshake(); err != nil {
			results <- smtpTestResult{err: err}
			return
		}
		connection = tlsConnection
		tlsActive = true
	}
	reader := bufio.NewReader(connection)
	writer := bufio.NewWriter(connection)
	writeResponse := func(response string) error {
		if _, err := writer.WriteString(response + "\r\n"); err != nil {
			return err
		}
		return writer.Flush()
	}
	if err := writeResponse("220 localhost ESMTP"); err != nil {
		results <- smtpTestResult{err: err}
		return
	}
	var message strings.Builder
	for {
		command, err := reader.ReadString('\n')
		if err != nil {
			results <- smtpTestResult{err: err}
			return
		}
		trimmed := strings.TrimRight(command, "\r\n")
		switch {
		case strings.HasPrefix(trimmed, "EHLO "):
			capabilities := "250-localhost\r\n250-AUTH PLAIN\r\n250 OK\r\n"
			if tlsMode == "starttls" && !tlsActive {
				capabilities = "250-localhost\r\n250-STARTTLS\r\n250 OK\r\n"
			}
			if _, err := writer.WriteString(capabilities); err != nil {
				results <- smtpTestResult{err: err}
				return
			}
			if err := writer.Flush(); err != nil {
				results <- smtpTestResult{err: err}
				return
			}
		case trimmed == "STARTTLS":
			if err := writeResponse("220 2.0.0 ready to start TLS"); err != nil {
				results <- smtpTestResult{err: err}
				return
			}
			tlsConnection := tls.Server(connection, serverTLSConfig)
			if err := tlsConnection.Handshake(); err != nil {
				results <- smtpTestResult{err: err}
				return
			}
			connection = tlsConnection
			reader = bufio.NewReader(connection)
			writer = bufio.NewWriter(connection)
			tlsActive = true
		case strings.HasPrefix(trimmed, "AUTH PLAIN"):
			encoded := strings.TrimSpace(strings.TrimPrefix(trimmed, "AUTH PLAIN"))
			if encoded == "" {
				if err := writeResponse("334 "); err != nil {
					results <- smtpTestResult{err: err}
					return
				}
				encoded, err = reader.ReadString('\n')
				if err != nil {
					results <- smtpTestResult{err: err}
					return
				}
				encoded = strings.TrimSpace(encoded)
			}
			credentials, decodeErr := base64.StdEncoding.DecodeString(encoded)
			parts := strings.Split(string(credentials), "\x00")
			authenticated := decodeErr == nil && len(parts) == 3 && parts[1] == "smtp-user@example.test" && parts[2] == "smtp-test-password"
			if rejectAuthentication || !authenticated {
				writeResponse("535 5.7.8 Authentication credentials invalid")
				results <- smtpTestResult{authenticated: false}
				return
			}
			if err := writeResponse("235 2.7.0 Authentication successful"); err != nil {
				results <- smtpTestResult{err: err}
				return
			}
		case strings.HasPrefix(trimmed, "MAIL FROM:") || strings.HasPrefix(trimmed, "RCPT TO:"):
			if err := writeResponse("250 2.1.0 accepted"); err != nil {
				results <- smtpTestResult{err: err}
				return
			}
		case trimmed == "DATA":
			if err := writeResponse("354 end with <CRLF>.<CRLF>"); err != nil {
				results <- smtpTestResult{err: err}
				return
			}
			for {
				line, err := reader.ReadString('\n')
				if err != nil {
					results <- smtpTestResult{err: err}
					return
				}
				if line == ".\r\n" {
					break
				}
				message.WriteString(strings.TrimPrefix(line, "."))
			}
			if err := writeResponse("250 2.0.0 queued"); err != nil {
				results <- smtpTestResult{err: err}
				return
			}
			results <- smtpTestResult{message: message.String(), authenticated: true}
		case trimmed == "QUIT":
			writeResponse("221 2.0.0 bye")
			return
		default:
			if err := writeResponse("250 OK"); err != nil {
				results <- smtpTestResult{err: err}
				return
			}
		}
	}
}

func configureEmailCodeAPI(t *testing.T, server smtpTestServer) *gin.Engine {
	t.Helper()
	t.Setenv("SERVER_ENV", "test")
	t.Setenv("DB_TYPE", "sqlite")
	t.Setenv("DB_NAME", fmt.Sprintf("file:email-auth-smtp-http-test-%s?mode=memory&cache=shared", server.port))
	t.Setenv("JWT_SECRET", "email-auth-test-secret")
	t.Setenv("SMTP_HOST", "localhost")
	t.Setenv("SMTP_PORT", server.port)
	t.Setenv("SMTP_USERNAME", "smtp-user@example.test")
	t.Setenv("SMTP_PASSWORD", "smtp-test-password")
	t.Setenv("SMTP_FROM_EMAIL", "login@example.test")
	t.Setenv("SMTP_TLS_MODE", server.tlsMode)
	t.Setenv("SMTP_TLS_CA_FILE", server.caFile)
	t.Setenv("POSTMARK_SERVER_TOKEN", "")
	t.Setenv("POSTMARK_FROM_EMAIL", "")
	if err := config.InitDB(config.Load()); err != nil {
		t.Fatal(err)
	}
	router := gin.New()
	router.POST("/api/auth/email/request", RequestEmailLoginCode)
	return router
}

func TestRequestEmailLoginCodeUsesSMTPViaAPI(t *testing.T) {
	gin.SetMode(gin.TestMode)
	server := startSMTPTestServer(t, "implicit", false)
	router := configureEmailCodeAPI(t, server)
	request := httptest.NewRequest(http.MethodPost, "/api/auth/email/request", strings.NewReader(`{"email":"recipient@example.com"}`))
	request.Header.Set("Content-Type", "application/json")
	response := httptest.NewRecorder()
	router.ServeHTTP(response, request)
	if response.Code != http.StatusOK || !strings.Contains(response.Body.String(), `"code":0`) {
		t.Fatalf("expected HTTP 200 code:0, got %d: %s", response.Code, response.Body.String())
	}
	result := <-server.results
	if result.err != nil {
		t.Fatalf("SMTP test server failed: %v", result.err)
	}
	if !result.authenticated || !strings.Contains(result.message, "To: <recipient@example.com>") || !strings.Contains(result.message, "10 分钟内有效") {
		t.Fatalf("unexpected SMTP message or authentication result: authenticated=%t", result.authenticated)
	}
}

func TestRequestEmailLoginCodeReturnsServiceUnavailableWhenSMTPAuthFails(t *testing.T) {
	gin.SetMode(gin.TestMode)
	server := startSMTPTestServer(t, "starttls", true)
	router := configureEmailCodeAPI(t, server)
	request := httptest.NewRequest(http.MethodPost, "/api/auth/email/request", strings.NewReader(`{"email":"recipient@example.com"}`))
	request.Header.Set("Content-Type", "application/json")
	response := httptest.NewRecorder()
	router.ServeHTTP(response, request)
	if response.Code != http.StatusServiceUnavailable || strings.Contains(response.Body.String(), `"code":0`) {
		t.Fatalf("expected HTTP 503 with nonzero code, got %d: %s", response.Code, response.Body.String())
	}
	result := <-server.results
	if result.err != nil {
		t.Fatalf("SMTP test server failed: %v", result.err)
	}
	if result.authenticated {
		t.Fatal("expected SMTP authentication to be rejected")
	}
}

func TestRequestEmailLoginCodeLiveQQSMTP(t *testing.T) {
	if os.Getenv("RUN_QQ_SMTP_LIVE_TEST") != "1" {
		t.Skip("set RUN_QQ_SMTP_LIVE_TEST=1 to send a real email through QQ SMTP")
	}
	recipient := os.Getenv("SMTP_TEST_TO")
	if recipient == "" {
		t.Fatal("SMTP_TEST_TO must be set to the test recipient")
	}
	testEnvironment, err := godotenv.Read("../../.env.test")
	if err != nil {
		t.Fatalf("load local SMTP test environment: %v", err)
	}
	for name, value := range testEnvironment {
		t.Setenv(name, value)
	}
	if err := config.InitDB(config.Load()); err != nil {
		t.Fatal(err)
	}
	router := gin.New()
	router.POST("/api/auth/email/request", RequestEmailLoginCode)
	requestBody := fmt.Sprintf(`{"email":%q}`, recipient)
	request := httptest.NewRequest(http.MethodPost, "/api/auth/email/request", strings.NewReader(requestBody))
	request.Header.Set("Content-Type", "application/json")
	response := httptest.NewRecorder()
	router.ServeHTTP(response, request)
	if response.Code != http.StatusOK || !strings.Contains(response.Body.String(), `"code":0`) {
		t.Fatalf("expected HTTP 200 code:0, got %d: %s", response.Code, response.Body.String())
	}
}