Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- func (app *application) rateLimit(next http.Handler) http.Handler {
- type client struct {
- limiter *rate.Limiter
- lastseen time.Time
- }
- var (
- mu sync.Mutex
- clients = make(map[string]*client)
- rps = app.config.limiter.rps
- burst = app.config.limiter.burst
- )
- const (
- cleanUpTime = 3 * time.Minute
- cleanUpTimeTG = 30 * time.Minute
- sleepDuration = 1 * time.Minute
- )
- go func() {
- for {
- time.Sleep(sleepDuration)
- mu.Lock()
- for ip, client := range clients {
- // default rate limiter
- if client.limiter.Limit() == rate.Limit(rps) && time.Since(client.lastseen) > cleanUpTime {
- delete(clients, ip)
- } else if time.Since(client.lastseen) > cleanUpTimeTG {
- delete(clients, ip)
- }
- }
- mu.Unlock()
- }
- }()
- return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
- if app.config.limiter.enabled {
- ip := realip.FromRequest(r)
- if r.URL.Path == "/send-form" && r.Method == http.MethodPost {
- rps = app.config.limiter.rpsTG
- burst = app.config.limiter.burstTG
- // Modifying the client's IP to ensure a separate rate limiter for form submissions
- ip = "tg_client" + ip
- }
- mu.Lock()
- if _, found := clients[ip]; !found {
- clients[ip] = &client{
- limiter: rate.NewLimiter(rate.Limit(rps), burst),
- }
- }
- clients[ip].lastseen = time.Now()
- if !clients[ip].limiter.Allow() {
- mu.Unlock()
- app.rateLimitExceededResponse(w, r)
- return
- }
- mu.Unlock()
- }
- next.ServeHTTP(w, r)
- })
- }
Advertisement
Add Comment
Please, Sign In to add comment