didip / tollbooth

Simple middleware to rate-limit HTTP requests.

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Different behavior on limiting based on header key for v4.0.2 and v7

Xinyu-bot opened this issue · comments

Hi there, let me get to the problem directly.

Code I wrote is like:

var limitHeaders = map[string][]string{ "user_id": {} }

rl := tollbooth.NewLimiter(1, nil)
rl.SetIPLookups([]string{"X-Forwarded-For", "RemoteAddr", "X-Real-IP"})
rl.SetHeaders(limitHeaders) // to limit the requests that contain the key "user_id" in header, regardless of its value 

r := gin.New()
r.POST("/user", LimitHandlerFunc(rl), SomeBusinessHandler)
...

while LimitHandlerFunc is defined as below. The code was taken from https://github.com/didip/tollbooth_gin/blob/master/tollbooth_gin.go, and I only changed the package imported to v7.

func LimitHandlerFunc(lmt *limiter.Limiter) gin.HandlerFunc {
	return func(c *gin.Context) {
		httpError := tollbooth.LimitByRequest(lmt, c.Writer, c.Request)
		if httpError != nil {
			c.Data(httpError.StatusCode, lmt.GetMessageContentType(), []byte(httpError.Message))
			c.Abort()
		} else {
			c.Next()
		}
	}
}

The original idea is to limit the request sent by every single user, so the expected behavior here is that as long as a request contains user_id key in its header, regardless of the value of it, the request will be counted for rate limiting. And that is what v4.0.2+incompatible behaves (downloaded by command go install github.com/didip/tollbooth@latest).

However, when I switched to github.com/didip/tollbooth/v7 and github.com/didip/tollbooth/v7/limiter, the rate limiting for "user_id" would not work at all. Unless I specify the headers map as below,

var limitHeaders = map[string][]string{ "user_id": {"user1", "user2"} }

and then the rate limiting will only work for the requests that contains user_id: user1 or user_id: user2 in their headers.

I compared the code of function tollbooth.LimitByRequest between v4.0.2 and v7, and I believe these following codes lead to the different behaviors, as they do not present in v4.0.2 but in v7.

		// If request contains the header key but not the values,
		// skip limiter
		requestHeadersDefinedInLimiter = false

		for headerKey, headerValues := range lmtHeaders {
			for _, headerValue := range headerValues {
				if r.Header.Get(headerKey) == headerValue {
					requestHeadersDefinedInLimiter = true
					break
				}
			}
		}

		if !requestHeadersDefinedInLimiter {
			return true
		}

https://github.com/didip/tollbooth/blob/master/tollbooth.go#L110~L126

Is this the desired behavior, that v7 is supposed to only check headers with pre-defined and definite numbers of values?
I would guess it is not, since I found these following lines actually handle a situation where headerValues is empty slice

			if len(headerValues) == 0 {
				// If header values are empty, rate-limit all request containing headerKey.
				headerValuesToLimit = append(headerValuesToLimit, []string{headerKey, reqHeaderValue})

https://github.com/didip/tollbooth/blob/master/tollbooth.go#L225~L227

If it is not desired behavior, I would propose a little check on length of headerValues so that the behaviors of v4.0.2 and of v7 can be the same:

		for headerKey, headerValues := range lmtHeaders {
			if len(headerValues) == 0 { 
				requestHeadersDefinedInLimiter = true
				continue 
			}  
			for _, headerValue := range headerValues {
				if r.Header.Get(headerKey) == headerValue {
					requestHeadersDefinedInLimiter = true
					break
				}
			}
		}

I could make a PR real quick (so that I can be a contributor!), but I think it would be better if maintainers can tell me if this is actually desired behavior or not 😅

Anyway, PR is here #104 😅 If I made any mistake, or the modification is unnecessary, please let me know.

Should be fixed by #104