pkg/cache/cache_redis.go

Rama master · réplica de lectura del repositorio original.

package cache

import (
	"context"
	"errors"
	"strconv"
	"time"

	rediscache "github.com/go-redis/cache/v9"
	"github.com/redis/go-redis/v9"
	"github.com/shellhub-io/shellhub/pkg/clock"
)

type redisCache struct {
	cache  *rediscache.Cache
	client *redis.Client
	cfg    Config
	clock  clock.Clock
}

var _ Cache = &redisCache{}

// NewRedisCache connects to Redis and returns a cache backed by it, configured by cfg and reading c
// to time account lockouts. It fails rather than falling back: a misconfigured URI should surface
// at startup, not as silent cache misses forever.
func NewRedisCache(uri string, pool int, cfg Config, c clock.Clock) (Cache, error) {
	opt, err := redis.ParseURL(uri)
	if err != nil {
		return nil, err
	}

	if pool > 0 {
		opt.PoolSize = pool
	}

	client := redis.NewClient(opt)

	return &redisCache{
		cfg:    cfg,
		client: client,
		clock:  c,
		cache: rediscache.New(&rediscache.Options{
			Redis: client,
		}),
	}, nil
}

// Get gets the cache value for the given key.
// NOTE: missing key is not an error.
func (c *redisCache) Get(ctx context.Context, key string, value any) error {
	err := c.cache.Get(ctx, key, value)
	if errors.Is(err, rediscache.ErrCacheMiss) {
		return nil
	}

	return err
}

// Set puts value into cache with key and expire time.
func (c *redisCache) Set(ctx context.Context, key string, value any, ttl time.Duration) error {
	return c.cache.Set(&rediscache.Item{Ctx: ctx, Key: key, Value: value, TTL: ttl})
}

// SetNX atomically sets key to value with the given ttl only if it does not
// already exist, reporting whether it was set. Unlike Set (which serializes
// through go-redis/cache), this goes straight to the Redis SETNX so it can back
// single-use reservations.
func (c *redisCache) SetNX(ctx context.Context, key string, value any, ttl time.Duration) (bool, error) {
	return c.client.SetNX(ctx, key, value, ttl).Result()
}

var compareAndDeleteScript = redis.NewScript(`
if redis.call("GET", KEYS[1]) == ARGV[1] then
	return redis.call("DEL", KEYS[1])
end
return 0
`)

func (c *redisCache) CompareAndDelete(ctx context.Context, key, value string) (bool, error) {
	deleted, err := compareAndDeleteScript.Run(ctx, c.client, []string{key}, value).Int()
	if err != nil {
		return false, err
	}

	return deleted == 1, nil
}

// Delete deletes cached value by given key.
func (c *redisCache) Delete(ctx context.Context, key string) error {
	if err := c.cache.Get(ctx, key, nil); errors.Is(err, rediscache.ErrCacheMiss) {
		return nil
	}

	return c.cache.Delete(ctx, key)
}

func (c *redisCache) HasAccountLockout(ctx context.Context, source, id string) (int64, int, error) {
	if c.cfg.MaximumAccountLockout <= 0 {
		return 0, 0, nil
	}

	lockoutDeadline := "0"
	if err := c.Get(ctx, "account-lockout="+source+":"+id, &lockoutDeadline); err != nil {
		return 0, 0, err
	}

	attemptSTR := "0"
	if err := c.Get(ctx, "login-attempt="+source+":"+id, &attemptSTR); err != nil {
		return 0, 0, err
	}

	lockout, _ := strconv.ParseInt(lockoutDeadline, 10, 0)
	attempt, _ := strconv.Atoi(attemptSTR)

	if c.clock.Now().Unix() >= lockout {
		return 0, attempt, nil
	}

	return lockout, attempt, nil
}

func (c *redisCache) StoreLoginAttempt(ctx context.Context, source, id string) (int64, int, error) {
	if c.cfg.MaximumAccountLockout <= 0 {
		return 0, 0, nil
	}

	attemptSTR := "0"
	if err := c.Get(ctx, "login-attempt="+source+":"+id, &attemptSTR); err != nil {
		return 0, 0, err
	}

	attempt, _ := strconv.Atoi(attemptSTR)
	attempt++

	now := c.clock.Now()
	lockoutTTL, attemptTTL := c.cfg.lockout(attempt)

	if err := c.Set(ctx, "login-attempt="+source+":"+id, strconv.Itoa(attempt), attemptTTL); err != nil {
		return 0, attempt, err
	}

	if attempt <= 2 {
		return 0, attempt, nil
	}

	lockoutDeadline := strconv.FormatInt(now.Add(lockoutTTL).Unix(), 10)
	if err := c.Set(ctx, "account-lockout="+source+":"+id, lockoutDeadline, lockoutTTL); err != nil {
		return 0, attempt, err
	}

	lockout, _ := strconv.ParseInt(lockoutDeadline, 10, 0)

	return lockout, attempt, nil
}

func (c *redisCache) ResetLoginAttempts(ctx context.Context, source, id string) error {
	if c.cfg.MaximumAccountLockout <= 0 {
		return nil
	}

	if err := c.Delete(ctx, "login-attempt="+source+":"+id); err != nil {
		return err
	}

	return c.Delete(ctx, "account-lockout="+source+":"+id)
}