基于 IP 限制 HTTP 访问频率的 Go 实现

2024-08-13 09:01:25 浏览数 (1)

在构建高并发的 HTTP 服务时，限制某个 IP 的访问频率是一个常见的需求。无论是为了防止恶意攻击，还是为了节约服务器资源，这种限制都能有效地保护服务的可用性。本文将详细介绍如何在 Go 语言中实现基于 IP 的 HTTP 访问频率限制。

1. 背景与意义

当我们部署一个公开的 API 服务时，常常会遇到一些恶意用户或爬虫，它们会对服务器发起大量请求。如果不加限制，服务器可能会被过多的请求拖垮，从而影响正常用户的访问体验。因此，为每个 IP 地址设置访问频率限制（即速率限制）是必要的。

速率限制可以防止以下几种情况：

  • 拒绝服务攻击（DoS）: 恶意用户通过高频率的请求导致服务器资源耗尽，从而无法响应正常用户的请求。
  • 滥用资源: 某些用户可能滥用 API，频繁调用接口，消耗大量资源。
  • 爬虫的过度抓取: 不受限制的爬虫可能会在短时间内抓取大量数据，影响服务器性能。

通过在服务端实现基于 IP 的访问频率限制，可以有效避免这些问题。

2. Go 中的速率限制概述

在 Go 中，速率限制可以通过多种方式实现，其中最常见的方法是使用令牌桶（Token Bucket）算法。令牌桶算法是一种经典的速率限制算法，它通过向桶中添加令牌来限制操作的频率。

每个请求到来时，服务器会检查桶中是否有可用的令牌。如果有可用的令牌，则允许请求通过，并从桶中移除一个令牌；如果没有令牌，则拒绝请求。令牌会以固定的速率不断加入桶中，确保请求频率不会超过预定的阈值。

Go 提供了丰富的标准库和第三方库，可以帮助我们实现速率限制。

3. 使用 golang.org/x/time/rate 实现 IP 限制

golang.org/x/time/rate 是 Go 提供的一个用于速率限制的包，它基于令牌桶算法实现。首先，我们可以为每个 IP 创建一个 rate.Limiter，并根据请求频率限制配置速率。

3.1 安装依赖

首先，确保你已经安装了 golang.org/x/time/rate 包。如果没有安装，可以通过以下命令安装：

代码语言：javascript复制
go get golang.org/x/time/rate
3.2 基本的限速实现

以下是一个简单的例子，展示如何使用 rate.Limiter 来限制 IP 地址的访问频率：

代码语言：javascript复制
package main

import (
    "net/http"
    "sync"
    "time"
    "golang.org/x/time/rate"
)

var visitors = make(map[string]*rate.Limiter)
var mu sync.Mutex

// 每秒允许5次请求，最多存储10个令牌
func getLimiter(ip string) *rate.Limiter {
    mu.Lock()
    defer mu.Unlock()

    limiter, exists := visitors[ip]
    if !exists {
        limiter = rate.NewLimiter(5, 10)
        visitors[ip] = limiter
    }

    return limiter
}

func limitMiddleware(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        ip := r.RemoteAddr

        limiter := getLimiter(ip)
        if !limiter.Allow() {
            http.Error(w, "Too Many Requests", http.StatusTooManyRequests)
            return
        }

        next.ServeHTTP(w, r)
    })
}

func main() {
    mux := http.NewServeMux()
    mux.Handle("/", limitMiddleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        w.Write([]byte("Hello, World!"))
    })))

    http.ListenAndServe(":8080", mux)
}

在这个例子中，每个 IP 地址都被限制为每秒最多发起 5 个请求。如果请求超出限制，服务器将返回 429 状态码。

3.3 清理过期的限制器

在上面的代码中，我们为每个 IP 地址都创建了一个 rate.Limiter，并将其保存在 visitors 映射中。随着时间的推移，映射中会积累大量的过期 IP 地址，导致内存占用增加。因此，我们需要定期清理这些不再使用的限制器。

我们可以通过以下方式实现定期清理：

代码语言：javascript复制
func cleanupVisitors() {
    for {
        time.Sleep(time.Minute)
        mu.Lock()
        for ip, limiter := range visitors {
            if limiter.Allow() { // 如果1分钟内没有请求到达
                delete(visitors, ip)
            }
        }
        mu.Unlock()
    }
}

在 main 函数中，我们可以启动一个 goroutine 来定期调用 cleanupVisitors 函数：

代码语言：javascript复制
func main() {
    go cleanupVisitors()

    mux := http.NewServeMux()
    mux.Handle("/", limitMiddleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        w.Write([]byte("Hello, World!"))
    })))

    http.ListenAndServe(":8080", mux)
}

通过这种方式，我们可以有效地管理内存，避免过期 IP 地址的积累。

4. 增强的限制策略

在实际应用中，速率限制的策略可能会更为复杂。例如，我们可能希望根据不同的路径、用户角色或时间段来调整限制。以下是一些常见的增强策略。

4.1 基于路径的限制

对于不同的 API 端点，我们可能希望设置不同的速率限制。例如，/login 路径的请求可能比普通的 GET 请求更为敏感，因此我们可能需要对其施加更严格的限制。

我们可以根据请求的 URL 路径来选择不同的 rate.Limiter：

代码语言：javascript复制
func getLimiter(ip, path string) *rate.Limiter {
    mu.Lock()
    defer mu.Unlock()

    key := ip   path
    limiter, exists := visitors[key]
    if !exists {
        var r rate.Limit
        if path == "/login" {
            r = 1 // 每秒1个请求
        } else {
            r = 5 // 每秒5个请求
        }
        limiter = rate.NewLimiter(r, 10)
        visitors[key] = limiter
    }

    return limiter
}
4.2 基于用户角色的限制

在某些应用中，不同用户角色可能拥有不同的访问频率。例如，管理员可能有更高的请求速率，而普通用户则受到更严格的限制。我们可以通过识别用户角色来应用不同的限制策略：

代码语言：javascript复制
func getLimiter(ip, role string) *rate.Limiter {
    mu.Lock()
    defer mu.Unlock()

    key := ip   role
    limiter, exists := visitors[key]
    if !exists {
        var r rate.Limit
        if role == "admin" {
            r = 10 // 每秒10个请求
        } else {
            r = 3 // 每秒3个请求
        }
        limiter = rate.NewLimiter(r, 10)
        visitors[key] = limiter
    }

    return limiter
}

在应用中，你可以根据用户认证信息（如 JWT Token）来识别用户角色，并使用相应的限制策略。

4.3 使用 Redis 分布式存储限制器状态

在分布式系统中，多个实例可能会同时处理请求，因此每个实例都需要共享限制器状态。此时，我们可以使用 Redis 来存储和管理 rate.Limiter 的状态。

通过 Redis，我们可以确保所有实例共享同一套速率限制数据，从而实现全局一致的限制策略。

首先，安装 Redis 客户端库：

代码语言：javascript复制
go get github.com/go-redis/redis/v8

然后，在代码中使用 Redis 存储限制器状态：

代码语言：javascript复制
import (
    "context"
    "github.com/go-redis/redis/v8"
    "time"
)

var rdb = redis.NewClient(&redis.Options{
    Addr: "localhost:6379",
    Password: "",
    DB: 0,
})

func getLimiter(ip string) *rate.Limiter {
    key := "limiter:"   ip
    result, err := rdb.Get(context.Background(), key).Result()

    var limiter *rate.Limiter
    if err == redis.Nil || result == "" {
        limiter = rate.NewLimiter(5, 10)
        rdb.Set(context.Background(), key, limiter, time.Minute)
    } else {
        // 反序列化 limiter
        // limiter = deserialize(result)
    }

    return limiter
}

Redis 提供了持久化和自动过期功能，这使得它非常适合用来管理限流器状态。

腾讯技术创作特训营S8

0 人点赞