Files
heicode-mananger/heicode/model/token.go
T
chenchenandClaude Opus 4.7 d8f61957f0 feat(manager): version 1.4.2 — H2 sk- hash + M9 audit drawer + M3 docs + M7 vault
Bundled release bumping Manager to 1.4.2 with four product-doc gap
closures lined up in a single deploy.

VERSION:
  - 1.2.0 → 1.4.2 (catches up after Sprints 1-5 shipped under 1.2.0)

H2 — sk- hash phase A (server-side, zero client impact):
  - tokens table: new key_hash varchar(64) index column
  - Token.Insert() dual-writes Key + KeyHash on every new token
  - BackfillTokenKeyHash() runs at startup, batches 500 rows at a
    time, idempotent. Fills legacy rows that pre-date the column
    without blocking app boot
  - 5 unit tests pin: sha256 correctness, dual-write on Insert,
    empty Key → empty hash, backfill behaviour, idempotency
  - Phase B (switch lookup index off plaintext + drop Key column)
    can ship later once telemetry shows key_hash IS NULL count is 0

M9 — task detail drawer with audit timeline:
  - Deployments page click → Sheet drawer with RunDetailPanel +
    new RunAuditTimeline component
  - Timeline pulls from existing /api/agnet/deployments/:id/events
    which Sprint 1 already wired to the persistent
    agnet_audit_events table — no new backend
  - Vertical timeline w/ coloured dots (primary / amber / rose by
    classifyEventLevel), occurred_at + correlation_id per row,
    max-height + overflow for long traces
  - 15s polling; empty/loading/error states all rendered

M3 — project_doc as a first-class binding step:
  - Resource binding wizard split "SK or project docs" into two
    distinct steps: "Connect project docs" + "Connect SK skill packs"
  - Each step's Connect button pre-selects the matching type in
    the advanced sheet so users don't accidentally tag a doc repo
    as Git or SK
  - Summary dialog still receives the combined skOrDocSources view
    to keep the recommendation-card contract unchanged

M7 — secret vault status (admin panel):
  - controller/secret_store.go: new GetSecretStoreStatus handler
    + fetchHealth() method. Hits OpenBao /sys/health (token-less
    upstream endpoint), maps to a sanitized response — NEVER
    returns secret names or values per product docs §13.9
  - Graceful degradation: env vars unset → "not configured" pill;
    network error → "unreachable"; sealed → amber warning; healthy
    → green
  - Mounted at GET /api/secret-store/status behind middleware.AdminAuth
  - New SecretStoreSection in system-settings/maintenance,
    registered before Performance. Read-only card with refresh
    button, 7 status fields, message line, "how to enable" hint

Verification:
  - go vet ./... clean
  - go test ./controller/... ./middleware/... ./model/... all green
  - tsc --noEmit clean
  - Backend M7 endpoint deliberately tolerant — production may not
    have OPENBAO_ADDR set yet, UI shows "not configured" instead of
    500ing

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-05-25 11:19:42 +08:00

792 lines
26 KiB
Go
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
package model
import (
"crypto/sha256"
"encoding/hex"
"errors"
"fmt"
"strings"
"github.com/heicode/manager/common"
"github.com/heicode/manager/setting"
"github.com/heicode/manager/setting/operation_setting"
"github.com/bytedance/gopkg/util/gopool"
"gorm.io/gorm"
)
type Token struct {
Id int `json:"id"`
UserId int `json:"user_id" gorm:"index"`
Key string `json:"key" gorm:"type:varchar(128);uniqueIndex"`
// KeyHash is the SHA-256 hex digest of the raw `Key` value. H2
// phase A: dual-write — every Insert and key-update path writes
// both Key and KeyHash. ValidateUserToken still looks up by Key
// for now; once telemetry shows zero-NULL coverage we can flip
// the lookup to KeyHash, drop the Key uniqueIndex, and stop
// storing plaintext bearers in the DB. This is the upgrade path
// the V2 device-binding work needs to complete the "no long-lived
// plaintext credential at rest" promise from the product docs §13.9.
//
// Indexed for forward compatibility — phase B switches lookups to
// this column. Nullable / empty during the rollout window:
// - new rows: filled by Token.Insert / UpdateKey paths
// - legacy rows: filled by BackfillTokenKeyHash on startup
// SHA-256 hex == 64 chars; varchar(64) is the natural size.
KeyHash string `json:"-" gorm:"type:varchar(64);index;column:key_hash;default:''"`
Status int `json:"status" gorm:"default:1"`
Name string `json:"name" gorm:"index" `
CreatedTime int64 `json:"created_time" gorm:"bigint"`
AccessedTime int64 `json:"accessed_time" gorm:"bigint"`
ExpiredTime int64 `json:"expired_time" gorm:"bigint;default:-1"` // -1 means never expired
RemainQuota int `json:"remain_quota" gorm:"default:0"`
UnlimitedQuota bool `json:"unlimited_quota"`
ModelLimitsEnabled bool `json:"model_limits_enabled"`
ModelLimits string `json:"model_limits" gorm:"type:text"`
AllowIps *string `json:"allow_ips" gorm:"default:''"`
UsedQuota int `json:"used_quota" gorm:"default:0"` // used quota
Group string `json:"group" gorm:"default:''"`
CrossGroupRetry bool `json:"cross_group_retry"` // 跨分组重试,仅auto分组有效
HideFromUserUI bool `json:"hide_from_user_ui" gorm:"default:false"` // 系统托管令牌:列表/API key 详情对用户隐藏(中继仍可用)
// Device-binding fields (P0 of device-signature work). When DevicePubkey
// is set, this token is bound to a specific Heicode client install and
// MUST be presented with X-Heicode-{Device-Id,Timestamp,Nonce,Signature}
// headers; bare bearer use of the key is rejected. When nil, the token
// is a legacy bare-bearer token (CLI/SDK back-compat path).
DeviceId *string `json:"device_id" gorm:"type:varchar(64);index"` // client-generated UUID v4
DevicePubkey *string `json:"device_pubkey" gorm:"type:text"` // base64 32-byte Ed25519 pubkey
DeviceFingerprint string `json:"device_fingerprint" gorm:"type:varchar(64);default:''"` // sha256 of HWID + hostname
DeviceName string `json:"device_name" gorm:"type:varchar(64);default:''"` // user-friendly: "Chen's MacBook"
DevicePlatform string `json:"device_platform" gorm:"type:varchar(16);default:''"` // darwin / windows / linux
DeviceAppVersion string `json:"device_app_version" gorm:"type:varchar(32);default:''"` // cc-haha version at pair time
DeviceBoundAt int64 `json:"device_bound_at" gorm:"bigint;default:0"`
DeviceLastSeenIp string `json:"device_last_seen_ip" gorm:"type:varchar(45);default:''"`
DeviceLastUsedAt int64 `json:"device_last_used_at" gorm:"bigint;default:0"`
RequireDeviceBinding bool `json:"require_device_binding" gorm:"default:false"` // per-token override: force signed only
RevokedAt int64 `json:"revoked_at" gorm:"bigint;default:0"`
RevokedReason string `json:"revoked_reason" gorm:"type:varchar(128);default:''"`
DeletedAt gorm.DeletedAt `gorm:"index"`
}
func (token *Token) Clean() {
token.Key = ""
}
func MaskTokenKey(key string) string {
if key == "" {
return ""
}
if len(key) <= 4 {
return strings.Repeat("*", len(key))
}
if len(key) <= 8 {
return key[:2] + "****" + key[len(key)-2:]
}
return key[:4] + "**********" + key[len(key)-4:]
}
func (token *Token) GetFullKey() string {
return token.Key
}
func (token *Token) GetMaskedKey() string {
return MaskTokenKey(token.Key)
}
func (token *Token) GetIpLimits() []string {
// delete empty spaces
//split with \n
ipLimits := make([]string, 0)
if token.AllowIps == nil {
return ipLimits
}
cleanIps := strings.ReplaceAll(*token.AllowIps, " ", "")
if cleanIps == "" {
return ipLimits
}
ips := strings.Split(cleanIps, "\n")
for _, ip := range ips {
ip = strings.TrimSpace(ip)
ip = strings.ReplaceAll(ip, ",", "")
if ip != "" {
ipLimits = append(ipLimits, ip)
}
}
return ipLimits
}
func GetAllUserTokens(userId int, startIdx int, num int) ([]*Token, error) {
var tokens []*Token
var err error
err = DB.Where("user_id = ?", userId).Order("id desc").Limit(num).Offset(startIdx).Find(&tokens).Error
return tokens, err
}
// GetAllUserTokensVisibleInUI 返回控制台「令牌」列表中应对用户展示的条目(排除系统托管、对用户隐藏的 key)。
func GetAllUserTokensVisibleInUI(userId int, startIdx int, num int) ([]*Token, error) {
var tokens []*Token
err := DB.Where("user_id = ? AND hide_from_user_ui = ?", userId, false).
Order("id desc").Limit(num).Offset(startIdx).Find(&tokens).Error
return tokens, err
}
// CountUserTokensVisibleInUI 与列表分页一致,仅统计对用户可见的令牌。
func CountUserTokensVisibleInUI(userId int) (int64, error) {
var total int64
err := DB.Model(&Token{}).
Where("user_id = ? AND hide_from_user_ui = ?", userId, false).
Count(&total).Error
return total, err
}
// FindTokenByDeviceId is the V2 token-lookup path used when the client
// is authenticated by Ed25519 signature instead of bare bearer. The
// device_id header replaces the role bare sk- played in legacy
// TokenAuth: it's a stable handle the Manager already issued at pair
// time, and the device_pubkey on the row is what verifies the request
// signature.
//
// Returns gorm.ErrRecordNotFound when the device_id doesn't match any
// row — caller should treat as 401.
func FindTokenByDeviceId(deviceId string) (*Token, error) {
if deviceId == "" {
return nil, gorm.ErrRecordNotFound
}
var t Token
err := DB.Where("device_id = ? AND device_pubkey IS NOT NULL AND device_pubkey <> ''", deviceId).
First(&t).Error
if err != nil {
return nil, err
}
return &t, nil
}
// GetUserDeviceBoundTokens returns the user's device-bound tokens for the
// "Devices" management page. Excludes:
// - legacy bare-bearer tokens (device_pubkey null OR empty string —
// both shapes show up from partial-migration rows / direct SQL)
// - revoked tokens (revoked_at > 0). Users complained that revoked
// rows lingered in the list as noise ("已撤销, 为什么还有记录"); we
// keep the soft-delete row in DB for audit + so re-pair can heal
// it (controller/device.go::PairDevice reactivate branch), but it
// no longer shows up in the user-facing list.
//
// Sort order: "freshness". A device that has been used recently sorts
// above one that hasn't been used in days. A device that has *never*
// been used (last_used_at = 0) falls back to its pair time, so the
// just-paired device the user expects to see at the top doesn't sink
// to the bottom of the list behind an old-but-actively-used machine.
//
// CASE WHEN is portable across SQLite / MySQL / PostgreSQL (per
// heicode/CLAUDE.md Rule 2). `COALESCE(NULLIF(...))` would also work
// but reads less obviously when the value is an int64 zero rather
// than NULL.
func GetUserDeviceBoundTokens(userId int) ([]*Token, error) {
var tokens []*Token
err := DB.Where(
"user_id = ? AND device_pubkey IS NOT NULL AND device_pubkey <> '' AND revoked_at = 0",
userId,
).
Order("CASE WHEN device_last_used_at = 0 THEN device_bound_at ELSE device_last_used_at END DESC, id DESC").
Find(&tokens).Error
return tokens, err
}
// CountUserDeviceBoundTokens counts ACTIVE (non-revoked) device-bound
// tokens for the per-user cap enforcement at pair time.
func CountUserDeviceBoundTokens(userId int) (int64, error) {
var total int64
err := DB.Model(&Token{}).
Where("user_id = ? AND device_pubkey IS NOT NULL AND device_pubkey <> '' AND revoked_at = 0", userId).
Count(&total).Error
return total, err
}
// UpdateTokenDeviceLastSeen records the most-recent IP and timestamp this
// device-bound token was authenticated with. Called from a goroutine after
// successful signature verification — should be best-effort, never block.
// Nil-guards DB so unit tests (or partial-init binaries) don't panic.
func UpdateTokenDeviceLastSeen(tokenId int, ip string, nowMs int64) error {
if DB == nil {
return nil
}
return DB.Model(&Token{}).Where("id = ?", tokenId).Updates(map[string]interface{}{
"device_last_seen_ip": ip,
"device_last_used_at": nowMs,
}).Error
}
// RevokeDevice soft-disables a device-bound token. We don't hard-delete
// so the audit trail (last seen IP, fingerprint) stays inspectable.
func RevokeDevice(tokenId int, reason string) error {
now := common.GetTimestamp() * 1000
return DB.Model(&Token{}).Where("id = ?", tokenId).Updates(map[string]interface{}{
"status": common.TokenStatusDisabled,
"revoked_at": now,
"revoked_reason": reason,
}).Error
}
// ReactivateDevice undoes RevokeDevice for the case where the SAME device
// (same device_id + same pubkey) re-pairs after being revoked. The pair
// controller has already verified user ownership and pubkey match; this
// only clears the soft-disable state. We deliberately keep the original
// device_fingerprint / device_bound_at — the device IS the same one, and
// rewriting bound_at would corrupt the audit timeline.
//
// `gorm.Updates` with a map skips zero values, but we want to explicitly
// set revoked_at back to 0 and revoked_reason back to '' — so we use
// `Select` to force them through. Same trick the rest of the codebase
// uses when zeroing GORM fields (e.g. ResetUserPassword in user.go).
func ReactivateDevice(tokenId int) error {
if DB == nil {
return nil
}
return DB.Model(&Token{}).Where("id = ?", tokenId).
Select("status", "revoked_at", "revoked_reason").
Updates(map[string]interface{}{
"status": common.TokenStatusEnabled,
"revoked_at": int64(0),
"revoked_reason": "",
}).Error
}
// EnsureUserRelayToken 在用户没有任何 relay 令牌时创建一枚系统托管令牌(对用户隐藏),用于登录后即可走网关。
func EnsureUserRelayToken(userId int, username string) {
if userId <= 0 {
return
}
count, err := CountUserTokens(userId)
if err != nil {
common.SysLog("EnsureUserRelayToken count: " + err.Error())
return
}
if count > 0 {
return
}
key, err := common.GenerateKey()
if err != nil {
common.SysLog("EnsureUserRelayToken generate key: " + err.Error())
return
}
token := Token{
UserId: userId,
Name: username + "的初始令牌",
Key: key,
CreatedTime: common.GetTimestamp(),
AccessedTime: common.GetTimestamp(),
ExpiredTime: -1,
RemainQuota: 500000,
UnlimitedQuota: true,
ModelLimitsEnabled: false,
HideFromUserUI: true,
}
if setting.DefaultUseAutoGroup {
token.Group = "auto"
}
if err := token.Insert(); err != nil {
common.SysLog("EnsureUserRelayToken insert: " + err.Error())
}
}
// sanitizeLikePattern 校验并清洗用户输入的 LIKE 搜索模式。
// 规则:
// 1. 转义 ! 和 _(使用 ! 作为 ESCAPE 字符,兼容 MySQL/PostgreSQL/SQLite)
// 2. 连续的 % 合并为单个 %
// 3. 最多允许 2 个 %
// 4. 含 % 时(模糊搜索),去掉 % 后关键词长度必须 >= 2
// 5. 不含 % 时按精确匹配
func sanitizeLikePattern(input string) (string, error) {
// 1. 先转义 ESCAPE 字符 ! 自身,再转义 _
// 使用 ! 而非 \ 作为 ESCAPE 字符,避免 MySQL 中反斜杠的字符串转义问题
input = strings.ReplaceAll(input, "!", "!!")
input = strings.ReplaceAll(input, `_`, `!_`)
// 2. 连续的 % 直接拒绝
if strings.Contains(input, "%%") {
return "", errors.New("搜索模式中不允许包含连续的 % 通配符")
}
// 3. 统计 % 数量,不得超过 2
count := strings.Count(input, "%")
if count > 2 {
return "", errors.New("搜索模式中最多允许包含 2 个 % 通配符")
}
// 4. 含 % 时,去掉 % 后关键词长度必须 >= 2
if count > 0 {
stripped := strings.ReplaceAll(input, "%", "")
if len(stripped) < 2 {
return "", errors.New("使用模糊搜索时,关键词长度至少为 2 个字符")
}
return input, nil
}
// 5. 无 % 时,精确全匹配
return input, nil
}
const searchHardLimit = 100
// SearchUserTokens 搜索用户令牌。restrictVisibleInUI 为 true 时排除系统托管、对用户隐藏的令牌(与普通用户控制台列表一致)。
func SearchUserTokens(userId int, keyword string, token string, offset int, limit int, restrictVisibleInUI bool) (tokens []*Token, total int64, err error) {
// model 层强制截断
if limit <= 0 || limit > searchHardLimit {
limit = searchHardLimit
}
if offset < 0 {
offset = 0
}
if token != "" {
token = strings.TrimPrefix(token, "sk-")
}
// 超量用户(令牌数超过上限)只允许精确搜索,禁止模糊搜索
maxTokens := operation_setting.GetMaxUserTokens()
hasFuzzy := strings.Contains(keyword, "%") || strings.Contains(token, "%")
if hasFuzzy {
var count int64
var err error
if restrictVisibleInUI {
count, err = CountUserTokensVisibleInUI(userId)
} else {
count, err = CountUserTokens(userId)
}
if err != nil {
common.SysLog("failed to count user tokens: " + err.Error())
return nil, 0, errors.New("获取令牌数量失败")
}
if int(count) > maxTokens {
return nil, 0, errors.New("令牌数量超过上限,仅允许精确搜索,请勿使用 % 通配符")
}
}
baseQuery := DB.Model(&Token{}).Where("user_id = ?", userId)
if restrictVisibleInUI {
baseQuery = baseQuery.Where("hide_from_user_ui = ?", false)
}
// 非空才加 LIKE 条件,空则跳过(不过滤该字段)
if keyword != "" {
keywordPattern, err := sanitizeLikePattern(keyword)
if err != nil {
return nil, 0, err
}
baseQuery = baseQuery.Where("name LIKE ? ESCAPE '!'", keywordPattern)
}
if token != "" {
tokenPattern, err := sanitizeLikePattern(token)
if err != nil {
return nil, 0, err
}
baseQuery = baseQuery.Where(commonKeyCol+" LIKE ? ESCAPE '!'", tokenPattern)
}
// 先查匹配总数(用于分页,受 maxTokens 上限保护,避免全表 COUNT)
err = baseQuery.Limit(maxTokens).Count(&total).Error
if err != nil {
common.SysError("failed to count search tokens: " + err.Error())
return nil, 0, errors.New("搜索令牌失败")
}
// 再分页查数据
err = baseQuery.Order("id desc").Offset(offset).Limit(limit).Find(&tokens).Error
if err != nil {
common.SysError("failed to search tokens: " + err.Error())
return nil, 0, errors.New("搜索令牌失败")
}
return tokens, total, nil
}
func ValidateUserToken(key string) (token *Token, err error) {
if key == "" {
return nil, ErrTokenNotProvided
}
token, err = GetTokenByKey(key, false)
if err == nil {
if token.Status == common.TokenStatusExhausted ||
token.Status == common.TokenStatusExpired ||
token.Status != common.TokenStatusEnabled {
return token, ErrTokenInvalid
}
if token.ExpiredTime != -1 && token.ExpiredTime < common.GetTimestamp() {
if !common.RedisEnabled {
token.Status = common.TokenStatusExpired
err := token.SelectUpdate()
if err != nil {
common.SysLog("failed to update token status" + err.Error())
}
}
return token, ErrTokenInvalid
}
if !token.UnlimitedQuota && token.RemainQuota <= 0 {
if !common.RedisEnabled {
token.Status = common.TokenStatusExhausted
err := token.SelectUpdate()
if err != nil {
common.SysLog("failed to update token status" + err.Error())
}
}
return token, ErrTokenInvalid
}
return token, nil
}
common.SysLog("ValidateUserToken: failed to get token: " + err.Error())
if errors.Is(err, gorm.ErrRecordNotFound) {
return nil, ErrTokenInvalid
}
return nil, fmt.Errorf("%w: %v", ErrDatabase, err)
}
func GetTokenByIds(id int, userId int) (*Token, error) {
if id == 0 || userId == 0 {
return nil, errors.New("id 或 userId 为空!")
}
token := Token{Id: id, UserId: userId}
var err error = nil
err = DB.First(&token, "id = ? and user_id = ?", id, userId).Error
return &token, err
}
func GetTokenById(id int) (*Token, error) {
if id == 0 {
return nil, errors.New("id 为空!")
}
token := Token{Id: id}
var err error = nil
err = DB.First(&token, "id = ?", id).Error
if shouldUpdateRedis(true, err) {
gopool.Go(func() {
if err := cacheSetToken(token); err != nil {
common.SysLog("failed to update user status cache: " + err.Error())
}
})
}
return &token, err
}
func GetTokenByKey(key string, fromDB bool) (token *Token, err error) {
defer func() {
// Update Redis cache asynchronously on successful DB read
if shouldUpdateRedis(fromDB, err) && token != nil {
gopool.Go(func() {
if err := cacheSetToken(*token); err != nil {
common.SysLog("failed to update user status cache: " + err.Error())
}
})
}
}()
if !fromDB && common.RedisEnabled {
// Try Redis first
token, err := cacheGetTokenByKey(key)
if err == nil {
return token, nil
}
// Don't return error - fall through to DB
}
fromDB = true
err = DB.Where(commonKeyCol+" = ?", key).First(&token).Error
return token, err
}
// BackfillTokenKeyHash fills tokens.key_hash for any rows that still
// have an empty hash but a non-empty plaintext Key. Idempotent — safe
// to call on every container startup. Batches rows so a large legacy
// `tokens` table (10k+ rows) doesn't blow up memory.
//
// This is the H2 phase-A backfill. Once telemetry shows
// COUNT(*) WHERE key_hash = '' AND key != '' equals zero across all
// environments, phase B switches ValidateUserToken to look up by hash
// and we can drop the Key column.
func BackfillTokenKeyHash() error {
if DB == nil {
return nil
}
const batchSize = 500
totalUpdated := 0
for {
var rows []Token
// Select only id + key, skip deleted-at rows, only the ones
// that still need a hash. The model's regular Find honours
// soft-delete; that's correct — we don't backfill tombstones.
err := DB.Select("id", "key").
Where("(key_hash IS NULL OR key_hash = '') AND key IS NOT NULL AND key <> ''").
Limit(batchSize).
Find(&rows).Error
if err != nil {
return err
}
if len(rows) == 0 {
break
}
for _, r := range rows {
hash := computeKeyHash(r.Key)
if hash == "" {
continue
}
if err := DB.Model(&Token{}).Where("id = ?", r.Id).
Update("key_hash", hash).Error; err != nil {
return err
}
totalUpdated++
}
if len(rows) < batchSize {
break
}
}
if totalUpdated > 0 {
common.SysLog(fmt.Sprintf("BackfillTokenKeyHash: filled %d rows", totalUpdated))
}
return nil
}
// computeKeyHash returns the SHA-256 hex digest of `key`. H2 phase A —
// stored in tokens.key_hash alongside the plaintext Key so we can move
// the auth lookup off plaintext in phase B without a downtime
// migration. Empty key returns empty string (zero-value, intentional).
func computeKeyHash(key string) string {
if key == "" {
return ""
}
sum := sha256.Sum256([]byte(key))
return hex.EncodeToString(sum[:])
}
func (token *Token) Insert() error {
// H2: always populate KeyHash on write so phase B can flip the
// lookup index without a one-off backfill round-trip. Safe to
// re-set on every Insert — value is deterministic from Key.
token.KeyHash = computeKeyHash(token.Key)
return DB.Create(token).Error
}
// Update Make sure your token's fields is completed, because this will update non-zero values
func (token *Token) Update() (err error) {
defer func() {
if shouldUpdateRedis(true, err) {
gopool.Go(func() {
err := cacheSetToken(*token)
if err != nil {
common.SysLog("failed to update token cache: " + err.Error())
}
})
}
}()
err = DB.Model(token).Select("name", "status", "expired_time", "remain_quota", "unlimited_quota",
"model_limits_enabled", "model_limits", "allow_ips", "group", "cross_group_retry").Updates(token).Error
return err
}
func (token *Token) SelectUpdate() (err error) {
defer func() {
if shouldUpdateRedis(true, err) {
gopool.Go(func() {
err := cacheSetToken(*token)
if err != nil {
common.SysLog("failed to update token cache: " + err.Error())
}
})
}
}()
// This can update zero values
return DB.Model(token).Select("accessed_time", "status").Updates(token).Error
}
func (token *Token) Delete() (err error) {
defer func() {
if shouldUpdateRedis(true, err) {
gopool.Go(func() {
err := cacheDeleteToken(token.Key)
if err != nil {
common.SysLog("failed to delete token cache: " + err.Error())
}
})
}
}()
err = DB.Delete(token).Error
return err
}
func (token *Token) IsModelLimitsEnabled() bool {
return token.ModelLimitsEnabled
}
func (token *Token) GetModelLimits() []string {
if token.ModelLimits == "" {
return []string{}
}
return strings.Split(token.ModelLimits, ",")
}
func (token *Token) GetModelLimitsMap() map[string]bool {
limits := token.GetModelLimits()
limitsMap := make(map[string]bool)
for _, limit := range limits {
limitsMap[limit] = true
}
return limitsMap
}
func DisableModelLimits(tokenId int) error {
token, err := GetTokenById(tokenId)
if err != nil {
return err
}
token.ModelLimitsEnabled = false
token.ModelLimits = ""
return token.Update()
}
func DeleteTokenById(id int, userId int) (err error) {
// Why we need userId here? In case user want to delete other's token.
if id == 0 || userId == 0 {
return errors.New("id 或 userId 为空!")
}
token := Token{Id: id, UserId: userId}
err = DB.Where(token).First(&token).Error
if err != nil {
return err
}
return token.Delete()
}
func IncreaseTokenQuota(tokenId int, key string, quota int) (err error) {
if quota < 0 {
return errors.New("quota 不能为负数!")
}
if common.RedisEnabled {
gopool.Go(func() {
err := cacheIncrTokenQuota(key, int64(quota))
if err != nil {
common.SysLog("failed to increase token quota: " + err.Error())
}
})
}
if common.BatchUpdateEnabled {
addNewRecord(BatchUpdateTypeTokenQuota, tokenId, quota)
return nil
}
return increaseTokenQuota(tokenId, quota)
}
func increaseTokenQuota(id int, quota int) (err error) {
err = DB.Model(&Token{}).Where("id = ?", id).Updates(
map[string]interface{}{
"remain_quota": gorm.Expr("remain_quota + ?", quota),
"used_quota": gorm.Expr("used_quota - ?", quota),
"accessed_time": common.GetTimestamp(),
},
).Error
return err
}
func DecreaseTokenQuota(id int, key string, quota int) (err error) {
if quota < 0 {
return errors.New("quota 不能为负数!")
}
if common.RedisEnabled {
gopool.Go(func() {
err := cacheDecrTokenQuota(key, int64(quota))
if err != nil {
common.SysLog("failed to decrease token quota: " + err.Error())
}
})
}
if common.BatchUpdateEnabled {
addNewRecord(BatchUpdateTypeTokenQuota, id, -quota)
return nil
}
return decreaseTokenQuota(id, quota)
}
func decreaseTokenQuota(id int, quota int) (err error) {
err = DB.Model(&Token{}).Where("id = ?", id).Updates(
map[string]interface{}{
"remain_quota": gorm.Expr("remain_quota - ?", quota),
"used_quota": gorm.Expr("used_quota + ?", quota),
"accessed_time": common.GetTimestamp(),
},
).Error
return err
}
// CountUserTokens returns total number of tokens for the given user, used for pagination
func CountUserTokens(userId int) (int64, error) {
var total int64
err := DB.Model(&Token{}).Where("user_id = ?", userId).Count(&total).Error
return total, err
}
// BatchDeleteTokens 删除指定用户的一组令牌,返回成功删除数量
func BatchDeleteTokens(ids []int, userId int) (int, error) {
if len(ids) == 0 {
return 0, errors.New("ids 不能为空!")
}
tx := DB.Begin()
var tokens []Token
if err := tx.Where("user_id = ? AND id IN (?)", userId, ids).Find(&tokens).Error; err != nil {
tx.Rollback()
return 0, err
}
if err := tx.Where("user_id = ? AND id IN (?)", userId, ids).Delete(&Token{}).Error; err != nil {
tx.Rollback()
return 0, err
}
if err := tx.Commit().Error; err != nil {
return 0, err
}
if common.RedisEnabled {
gopool.Go(func() {
for _, t := range tokens {
_ = cacheDeleteToken(t.Key)
}
})
}
return len(tokens), nil
}
func GetTokenKeysByIds(ids []int, userId int) ([]Token, error) {
var tokens []Token
err := DB.Select("id", commonKeyCol).
Where("user_id = ? AND id IN (?)", userId, ids).
Find(&tokens).Error
return tokens, err
}
// InvalidateUserTokensCache 清理指定用户所有令牌在 Redis 中的缓存,
// 配合 InvalidateUserCache 使用,可在用户被禁用/删除时立即阻断其令牌的请求。
// 下一次请求将从数据库重新加载令牌及用户状态,从而立即识别出被禁用的用户。
func InvalidateUserTokensCache(userId int) error {
if !common.RedisEnabled {
return nil
}
if userId <= 0 {
return errors.New("userId 无效")
}
var tokens []Token
if err := DB.Unscoped().
Select("id", commonKeyCol).
Where("user_id = ?", userId).
Find(&tokens).Error; err != nil {
return err
}
var firstErr error
for _, t := range tokens {
if t.Key == "" {
continue
}
if err := cacheDeleteToken(t.Key); err != nil && firstErr == nil {
firstErr = err
}
}
return firstErr
}