go - 速率限制特定端点

Question

我是 GoLang 新手，正在开发我的第一个 API。我有两个端点，我只想对其中之一进行速率限制。我找到了helpful tutorial为了让我开始，我的方法基于教程，认识到这种方法会限制我的两个端点的速率:

var limiter = rate.NewLimiter(rate.Every((1*time.Hour)/3), 1)

func limit(next http.Handler) http.Handler {
    return http.HandlerFunc(func(res http.ResponseWriter, req *http.Request) {
        if limiter.Allow() == false {
            http.Error(res, http.StatusText(429), http.StatusTooManyRequests)
            return
        }
        next.ServeHTTP(res, req)
    })
}

func main() {
    mux := http.NewServeMux()
    mux.HandleFunc("/", createNewToken)
    mux.HandleFunc("/notify", sendPushNotificationToAllTokens)

    log.Fatal(http.ListenAndServeTLS(":5050", "localhost.crt", "localhost.key", limit(mux)))
}

我研究了 http.Handle and http.HandleFunc 之间的区别并且天真地相信我可以用 http.HandleFunc 替换 http.Handle。这种方法是完全有缺陷的，因为 HandlerFunc 中包含的逻辑永远不会执行:

var limiter = rate.NewLimiter(rate.Every(1*time.Hour/3), 1)

func limit(next http.HandlerFunc) http.HandlerFunc {
    return func(res http.ResponseWriter, req *http.Request) {
        if limiter.Allow() == false {
            http.Error(res, http.StatusText(429), http.StatusTooManyRequests)
            return
        }
        next.ServeHTTP(res, req)
    }
}

func main() {
    //mux := http.NewServeMux()
    http.HandleFunc("/", createNewToken)
    http.HandleFunc("/notify", sendPushNotificationToAllTokens)

    // attempt to only rate limit the /notify endpoint 
    log.Fatal(http.ListenAndServeTLS(":5050", "localhost.crt", "localhost.key", limit(sendPushNotificationToAllTokens)))

任何人都可以解释为什么这不起作用，以及我如何解决这个问题以仅对特定端点进行速率限制？

Answer 1

使用普通的http.Handler和http.HanlderFunc之间的区别在这里并不重要。 http.HandleFunc 只是将常规函数转换为 http.Handler 的一种方法 - 它本质上与原始版本的 limit 执行相同的操作>.

您的 limit 实现看起来都很好；可能第二个更好，因为它更简单。相反，问题出在 main 中。当您调用 http.ListenAndServeTLS 并为最终参数提供值时，它会请求仅将您作为最终参数传入的处理程序用作根请求处理程序。除非您传入 nil 作为最后一个参数，否则对 http.Handle() 或 http.HandleFunc() 的任何调用都会被忽略。

您想要做的是将 limit 应用于您想要限制的特定处理程序。为此，您有两种选择。首先，您可以像第一个代码片段中那样使用 ServeMux:

func main() {
    mux := http.NewServeMux()
    mux.HandleFunc("/", createNewToken)
    // Limit only the handler for "/notify".
    mux.HandleFunc("/notify", limit(sendPushNotificationToAllTokens))

    // Don't limit the whole mux.
    log.Fatal(http.ListenAndServeTLS(":5050", "localhost.crt", "localhost.key", mux))
}

或者，您可以执行类似于第二个代码片段的操作，但将 nil 作为最后一个参数传递给 http.ListenAndServeTLS，以便默认的 http使用.ServeMux，这意味着对http.HandleFunc()的调用将受到尊重:

func main() {
    http.HandleFunc("/", createNewToken)
    // Limit only the handler for "/notify".
    http.HandleFunc("/notify", limit(sendPushNotificationToAllTokens))

    // Pass in nil here so that http.DefaultServeMux is used.
    log.Fatal(http.ListenAndServeTLS(":5050", "localhost.crt", "localhost.key", nil))
}

对于简单的应用程序，第一种方法可能就很好。对于更复杂的事情，我建议使用后面的方法，因为如果您打开多个服务器或执行其他更复杂的事情，它会起作用。