Reputation: 709
I have developed a HTTP server using Go. Now I would like to implement a rate limiter, so that I can check if a HTTP request from a particular IP is sending more than 10 HTTP requests in 1 min, I can put that IP in a block list for some time say (1 hr) and in the meantime is the same subscriber is sending requests while he/she was in blocking period I would be sending 429 error response from HTTP server.
I have written a code for this, but in this am able to block the IP addresses, but in that its unblocking all the IP's after 1 hr of time. I am expecting first come first unblock.
Package main
import (
"log"
"net/http"
"strings"
"time"
)
func main() {
fs := http.FileServer(http.Dir("./html/"))
http.Handle("/", fs)
log.Println("Listening..")
go clearLastRequestsIPs()
go clearBlockedIPs()
err := http.ListenAndServe(":8080", middleware(nil))
if err != nil {
log.Fatalln(err)
}
}
// Stores last requests IPs
var lastRequestsIPs []string
// Block IP for 1 hours
var blockedIPs []string
func middleware(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
ipAddr := strings.Split(r.RemoteAddr, ":")[0]
if existsBlockedIP(ipAddr) {
http.Error(w, "", http.StatusTooManyRequests)
return
}
// how many requests the current IP made in last 5 mins
requestCounter := 0
for _, ip := range lastRequestsIPs {
if ip == ipAddr {
requestCounter++
}
}
if requestCounter >= 1000 {
blockedIPs = append(blockedIPs, ipAddr)
http.Error(w, "", http.StatusTooManyRequests)
return
}
lastRequestsIPs = append(lastRequestsIPs, ipAddr)
if next == nil {
http.DefaultServeMux.ServeHTTP(w, r)
return
}
next.ServeHTTP(w, r)
})
}
func existsBlockedIP(ipAddr string) bool {
for _, ip := range blockedIPs {
if ip == ipAddr {
return true
}
}
return false
}
func existsLastRequest(ipAddr string) bool {
for _, ip := range lastRequestsIPs {
if ip == ipAddr {
return true
}
}
return false
}
// Clears lastRequestsIPs array every 1 hrs
func clearLastRequestsIPs() {
for {
lastRequestsIPs = []string{}
time.Sleep(time.Hour * 1)
}
}
// Clears blockedIPs array every 1 hours
func clearBlockedIPs() {
for {
blockedIPs = []string{}
time.Sleep(time.Hour * 1)
}
}
Upvotes: 0
Views: 1374
Reputation: 64657
You could use some middleware like this:
type Limiter struct {
ipCount map[string]int
sync.Mutex
}
var limiter Limiter
func init() {
limiter.ipCount = make(map[string]int)
}
func limit(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
// Get the IP address for the current user.
ip, _, err := net.SplitHostPort(r.RemoteAddr)
if err != nil {
log.Println(err.Error())
http.Error(w, "Internal Server Error", http.StatusInternalServerError)
return
}
// Get the # of times the visitor has visited in the last 60 seconds
limiter.Lock()
count, ok := limiter.ipCount[ip]
if !ok {
limiter.ipCount[ip] = 0
}
if count > 10 {
limiter.Unlock()
http.Error(w, http.StatusText(429), http.StatusTooManyRequests)
return
} else {
limiter.ipCount[ip]++
}
time.AfterFunc(time.Second * 60, func() {
limiter.Lock()
limiter.ipCount[ip]--
limiter.Unlock()
})
if limiter.ipCount[ip] == 10 {
// set it to 20 so the decrement timers will only decrease it to
// 10, and they stay blocked until the next timer resets it to 0
limiter.ipCount[ip] = 20
time.AfterFunc(time.Hour, func() {
limiter.Lock()
limiter.ipCount[ip] = 0
limiter.Unlock()
})
}
limiter.Unlock()
next.ServeHTTP(w, r)
})
}
Upvotes: 2