iotlcm/internal/token/token.go

// Package token provides the process-local shared token core used by OneLink
// clients. It deliberately does not own channel configuration or generated
// database types; callers seed it from their channel cache.
package token

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"sync"
	"time"

	"git.5iot.com/iotlcm-public/iotlcm-upstream-chan/internal/model"
	"git.5iot.com/iotlcm-public/iotlcm-upstream-chan/internal/model/entity"
	"git.5iot.com/iotlcm-public/iotlcm-upstream-chan/internal/service"
	"github.com/gogf/gf/v2/frame/g"
	"github.com/gogf/gf/v2/os/gtime"
	"golang.org/x/sync/singleflight"
)

var (
	ErrEmptyToken       = errors.New("token fetcher returned an empty token")
	ErrTokenUnavailable = errors.New("token unavailable")
)

// OneLinkToken is the value cached for a channel. Times are in the local
// process clock and are persisted as gtime values by the manager.
type OneLinkToken struct {
	Value       string
	DeviceValue string
	CreatedAt   time.Time
	ExpiresAt   time.Time
}

func (t OneLinkToken) Valid(now time.Time) bool {
	return strings.TrimSpace(t.Value) != "" && !t.ExpiresAt.IsZero() && now.Before(t.ExpiresAt)
}

// TokenFetcher obtains a token from the upstream provider. FetchToken must
// honor ctx cancellation.
type TokenFetcher interface {
	FetchToken(ctx context.Context, chanKey string) (OneLinkToken, error)
}

// TokenProvider is the consumer-facing provider interface.
type TokenProvider interface {
	GetToken(ctx context.Context, chanKey string) (string, error)
}

// TokenPersister is injectable for tests and alternate storage. The default
// manager uses service.ChannelUpdate().UpdateAccessToken.
type TokenPersister interface {
	UpdateAccessToken(ctx context.Context, input *model.UpdateAccessTokenInput) error
}

type TokenFetcherFunc func(context.Context, string) (OneLinkToken, error)

func (f TokenFetcherFunc) FetchToken(ctx context.Context, chanKey string) (OneLinkToken, error) {
	return f(ctx, chanKey)
}

type TokenPersisterFunc func(context.Context, *model.UpdateAccessTokenInput) error

func (f TokenPersisterFunc) UpdateAccessToken(ctx context.Context, input *model.UpdateAccessTokenInput) error {
	return f(ctx, input)
}

type Option func(*manager)

func WithRefreshAhead(d time.Duration) Option {
	return func(m *manager) {
		if d >= 0 {
			m.refreshAhead = d
		}
	}
}

// WithUntilInvalidated allows tokens without an expiry timestamp. They remain
// valid until the consumer explicitly invalidates the rejected value.
func WithUntilInvalidated() Option {
	return func(m *manager) {
		m.untilInvalidated = true
	}
}

func WithClock(now func() time.Time) Option {
	return func(m *manager) {
		if now != nil {
			m.now = now
		}
	}
}
func WithPersister(p TokenPersister) Option {
	return func(m *manager) {
		if p != nil {
			m.persister = p
		}
	}
}

type tokenEntry struct {
	mu      sync.RWMutex
	token   OneLinkToken
	version uint64
}

// TokenManager is the shared-token contract consumed by channel clients.
// Implementations are safe for concurrent use.
type TokenManager interface {
	TokenProvider
	Get(context.Context, string) (OneLinkToken, error)
	Seed(string, OneLinkToken)
	SeedFromChannel(*entity.TSimChannels)
	SeedFromConfig(*model.SimChannelConfigItemCache)
	InvalidateIfCurrent(string, string) bool
	Current(string) (OneLinkToken, bool)
}

// TokenManager caches one token per channel and coalesces refreshes per
// channel. A failed refresh never overwrites a cached token.
type manager struct {
	fetcher          TokenFetcher
	persister        TokenPersister
	refreshAhead     time.Duration
	untilInvalidated bool
	now              func() time.Time
}

// Runtime and request-scoped test instances share token state by chanKey.
// This keeps their refreshes in one flight without coupling either instance
// to the runtime channel pool.
var (
	sharedTokenEntries sync.Map // map[string]*tokenEntry
	sharedTokenFlights singleflight.Group
)

func NewTokenManager(fetcher TokenFetcher, options ...Option) TokenManager {
	m := &manager{fetcher: fetcher, now: time.Now}
	for _, option := range options {
		if option != nil {
			option(m)
		}
	}
	return m
}

func (m *manager) entry(chanKey string) *tokenEntry {
	e, _ := sharedTokenEntries.LoadOrStore(chanKey, &tokenEntry{})
	return e.(*tokenEntry)
}

func (m *manager) valid(t OneLinkToken, now time.Time) bool {
	if strings.TrimSpace(t.Value) == "" {
		return false
	}
	if t.ExpiresAt.IsZero() {
		return m.untilInvalidated
	}
	return now.Before(t.ExpiresAt)
}

func (m *manager) reusable(t OneLinkToken, now time.Time) bool {
	if !m.valid(t, now) {
		return false
	}
	return t.ExpiresAt.IsZero() || m.refreshAhead <= 0 || now.Before(t.ExpiresAt.Add(-m.refreshAhead))
}

// Seed initializes or replaces the in-process value. It is safe to call when
// channel configuration is reloaded.
func (m *manager) Seed(chanKey string, value OneLinkToken) {
	e := m.entry(chanKey)
	e.mu.Lock()
	e.token = value
	e.version++
	e.mu.Unlock()
}

// SeedFromChannel copies access_token/access_token_expire from channel config.
func (m *manager) SeedFromChannel(channel *entity.TSimChannels) {
	if channel == nil {
		return
	}
	var created, expires time.Time
	if channel.AccessTokenCreated != nil {
		created = channel.AccessTokenCreated.Time
	}
	if channel.AccessTokenExpire != nil {
		expires = channel.AccessTokenExpire.Time
	}
	incoming := OneLinkToken{Value: channel.AccessToken, DeviceValue: channel.RefreshToken, CreatedAt: created, ExpiresAt: expires}
	e := m.entry(channel.ChanKey)
	e.mu.Lock()
	current := e.token
	if strings.TrimSpace(current.Value) == "" || incoming.CreatedAt.After(current.CreatedAt) ||
		(incoming.Value == current.Value && incoming.DeviceValue != current.DeviceValue) {
		e.token = incoming
		e.version++
	}
	e.mu.Unlock()
}

// SeedFromConfig is a convenience for the channel service's cached config.
func (m *manager) SeedFromConfig(config *model.SimChannelConfigItemCache) {
	if config == nil {
		return
	}
	m.SeedFromChannel(config.Channel)
}

// GetToken returns a valid token, refreshing when it is expired or within the
// configured refresh-ahead window. A still-valid old token is returned when
// refresh or persistence fails.
func (m *manager) GetToken(ctx context.Context, chanKey string) (string, error) {
	t, err := m.Get(ctx, chanKey)
	return t.Value, err
}

func (m *manager) Get(ctx context.Context, chanKey string) (OneLinkToken, error) {
	if err := ctx.Err(); err != nil {
		return OneLinkToken{}, err
	}
	if strings.TrimSpace(chanKey) == "" {
		return OneLinkToken{}, fmt.Errorf("chanKey is required")
	}
	e := m.entry(chanKey)
	now := m.now()
	e.mu.RLock()
	current := e.token
	e.mu.RUnlock()
	if m.reusable(current, now) {
		return current, nil
	}

	resultCh := sharedTokenFlights.DoChan(chanKey, func() (any, error) {
		// Re-check after joining the per-channel flight.
		e.mu.RLock()
		current := e.token
		e.mu.RUnlock()
		now := m.now()
		if m.reusable(current, now) {
			return current, nil
		}
		if m.fetcher == nil {
			return nil, ErrTokenUnavailable
		}
		fresh, err := m.fetcher.FetchToken(ctx, chanKey)
		if err != nil {
			return nil, err
		}
		fresh.Value = strings.TrimSpace(fresh.Value)
		if fresh.Value == "" {
			return nil, ErrEmptyToken
		}
		if fresh.ExpiresAt.IsZero() && !m.untilInvalidated {
			return nil, fmt.Errorf("token expiry must be in the future")
		}
		if !fresh.ExpiresAt.IsZero() && !fresh.ExpiresAt.After(now) {
			return nil, fmt.Errorf("token expiry must be in the future")
		}
		if fresh.CreatedAt.IsZero() {
			fresh.CreatedAt = now
		}
		if err := m.persist(ctx, chanKey, fresh); err != nil {
			return nil, err
		}
		g.Log().Infof(ctx, "shared token refreshed and persisted, chanKey: %s", chanKey)
		e.mu.Lock()
		e.token = fresh
		e.version++
		e.mu.Unlock()
		return fresh, nil
	})
	select {
	case <-ctx.Done():
		return OneLinkToken{}, ctx.Err()
	case result := <-resultCh:
		if result.Err == nil {
			return result.Val.(OneLinkToken), nil
		}
		now = m.now()
		e.mu.RLock()
		current = e.token
		e.mu.RUnlock()
		if m.valid(current, now) {
			return current, nil
		}
		return OneLinkToken{}, result.Err
	}
}

func (m *manager) persist(ctx context.Context, chanKey string, t OneLinkToken) error {
	p := m.persister
	if p == nil {
		p = channelUpdatePersister{}
	}
	var expiresAt *gtime.Time
	if !t.ExpiresAt.IsZero() {
		expiresAt = gtime.New(t.ExpiresAt)
	}
	return p.UpdateAccessToken(ctx, &model.UpdateAccessTokenInput{
		ChanKey: chanKey, AccessToken: t.Value, RefreshToken: t.DeviceValue,
		AccessTokenCreated: gtime.New(t.CreatedAt), AccessTokenExpire: expiresAt, RefreshTokenExpire: expiresAt,
	})
}

type channelUpdatePersister struct{}

func (channelUpdatePersister) UpdateAccessToken(ctx context.Context, input *model.UpdateAccessTokenInput) error {
	return service.ChannelUpdate().UpdateAccessToken(ctx, input)
}

// InvalidateIfCurrent atomically removes value only when it still matches.
// It returns false when another request has already installed a new token.
func (m *manager) InvalidateIfCurrent(chanKey, value string) bool {
	e := m.entry(chanKey)
	e.mu.Lock()
	defer e.mu.Unlock()
	if e.token.Value != value {
		return false
	}
	e.token = OneLinkToken{}
	e.version++
	return true
}

// Current returns the cached value without triggering a refresh.
func (m *manager) Current(chanKey string) (OneLinkToken, bool) {
	e, ok := sharedTokenEntries.Load(chanKey)
	if !ok {
		return OneLinkToken{}, false
	}
	e.(*tokenEntry).mu.RLock()
	t := e.(*tokenEntry).token
	e.(*tokenEntry).mu.RUnlock()
	return t, t.Value != ""
}