127 lines
2.9 KiB
Go
127 lines
2.9 KiB
Go
package mysql
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"strings"
|
|
"time"
|
|
|
|
"gorm.io/gorm"
|
|
"gorm.io/gorm/clause"
|
|
|
|
"wucher/internal/domain/transient"
|
|
)
|
|
|
|
type TokenStore struct {
|
|
db *gorm.DB
|
|
now func() time.Time
|
|
}
|
|
|
|
func NewTokenStore(db *gorm.DB) *TokenStore {
|
|
return &TokenStore{
|
|
db: db,
|
|
now: func() time.Time { return time.Now().UTC() },
|
|
}
|
|
}
|
|
|
|
func (s *TokenStore) Set(ctx context.Context, key string, value []byte, ttl time.Duration) error {
|
|
if s == nil || s.db == nil {
|
|
return errors.New("token store db is required")
|
|
}
|
|
key = strings.TrimSpace(key)
|
|
if key == "" {
|
|
return errors.New("token store key is required")
|
|
}
|
|
|
|
now := s.now().UTC()
|
|
var expiresAt *time.Time
|
|
if ttl > 0 {
|
|
ts := now.Add(ttl)
|
|
expiresAt = &ts
|
|
}
|
|
|
|
entry := transient.TokenEntry{
|
|
StoreKey: key,
|
|
Value: append([]byte(nil), value...),
|
|
ExpiresAt: expiresAt,
|
|
}
|
|
return s.db.WithContext(ctx).Clauses(clause.OnConflict{
|
|
Columns: []clause.Column{{Name: "store_key"}},
|
|
DoUpdates: clause.Assignments(map[string]any{
|
|
"value": entry.Value,
|
|
"expires_at": entry.ExpiresAt,
|
|
"updated_at": now,
|
|
}),
|
|
}).Create(&entry).Error
|
|
}
|
|
|
|
func (s *TokenStore) SetIfNotExists(ctx context.Context, key string, value []byte, ttl time.Duration) (bool, error) {
|
|
if s == nil || s.db == nil {
|
|
return false, errors.New("token store db is required")
|
|
}
|
|
key = strings.TrimSpace(key)
|
|
if key == "" {
|
|
return false, errors.New("token store key is required")
|
|
}
|
|
|
|
now := s.now().UTC()
|
|
var expiresAt *time.Time
|
|
if ttl > 0 {
|
|
ts := now.Add(ttl)
|
|
expiresAt = &ts
|
|
if err := s.db.WithContext(ctx).
|
|
Delete(&transient.TokenEntry{}, "store_key = ? AND expires_at IS NOT NULL AND expires_at <= ?", key, now).Error; err != nil {
|
|
return false, err
|
|
}
|
|
}
|
|
|
|
entry := transient.TokenEntry{
|
|
StoreKey: key,
|
|
Value: append([]byte(nil), value...),
|
|
ExpiresAt: expiresAt,
|
|
}
|
|
tx := s.db.WithContext(ctx).Clauses(clause.OnConflict{
|
|
Columns: []clause.Column{{Name: "store_key"}},
|
|
DoNothing: true,
|
|
}).Create(&entry)
|
|
if tx.Error != nil {
|
|
return false, tx.Error
|
|
}
|
|
return tx.RowsAffected > 0, nil
|
|
}
|
|
|
|
func (s *TokenStore) Get(ctx context.Context, key string) ([]byte, error) {
|
|
if s == nil || s.db == nil {
|
|
return nil, errors.New("token store db is required")
|
|
}
|
|
key = strings.TrimSpace(key)
|
|
if key == "" {
|
|
return nil, nil
|
|
}
|
|
|
|
var entry transient.TokenEntry
|
|
tx := s.db.WithContext(ctx).Where("store_key = ?", key).Limit(1).Find(&entry)
|
|
if tx.Error != nil {
|
|
return nil, tx.Error
|
|
}
|
|
if tx.RowsAffected == 0 {
|
|
return nil, nil
|
|
}
|
|
if entry.ExpiresAt != nil && !entry.ExpiresAt.After(s.now().UTC()) {
|
|
_ = s.Delete(ctx, key)
|
|
return nil, nil
|
|
}
|
|
return append([]byte(nil), entry.Value...), nil
|
|
}
|
|
|
|
func (s *TokenStore) Delete(ctx context.Context, key string) error {
|
|
if s == nil || s.db == nil {
|
|
return errors.New("token store db is required")
|
|
}
|
|
key = strings.TrimSpace(key)
|
|
if key == "" {
|
|
return nil
|
|
}
|
|
return s.db.WithContext(ctx).Delete(&transient.TokenEntry{}, "store_key = ?", key).Error
|
|
}
|