stock-analysis/internal/service/auth.go

// Package service 提供用户注册/登录与验证码服务。
package service

import (
	"context"
	"crypto/rand"
	"crypto/subtle"
	"encoding/hex"
	"fmt"
	"log"
	"math/big"
	"net/mail"
	"regexp"
	"strings"
	"sync"
	"time"

	"golang.org/x/crypto/bcrypt"

	"stock-analysis/internal/model"
)

// AuthService 用户认证服务。
type AuthService struct {
	store AuthStore
	// 邮件发送器，未配置时使用控制台打印。
	mailer Mailer
	// 验证码有效期。
	codeTTL time.Duration

	tokenMu sync.RWMutex
	tokens  map[string]int64 // token -> userID
}

// AuthStore 定义认证所需的持久化接口。
type AuthStore interface {
	SaveUser(ctx context.Context, u model.User) (int64, error)
	GetUserByID(ctx context.Context, id int64) (*model.User, error)
	GetUserByEmail(ctx context.Context, email string) (*model.User, error)
	GetUserByUsername(ctx context.Context, username string) (*model.User, error)
	SaveVerifyCode(ctx context.Context, v model.VerifyCode) error
	GetLatestVerifyCode(ctx context.Context, email, purpose string) (*model.VerifyCode, error)
	MarkVerifyCodeUsed(ctx context.Context, email, code, purpose string) error
}

// Mailer 邮件发送接口。
type Mailer interface {
	SendCode(email, code string) error
}

// consoleMailer 默认邮件发送器，仅打印到控制台。
type consoleMailer struct{}

func (consoleMailer) SendCode(email, code string) error {
	fmt.Printf("[验证码] 邮箱: %s 验证码: %s\n", email, code)
	return nil
}

// NewAuthService 创建认证服务。
func NewAuthService(store AuthStore) *AuthService {
	return &AuthService{
		store:   store,
		mailer:  consoleMailer{},
		codeTTL: 10 * time.Minute,
		tokens:  make(map[string]int64),
	}
}

// SetMailer 设置外部邮件发送器。
func (s *AuthService) SetMailer(m Mailer) { s.mailer = m }

// SendVerifyCode 发送邮箱验证码。
func (s *AuthService) SendVerifyCode(ctx context.Context, email, purpose string) error {
	email = normalizeEmail(email)
	if !isValidEmail(email) {
		return fmt.Errorf("邮箱格式不正确")
	}
	if purpose != "register" && purpose != "login" && purpose != "reset" {
		purpose = "register"
	}

	code := randomCode(6)
	v := model.VerifyCode{
		Email:     email,
		Code:      code,
		Purpose:   purpose,
		ExpiresAt: time.Now().Add(s.codeTTL).Unix(),
	}
	if err := s.store.SaveVerifyCode(ctx, v); err != nil {
		return fmt.Errorf("保存验证码失败: %w", err)
	}
	if s.mailer != nil {
		if err := s.mailer.SendCode(email, code); err != nil {
			return fmt.Errorf("发送验证码失败: %w", err)
		}
	}
	return nil
}

// Register 邮箱+验证码注册，包含用户名。
func (s *AuthService) Register(ctx context.Context, username, email, code, password string) (*model.User, error) {
	email = normalizeEmail(email)
	if !isValidEmail(email) {
		return nil, fmt.Errorf("邮箱格式不正确")
	}
	if len(password) < 6 {
		return nil, fmt.Errorf("密码长度不能少于6位")
	}
	username = strings.TrimSpace(username)
	if username == "" {
		return nil, fmt.Errorf("用户名不能为空")
	}

	if err := s.verifyCode(ctx, email, code, "register"); err != nil {
		return nil, err
	}

	existing, err := s.store.GetUserByEmail(ctx, email)
	if err != nil {
		return nil, fmt.Errorf("查询用户失败: %w", err)
	}
	if existing != nil {
		return nil, fmt.Errorf("该邮箱已注册")
	}

	hash, err := bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost)
	if err != nil {
		return nil, fmt.Errorf("密码加密失败: %w", err)
	}

	now := time.Now().Format("2006-01-02 15:04:05")
	u := model.User{
		Username:     username,
		Email:        email,
		PasswordHash: string(hash),
		CreatedAt:    now,
		UpdatedAt:    now,
	}
	id, err := s.store.SaveUser(ctx, u)
	if err != nil {
		return nil, fmt.Errorf("保存用户失败: %w", err)
	}
	u.ID = id

	// 标记验证码已使用。
	_ = s.store.MarkVerifyCodeUsed(ctx, email, code, "register")
	return &u, nil
}

// Login 邮箱+验证码登录。
func (s *AuthService) Login(ctx context.Context, email, code string) (*model.User, string, error) {
	email = normalizeEmail(email)
	if !isValidEmail(email) {
		return nil, "", fmt.Errorf("邮箱格式不正确")
	}

	if err := s.verifyCode(ctx, email, code, "login"); err != nil {
		return nil, "", err
	}

	u, err := s.store.GetUserByEmail(ctx, email)
	if err != nil {
		return nil, "", fmt.Errorf("查询用户失败: %w", err)
	}
	if u == nil {
		return nil, "", fmt.Errorf("该邮箱未注册")
	}

	_ = s.store.MarkVerifyCodeUsed(ctx, email, code, "login")
	token := s.issueToken(u.ID)
	return u, token, nil
}

// LoginByUsername 用户名+密码登录。
func (s *AuthService) LoginByUsername(ctx context.Context, username, password string) (*model.User, string, error) {
	username = strings.TrimSpace(username)
	if username == "" {
		return nil, "", fmt.Errorf("用户名不能为空")
	}
	u, err := s.store.GetUserByUsername(ctx, username)
	if err != nil {
		return nil, "", fmt.Errorf("查询用户失败: %w", err)
	}
	if u == nil {
		return nil, "", fmt.Errorf("用户名或密码错误")
	}
	if err := bcrypt.CompareHashAndPassword([]byte(u.PasswordHash), []byte(password)); err != nil {
		return nil, "", fmt.Errorf("用户名或密码错误")
	}
	token := s.issueToken(u.ID)
	return u, token, nil
}

// EnsureAdmin 确保 admin 用户存在，不存在则创建。
func (s *AuthService) EnsureAdmin(ctx context.Context) error {
	u, err := s.store.GetUserByUsername(ctx, "admin")
	if err != nil {
		return fmt.Errorf("查询admin: %w", err)
	}
	if u != nil {
		return nil // 已存在
	}
	hash, err := bcrypt.GenerateFromPassword([]byte("123456"), bcrypt.DefaultCost)
	if err != nil {
		return fmt.Errorf("加密admin密码: %w", err)
	}
	now := time.Now().Format("2006-01-02 15:04:05")
	admin := model.User{
		Username:     "admin",
		Email:        "admin@stock.local",
		PasswordHash: string(hash),
		CreatedAt:    now,
		UpdatedAt:    now,
	}
	if _, err := s.store.SaveUser(ctx, admin); err != nil {
		return fmt.Errorf("创建admin: %w", err)
	}
	log.Printf("[auth] admin 管理员账户已创建（用户名: admin, 密码: 123456）")
	return nil
}

// issueToken 生成并保存会话 token。
func (s *AuthService) issueToken(userID int64) string {
	token, _ := RandomToken(32)
	s.tokenMu.Lock()
	s.tokens[token] = userID
	s.tokenMu.Unlock()
	return token
}

// ValidateToken 校验 token 并返回对应用户。
func (s *AuthService) ValidateToken(ctx context.Context, token string) (*model.User, error) {
	s.tokenMu.RLock()
	userID, ok := s.tokens[token]
	s.tokenMu.RUnlock()
	if !ok {
		return nil, fmt.Errorf("未登录或登录已过期")
	}
	u, err := s.store.GetUserByID(ctx, userID)
	if err != nil {
		return nil, fmt.Errorf("查询用户失败: %w", err)
	}
	if u == nil {
		return nil, fmt.Errorf("用户不存在")
	}
	return u, nil
}

// Logout 注销指定 token。
func (s *AuthService) Logout(token string) {
	s.tokenMu.Lock()
	delete(s.tokens, token)
	s.tokenMu.Unlock()
}

// verifyCode 校验验证码。
func (s *AuthService) verifyCode(ctx context.Context, email, code, purpose string) error {
	stored, err := s.store.GetLatestVerifyCode(ctx, email, purpose)
	if err != nil {
		return fmt.Errorf("查询验证码失败: %w", err)
	}
	if stored == nil {
		return fmt.Errorf("验证码不存在，请先获取")
	}
	if stored.Used {
		return fmt.Errorf("验证码已使用")
	}
	if time.Now().Unix() > stored.ExpiresAt {
		return fmt.Errorf("验证码已过期")
	}
	if subtle.ConstantTimeCompare([]byte(stored.Code), []byte(code)) != 1 {
		return fmt.Errorf("验证码错误")
	}
	return nil
}

func normalizeEmail(email string) string {
	return strings.ToLower(strings.TrimSpace(email))
}

var emailRegex = regexp.MustCompile(`^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$`)

func isValidEmail(email string) bool {
	if !emailRegex.MatchString(email) {
		return false
	}
	_, err := mail.ParseAddress(email)
	return err == nil
}

// randomCode 生成指定长度的数字验证码。
func randomCode(length int) string {
	const digits = "0123456789"
	var b strings.Builder
	for i := 0; i < length; i++ {
		n, _ := rand.Int(rand.Reader, big.NewInt(int64(len(digits))))
		b.WriteByte(digits[n.Int64()])
	}
	return b.String()
}

// RandomToken 生成随机 token。
func RandomToken(length int) (string, error) {
	buf := make([]byte, length)
	if _, err := rand.Read(buf); err != nil {
		return "", err
	}
	return hex.EncodeToString(buf), nil
}