Guest User

Untitled

a guest
Nov 14th, 2023
96
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Go 1.53 KB | None | 0 0
  1. func (app *application) rateLimit(next http.Handler) http.Handler {
  2.     type client struct {
  3.         limiter  *rate.Limiter
  4.         lastseen time.Time
  5.     }
  6.  
  7.     var (
  8.         mu      sync.Mutex
  9.         clients = make(map[string]*client)
  10.         rps     = app.config.limiter.rps
  11.         burst   = app.config.limiter.burst
  12.     )
  13.  
  14.     const (
  15.         cleanUpTime   = 3 * time.Minute
  16.         cleanUpTimeTG = 30 * time.Minute
  17.         sleepDuration = 1 * time.Minute
  18.     )
  19.     go func() {
  20.         for {
  21.             time.Sleep(sleepDuration)
  22.             mu.Lock()
  23.             for ip, client := range clients {
  24.                 // default rate limiter
  25.                 if client.limiter.Limit() == rate.Limit(rps) && time.Since(client.lastseen) > cleanUpTime {
  26.                     delete(clients, ip)
  27.                 } else if time.Since(client.lastseen) > cleanUpTimeTG {
  28.                     delete(clients, ip)
  29.                 }
  30.             }
  31.             mu.Unlock()
  32.         }
  33.     }()
  34.  
  35.     return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  36.         if app.config.limiter.enabled {
  37.             ip := realip.FromRequest(r)
  38.  
  39.             if r.URL.Path == "/send-form" && r.Method == http.MethodPost {
  40.                 rps = app.config.limiter.rpsTG
  41.                 burst = app.config.limiter.burstTG
  42.                 // Modifying the client's IP to ensure a separate rate limiter for form submissions
  43.                 ip = "tg_client" + ip
  44.  
  45.             }
  46.             mu.Lock()
  47.             if _, found := clients[ip]; !found {
  48.                 clients[ip] = &client{
  49.                     limiter: rate.NewLimiter(rate.Limit(rps), burst),
  50.                 }
  51.             }
  52.             clients[ip].lastseen = time.Now()
  53.             if !clients[ip].limiter.Allow() {
  54.                 mu.Unlock()
  55.                 app.rateLimitExceededResponse(w, r)
  56.                 return
  57.             }
  58.             mu.Unlock()
  59.         }
  60.         next.ServeHTTP(w, r)
  61.     })
  62.  
  63. }
Advertisement
Add Comment
Please, Sign In to add comment