golang http 服务基于 ip 的访问频率限制

编程 评论

转自learnku

简介

  • 如果你运行 HTTP 服务，并且希望限制 HTTP 的访问频率，那么你可以借助一些比较稳定的工具，例如： tollbooth。不过如果你构建的应用比较简单，也可以自己来实现。

  • 我们可以使用一个现有的 Go 包 x/time/rate。

  • 本课程，我们将创建一个简单的中间件实现基于 IP 限制 HTTP 访问频率。

简单的 HTTP 服务

  • 让我们从创建一个简单的 HTTP 服务开始，它有非常简单的终端。 但是，因为它的访问频率可能非常高，因此我们要为它添加频率限制。

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    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) {
        // 某些消耗很高的数据库请求
        w.Write([]byte("alles gut"))
    }

  • 通过 main.go 我们启动服务，监听 :8888 端口，这样我们就有了一个简单的终端 /。

golang.org/x/time/rate

  • 我们将使用名为 x/time/rate 的 Go 包，它提供了一个令牌桶速率限制器算法。rate#Limiter 控制允许事件发生的频率。它实现了一个大小为 b 的「令牌桶」，最初是满的，并以每秒 r 的速度重新填充令牌。通俗地讲，就是在任何足够大的时间间隔内，限制器将速率限制为每秒 r 个令牌，最大突发大小为 b 个事件。

  • 由于我们希望实现每个 IP 地址的速率限制器，我们还需要维护一个限制器映射。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
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 创建了一个新的速率限制器，并将其添加到 ips 映射中,
// 使用 IP地址作为密钥
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 返回所提供的IP地址的速率限制器(如果存在的话).
// 否则调用 AddIP 将 IP 地址添加到映射中
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 服务并将中间件添加到所有端点，如果 IP 达到限制，它将响应 429 Too Many Requests，否则，它将继续该请求。

  • 每一个经过中间件的请求，我们都会调用 limitMiddleware 函数中的全局方法 Allow()。如果存储桶中没有令牌了，该方法会返回 false，该请求会收到 429 Too Many Requests 的响应。否则 Allow() 方法将消耗一个令牌，并将请求传递给下一个程序。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
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) {
    // 非常重要的数据请求（译者注：这句话没理解到位）
    w.Write([]byte("alles gut"))
}

编译 & 执行

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

测试

  • 这是我喜欢使用的一个非常好的来进行 HTTP 负载测试的工具，它叫做 vegeta (它也是用 Go 编写的)。

    1
    brew install vegeta

  • 我们需要创建一个简单的配置文件，来展示我们希望生成的请求。

    1
    GET http://localhost:8888/

  • 然后运行攻击 10 秒，每个时间单位 100 个请求。

    1
    vegeta attack -duration=10s -rate=100 -targets=vegeta.conf | vegeta report

  • 结果，您将看到一些请求返回了 200，但是大多数都返回了 429。

QinHanGame Developer & Designer

当你发现自己的才华撑不起野心时，就请安静下来学习吧！