pangolin/server/internal/auth/integration_test.go

//go:build integration

package auth

import (
	"context"
	"database/sql"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	_ "github.com/go-sql-driver/mysql"
	"github.com/redis/go-redis/v9"
	"github.com/testcontainers/testcontainers-go"
	tcmysql "github.com/testcontainers/testcontainers-go/modules/mysql"
	tcredis "github.com/testcontainers/testcontainers-go/modules/redis"

	"github.com/wangjia/pangolin/server/internal/codes"
)

func setupMySQL(t *testing.T) *sql.DB {
	t.Helper()
	ctx := context.Background()
	ctr, err := tcmysql.Run(ctx, "mysql:8.0",
		tcmysql.WithDatabase("pangolin_test"),
		tcmysql.WithUsername("root"),
		tcmysql.WithPassword("test"),
	)
	testcontainers.CleanupContainer(t, ctr)
	if err != nil {
		t.Fatalf("mysql container: %v", err)
	}
	dsn, err := ctr.ConnectionString(ctx, "parseTime=true", "loc=UTC", "time_zone=%27%2B00%3A00%27")
	if err != nil {
		t.Fatalf("mysql dsn: %v", err)
	}
	db, err := sql.Open("mysql", dsn)
	if err != nil {
		t.Fatalf("open mysql: %v", err)
	}
	t.Cleanup(func() { db.Close() })
	if err := applyAuthSchema(db); err != nil {
		t.Fatalf("schema: %v", err)
	}
	return db
}

func setupRedis(t *testing.T) *redis.Client {
	t.Helper()
	ctx := context.Background()
	ctr, err := tcredis.Run(ctx, "redis:7-alpine")
	testcontainers.CleanupContainer(t, ctr)
	if err != nil {
		t.Fatalf("redis container: %v", err)
	}
	addr, err := ctr.ConnectionString(ctx)
	if err != nil {
		t.Fatalf("redis addr: %v", err)
	}
	for _, p := range []string{"redis://", "rediss://"} {
		if len(addr) > len(p) && addr[:len(p)] == p {
			addr = addr[len(p):]
		}
	}
	rdb := redis.NewClient(&redis.Options{Addr: addr})
	t.Cleanup(func() { rdb.Close() })
	return rdb
}

func applyAuthSchema(db *sql.DB) error {
	stmts := []string{
		`CREATE TABLE IF NOT EXISTS plans (
			id            BIGINT UNSIGNED AUTO_INCREMENT PRIMARY KEY,
			code          ENUM('free','pro','team') NOT NULL UNIQUE,
			max_devices   INT NOT NULL DEFAULT 1,
			daily_minutes INT NULL,
			ad_gate       BOOLEAN NOT NULL DEFAULT FALSE
		) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4`,
		`CREATE TABLE IF NOT EXISTS users (
			id          BIGINT UNSIGNED AUTO_INCREMENT PRIMARY KEY,
			uuid        CHAR(36) NOT NULL UNIQUE,
			email       VARCHAR(255) NOT NULL UNIQUE,
			pw_hash     VARCHAR(255) NOT NULL,
			dp_uuid     CHAR(36) NOT NULL,
			status      ENUM('active','banned') NOT NULL DEFAULT 'active',
			totp_enabled BOOLEAN NOT NULL DEFAULT FALSE,
			created_at  DATETIME(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6)
		) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4`,
		`CREATE TABLE IF NOT EXISTS subscriptions (
			id         BIGINT UNSIGNED AUTO_INCREMENT PRIMARY KEY,
			user_id    BIGINT UNSIGNED NOT NULL,
			plan_id    BIGINT UNSIGNED NOT NULL,
			expires_at DATETIME(6) NOT NULL,
			source     ENUM('trial','code') NOT NULL,
			created_at DATETIME(6) NOT NULL DEFAULT CURRENT_TIMESTAMP(6),
			FOREIGN KEY (plan_id) REFERENCES plans(id),
			INDEX idx_user_exp (user_id, expires_at)
		) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4`,
		`INSERT IGNORE INTO plans (code, max_devices, daily_minutes, ad_gate)
		 VALUES ('free',1,10,TRUE),('pro',5,NULL,FALSE),('team',10,NULL,FALSE)`,
	}
	for _, s := range stmts {
		if _, err := db.Exec(s); err != nil {
			return err
		}
	}
	return nil
}

func newIntegrationService(t *testing.T, db *sql.DB, rdb *redis.Client) *Service {
	t.Helper()
	store := NewSQLStore(db)
	rl := NewRateLimiter(rdb, nil)
	key := newRSAKey(t)
	tm, err := NewTokenManager(rdb, TokenConfig{SignKey: key, SignKID: "k1"})
	if err != nil {
		t.Fatalf("token manager: %v", err)
	}
	return NewService(store, rdb, rl, tm, NewLogMailer(nil), ServiceConfig{}, nil)
}

// TestIntegration_FullChain exercises register → login → refresh → protected
// route against real MySQL 8 and Redis containers.
func TestIntegration_FullChain(t *testing.T) {
	db := setupMySQL(t)
	rdb := setupRedis(t)
	svc := newIntegrationService(t, db, rdb)
	ctx := context.Background()
	const email = "integration@example.com"
	const pw = "password-integration"

	// 1. Send code (read it back from Redis to simulate the user).
	if _, e := svc.SendCode(ctx, email, "198.51.100.7"); e != nil {
		t.Fatalf("SendCode: %v", e)
	}
	code, err := rdb.Get(ctx, codeKey(email)).Result()
	if err != nil {
		t.Fatalf("code not stored: %v", err)
	}

	// 2. Register → trial subscription must exist for 7 days.
	pair, apiErr := svc.Register(ctx, email, code, pw, "203.0.113.10", DeviceMeta{})
	if apiErr != nil {
		t.Fatalf("Register: %v", apiErr)
	}

	var plan, source string
	var expires time.Time
	err = db.QueryRowContext(ctx,
		`SELECT p.code, s.source, s.expires_at
		   FROM subscriptions s JOIN plans p ON p.id = s.plan_id
		   JOIN users u ON u.id = s.user_id
		  WHERE u.email = ?`, email).Scan(&plan, &source, &expires)
	if err != nil {
		t.Fatalf("query trial: %v", err)
	}
	if plan != "pro" || source != "trial" {
		t.Fatalf("trial = %s/%s, want pro/trial", plan, source)
	}
	days := time.Until(expires).Hours() / 24
	if days < 6.5 || days > 7.1 {
		t.Fatalf("trial length = %.2f days, want ~7", days)
	}

	// 3. Duplicate email — anti-enumeration: registered email gets no code, and
	// even a forced code yields the generic code error (not "email exists").
	if _, e := svc.SendCode(ctx, email, ""); e != nil && e.Code != ErrRateLimited.Code {
		t.Fatalf("second SendCode: %v", e)
	}
	// Force a fresh code regardless of rate limit.
	_ = rdb.Set(ctx, codeKey(email), code, 10*time.Minute).Err()
	if _, e := svc.Register(ctx, email, code, pw, "203.0.113.10", DeviceMeta{}); e == nil || e.Code != ErrCodeInvalid.Code {
		t.Fatalf("want code_invalid (anti-enumeration), got %v", e)
	}

	// 4. Login.
	loginPair, _, apiErr := svc.Login(ctx, email, pw, "198.51.100.7", DeviceMeta{})
	if apiErr != nil {
		t.Fatalf("Login: %v", apiErr)
	}

	// 5. Refresh rotates.
	rotated, apiErr := svc.Refresh(ctx, loginPair.Tokens.RefreshToken)
	if apiErr != nil {
		t.Fatalf("Refresh: %v", apiErr)
	}
	if _, e := svc.Refresh(ctx, loginPair.Tokens.RefreshToken); e == nil {
		t.Fatal("old refresh token must be rejected after rotation")
	}

	// 6. Access a protected route with the rotated access token.
	protected := RequireAuth(svc.tokens)(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		uid, ok := UserIDFromContext(r.Context())
		if !ok || uid == 0 {
			w.WriteHeader(http.StatusUnauthorized)
			return
		}
		// Confirm interop with the codes module's context key.
		if _, ok := r.Context().Value(codes.CtxKeyUserID).(int64); !ok {
			w.WriteHeader(http.StatusInternalServerError)
			return
		}
		w.WriteHeader(http.StatusOK)
	}))

	req := httptest.NewRequest(http.MethodGet, "/me", nil)
	req.Header.Set("Authorization", "Bearer "+rotated.AccessToken)
	rec := httptest.NewRecorder()
	protected.ServeHTTP(rec, req)
	if rec.Code != http.StatusOK {
		t.Fatalf("protected route status = %d, want 200", rec.Code)
	}

	// Pair returned at registration is also a valid access token.
	if _, e := svc.tokens.ParseAccess(pair.AccessToken); e != nil {
		t.Fatalf("register access token invalid: %v", e)
	}
}