new-api/model/user_access_token.go

package model

import (
	"errors"
	"strconv"
	"strings"
	"sync/atomic"
	"time"

	"github.com/QuantumNous/new-api/common"

	"gorm.io/gorm"
	"gorm.io/gorm/clause"
)

const (
	AccessTokenPrefix                = "nap_"
	AccessTokenSessionPrefix         = "pat:"
	AccessTokenSessionVersion        = 1
	legacyAccessTokenRetireAtKey     = "LegacyAccessTokenRetireAt"
	legacyAccessTokenTransitionDays  = 30
	userAccessTokenCleanupBatchSize  = 500
	userAccessTokenTouchIntervalSecs = 60
)

var (
	ErrAccessTokenNotFound      = errors.New("access token not found")
	ErrAccessTokenLimit         = errors.New("access token limit reached")
	ErrAccessTokenExpired       = errors.New("access token has expired")
	ErrLegacyAccessTokenRetired = errors.New("legacy access token has been retired")
	errLegacyRetireAtReadOnly   = errors.New("LegacyAccessTokenRetireAt is managed by the server")
)

// legacyAccessTokenRetireAt caches the deployment-wide deadline loaded at
// startup. Zero means it was never loaded, which is treated as retired.
var legacyAccessTokenRetireAt atomic.Int64

// UserAccessToken is a scoped dashboard personal access token. Only the
// SHA-256 fingerprint of the plaintext is stored; it doubles as the audit
// token_ref.
type UserAccessToken struct {
	Id         int    `json:"id"`
	UserId     int    `json:"-" gorm:"index"`
	Name       string `json:"name" gorm:"type:varchar(64)"`
	TokenHash  string `json:"token_ref" gorm:"type:varchar(64);uniqueIndex"`
	TokenHint  string `json:"token_hint" gorm:"type:varchar(8)"`
	Scopes     string `json:"-" gorm:"type:varchar(2048)"`
	ExpiresAt  int64  `json:"expires_at" gorm:"type:bigint;index"`
	LastUsedAt int64  `json:"last_used_at" gorm:"type:bigint"`
	LastUsedIp string `json:"last_used_ip" gorm:"type:varchar(64)"`
	CreatedAt  int64  `json:"created_at" gorm:"type:bigint"`
}

func (UserAccessToken) TableName() string {
	return "user_access_tokens"
}

// GetScopes returns the stored scope list; a malformed value grants nothing.
func (token *UserAccessToken) GetScopes() []string {
	if token == nil || token.Scopes == "" {
		return nil
	}
	var scopes []string
	if err := common.UnmarshalJsonStr(token.Scopes, &scopes); err != nil {
		common.SysError("invalid access token scopes: " + err.Error())
		return nil
	}
	return scopes
}

func (token *UserAccessToken) SetScopes(scopes []string) error {
	data, err := common.Marshal(scopes)
	if err != nil {
		return err
	}
	token.Scopes = string(data)
	return nil
}

// Expired reports whether the token has an absolute expiry at or before now.
func (token *UserAccessToken) Expired(now int64) bool {
	return token.ExpiresAt > 0 && token.ExpiresAt <= now
}

func FindUserAccessTokenByHash(hash string) (*UserAccessToken, error) {
	if hash == "" {
		return nil, nil
	}
	var token UserAccessToken
	err := DB.Where("token_hash = ?", hash).First(&token).Error
	if errors.Is(err, gorm.ErrRecordNotFound) {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}
	return &token, nil
}

func GetUserAccessToken(userID, id int) (*UserAccessToken, error) {
	var token UserAccessToken
	err := DB.Where("id = ? AND user_id = ?", id, userID).First(&token).Error
	if errors.Is(err, gorm.ErrRecordNotFound) {
		return nil, ErrAccessTokenNotFound
	}
	if err != nil {
		return nil, err
	}
	return &token, nil
}

func ListUserAccessTokens(userID int) ([]UserAccessToken, error) {
	tokens := make([]UserAccessToken, 0)
	err := DB.Where("user_id = ?", userID).Order("created_at DESC").Order("id DESC").Find(&tokens).Error
	return tokens, err
}

// CreateUserAccessToken inserts a token while holding the owner's users row, so
// concurrent creates cannot exceed limit. Expired rows awaiting cleanup count.
func CreateUserAccessToken(userID int, token *UserAccessToken, limit int) error {
	if userID <= 0 || token == nil || token.TokenHash == "" {
		return ErrAccessTokenNotFound
	}
	return DB.Transaction(func(tx *gorm.DB) error {
		var user User
		if err := lockForUpdate(tx).Select("id").First(&user, userID).Error; err != nil {
			return err
		}
		var count int64
		if err := tx.Model(&UserAccessToken{}).Where("user_id = ?", userID).Count(&count).Error; err != nil {
			return err
		}
		if count >= int64(limit) {
			return ErrAccessTokenLimit
		}
		token.Id = 0
		token.UserId = userID
		if token.CreatedAt == 0 {
			token.CreatedAt = common.GetTimestamp()
		}
		return tx.Create(token).Error
	})
}

// UpdateUserAccessToken renames one token owned by userID and, when scopes is
// not nil, replaces its grant.
func UpdateUserAccessToken(userID, id int, name string, scopes []string) (*UserAccessToken, error) {
	updates := map[string]any{"name": name}
	if scopes != nil {
		var grant UserAccessToken
		if err := grant.SetScopes(scopes); err != nil {
			return nil, err
		}
		updates["scopes"] = grant.Scopes
	}
	result := DB.Model(&UserAccessToken{}).Where("id = ? AND user_id = ?", id, userID).Updates(updates)
	if result.Error != nil {
		return nil, result.Error
	}
	token, err := GetUserAccessToken(userID, id)
	if err != nil {
		return nil, err
	}
	return token, nil
}

// DeleteUserAccessToken hard-deletes one token owned by userID and returns the
// deleted row for auditing.
func DeleteUserAccessToken(userID, id int) (*UserAccessToken, error) {
	var token UserAccessToken
	err := DB.Transaction(func(tx *gorm.DB) error {
		if err := lockForUpdate(tx).Where("id = ? AND user_id = ?", id, userID).First(&token).Error; err != nil {
			if errors.Is(err, gorm.ErrRecordNotFound) {
				return ErrAccessTokenNotFound
			}
			return err
		}
		return tx.Where("id = ? AND user_id = ?", id, userID).Delete(&UserAccessToken{}).Error
	})
	if err != nil {
		return nil, err
	}
	return &token, nil
}

func DeleteUserAccessTokensWithTx(tx *gorm.DB, userID int) (int64, error) {
	result := tx.Where("user_id = ?", userID).Delete(&UserAccessToken{})
	return result.RowsAffected, result.Error
}

// TouchUserAccessToken records the last use at most once per interval per token.
func TouchUserAccessToken(id int, ip string, now int64) error {
	if len(ip) > 64 {
		ip = ip[:64]
	}
	return DB.Model(&UserAccessToken{}).
		Where("id = ? AND last_used_at < ?", id, now-userAccessTokenTouchIntervalSecs).
		Updates(map[string]any{"last_used_at": now, "last_used_ip": ip}).Error
}

// DeleteExpiredUserAccessTokens removes tokens that expired before the cutoff.
// Tokens that never expire are kept.
func DeleteExpiredUserAccessTokens(before int64) error {
	for {
		var ids []int
		if err := DB.Model(&UserAccessToken{}).
			Where("expires_at > 0 AND expires_at < ?", before).
			Order("id").Limit(userAccessTokenCleanupBatchSize).Pluck("id", &ids).Error; err != nil {
			return err
		}
		if len(ids) == 0 {
			return nil
		}
		if err := DB.Where("id IN ?", ids).Delete(&UserAccessToken{}).Error; err != nil {
			return err
		}
		if len(ids) < userAccessTokenCleanupBatchSize {
			return nil
		}
	}
}

// AccessTokenHint is the last four characters of a token, shown so users can
// tell their tokens apart. PostgreSQL pads legacy CHAR(32) values with spaces.
func AccessTokenHint(token string) string {
	token = strings.TrimRight(token, " ")
	if len(token) <= 4 {
		return ""
	}
	return token[len(token)-4:]
}

// AccessTokenSessionID is the step-up session identifier bound to a token.
// Browser session IDs are UUIDs and can never carry this prefix.
func AccessTokenSessionID(id int) string {
	return AccessTokenSessionPrefix + strconv.Itoa(id)
}

func ParseAccessTokenSessionID(sid string) (int, bool) {
	raw, ok := strings.CutPrefix(sid, AccessTokenSessionPrefix)
	if !ok {
		return 0, false
	}
	id, err := strconv.Atoi(raw)
	if err != nil || id <= 0 || strconv.Itoa(id) != raw {
		return 0, false
	}
	return id, true
}

// ValidateAccessTokenIdentity checks a token-bound step-up identity without
// taking locks, for use before issuing a verification proof.
func ValidateAccessTokenIdentity(identity AuthSessionIdentity) error {
	tokenID, ok := ParseAccessTokenSessionID(identity.SessionID)
	if !ok {
		return ErrUserSessionInactive
	}
	return checkAccessTokenIdentity(DB, tokenID, identity, false)
}

func validateAccessTokenIdentityWithTx(tx *gorm.DB, tokenID int, identity AuthSessionIdentity) error {
	return checkAccessTokenIdentity(tx, tokenID, identity, true)
}

// checkAccessTokenIdentity locks users before user_access_tokens, matching
// CreateUserAccessToken, and reports every mismatch as an inactive session.
func checkAccessTokenIdentity(tx *gorm.DB, tokenID int, identity AuthSessionIdentity, lock bool) error {
	if identity.UserID <= 0 || identity.UserAuthVersion <= 0 || identity.SessionVersion != AccessTokenSessionVersion {
		return ErrUserSessionInactive
	}
	query := tx
	if lock {
		query = lockForUpdate(tx)
	}
	var user User
	if err := query.Select("id", "status", "auth_version").First(&user, identity.UserID).Error; err != nil {
		if errors.Is(err, gorm.ErrRecordNotFound) {
			return ErrUserSessionInactive
		}
		return err
	}
	if user.Status != common.UserStatusEnabled || user.AuthVersion != identity.UserAuthVersion {
		return ErrUserSessionInactive
	}
	query = tx
	if lock {
		query = lockForUpdate(tx)
	}
	var token UserAccessToken
	if err := query.Where("id = ? AND user_id = ?", tokenID, identity.UserID).First(&token).Error; err != nil {
		if errors.Is(err, gorm.ErrRecordNotFound) {
			return ErrUserSessionInactive
		}
		return err
	}
	if token.Expired(time.Now().Unix()) {
		return ErrUserSessionInactive
	}
	return nil
}

// EnsureLegacyAccessTokenRetireAt records the legacy token deadline on first
// start (now + 30 days) and loads the stored value. The first node to start
// wins; later starts and other nodes only read it.
func EnsureLegacyAccessTokenRetireAt(now int64) error {
	deadline := strconv.FormatInt(now+legacyAccessTokenTransitionDays*24*60*60, 10)
	if err := DB.Clauses(clause.OnConflict{DoNothing: true}).Create(&Option{Key: legacyAccessTokenRetireAtKey, Value: deadline}).Error; err != nil {
		return err
	}
	var option Option
	if err := DB.Where(&Option{Key: legacyAccessTokenRetireAtKey}).First(&option).Error; err != nil {
		return err
	}
	retireAt, err := strconv.ParseInt(strings.TrimSpace(option.Value), 10, 64)
	if err != nil || retireAt <= 0 {
		return errors.New("invalid LegacyAccessTokenRetireAt option value")
	}
	legacyAccessTokenRetireAt.Store(retireAt)
	return nil
}

func LegacyAccessTokenRetireAt() int64 {
	return legacyAccessTokenRetireAt.Load()
}

func LegacyAccessTokensRetired(now int64) bool {
	retireAt := legacyAccessTokenRetireAt.Load()
	return retireAt <= 0 || now >= retireAt
}