157 lines
3.3 KiB
Go
157 lines
3.3 KiB
Go
package system
|
||
|
||
import (
|
||
"sync"
|
||
"time"
|
||
|
||
"github.com/go-kratos/kratos/v2/transport/http"
|
||
"github.com/mojocn/base64Captcha"
|
||
)
|
||
|
||
// CaptchaService 验证码服务
|
||
type CaptchaService struct {
|
||
config CaptchaConfig
|
||
store base64Captcha.Store
|
||
blackCache *BlackCache
|
||
}
|
||
|
||
// CaptchaConfig 验证码配置
|
||
type CaptchaConfig struct {
|
||
KeyLong int
|
||
ImgWidth int
|
||
ImgHeight int
|
||
OpenCaptcha int // 防爆次数，0表示关闭
|
||
OpenCaptchaTimeOut int // 缓存超时时间(秒)
|
||
}
|
||
|
||
// BlackCache 黑名单缓存
|
||
type BlackCache struct {
|
||
mu sync.RWMutex
|
||
items map[string]*cacheItem
|
||
}
|
||
|
||
type cacheItem struct {
|
||
value int
|
||
expiration time.Time
|
||
}
|
||
|
||
// NewBlackCache 创建黑名单缓存
|
||
func NewBlackCache() *BlackCache {
|
||
bc := &BlackCache{
|
||
items: make(map[string]*cacheItem),
|
||
}
|
||
go bc.cleanup()
|
||
return bc
|
||
}
|
||
|
||
func (c *BlackCache) Set(key string, value int, duration time.Duration) {
|
||
c.mu.Lock()
|
||
defer c.mu.Unlock()
|
||
c.items[key] = &cacheItem{
|
||
value: value,
|
||
expiration: time.Now().Add(duration),
|
||
}
|
||
}
|
||
|
||
func (c *BlackCache) Get(key string) (int, bool) {
|
||
c.mu.RLock()
|
||
defer c.mu.RUnlock()
|
||
item, ok := c.items[key]
|
||
if !ok || time.Now().After(item.expiration) {
|
||
return 0, false
|
||
}
|
||
return item.value, true
|
||
}
|
||
|
||
func (c *BlackCache) cleanup() {
|
||
ticker := time.NewTicker(time.Minute)
|
||
for range ticker.C {
|
||
c.mu.Lock()
|
||
now := time.Now()
|
||
for k, v := range c.items {
|
||
if now.After(v.expiration) {
|
||
delete(c.items, k)
|
||
}
|
||
}
|
||
c.mu.Unlock()
|
||
}
|
||
}
|
||
|
||
// CaptchaResponse 验证码响应
|
||
type CaptchaResponse struct {
|
||
CaptchaId string `json:"captchaId"`
|
||
PicPath string `json:"picPath"`
|
||
CaptchaLength int `json:"captchaLength"`
|
||
OpenCaptcha bool `json:"openCaptcha"`
|
||
}
|
||
|
||
// NewCaptchaService 创建验证码服务
|
||
func NewCaptchaService(config CaptchaConfig, store base64Captcha.Store) *CaptchaService {
|
||
if store == nil {
|
||
store = base64Captcha.DefaultMemStore
|
||
}
|
||
return &CaptchaService{
|
||
config: config,
|
||
store: store,
|
||
blackCache: NewBlackCache(),
|
||
}
|
||
}
|
||
|
||
// RegisterRoutes 注册路由
|
||
func (s *CaptchaService) RegisterRoutes(srv *http.Server) {
|
||
r := srv.Route("/")
|
||
r.POST("/base/captcha", s.handleCaptcha)
|
||
}
|
||
|
||
func (s *CaptchaService) handleCaptcha(ctx http.Context) error {
|
||
// 获取客户端IP
|
||
clientIP := ctx.Request().RemoteAddr
|
||
|
||
// 判断验证码是否开启
|
||
openCaptcha := s.config.OpenCaptcha
|
||
openCaptchaTimeOut := s.config.OpenCaptchaTimeOut
|
||
|
||
v, ok := s.blackCache.Get(clientIP)
|
||
if !ok {
|
||
s.blackCache.Set(clientIP, 1, time.Second*time.Duration(openCaptchaTimeOut))
|
||
}
|
||
|
||
var oc bool
|
||
if openCaptcha == 0 || openCaptcha < v {
|
||
oc = true
|
||
}
|
||
|
||
// 生成验证码
|
||
driver := base64Captcha.NewDriverDigit(
|
||
s.config.ImgHeight,
|
||
s.config.ImgWidth,
|
||
s.config.KeyLong,
|
||
0.7,
|
||
80,
|
||
)
|
||
cp := base64Captcha.NewCaptcha(driver, s.store)
|
||
id, b64s, _, err := cp.Generate()
|
||
if err != nil {
|
||
return ctx.Result(200, map[string]any{
|
||
"code": 7,
|
||
"msg": "验证码获取失败",
|
||
})
|
||
}
|
||
|
||
return ctx.Result(200, map[string]any{
|
||
"code": 0,
|
||
"msg": "验证码获取成功",
|
||
"data": CaptchaResponse{
|
||
CaptchaId: id,
|
||
PicPath: b64s,
|
||
CaptchaLength: s.config.KeyLong,
|
||
OpenCaptcha: oc,
|
||
},
|
||
})
|
||
}
|
||
|
||
// Verify 验证验证码
|
||
func (s *CaptchaService) Verify(id, answer string) bool {
|
||
return s.store.Verify(id, answer, true)
|
||
}
|