spark-server/internal/module/auth/logic/refresh.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, pair TokenPair, ttl time.Duration) error
	Get(ctx context.Context, userID int64, deviceID uint64) (TokenPair, error)
	Delete(ctx context.Context, userID int64, deviceID uint64) error
	ListByUser(ctx context.Context, userID int64) ([]DevicePair, error)
}

// DevicePair is the live token pair for one device of a user.
type DevicePair struct {
	DeviceID uint64
	Pair     TokenPair
}

type MemoryPairs struct {
	mu    sync.Mutex
	items map[string]storedPair
}

type storedPair struct {
	pair TokenPair
	exp  int64
}

func NewMemoryPairs() *MemoryPairs {
	return &MemoryPairs{items: make(map[string]storedPair)}
}

func pairKey(userID int64, deviceID uint64) string {
	return strconv.FormatInt(userID, 10) + ":" + strconv.FormatUint(deviceID, 10)
}

func (s *MemoryPairs) Put(_ context.Context, userID int64, deviceID uint64, pair TokenPair, ttl time.Duration) error {
	if userID <= 0 || deviceID == 0 || ttl <= 0 {
		return nil
	}
	s.mu.Lock()
	s.items[pairKey(userID, deviceID)] = storedPair{pair: pair, exp: time.Now().Add(ttl).UnixNano()}
	s.mu.Unlock()
	return nil
}

func (s *MemoryPairs) Get(_ context.Context, userID int64, deviceID uint64) (TokenPair, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	row, ok := s.items[pairKey(userID, deviceID)]
	if !ok {
		return TokenPair{}, nil
	}
	if time.Now().UnixNano() >= row.exp {
		delete(s.items, pairKey(userID, deviceID))
		return TokenPair{}, nil
	}
	return row.pair, nil
}

func (s *MemoryPairs) Delete(_ context.Context, userID int64, deviceID uint64) error {
	s.mu.Lock()
	delete(s.items, pairKey(userID, deviceID))
	s.mu.Unlock()
	return nil
}

func (s *MemoryPairs) ListByUser(_ context.Context, userID int64) ([]DevicePair, error) {
	if userID <= 0 {
		return nil, nil
	}
	prefix := strconv.FormatInt(userID, 10) + ":"
	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, err := strconv.ParseUint(key[len(prefix):], 10, 64)
		if err != nil || did == 0 {
			continue
		}
		out = append(out, DevicePair{DeviceID: did, Pair: row.pair})
	}
	return out, nil
}

type RedisPairs struct {
	rdb    *redis.Client
	prefix string
}

func NewRedisPairs(rdb *redis.Client, prefix string) *RedisPairs {
	if prefix == "" {
		prefix = "refresh:"
	}
	return &RedisPairs{rdb: rdb, prefix: prefix}
}

func (s *RedisPairs) key(userID int64, deviceID uint64) string {
	return s.prefix + pairKey(userID, deviceID)
}

func (s *RedisPairs) Put(ctx context.Context, userID int64, deviceID uint64, 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), raw, ttl).Err()
}

func (s *RedisPairs) Get(ctx context.Context, userID int64, deviceID uint64) (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)).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) 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)).Err()
}

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 + strconv.FormatInt(userID, 10) + ":"
	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 {
			did, err := strconv.ParseUint(strings.TrimPrefix(key, prefix), 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 pair TokenPair
			if err := json.Unmarshal(raw, &pair); err != nil {
				continue
			}
			out = append(out, DevicePair{DeviceID: did, Pair: pair})
		}
		cursor = next
		if cursor == 0 {
			break
		}
	}
	return out, nil
}