Rate limit specific endpoints-CodePudding

I am new to GoLang and working on my first API. I have two endpoints, and I want to rate limit only one of them. I found a helpful tutorial to get me started, and I've based my approach off of the tutorial, recognizing that this approach will rate limit both of my endpoints:

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)))
}

I researched the difference between http.Handle and http.HandleFunc and naively believed that I could substitute http.HandleFunc for http.Handle. This approach is completely flawed as the logic contained in the HandlerFunc never executes:

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)))

Can anyone explain why this does not work, and how I could approach this problem to only rate limit a specific endpoint?

CodePudding user response：

The distinction between using a plain http.Handler and a http.HanlderFunc doesn't really matter here. http.HandleFunc is just a way to convert a regular function into a http.Handler - it essentially does the same thing as your original version of limit.

Your implementations of limit both look fine; probably the second is better because it's simpler. Instead, the issue is in main. When you call http.ListenAndServeTLS and provide a value for the final argument, it requests that only the handler you pass in as that final argument be used as the root request handler. Any calls to http.Handle() or http.HandleFunc() are ignored unless you pass in nil as this final argument.

What you want to do instead is apply limit to the specific handler you want to limit. You have two options for this. First, you can use a ServeMux like in your first code snippet:

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))
}

Alternatively, you can do something more like your second code snippet, but upass in nil for the final argument to http.ListenAndServeTLS so that the default http.ServeMux is used, meaning that the calls to http.HandleFunc() will be respected:

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))
}

For a simple application, the first approach is probably fine. For anything more complex, I'd recommend the later approach because it will work if you open multiple servers or do other more complex things.