53 lines
875 B
Go
53 lines
875 B
Go
package middleware
|
|
|
|
import (
|
|
"sync"
|
|
"time"
|
|
"yintai-api/base/config"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
func LoginLimit() gin.HandlerFunc {
|
|
const maxPerSecond = 2
|
|
var (
|
|
mu sync.Mutex
|
|
hits = map[string][]time.Time{}
|
|
)
|
|
return func(c *gin.Context) {
|
|
ip := c.ClientIP()
|
|
now := time.Now()
|
|
cutoff := now.Add(-time.Second)
|
|
|
|
mu.Lock()
|
|
recent := make([]time.Time, 0, maxPerSecond)
|
|
for _, t := range hits[ip] {
|
|
if t.After(cutoff) {
|
|
recent = append(recent, t)
|
|
}
|
|
}
|
|
if len(recent) >= maxPerSecond {
|
|
hits[ip] = recent
|
|
mu.Unlock()
|
|
panic(config.ErrLimit.New())
|
|
}
|
|
hits[ip] = append(recent, now)
|
|
if len(hits) > 1024 {
|
|
for k, v := range hits {
|
|
alive := false
|
|
for _, t := range v {
|
|
if t.After(cutoff) {
|
|
alive = true
|
|
break
|
|
}
|
|
}
|
|
if !alive {
|
|
delete(hits, k)
|
|
}
|
|
}
|
|
}
|
|
mu.Unlock()
|
|
c.Next()
|
|
}
|
|
}
|