126 lines
2.6 KiB
Go
126 lines
2.6 KiB
Go
package smsstore
|
|
|
|
import (
|
|
"crypto/rand"
|
|
"fmt"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
type Purpose string
|
|
|
|
const (
|
|
PurposeLogin Purpose = "login"
|
|
PurposeBind Purpose = "bind"
|
|
)
|
|
|
|
type Record struct {
|
|
Code string
|
|
ExpiresAt time.Time
|
|
SentAt time.Time
|
|
Attempts int
|
|
}
|
|
|
|
type Store struct {
|
|
mu sync.Mutex
|
|
byKey map[string]*Record
|
|
ttl time.Duration
|
|
resendGap time.Duration
|
|
maxAttempts int
|
|
fixedCode string // 开发固定码；空则随机 6 位
|
|
}
|
|
|
|
func New(ttl, resendGap time.Duration, fixedCode string) *Store {
|
|
if ttl <= 0 {
|
|
ttl = 5 * time.Minute
|
|
}
|
|
if resendGap <= 0 {
|
|
resendGap = 60 * time.Second
|
|
}
|
|
return &Store{
|
|
byKey: map[string]*Record{},
|
|
ttl: ttl,
|
|
resendGap: resendGap,
|
|
maxAttempts: 5,
|
|
fixedCode: fixedCode,
|
|
}
|
|
}
|
|
|
|
func key(purpose Purpose, phone string) string {
|
|
return string(purpose) + ":" + phone
|
|
}
|
|
|
|
// Issue 生成并保存验证码；若未到重发间隔返回 retryAfter>0。
|
|
func (s *Store) Issue(purpose Purpose, phone string) (code string, expiresIn, retryAfter int, err error) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
s.gcLocked()
|
|
k := key(purpose, phone)
|
|
now := time.Now()
|
|
if old, ok := s.byKey[k]; ok && now.Sub(old.SentAt) < s.resendGap {
|
|
left := int((s.resendGap - now.Sub(old.SentAt)).Seconds())
|
|
if left < 1 {
|
|
left = 1
|
|
}
|
|
return "", 0, left, fmt.Errorf("发送过于频繁，请 %d 秒后再试", left)
|
|
}
|
|
code = s.fixedCode
|
|
if code == "" {
|
|
var e error
|
|
code, e = randomDigits(6)
|
|
if e != nil {
|
|
return "", 0, 0, e
|
|
}
|
|
}
|
|
s.byKey[k] = &Record{
|
|
Code: code,
|
|
ExpiresAt: now.Add(s.ttl),
|
|
SentAt: now,
|
|
}
|
|
return code, int(s.ttl.Seconds()), 0, nil
|
|
}
|
|
|
|
// Consume 校验并消费验证码（成功后删除）。
|
|
func (s *Store) Consume(purpose Purpose, phone, code string) error {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
s.gcLocked()
|
|
k := key(purpose, phone)
|
|
rec, ok := s.byKey[k]
|
|
if !ok || time.Now().After(rec.ExpiresAt) {
|
|
delete(s.byKey, k)
|
|
return fmt.Errorf("验证码无效或已过期")
|
|
}
|
|
if rec.Code != code {
|
|
rec.Attempts++
|
|
if rec.Attempts >= s.maxAttempts {
|
|
delete(s.byKey, k)
|
|
return fmt.Errorf("验证码错误次数过多，请重新获取")
|
|
}
|
|
return fmt.Errorf("验证码错误")
|
|
}
|
|
delete(s.byKey, k)
|
|
return nil
|
|
}
|
|
|
|
func (s *Store) gcLocked() {
|
|
now := time.Now()
|
|
for k, r := range s.byKey {
|
|
if now.After(r.ExpiresAt.Add(time.Minute)) {
|
|
delete(s.byKey, k)
|
|
}
|
|
}
|
|
}
|
|
|
|
func randomDigits(n int) (string, error) {
|
|
b := make([]byte, n)
|
|
if _, err := rand.Read(b); err != nil {
|
|
return "", err
|
|
}
|
|
out := make([]byte, n)
|
|
for i := range b {
|
|
out[i] = '0' + b[i]%10
|
|
}
|
|
return string(out), nil
|
|
}
|