spark-server/internal/module/auth/logic/refresh.go

package logic

import (
	"context"
	"encoding/json"
	"strconv"
	"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
}

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
}

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()
}