59 lines
1.3 KiB
Go
59 lines
1.3 KiB
Go
package model
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
type MemoryStore struct {
|
|
mu sync.Mutex
|
|
codes map[string]codeEntry
|
|
cooldowns map[string]time.Time
|
|
}
|
|
|
|
type codeEntry struct {
|
|
code string
|
|
exp time.Time
|
|
}
|
|
|
|
func NewMemoryStore() *MemoryStore {
|
|
return &MemoryStore{codes: make(map[string]codeEntry), cooldowns: make(map[string]time.Time)}
|
|
}
|
|
|
|
func (m *MemoryStore) Put(_ context.Context, key, code string, ttlSeconds int) error {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
m.codes[key] = codeEntry{code: code, exp: time.Now().Add(time.Duration(ttlSeconds) * time.Second)}
|
|
return nil
|
|
}
|
|
|
|
func (m *MemoryStore) Get(_ context.Context, key string) (string, error) {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
e, ok := m.codes[key]
|
|
if !ok || time.Now().After(e.exp) {
|
|
delete(m.codes, key)
|
|
return "", nil
|
|
}
|
|
return e.code, nil
|
|
}
|
|
|
|
func (m *MemoryStore) Del(_ context.Context, key string) error {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
delete(m.codes, key)
|
|
delete(m.cooldowns, key)
|
|
return nil
|
|
}
|
|
|
|
func (m *MemoryStore) SetCooldown(_ context.Context, key string, ttlSeconds int) (bool, error) {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
if exp, ok := m.cooldowns[key]; ok && time.Now().Before(exp) {
|
|
return false, nil
|
|
}
|
|
m.cooldowns[key] = time.Now().Add(time.Duration(ttlSeconds) * time.Second)
|
|
return true, nil
|
|
}
|