298 lines
7.4 KiB
Go
298 lines
7.4 KiB
Go
package logic
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"strconv"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/redis/go-redis/v9"
|
|
|
|
"spark-server/internal/pkg/errs"
|
|
)
|
|
|
|
const (
|
|
tokenUseAccess = "access"
|
|
tokenUseRefresh = "refresh"
|
|
)
|
|
|
|
// TokenPair is the current access+refresh jti for one user device.
|
|
type TokenPair struct {
|
|
Access string `json:"a"`
|
|
Refresh string `json:"r"`
|
|
}
|
|
|
|
// PairStore holds the live token pair per device so refresh can rotate and logout can revoke both.
|
|
type PairStore interface {
|
|
Put(ctx context.Context, userID int64, deviceID uint64, appID uint32, pair TokenPair, ttl time.Duration) error
|
|
Get(ctx context.Context, userID int64, deviceID uint64, appID uint32) (TokenPair, error)
|
|
Delete(ctx context.Context, userID int64, deviceID uint64, appID uint32) error
|
|
ListByUser(ctx context.Context, userID int64) ([]DevicePair, error)
|
|
DeleteAllOnDevice(ctx context.Context, userID int64, deviceID uint64) error
|
|
}
|
|
|
|
// DevicePair is the live token pair for one device of a user.
|
|
type DevicePair struct {
|
|
DeviceID uint64
|
|
AppID uint32
|
|
Pair TokenPair
|
|
}
|
|
|
|
type MemoryPairs struct {
|
|
mu sync.Mutex
|
|
mode PairKeyMode
|
|
items map[string]storedPair
|
|
}
|
|
|
|
type storedPair struct {
|
|
pair TokenPair
|
|
exp int64
|
|
}
|
|
|
|
func NewMemoryPairs() *MemoryPairs {
|
|
return NewMemoryPairsMode(PairKeyTriple)
|
|
}
|
|
|
|
func NewMemoryPairsMode(mode PairKeyMode) *MemoryPairs {
|
|
return &MemoryPairs{mode: normalizePairKeyMode(mode), items: make(map[string]storedPair)}
|
|
}
|
|
|
|
func (s *MemoryPairs) Put(_ context.Context, userID int64, deviceID uint64, appID uint32, pair TokenPair, ttl time.Duration) error {
|
|
if userID <= 0 || deviceID == 0 || ttl <= 0 {
|
|
return nil
|
|
}
|
|
key := storageKey(s.mode, userID, deviceID, appID)
|
|
s.mu.Lock()
|
|
s.items[key] = storedPair{pair: pair, exp: time.Now().Add(ttl).UnixNano()}
|
|
s.mu.Unlock()
|
|
return nil
|
|
}
|
|
|
|
func (s *MemoryPairs) Get(_ context.Context, userID int64, deviceID uint64, appID uint32) (TokenPair, error) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
key := storageKey(s.mode, userID, deviceID, appID)
|
|
row, ok := s.items[key]
|
|
if !ok {
|
|
return TokenPair{}, nil
|
|
}
|
|
if time.Now().UnixNano() >= row.exp {
|
|
delete(s.items, key)
|
|
return TokenPair{}, nil
|
|
}
|
|
return row.pair, nil
|
|
}
|
|
|
|
func (s *MemoryPairs) Delete(_ context.Context, userID int64, deviceID uint64, appID uint32) error {
|
|
s.mu.Lock()
|
|
delete(s.items, storageKey(s.mode, userID, deviceID, appID))
|
|
s.mu.Unlock()
|
|
return nil
|
|
}
|
|
|
|
func (s *MemoryPairs) DeleteAllOnDevice(_ context.Context, userID int64, deviceID uint64) error {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
if s.mode == PairKeyPair {
|
|
delete(s.items, storageKey(s.mode, userID, deviceID, 0))
|
|
return nil
|
|
}
|
|
prefix := deviceKeyPrefix(userID, deviceID) + ":"
|
|
for key := range s.items {
|
|
if strings.HasPrefix(key, prefix) {
|
|
delete(s.items, key)
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (s *MemoryPairs) ListByUser(_ context.Context, userID int64) ([]DevicePair, error) {
|
|
if userID <= 0 {
|
|
return nil, nil
|
|
}
|
|
prefix := userKeyPrefix(userID)
|
|
now := time.Now().UnixNano()
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
out := make([]DevicePair, 0)
|
|
for key := range s.items {
|
|
row := s.items[key]
|
|
if !strings.HasPrefix(key, prefix) {
|
|
continue
|
|
}
|
|
if now >= row.exp {
|
|
delete(s.items, key)
|
|
continue
|
|
}
|
|
did, aid, ok := parsePairKey(key, prefix)
|
|
if !ok || did == 0 {
|
|
continue
|
|
}
|
|
_ = aid
|
|
out = append(out, DevicePair{DeviceID: did, AppID: aid, Pair: row.pair})
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
func parsePairKey(key, userPrefix string) (deviceID uint64, appID uint32, ok bool) {
|
|
rest := strings.TrimPrefix(key, userPrefix)
|
|
parts := strings.Split(rest, ":")
|
|
if len(parts) < 2 {
|
|
return 0, 0, false
|
|
}
|
|
did, err := strconv.ParseUint(parts[0], 10, 64)
|
|
if err != nil || did == 0 {
|
|
return 0, 0, false
|
|
}
|
|
if len(parts) == 2 {
|
|
return did, 0, true
|
|
}
|
|
aid64, err := strconv.ParseUint(parts[1], 10, 32)
|
|
if err != nil {
|
|
return did, 0, true
|
|
}
|
|
return did, uint32(aid64), true
|
|
}
|
|
|
|
type RedisPairs struct {
|
|
rdb *redis.Client
|
|
prefix string
|
|
mode PairKeyMode
|
|
}
|
|
|
|
func NewRedisPairs(rdb *redis.Client, prefix string) *RedisPairs {
|
|
return NewRedisPairsMode(rdb, prefix, PairKeyTriple)
|
|
}
|
|
|
|
func NewRedisPairsMode(rdb *redis.Client, prefix string, mode PairKeyMode) *RedisPairs {
|
|
if prefix == "" {
|
|
prefix = "refresh:"
|
|
}
|
|
return &RedisPairs{rdb: rdb, prefix: prefix, mode: normalizePairKeyMode(mode)}
|
|
}
|
|
|
|
func (s *RedisPairs) key(userID int64, deviceID uint64, appID uint32) string {
|
|
return s.prefix + storageKey(s.mode, userID, deviceID, appID)
|
|
}
|
|
|
|
func (s *RedisPairs) Put(ctx context.Context, userID int64, deviceID uint64, appID uint32, pair TokenPair, ttl time.Duration) error {
|
|
if s == nil || s.rdb == nil {
|
|
return errs.ErrAvailable.WithInnerMsg("jwt refresh redis not ready")
|
|
}
|
|
if userID <= 0 || deviceID == 0 || ttl <= 0 {
|
|
return nil
|
|
}
|
|
raw, err := json.Marshal(pair)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return s.rdb.Set(ctx, s.key(userID, deviceID, appID), raw, ttl).Err()
|
|
}
|
|
|
|
func (s *RedisPairs) Get(ctx context.Context, userID int64, deviceID uint64, appID uint32) (TokenPair, error) {
|
|
if s == nil || s.rdb == nil {
|
|
return TokenPair{}, errs.ErrAvailable.WithInnerMsg("jwt refresh redis not ready")
|
|
}
|
|
raw, err := s.rdb.Get(ctx, s.key(userID, deviceID, appID)).Bytes()
|
|
if err == redis.Nil {
|
|
return TokenPair{}, nil
|
|
}
|
|
if err != nil {
|
|
return TokenPair{}, err
|
|
}
|
|
var pair TokenPair
|
|
if err := json.Unmarshal(raw, &pair); err != nil {
|
|
return TokenPair{}, err
|
|
}
|
|
return pair, nil
|
|
}
|
|
|
|
func (s *RedisPairs) Delete(ctx context.Context, userID int64, deviceID uint64, appID uint32) error {
|
|
if s == nil || s.rdb == nil {
|
|
return errs.ErrAvailable.WithInnerMsg("jwt refresh redis not ready")
|
|
}
|
|
return s.rdb.Del(ctx, s.key(userID, deviceID, appID)).Err()
|
|
}
|
|
|
|
func (s *RedisPairs) DeleteAllOnDevice(ctx context.Context, userID int64, deviceID uint64) error {
|
|
if s == nil || s.rdb == nil {
|
|
return errs.ErrAvailable.WithInnerMsg("jwt refresh redis not ready")
|
|
}
|
|
if s.mode == PairKeyPair {
|
|
return s.rdb.Del(ctx, s.key(userID, deviceID, 0)).Err()
|
|
}
|
|
pattern := s.prefix + deviceKeyPrefix(userID, deviceID) + ":*"
|
|
var cursor uint64
|
|
for {
|
|
keys, next, err := s.rdb.Scan(ctx, cursor, pattern, 100).Result()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if len(keys) > 0 {
|
|
if err := s.rdb.Del(ctx, keys...).Err(); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
cursor = next
|
|
if cursor == 0 {
|
|
break
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (s *RedisPairs) ListByUser(ctx context.Context, userID int64) ([]DevicePair, error) {
|
|
if s == nil || s.rdb == nil {
|
|
return nil, errs.ErrAvailable.WithInnerMsg("jwt refresh redis not ready")
|
|
}
|
|
if userID <= 0 {
|
|
return nil, nil
|
|
}
|
|
prefix := s.prefix + userKeyPrefix(userID)
|
|
var cursor uint64
|
|
out := make([]DevicePair, 0)
|
|
for {
|
|
keys, next, err := s.rdb.Scan(ctx, cursor, prefix+"*", 100).Result()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
for _, key := range keys {
|
|
rest := strings.TrimPrefix(key, prefix)
|
|
parts := strings.Split(rest, ":")
|
|
if len(parts) < 1 {
|
|
continue
|
|
}
|
|
did, err := strconv.ParseUint(parts[0], 10, 64)
|
|
if err != nil || did == 0 {
|
|
continue
|
|
}
|
|
raw, err := s.rdb.Get(ctx, key).Bytes()
|
|
if err == redis.Nil {
|
|
continue
|
|
}
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
var aid uint32
|
|
if len(parts) >= 2 {
|
|
aid64, err := strconv.ParseUint(parts[1], 10, 32)
|
|
if err == nil {
|
|
aid = uint32(aid64)
|
|
}
|
|
}
|
|
var pair TokenPair
|
|
if err := json.Unmarshal(raw, &pair); err != nil {
|
|
continue
|
|
}
|
|
out = append(out, DevicePair{DeviceID: did, AppID: aid, Pair: pair})
|
|
}
|
|
cursor = next
|
|
if cursor == 0 {
|
|
break
|
|
}
|
|
}
|
|
return out, nil
|
|
}
|