pangolin/server/internal/auth/store.go

package auth

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/google/uuid"
)

// User is the subset of the users row the auth module needs.
type User struct {
	ID          int64
	UUID        string
	Email       string
	PwHash      string
	DpUUID      string
	Status      string // "active" | "banned"
	TOTPEnabled bool   // two-factor enabled → login requires a second TOTP step
}

// Sentinel store errors. Service maps these to API errors.
var (
	// ErrEmailTaken is returned by CreateUserWithTrial on a duplicate email.
	ErrEmailTaken = errors.New("auth: email already registered")
	// ErrNotFound is returned when a user lookup yields no row.
	ErrNotFound = errors.New("auth: user not found")
)

// UserStore is the persistence contract for the auth module. The MySQL
// implementation lives in this file; tests substitute an in-memory fake.
type UserStore interface {
	// CreateUserWithTrial atomically inserts a new user and a 7-day PRO trial
	// subscription (source='trial') in a single transaction. The email UNIQUE
	// constraint guarantees a single trial per address; a duplicate returns
	// ErrEmailTaken. trialDays controls the trial length.
	CreateUserWithTrial(ctx context.Context, email, pwHash string, trialDays int) (*User, error)
	// GetUserByEmail returns the user for login. ErrNotFound when absent.
	GetUserByEmail(ctx context.Context, email string) (*User, error)
}

// Noticer 抽象 notices.Store 的事务内插入入口(测试可替身;生产传 notices.NewStore(db))。
// auth → notices 单向依赖,不引入 import cycle,故本地声明接口、不直接 import notices
// (参照 reward 包 Noticer 的做法)。
type Noticer interface {
	InsertNoticeTx(ctx context.Context, tx *sql.Tx, userID int64, typ, titleZH, titleEN, bodyZH, bodyEN, link string, now time.Time) error
}

// SQLStore is the MySQL-backed UserStore.
type SQLStore struct {
	db      *sql.DB
	noticer Noticer
}

// NewSQLStore builds a SQLStore.
func NewSQLStore(db *sql.DB) *SQLStore { return &SQLStore{db: db} }

// SetNoticer 注入注册欢迎通知钩子(notices.Store 满足此接口);为 nil 时跳过(装配前兼容)。
func (s *SQLStore) SetNoticer(n Noticer) { s.noticer = n }

// CreateUserWithTrial implements UserStore.
func (s *SQLStore) CreateUserWithTrial(ctx context.Context, email, pwHash string, trialDays int) (*User, error) {
	userUUID := uuid.NewString()
	dpUUID := uuid.NewString()

	tx, err := s.db.BeginTx(ctx, &sql.TxOptions{Isolation: sql.LevelReadCommitted})
	if err != nil {
		return nil, fmt.Errorf("auth: begin tx: %w", err)
	}
	committed := false
	defer func() {
		if !committed {
			_ = tx.Rollback()
		}
	}()

	now := time.Now().UTC()
	res, err := tx.ExecContext(ctx,
		`INSERT INTO users (uuid, email, pw_hash, dp_uuid, status, created_at)
		 VALUES (?, ?, ?, ?, 'active', ?)`,
		userUUID, email, pwHash, dpUUID, now)
	if err != nil {
		if isDuplicateKey(err) {
			return nil, ErrEmailTaken
		}
		return nil, fmt.Errorf("auth: insert user: %w", err)
	}
	userID, err := res.LastInsertId()
	if err != nil {
		return nil, fmt.Errorf("auth: user last id: %w", err)
	}

	// Resolve the PRO plan id and insert the trial subscription.
	var proID int64
	if err := tx.QueryRowContext(ctx, `SELECT id FROM plans WHERE code='pro'`).Scan(&proID); err != nil {
		return nil, fmt.Errorf("auth: lookup pro plan: %w", err)
	}
	expires := now.AddDate(0, 0, trialDays)
	if _, err := tx.ExecContext(ctx,
		`INSERT INTO subscriptions (user_id, plan_id, expires_at, source, created_at)
		 VALUES (?, ?, ?, 'trial', ?)`,
		userID, proID, expires, now); err != nil {
		return nil, fmt.Errorf("auth: insert trial subscription: %w", err)
	}

	// 欢迎站内通知:同事务插入,注册回滚则通知不留(零孤儿),与 reward/pay 到账接缝同语义。
	if s.noticer != nil {
		titleZH := "欢迎加入 Pangolin 🎉"
		titleEN := "Welcome to Pangolin 🎉"
		bodyZH := fmt.Sprintf("已为你开通 %d 天 PRO 会员,现在就选节点连接,畅享全球网络。", trialDays)
		bodyEN := fmt.Sprintf("Your %d-day PRO trial is active. Pick a node and connect to enjoy unrestricted access.", trialDays)
		if err := s.noticer.InsertNoticeTx(ctx, tx, userID, "reward", titleZH, titleEN, bodyZH, bodyEN, "", now); err != nil {
			return nil, fmt.Errorf("auth: insert welcome notice: %w", err)
		}
	}

	if err := tx.Commit(); err != nil {
		return nil, fmt.Errorf("auth: commit: %w", err)
	}
	committed = true

	return &User{
		ID:     userID,
		UUID:   userUUID,
		Email:  email,
		PwHash: pwHash,
		DpUUID: dpUUID,
		Status: "active",
	}, nil
}

// GetUserByEmail implements UserStore.
func (s *SQLStore) GetUserByEmail(ctx context.Context, email string) (*User, error) {
	var u User
	err := s.db.QueryRowContext(ctx,
		`SELECT id, uuid, email, pw_hash, dp_uuid, status, totp_enabled FROM users WHERE email = ?`,
		email).Scan(&u.ID, &u.UUID, &u.Email, &u.PwHash, &u.DpUUID, &u.Status, &u.TOTPEnabled)
	if err == sql.ErrNoRows {
		return nil, ErrNotFound
	}
	if err != nil {
		return nil, fmt.Errorf("auth: get user by email: %w", err)
	}
	return &u, nil
}

// isDuplicateKey reports whether err is a duplicate-key/unique-constraint
// violation, dialect-neutral: MySQL ("Duplicate entry" / 1062) and sqlite
// (modernc.org/sqlite "UNIQUE constraint failed" / SQLITE_CONSTRAINT_UNIQUE
// 2067) both matched by message text.
func isDuplicateKey(err error) bool {
	if err == nil {
		return false
	}
	msg := err.Error()
	return strings.Contains(msg, "Duplicate entry") || strings.Contains(msg, "1062") ||
		strings.Contains(msg, "UNIQUE constraint failed")
}