145 lines
3.7 KiB
Go
145 lines
3.7 KiB
Go
package system
|
||
|
||
import (
|
||
"crypto/rand"
|
||
"crypto/sha256"
|
||
"encoding/base32"
|
||
"encoding/hex"
|
||
"errors"
|
||
"strings"
|
||
"time"
|
||
|
||
"git.echol.cn/loser/Go-Web-Template/server/global"
|
||
sysModel "git.echol.cn/loser/Go-Web-Template/server/model/system"
|
||
"gorm.io/gorm"
|
||
"gorm.io/gorm/clause"
|
||
)
|
||
|
||
type InviteCodeService struct{}
|
||
|
||
var InviteCodeServiceApp = new(InviteCodeService)
|
||
|
||
type InviteCodeCurrent struct {
|
||
HasUnused bool `json:"hasUnused"`
|
||
Record *sysModel.InviteCode `json:"record,omitempty"`
|
||
}
|
||
|
||
type InviteCodeGenerated struct {
|
||
Code string `json:"code"`
|
||
CodeLast4 string `json:"codeLast4"`
|
||
ExpiresAt *time.Time `json:"expiresAt,omitempty"`
|
||
}
|
||
|
||
func sha256Hex(s string) string {
|
||
sum := sha256.Sum256([]byte(s))
|
||
return hex.EncodeToString(sum[:])
|
||
}
|
||
|
||
func generateCode() (string, error) {
|
||
buf := make([]byte, 32)
|
||
if _, err := rand.Read(buf); err != nil {
|
||
return "", err
|
||
}
|
||
enc := base32.StdEncoding.WithPadding(base32.NoPadding)
|
||
return strings.ToLower(enc.EncodeToString(buf)), nil
|
||
}
|
||
|
||
func (s *InviteCodeService) GetCurrentUnused(userID uint) (*sysModel.InviteCode, error) {
|
||
var record sysModel.InviteCode
|
||
now := time.Now()
|
||
err := global.GVA_DB.
|
||
Where("created_by_user_id = ? AND status = ?", userID, sysModel.InviteCodeStatusUnused).
|
||
Where(global.GVA_DB.Where("expires_at IS NULL").Or("expires_at > ?", now)).
|
||
Order("id desc").
|
||
First(&record).Error
|
||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||
return nil, nil
|
||
}
|
||
return &record, err
|
||
}
|
||
|
||
func (s *InviteCodeService) Generate(userID uint, expireHours int) (*InviteCodeGenerated, error) {
|
||
if expireHours <= 0 {
|
||
expireHours = 24
|
||
}
|
||
|
||
// 单活跃未使用码：存在则拒绝生成（用户已选择该策略）
|
||
existing, err := s.GetCurrentUnused(userID)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if existing != nil {
|
||
return nil, errors.New("已有未使用邀请码")
|
||
}
|
||
|
||
code, err := generateCode()
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
last4 := ""
|
||
if len(code) >= 4 {
|
||
last4 = code[len(code)-4:]
|
||
}
|
||
|
||
var expiresAt *time.Time
|
||
if expireHours > 0 {
|
||
t := time.Now().Add(time.Duration(expireHours) * time.Hour)
|
||
expiresAt = &t
|
||
}
|
||
|
||
record := sysModel.InviteCode{
|
||
CreatedByUserID: userID,
|
||
CodeHash: sha256Hex(code),
|
||
CodeLast4: last4,
|
||
Status: sysModel.InviteCodeStatusUnused,
|
||
ExpiresAt: expiresAt,
|
||
}
|
||
if err := global.GVA_DB.Create(&record).Error; err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
return &InviteCodeGenerated{Code: code, CodeLast4: last4, ExpiresAt: expiresAt}, nil
|
||
}
|
||
|
||
// ConsumeByPlainCode 校验并消耗邀请码（用于注册邀请码模式，需事务调用）
|
||
func (s *InviteCodeService) ConsumeByPlainCode(tx *gorm.DB, code string, inviteeUserID uint) (*sysModel.InviteCode, error) {
|
||
code = strings.TrimSpace(code)
|
||
if code == "" {
|
||
return nil, errors.New("邀请码不能为空")
|
||
}
|
||
|
||
now := time.Now()
|
||
var record sysModel.InviteCode
|
||
err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).
|
||
Where("code_hash = ? AND status = ?", sha256Hex(code), sysModel.InviteCodeStatusUnused).
|
||
Where(tx.Where("expires_at IS NULL").Or("expires_at > ?", now)).
|
||
First(&record).Error
|
||
if err != nil {
|
||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||
return nil, errors.New("邀请码无效")
|
||
}
|
||
return nil, err
|
||
}
|
||
|
||
usedAt := time.Now()
|
||
record.Status = sysModel.InviteCodeStatusUsed
|
||
record.UsedByUserID = &inviteeUserID
|
||
record.UsedAt = &usedAt
|
||
if err := tx.Model(&record).Select("status", "used_by_user_id", "used_at").Updates(&record).Error; err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
// 写邀请关系（invitee 唯一）
|
||
rel := sysModel.UserInviteRelation{
|
||
InviterUserID: record.CreatedByUserID,
|
||
InviteeUserID: inviteeUserID,
|
||
InviteCodeID: record.ID,
|
||
CreatedAt: time.Now(),
|
||
}
|
||
if err := tx.Create(&rel).Error; err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
return &record, nil
|
||
}
|