Ratelimit by header value creates keys using cached values

[caalberts]

Hello! First of all, thanks for the great library. It worked great, until I had to add rate limiting by user id in the request header.

I came across #43, however I found that the function tollbooth.BuildKeys build keys using existing header values in the cache. Specifically this line:

tollbooth/tollbooth.go

Line 89 in b10a036

for _, headerValue := range headerValues {

Should the key be created this way instead of looping through existing values?

sliceKeys = append(sliceKeys, []string{remoteIP, path, r.Method, headerKey, r.Header.Get(headerKey)})

As I understand it, this sliceKeys is then used to lookup against the rate limit cache. If sliceKeys contain existing headers, the current request would get rate limited due to one of the existing headers.

If I have used the package wrongly, please help me to understand how I should use it.

Here are the middleware I wrote and the test.

func RateLimit(h http.HandlerFunc) http.HandlerFunc {
	allocationLimiter := tollbooth.NewLimiter(1, &limiter.ExpirableOptions{DefaultExpirationTTL: time.Minute}).
		SetMethods([]string{"POST"})

	handler := func(w http.ResponseWriter, r *http.Request) {
		customerID := r.Header.Get("X_Owner_ID")
		fmt.Printf("customerID: %s\n", customerID)
		allocationLimiter.SetHeader("X_Owner_ID", []string{customerID})

		tollbooth.LimitFuncHandler(allocationLimiter, h).ServeHTTP(w, r)
	}

	return handler
}

func TestRateLimit(t *testing.T) {
	customerID1 := "1234"
	customerID2 := "5678"

	tests := []struct {
		name                string
		secondRequestStatus int
		customerIDs         []string
	}{
		{"different customer id", http.StatusOK, []string{customerID1, customerID2}},
		{"same customer id", http.StatusTooManyRequests, []string{customerID1, customerID1}},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			h := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
			testServer := httptest.NewServer(RateLimit(h))
			defer testServer.Close()
			client := &http.Client{}

			request1, err := http.NewRequest("POST", testServer.URL, nil)
			assert.Nil(t, err)
			request1.Header.Add("CustomerID", tt.customerIDs[0])

			response, err := client.Do(request1)
			assert.Nil(t, err)
			assert.Equal(t, http.StatusOK, response.StatusCode)

			request2, err := http.NewRequest("POST", testServer.URL, nil)
			assert.Nil(t, err)
			request2.Header.Add("CustomerID", tt.customerIDs[1])

			response, err = client.Do(request2)
			assert.Nil(t, err)
			assert.Equal(t, tt.secondRequestStatus, response.StatusCode)
		})
	}
}

[rvdwijngaard]

I was actually thinking the same; @didip can you please give your comments on this issue?

[jack-chung]

Agreed the headerValues part of the code looks not right. It's not checking the configured values against the request value, and also not using the request value to build the key.

[didip]

Man, I keep not getting notified over email.

This does look like a bug. I will take a look at it closer, it has been a while.

[didip]

This is a legitimate bug, see fix SHA above.

[jack-chung]

@didip For the case of a configured header name with no value, the new code puts the header name into the key but no value. I thought it's supposed to include the header value in the request as part of the key? Or did I misunderstand the intended use case?

Example:
If I want to throttle on a header called X-USER-ID, so that each user is bound by a request limit, this is not going to work with the new code (nor the old code). Only the header name is in the key, so different users coming from the same IP (e.g. behind corporate firewall) would all be subject to the same aggregated limit, not individual limits.

[didip]

@jack-chung If you only define request header (without values) and if everyone visiting has X-USER-ID, then yes, they all limited under the same bucket.

Do you think they should all be limited as individual limit?

[didip]

v4.0.2

[jack-chung]

@didip Yes that's how I see this feature should work, based on the use case I outlined above. Also, this new code still supports the "aggregated limit" use case if needed. Just set the header value to a static value for all users, then it works just like before.