kra-new/internal/data/api_token.go

package data

import (
	"context"
	"errors"
	"time"

	"kra/internal/biz"

	"gorm.io/gorm"
)

type apiTokenRepo struct{ data *Data }

func NewAPITokenRepo(data *Data) biz.APITokenRepo { return &apiTokenRepo{data: data} }

type apiTokenPO struct {
	ID          uint `gorm:"primaryKey"`
	CreatedAt   time.Time
	UpdatedAt   time.Time
	DeletedAt   gorm.DeletedAt `gorm:"index"`
	UserID      uint
	AuthorityID uint
	Token       string `gorm:"type:text"`
	Status      bool
	ExpiresAt   time.Time
	Remark      string
}

func (apiTokenPO) TableName() string { return "sys_api_tokens" }

type jwtBlacklistPO struct {
	ID        uint `gorm:"primaryKey"`
	CreatedAt time.Time
	UpdatedAt time.Time
	DeletedAt gorm.DeletedAt `gorm:"index"`
	JWT       string         `gorm:"column:jwt;type:text"`
}

func (jwtBlacklistPO) TableName() string { return "jwt_blacklists" }

func (r *apiTokenRepo) UserHasAuthority(ctx context.Context, userID, authorityID uint) (*biz.User, bool, error) {
	var po userPO
	if err := r.data.gormDB.WithContext(ctx).First(&po, userID).Error; err != nil {
		if errors.Is(err, gorm.ErrRecordNotFound) {
			return nil, false, errors.New("用户不存在")
		}
		return nil, false, err
	}
	var count int64
	err := r.data.gormDB.WithContext(ctx).Model(&userAuthorityPO{}).
		Joins("JOIN sys_authorities ON sys_authorities.authority_id = sys_user_authority.sys_authority_authority_id").
		Where("sys_user_authority.sys_user_id = ? AND sys_user_authority.sys_authority_authority_id = ? AND sys_authorities.deleted_at IS NULL", userID, authorityID).
		Count(&count).Error
	if err != nil {
		return nil, false, err
	}
	user, err := (&userRepo{data: r.data}).loadUser(ctx, &po)
	return user, count > 0 || po.AuthorityID == authorityID, err
}
func (r *apiTokenRepo) CreateAPIToken(ctx context.Context, v *biz.APIToken) error {
	po := apiTokenPO{UserID: v.UserID, AuthorityID: v.AuthorityID, Token: v.Token, Status: v.Status, ExpiresAt: v.ExpiresAt, Remark: v.Remark}
	if err := r.data.gormDB.WithContext(ctx).Create(&po).Error; err != nil {
		return err
	}
	v.ID = po.ID
	v.CreatedAt = po.CreatedAt
	return nil
}
func (r *apiTokenRepo) ListAPITokens(ctx context.Context, page, size int, userID uint, status *bool) ([]*biz.APIToken, int64, error) {
	db := r.data.gormDB.WithContext(ctx).Model(&apiTokenPO{})
	if userID != 0 {
		db = db.Where("user_id = ?", userID)
	}
	if status != nil {
		db = db.Where("status = ?", *status)
	}
	var total int64
	if err := db.Count(&total).Error; err != nil {
		return nil, 0, err
	}
	var pos []apiTokenPO
	if err := applyPagination(db.Order("created_at desc"), page, size, 100).Find(&pos).Error; err != nil {
		return nil, 0, err
	}
	userIDs := make([]uint, 0, len(pos))
	for _, po := range pos {
		userIDs = append(userIDs, po.UserID)
	}
	var userPOs []userPO
	if len(userIDs) > 0 {
		if err := r.data.gormDB.WithContext(ctx).Where("id IN ?", userIDs).Find(&userPOs).Error; err != nil {
			return nil, 0, err
		}
	}
	users := make(map[uint]*biz.User, len(userPOs))
	for i := range userPOs {
		users[userPOs[i].ID] = baseBizUser(&userPOs[i])
	}
	out := make([]*biz.APIToken, 0, len(pos))
	for _, po := range pos {
		v := &biz.APIToken{ID: po.ID, CreatedAt: po.CreatedAt, UpdatedAt: po.UpdatedAt, UserID: po.UserID, AuthorityID: po.AuthorityID, Token: po.Token, Status: po.Status, ExpiresAt: po.ExpiresAt, Remark: po.Remark, User: users[po.UserID]}
		out = append(out, v)
	}
	return out, total, nil
}
func (r *apiTokenRepo) DisableAPIToken(ctx context.Context, id uint) (string, error) {
	var po apiTokenPO
	if err := r.data.gormDB.WithContext(ctx).First(&po, id).Error; err != nil {
		return "", err
	}
	return po.Token, r.data.gormDB.WithContext(ctx).Model(&po).Update("status", false).Error
}
func (r *apiTokenRepo) DisableAndBlacklistAPIToken(ctx context.Context, id uint) error {
	// Persist the blacklist row first, then mark the
	// API-token record disabled. The two writes are intentionally separate in
	// the reference implementation, so a storage error has the same observable
	// result here.
	var po apiTokenPO
	if err := r.data.gormDB.WithContext(ctx).First(&po, id).Error; err != nil {
		return err
	}
	if err := r.BlacklistToken(ctx, po.Token); err != nil {
		return err
	}
	return r.data.gormDB.WithContext(ctx).Model(&po).Update("status", false).Error
}
func (r *apiTokenRepo) BlacklistToken(ctx context.Context, token string) error {
	return r.data.gormDB.WithContext(ctx).Create(&jwtBlacklistPO{JWT: token}).Error
}
func (r *apiTokenRepo) IsTokenDisabled(ctx context.Context, token string) (bool, error) {
	var blacklistCount int64
	if err := r.data.gormDB.WithContext(ctx).Model(&jwtBlacklistPO{}).Where("jwt = ?", token).Count(&blacklistCount).Error; err != nil {
		return false, err
	}
	// The JWT middleware consults only jwt_blacklists and lets JWT parsing
	// report expiry. API-token deletion writes this row before changing the
	// token status, so consulting status/ExpiresAt here would produce a
	// different error message and would make manually-disabled tokens behave
	// unlike tokens issued through the regular login flow.
	return blacklistCount > 0, nil
}