Topics

Golang IP Rate limiter with automated proxy ranges and bot reverse DNS check

package main

package main

import (
	"bufio"
	"github.com/astaxie/beego"
	"github.com/astaxie/beego/context"
	"github.com/astaxie/beego/logs"
	"github.com/ulule/limiter"
	"github.com/ulule/limiter/drivers/store/memory"
	"net"
	"net/http"
	"strconv"
	"strings"
	"sync"
)

type rateLimiter struct {
	generalLimiter *limiter.Limiter
	loginLimiter   *limiter.Limiter
	proxyBlocks    []*net.IPNet
	botIpCacheMu   sync.RWMutex
	botIpCache     map[string]bool
}

var (
	proxyUrls = []string{
		"https://www.cloudflare.com/ips-v4",
		"https://www.cloudflare.com/ips-v6",
	}
	proxyCidrs = []string{
		"127.0.0.1/32",
		"::1/128",
	}
)

func main() {
	r := &rateLimiter{
		botIpCache: make(map[string]bool),
	}
	r.botIpCache["127.0.0.1"] = true
	r.botIpCache["::1"] = true

	r.generalLimiter = newLimiter("2-S")
	r.loginLimiter = newLimiter("2-M")

	r.initProxyBlocks()

	//More on Beego filters here https://beego.me/docs/mvc/controller/filter.md
	beego.InsertFilter("/*", beego.BeforeRouter, func(c *context.Context) {
		rateLimit(r, c)
	}, true)

	//refer to https://beego.me/docs/mvc/controller/errors.md for error handling
	beego.ErrorHandler("429", func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusTooManyRequests)
		w.Write([]byte("Too Many Requests"))
		return
	})
	beego.Run()
}

func rateLimit(r *rateLimiter, ctx *context.Context) {
	var (
		limiterCtx   limiter.Context
		err          error
		req          = ctx.Request
		trustedProxy = false
	)

	ip := r.generalLimiter.GetIP(req)
	for _, block := range r.proxyBlocks {
		if block.Contains(ip) {
			trustedProxy = true
			ip = r.generalLimiter.GetIP(req)
			break
		}
	}
	//remove Forwarded for header if unknown proxy
	if !trustedProxy {
		req.Header.Del("X-Forwarded-For")
	}

	if strings.HasPrefix(ctx.Input.URL(), "/login") {
		limiterCtx, err = r.loginLimiter.Get(req.Context(), ip.String())
	} else {
		limiterCtx, err = r.generalLimiter.Get(req.Context(), ip.String())
	}
	if err != nil {
		ctx.Abort(http.StatusInternalServerError, err.Error())
		return
	}

	h := ctx.ResponseWriter.Header()
	h.Add("X-RateLimit-Limit", strconv.FormatInt(limiterCtx.Limit, 10))
	h.Add("X-RateLimit-Remaining", strconv.FormatInt(limiterCtx.Remaining, 10))
	h.Add("X-RateLimit-Reset", strconv.FormatInt(limiterCtx.Reset, 10))

	if limiterCtx.Reached {
		if r.isTrustedBot(ip.String()) {
			//beetak abasha...
			return
		}
		logs.Debug("Too Many Requests from %s on %s", ip, ctx.Input.URL())
		//refer to https://beego.me/docs/mvc/controller/errors.md for error handling
		ctx.Abort(http.StatusTooManyRequests, "429")
		return
	}

}

func (r *rateLimiter) initProxyBlocks() {
	var blocks []*net.IPNet
	for _, u := range proxyUrls {
		blocks = r.fetchProxyBlocks(u)
		if blocks != nil {
			r.proxyBlocks = append(r.proxyBlocks, blocks...)
		}
	}
	var block *net.IPNet
	for _, c := range proxyCidrs {
		if block = parseCidr(c); block != nil {
			r.proxyBlocks = append(r.proxyBlocks, block)
		}
	}

}

func (r *rateLimiter) fetchProxyBlocks(urlString string) []*net.IPNet {
	response, err := http.Get(urlString)
	if err != nil {
		logs.Error("http.Get => %v", err)
		return nil
	}
	defer response.Body.Close()

	var (
		line   string
		block  *net.IPNet
		blocks []*net.IPNet
	)
	scanner := bufio.NewScanner(response.Body)
	for scanner.Scan() {
		line = strings.TrimSpace(scanner.Text())
		if line != "" {
			if block = parseCidr(line); block != nil {
				blocks = append(blocks, block)
			}
		}
	}
	if scanner.Err() != nil {
		logs.Error("Error scanning proxy list: %s", scanner.Err())
		return nil
	}
	return blocks
}

func (r *rateLimiter) isTrustedBot(ip string) bool {
	r.botIpCacheMu.RLock()
	if _, exists := r.botIpCache[ip]; exists {
		r.botIpCacheMu.RUnlock()
		return true
	}
	r.botIpCacheMu.RUnlock()
	addr, err := net.LookupAddr(ip)
	if err != nil {
		logs.Error("Error getting host for ip %s Error: %v", ip, err)
		return false
	}
	for _, h := range addr {
		if strings.HasSuffix(h, "google.com.") ||
			strings.HasSuffix(h, "googlebot.com.") ||
			strings.HasSuffix(h, "msn.com.") ||
			strings.HasSuffix(h, "yandex.com.") ||
			strings.HasSuffix(h, "yandex.net.") ||
			strings.HasSuffix(h, "baidu.com.") ||
			strings.HasSuffix(h, "yahoo.com.") {
			ips, err := net.LookupHost(h)
			if err != nil {
				logs.Error("Error getting host for ip %s Error: %v", ip, err)
				return false
			}
			for _, _ip := range ips {
				if ip == _ip {
					r.botIpCacheMu.Lock()
					r.botIpCache[ip] = true
					r.botIpCacheMu.Unlock()
					return true
				}
			}
		}
	}
	return false
}

func parseCidr(s string) *net.IPNet {
	_, block, err := net.ParseCIDR(s)
	if err != nil {
		logs.Error("Error parsing proxy block: %v", err)
		return nil
	}
	return block
}

func PanicOnError(e error) {
	if e != nil {
		panic(e)
	}
}

func newLimiter(rate string) *limiter.Limiter {
	r, err := limiter.NewRateFromFormatted(rate)
	PanicOnError(err)
	return limiter.New(memory.NewStore(), r, limiter.WithTrustForwardHeader(true))
}

Fork it

By continuing to use the site, you agree to the use of cookies. more information

The cookie settings on this website are set to "allow cookies" to give you the best browsing experience possible. If you continue to use this website without changing your cookie settings or you click "Accept" below then you are consenting to this.

Close