github.com/argoproj/argo-cd/v3@v3.2.1/util/session/state.go (about)

     1  package session
     2  
     3  import (
     4  	"context"
     5  	"strings"
     6  	"sync"
     7  	"time"
     8  
     9  	"github.com/redis/go-redis/v9"
    10  	log "github.com/sirupsen/logrus"
    11  
    12  	utilio "github.com/argoproj/argo-cd/v3/util/io"
    13  )
    14  
    15  const (
    16  	revokedTokenPrefix = "revoked-token|"
    17  	newRevokedTokenKey = "new-revoked-token"
    18  )
    19  
    20  type userStateStorage struct {
    21  	attempts            map[string]LoginAttempts
    22  	redis               *redis.Client
    23  	revokedTokens       map[string]bool
    24  	recentRevokedTokens map[string]bool
    25  	lock                sync.RWMutex
    26  	resyncDuration      time.Duration
    27  }
    28  
    29  var _ UserStateStorage = &userStateStorage{}
    30  
    31  func NewUserStateStorage(redis *redis.Client) *userStateStorage {
    32  	return &userStateStorage{
    33  		attempts:            map[string]LoginAttempts{},
    34  		revokedTokens:       map[string]bool{},
    35  		recentRevokedTokens: map[string]bool{},
    36  		resyncDuration:      time.Second * 15,
    37  		redis:               redis,
    38  	}
    39  }
    40  
    41  // Init sets up watches on the revoked tokens and starts a ticker to periodically resync the revoked tokens from Redis.
    42  // Don't call this until after setting up all hooks on the Redis client, or you might encounter race conditions.
    43  func (storage *userStateStorage) Init(ctx context.Context) {
    44  	go storage.watchRevokedTokens(ctx)
    45  	ticker := time.NewTicker(storage.resyncDuration)
    46  	go func() {
    47  		storage.loadRevokedTokensSafe()
    48  		for range ticker.C {
    49  			storage.loadRevokedTokensSafe()
    50  		}
    51  	}()
    52  	go func() {
    53  		<-ctx.Done()
    54  		ticker.Stop()
    55  	}()
    56  }
    57  
    58  func (storage *userStateStorage) watchRevokedTokens(ctx context.Context) {
    59  	pubsub := storage.redis.Subscribe(ctx, newRevokedTokenKey)
    60  	defer utilio.Close(pubsub)
    61  
    62  	ch := pubsub.Channel()
    63  	for {
    64  		select {
    65  		case <-ctx.Done():
    66  			return
    67  		case val := <-ch:
    68  			storage.lock.Lock()
    69  			storage.revokedTokens[val.Payload] = true
    70  			storage.recentRevokedTokens[val.Payload] = true
    71  			storage.lock.Unlock()
    72  		}
    73  	}
    74  }
    75  
    76  func (storage *userStateStorage) loadRevokedTokensSafe() {
    77  	err := storage.loadRevokedTokens()
    78  	for err != nil {
    79  		log.Warnf("Failed to resync revoked tokens. retrying again in 1 minute: %v", err)
    80  		time.Sleep(time.Minute)
    81  		err = storage.loadRevokedTokens()
    82  	}
    83  }
    84  
    85  func (storage *userStateStorage) loadRevokedTokens() error {
    86  	redisRevokedTokens := map[string]bool{}
    87  	iterator := storage.redis.Scan(context.Background(), 0, revokedTokenPrefix+"*", 10000).Iterator()
    88  	for iterator.Next(context.Background()) {
    89  		parts := strings.Split(iterator.Val(), "|")
    90  		if len(parts) != 2 {
    91  			log.Warnf("Unexpected redis key prefixed with '%s'. Must have token id after the prefix but got: '%s'.",
    92  				revokedTokenPrefix,
    93  				iterator.Val())
    94  			continue
    95  		}
    96  		redisRevokedTokens[parts[1]] = true
    97  	}
    98  	if iterator.Err() != nil {
    99  		return iterator.Err()
   100  	}
   101  
   102  	storage.lock.Lock()
   103  	defer storage.lock.Unlock()
   104  	storage.revokedTokens = redisRevokedTokens
   105  	for recentRevokedToken := range storage.recentRevokedTokens {
   106  		storage.revokedTokens[recentRevokedToken] = true
   107  	}
   108  	storage.recentRevokedTokens = map[string]bool{}
   109  
   110  	return nil
   111  }
   112  
   113  func (storage *userStateStorage) GetLoginAttempts() map[string]LoginAttempts {
   114  	return storage.attempts
   115  }
   116  
   117  func (storage *userStateStorage) SetLoginAttempts(attempts map[string]LoginAttempts) error {
   118  	storage.attempts = attempts
   119  	return nil
   120  }
   121  
   122  func (storage *userStateStorage) RevokeToken(ctx context.Context, id string, expiringAt time.Duration) error {
   123  	storage.lock.Lock()
   124  	storage.revokedTokens[id] = true
   125  	storage.recentRevokedTokens[id] = true
   126  	storage.lock.Unlock()
   127  	if err := storage.redis.Set(ctx, revokedTokenPrefix+id, "", expiringAt).Err(); err != nil {
   128  		return err
   129  	}
   130  	return storage.redis.Publish(ctx, newRevokedTokenKey, id).Err()
   131  }
   132  
   133  func (storage *userStateStorage) IsTokenRevoked(id string) bool {
   134  	storage.lock.RLock()
   135  	defer storage.lock.RUnlock()
   136  	return storage.revokedTokens[id]
   137  }
   138  
   139  func (storage *userStateStorage) GetLockObject() *sync.RWMutex {
   140  	return &storage.lock
   141  }
   142  
   143  type UserStateStorage interface {
   144  	Init(ctx context.Context)
   145  	// GetLoginAttempts return number of concurrent login attempts
   146  	GetLoginAttempts() map[string]LoginAttempts
   147  	// SetLoginAttempts sets number of concurrent login attempts
   148  	SetLoginAttempts(attempts map[string]LoginAttempts) error
   149  	// RevokeToken revokes token with given id (information about revocation expires after specified timeout)
   150  	RevokeToken(ctx context.Context, id string, expiringAt time.Duration) error
   151  	// IsTokenRevoked checks if given token is revoked
   152  	IsTokenRevoked(id string) bool
   153  	// GetLockObject returns a lock used by the storage
   154  	GetLockObject() *sync.RWMutex
   155  }