Golang 基于IP地址的HTTP限速请求

Golang 基于IP地址的HTTP限速请求

使用 golang.org/x/time

在本教程中，我们将基于用户的IP地址创建一个简单的中间件来进行速率限制。

纯HTTP服务器
让我们从构建一个简单的HTTP服务器开始，该服务器具有非常简单的终结点。这可能是一个沉重的端点，这就是为什么我们要在那里添加速率限制。

package main

import (
    "log"
    "net/http"
)

func main() {
    mux := http.NewServeMux()
    mux.HandleFunc("/", okHandler)

    if err := http.ListenAndServe(":8888", mux); err != nil {
        log.Fatalf("unable to start server: %s", err.Error())
    }
}

func okHandler(w http.ResponseWriter, r *http.Request) {
    // Some very expensive database call
    w.Write([]byte("alles gut"))
}

在main.go我们启动服务器时`:8888｀。

我们将使用golang.org/x/time/rate，它提供令牌桶速率限制器算法。rate＃Limiter控制事件发生的频率。它实现了一个“令牌桶”，该令牌桶的大小b最初为满，然后以每秒r令牌数的速率重新填充。非正式地，在足够大的时间间隔内，限制器将速率限制为每秒r个令牌，最大突发大小为b个事件。

由于我们要对每个IP地址实施速率限制器，因此我们还需要维护一个限制器图。

package main

import (
    "sync"

    "golang.org/x/time/rate"
)

// IPRateLimiter .
type IPRateLimiter struct {
    ips map[string]*rate.Limiter
    mu  *sync.RWMutex
    r   rate.Limit
    b   int
}

// NewIPRateLimiter .
func NewIPRateLimiter(r rate.Limit, b int) *IPRateLimiter {
    i := &IPRateLimiter{
        ips: make(map[string]*rate.Limiter),
        mu:  &sync.RWMutex{},
        r:   r,
        b:   b,
    }

    return i
}

// AddIP creates a new rate limiter and adds it to the ips map,
// using the IP address as the key
func (i *IPRateLimiter) AddIP(ip string) *rate.Limiter {
    i.mu.Lock()
    defer i.mu.Unlock()

    limiter := rate.NewLimiter(i.r, i.b)

    i.ips[ip] = limiter

    return limiter
}

// GetLimiter returns the rate limiter for the provided IP address if it exists.
// Otherwise calls AddIP to add IP address to the map
func (i *IPRateLimiter) GetLimiter(ip string) *rate.Limiter {
    i.mu.Lock()
    limiter, exists := i.ips[ip]

    if !exists {
        i.mu.Unlock()
        return i.AddIP(ip)
    }

    i.mu.Unlock()

    return limiter
}

NewIPRateLimiter创建一个IP限制器实例，HTTP服务器将不得不调用GetLimiter该实例以获取指定IP的限制器（从映射或生成一个新的）。

中间件

让我们升级我们的HTTP Server并将中间件添加到所有端点，因此，如果IP达到限制，它将响应429 Too Many Requests，否则，它将继续该请求。

在该limitMiddleware函数中，Allow()每次收到HTTP请求时，我们都调用全局限制器的方法。如果存储桶中没有剩余令牌，Allow()则返回false，我们将向用户发送429 Too Many Requests响应。否则，调用Allow()将只消耗存储桶中的一个令牌，然后我们将控制权传递给链中的下一个处理程序。

package main

import (
    "log"
    "net/http"
)

var limiter = NewIPRateLimiter(1, 5)

func main() {
    mux := http.NewServeMux()
    mux.HandleFunc("/", okHandler)

    if err := http.ListenAndServe(":8888", limitMiddleware(mux)); err != nil {
        log.Fatalf("unable to start server: %s", err.Error())
    }
}

func limitMiddleware(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        limiter := limiter.GetLimiter(r.RemoteAddr)
        if !limiter.Allow() {
            http.Error(w, http.StatusText(http.StatusTooManyRequests), http.StatusTooManyRequests)
            return
        }

        next.ServeHTTP(w, r)
    })
}

func okHandler(w http.ResponseWriter, r *http.Request) {
    // Some very expensive database call
    w.Write([]byte("alles gut"))
}

生成并运行

go get golang.org/x/time/rate
go build -o server .
./server

测试

ab -n 1000 -c 100 http://localhost:8888/

其中－n表示请求数，－c表示并发数

其他

• uber-go漏桶限流器使用与原理分析