becomer
a51501909a
ci/woodpecker/tag/build Pipeline was successful
1. 预校验（避免无用查库）
原先任何非空字符串都会直接拿去查库。邀请码固定长度且只用
inviteCodeChars（去掉 0/O/1/L 后的 31 个字符），不合规的输入
可以无需查库就拒绝；用统一文案「邀请码无效」，不泄露额外信息。
2. 存量兼容（关键，差点造成线上事故）
最初把校验写成「必须 8 位」。查生产库发现：
6 位：551 个用户
8 位：1 个用户
即绝大多数用户的邀请码仍是 6 位旧码。若只接受 8 位，
这 551 人的码会被判无效，他们将无法被新用户填写邀请——
影响面很大。已改为同时接受 6 位与 8 位，字符集校验不变。
测试覆盖：8 位正常、6 位存量、5/7 位、空串、含易混淆字符
(0/O/1/I/L)、含小写、含符号、含中文。
232 lines
7.8 KiB
Go
232 lines
7.8 KiB
Go
package handler
|
||
|
||
import (
|
||
"crypto/rand"
|
||
"errors"
|
||
"fmt"
|
||
"log/slog"
|
||
"math/big"
|
||
"strings"
|
||
|
||
"cpdaily-new/internal/database"
|
||
"cpdaily-new/internal/model"
|
||
"cpdaily-new/internal/utils"
|
||
|
||
"github.com/gofiber/fiber/v2"
|
||
)
|
||
|
||
// 邀请码规则：每个用户在注册（首次登录）时自动获得一个邀请码，不再提供手动生成接口。
|
||
// - 字符集去掉易混淆的 0/O、1/I/L，便于口头/截图分享；
|
||
// - 长度 8，字符集 31 个 → 约 8.5 万亿种组合，配合唯一索引冲突重试，碰撞概率可忽略；
|
||
// - 使用 crypto/rand：旧实现以纳秒时间戳取模，可预测且并发注册时易生成相同码，
|
||
// 在 invite_code 唯一索引下会直接导致用户创建（登录）失败。
|
||
const (
|
||
inviteCodeChars = "ABCDEFGHJKMNPQRSTUVWXYZ23456789"
|
||
inviteCodeLength = 8
|
||
inviteCodeAttempts = 8
|
||
)
|
||
|
||
// newInviteCode 生成一个随机邀请码（不保证唯一，唯一性由调用方的写库重试保证）
|
||
func newInviteCode() (string, error) {
|
||
return randomFromCharset(inviteCodeChars, inviteCodeLength)
|
||
}
|
||
|
||
// randomFromCharset 用 crypto/rand 从字符集中均匀取 n 个字符
|
||
func randomFromCharset(charset string, n int) (string, error) {
|
||
max := big.NewInt(int64(len(charset)))
|
||
var sb strings.Builder
|
||
sb.Grow(n)
|
||
for i := 0; i < n; i++ {
|
||
idx, err := rand.Int(rand.Reader, max)
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
sb.WriteByte(charset[idx.Int64()])
|
||
}
|
||
return sb.String(), nil
|
||
}
|
||
|
||
// isInviteCodeConflict 判断是否为「邀请码」唯一索引冲突。
|
||
// 只认 invite_code 索引：用户名等其他唯一键冲突也会报 Duplicate entry，换邀请码重试毫无意义，
|
||
// 必须原样返回给调用方。MySQL 报错形如：
|
||
//
|
||
// Error 1062 (23000): Duplicate entry 'ABC123' for key 'sys_user.idx_sys_user_invite_code'
|
||
func isInviteCodeConflict(err error) bool {
|
||
if err == nil {
|
||
return false
|
||
}
|
||
msg := err.Error()
|
||
return strings.Contains(msg, "Duplicate entry") && strings.Contains(msg, "invite_code")
|
||
}
|
||
|
||
// withUniqueInviteCode 生成邀请码并执行写库操作 write(code)；
|
||
// 若写库因唯一索引冲突失败，则换一个新码重试，最多 inviteCodeAttempts 次。
|
||
// 冲突判定交给数据库唯一索引（而非「先查后写」），并发下依然正确。
|
||
func withUniqueInviteCode(write func(code string) error) (string, error) {
|
||
var lastErr error
|
||
for i := 0; i < inviteCodeAttempts; i++ {
|
||
code, err := newInviteCode()
|
||
if err != nil {
|
||
return "", fmt.Errorf("生成邀请码失败: %w", err)
|
||
}
|
||
err = write(code)
|
||
if err == nil {
|
||
return code, nil
|
||
}
|
||
if !isInviteCodeConflict(err) {
|
||
return "", err
|
||
}
|
||
lastErr = err
|
||
}
|
||
return "", fmt.Errorf("邀请码重复次数过多: %w", lastErr)
|
||
}
|
||
|
||
// ensureInviteCode 确保用户拥有邀请码：已有则直接返回；没有则生成并写库。
|
||
// 写库条件带 invite_code 为空，避免并发登录时把已生成的码覆盖掉；
|
||
// 若条件未命中（被并发请求抢先写入），重新读取最新值。
|
||
func ensureInviteCode(user *model.SysUser) error {
|
||
if strings.TrimSpace(user.InviteCode) != "" {
|
||
return nil
|
||
}
|
||
code, err := withUniqueInviteCode(func(code string) error {
|
||
res := database.DB.Model(&model.SysUser{}).
|
||
Where("id = ? AND (invite_code IS NULL OR invite_code = '')", user.ID).
|
||
Update("invite_code", code)
|
||
if res.Error != nil {
|
||
return res.Error
|
||
}
|
||
if res.RowsAffected == 0 {
|
||
return errInviteCodeAlreadySet
|
||
}
|
||
return nil
|
||
})
|
||
if errors.Is(err, errInviteCodeAlreadySet) {
|
||
var fresh model.SysUser
|
||
if e := database.DB.Select("id", "invite_code").First(&fresh, user.ID).Error; e != nil {
|
||
return e
|
||
}
|
||
user.InviteCode = fresh.InviteCode
|
||
return nil
|
||
}
|
||
if err != nil {
|
||
return err
|
||
}
|
||
user.InviteCode = code
|
||
slog.Info("已为用户补发邀请码", "userId", user.ID, "inviteCode", code)
|
||
return nil
|
||
}
|
||
|
||
var errInviteCodeAlreadySet = errors.New("invite code already set")
|
||
|
||
// BackfillInviteCodes 启动时为所有缺少邀请码的存量用户补齐邀请码（幂等，无缺失时只执行一次计数查询）。
|
||
func BackfillInviteCodes() {
|
||
var users []model.SysUser
|
||
if err := database.DB.Select("id", "invite_code").
|
||
Where("invite_code IS NULL OR invite_code = ''").
|
||
Find(&users).Error; err != nil {
|
||
slog.Error("查询缺少邀请码的用户失败", "error", err)
|
||
return
|
||
}
|
||
if len(users) == 0 {
|
||
return
|
||
}
|
||
ok := 0
|
||
for i := range users {
|
||
if err := ensureInviteCode(&users[i]); err != nil {
|
||
slog.Error("补发邀请码失败", "userId", users[i].ID, "error", err)
|
||
continue
|
||
}
|
||
ok++
|
||
}
|
||
slog.Info("存量用户邀请码补齐完成", "total", len(users), "success", ok)
|
||
}
|
||
|
||
// resetInviteCode 为指定用户重新生成邀请码（旧码立即失效）。
|
||
//
|
||
// 唯一性由 invite_code 唯一索引 + 冲突重试保证（与首次生成同一套逻辑）。
|
||
// 用带条件的方式写库：仅当当前值仍是调用方读到的旧码时才更新，
|
||
// 避免两个并发重置请求互相覆盖、导致其中一个返回的码实际已失效。
|
||
func resetInviteCode(user *model.SysUser) (string, error) {
|
||
if user == nil {
|
||
return "", errors.New("用户不存在")
|
||
}
|
||
oldCode := strings.TrimSpace(user.InviteCode)
|
||
|
||
code, err := withUniqueInviteCode(func(code string) error {
|
||
q := database.DB.Model(&model.SysUser{}).Where("id = ?", user.ID)
|
||
if oldCode != "" {
|
||
// 乐观锁：旧码未被他人改动才允许替换
|
||
q = q.Where("invite_code = ?", oldCode)
|
||
}
|
||
res := q.Update("invite_code", code)
|
||
if res.Error != nil {
|
||
return res.Error
|
||
}
|
||
if res.RowsAffected == 0 {
|
||
return errInviteCodeChanged
|
||
}
|
||
return nil
|
||
})
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
user.InviteCode = code
|
||
return code, nil
|
||
}
|
||
|
||
var errInviteCodeChanged = errors.New("invite code changed concurrently")
|
||
|
||
// ResetInviteCode 用户主动重置自己的邀请码。
|
||
//
|
||
// 影响：旧邀请码立即失效——已用旧码注册的邀请关系不受影响（关系已落库），
|
||
// 但旧码不能再被新用户使用。前端需在操作前明确提示这一点。
|
||
func ResetInviteCode(c *fiber.Ctx) error {
|
||
userID := c.Locals("userId").(uint)
|
||
var user model.SysUser
|
||
if err := database.DB.First(&user, userID).Error; err != nil {
|
||
return utils.FailAuth(c, "用户不存在")
|
||
}
|
||
// 没有码时先补齐，再重置（否则 oldCode 为空，乐观锁条件失效）
|
||
if strings.TrimSpace(user.InviteCode) == "" {
|
||
if err := ensureInviteCode(&user); err != nil {
|
||
return utils.FailServer(c, "生成邀请码失败，请稍后重试")
|
||
}
|
||
}
|
||
|
||
oldCode := user.InviteCode
|
||
code, err := resetInviteCode(&user)
|
||
if err != nil {
|
||
if errors.Is(err, errInviteCodeChanged) {
|
||
return utils.Fail(c, 409, "邀请码已被重置，请刷新页面后重试")
|
||
}
|
||
slog.Error("重置邀请码失败", "userId", userID, "error", err)
|
||
return utils.FailServer(c, "重置邀请码失败，请稍后重试")
|
||
}
|
||
slog.Info("用户重置邀请码", "userId", userID, "oldCode", oldCode, "newCode", code)
|
||
return utils.Success(c, fiber.Map{"inviteCode": code})
|
||
}
|
||
|
||
// inviteCodeLegacyLength 旧版邀请码长度。
|
||
//
|
||
// 2026-10 之前生成为 6 位，生产库中仍有大量存量用户沿用旧码
|
||
// （实测 551 个 6 位 / 1 个 8 位）。校验必须同时接受两种长度，
|
||
// 否则存量用户的码会被判为无效，导致他们无法被邀请。
|
||
const inviteCodeLegacyLength = 6
|
||
|
||
// isValidInviteCodeShape 校验邀请码的形态（长度 + 字符集），不查库。
|
||
//
|
||
// 长度接受 6（存量）与 8（当前）两种，字符集为 inviteCodeChars
|
||
// （已去掉易混淆的 0/O/1/L）。不合规的输入无需查库即可拒绝，
|
||
// 也避免把用户的任意长字符串当码去查。
|
||
func isValidInviteCodeShape(code string) bool {
|
||
if len(code) != inviteCodeLength && len(code) != inviteCodeLegacyLength {
|
||
return false
|
||
}
|
||
for i := 0; i < len(code); i++ {
|
||
if !strings.ContainsRune(inviteCodeChars, rune(code[i])) {
|
||
return false
|
||
}
|
||
}
|
||
return true
|
||
}
|