Skip to content

Commit

Permalink
We don’t need the ability to pick which header to use.
Browse files Browse the repository at this point in the history
  • Loading branch information
@didip
didip committed Mar 5, 2022
1 parent 983f008 commit 4a4cde1
Show file tree
Hide file tree
Showing 6 changed files with 49 additions and 101 deletions.
22 changes: 10 additions & 12 deletions README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -47,9 +47,8 @@ func main() {

// New in version >= 7, you must explicitly define how to pick the IP address.
lmt.SetIPLookup(limiter.IPLookup{
Name: "X-Real-IP",
HeaderIndexFromRight: 0,
IndexFromRight: 0,
Name: "X-Real-IP",
IndexFromRight: 0,
})

http.Handle("/", tollbooth.LimitFuncHandler(lmt, HelloHandler))
Expand Down Expand Up @@ -82,15 +81,15 @@ func main() {
// The name of lookup method.
// Possible options are: RemoteAddr, X-Forwarded-For, X-Real-IP, CF-Connecting-IP
// All other headers are considered unknown and will be ignored.
Name: "X-Real-IP",

// If there are multiple of the same header, this index determines which one to use.
// The index goes from right to left.
HeaderIndexFromRight: 0,
Name: "X-Real-IP",

// The index position to pick the ip address from a comma separated list.
// The index goes from right to left.
IndexFromRight: 0,
//
// When there are multiple of the same headers,
// we will concat them together in the order of first to last seen.
// And then we pick the IP using this index position.
IndexFromRight: 0,
})

// In version >= 7, lmt.SetIPLookups and lmt.GetIPLookups are removed.
Expand Down Expand Up @@ -155,9 +154,8 @@ func main() {
// New in version >= 7, you must explicitly define how to pick the IP address.
lmt.SetIPLookup(limiter.IPLookup{
Name: "X-Forwarded-For",
HeaderIndexFromRight: 0,
IndexFromRight: 0,
Name: "X-Forwarded-For",
IndexFromRight: 0,
})
// Set a custom message.
Expand Down
7 changes: 1 addition & 6 deletions libstring/libstring.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,12 +35,7 @@ func RemoteIPFromIPLookup(ipLookup limiter.IPLookup, r *http.Request) string {
case "X-Forwarded-For", "X-Real-IP", "CF-Connecting-IP":
ipAddrListCommaSeparated := r.Header.Values(ipLookup.Name)

headerIndex := len(ipAddrListCommaSeparated) - 1 - ipLookup.HeaderIndexFromRight
if headerIndex < 0 {
headerIndex = 0
}

ipAddrCommaSeparated := ipAddrListCommaSeparated[headerIndex]
ipAddrCommaSeparated := strings.Join(ipAddrListCommaSeparated, ",")

ips := strings.Split(ipAddrCommaSeparated, ",")
for i, p := range ips {
Expand Down
52 changes: 12 additions & 40 deletions libstring/libstring_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,9 +26,8 @@ func TestRemoteIPForwardedFor(t *testing.T) {
request.Header.Set("X-Real-IP", ipv6)

ip := RemoteIPFromIPLookup(limiter.IPLookup{
Name: "X-Forwarded-For",
HeaderIndexFromRight: 0,
IndexFromRight: 0,
Name: "X-Forwarded-For",
IndexFromRight: 0,
}, request)

if ip != "10.10.10.10" {
Expand All @@ -51,9 +50,8 @@ func TestRemoteIPRealIP(t *testing.T) {
request.Header.Set("X-Real-IP", ipv6)

ip := RemoteIPFromIPLookup(limiter.IPLookup{
Name: "X-Real-IP",
HeaderIndexFromRight: 0,
IndexFromRight: 0,
Name: "X-Real-IP",
IndexFromRight: 0,
}, request)

if ip != ipv6 {
Expand All @@ -73,9 +71,8 @@ func TestRemoteIPMultipleForwardedForIPAddresses(t *testing.T) {
request.Header.Set("X-Forwarded-For", "10.10.10.10,10.10.10.11")

ip := RemoteIPFromIPLookup(limiter.IPLookup{
Name: "X-Forwarded-For",
HeaderIndexFromRight: 0,
IndexFromRight: 0,
Name: "X-Forwarded-For",
IndexFromRight: 0,
}, request)

// Should get the last one
Expand All @@ -84,9 +81,8 @@ func TestRemoteIPMultipleForwardedForIPAddresses(t *testing.T) {
}

ip = RemoteIPFromIPLookup(limiter.IPLookup{
Name: "X-Forwarded-For",
HeaderIndexFromRight: 0,
IndexFromRight: 1,
Name: "X-Forwarded-For",
IndexFromRight: 1,
}, request)

// Should get the 2nd from last
Expand All @@ -96,9 +92,8 @@ func TestRemoteIPMultipleForwardedForIPAddresses(t *testing.T) {

// What about index out of bound? RemoteIP should simply choose index 0.
ip = RemoteIPFromIPLookup(limiter.IPLookup{
Name: "X-Forwarded-For",
HeaderIndexFromRight: 0,
IndexFromRight: 2,
Name: "X-Forwarded-For",
IndexFromRight: 2,
}, request)

if ip != "10.10.10.10" {
Expand All @@ -116,37 +111,14 @@ func TestRemoteIPMultipleForwardedForHeaders(t *testing.T) {
request.Header.Add("X-Forwarded-For", "10.10.10.10,10.10.10.11")

ip := RemoteIPFromIPLookup(limiter.IPLookup{
Name: "X-Forwarded-For",
HeaderIndexFromRight: 0,
IndexFromRight: 0,
Name: "X-Forwarded-For",
IndexFromRight: 0,
}, request)

// Should get the last header and the last IP
if ip != "10.10.10.11" {
t.Errorf("Did not get the right IP. IP: %v", ip)
}

ip = RemoteIPFromIPLookup(limiter.IPLookup{
Name: "X-Forwarded-For",
HeaderIndexFromRight: 1,
IndexFromRight: 0,
}, request)

// Should get the last IP from the first header
if ip != "8.8.4.4" {
t.Errorf("Did not get the right IP. IP: %v", ip)
}

ip = RemoteIPFromIPLookup(limiter.IPLookup{
Name: "X-Forwarded-For",
HeaderIndexFromRight: 1,
IndexFromRight: 1,
}, request)

// Should get the first IP from the first header
if ip != "8.8.8.8" {
t.Errorf("Did not get the right IP. IP: %v", ip)
}
}
func TestCanonicalizeIP(t *testing.T) {
tests := []struct {
Expand Down
4 changes: 0 additions & 4 deletions limiter/limiter.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -48,10 +48,6 @@ type IPLookup struct {
// All other headers are considered unknown and will be ignored.
Name string

// If there are multiple of the same header, this index determines which one to use.
// The index goes from right to left.
HeaderIndexFromRight int

// The index position to pick the ip address from a comma separated list.
// The index goes from right to left.
IndexFromRight int
Expand Down
5 changes: 2 additions & 3 deletions tollbooth_benchmark_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,9 +31,8 @@ func BenchmarkLimitByKeysWithExpiringBuckets(b *testing.B) {
func BenchmarkBuildKeys(b *testing.B) {
lmt := limiter.New(nil).SetMax(1) // Only 1 request per second is allowed.
lmt.SetIPLookup(limiter.IPLookup{
Name: "X-Real-IP",
HeaderIndexFromRight: 0,
IndexFromRight: 0,
Name: "X-Real-IP",
IndexFromRight: 0,
}).
SetHeaders(make(map[string][]string)).
SetHeader("X-Real-IP", []string{"2601:7:1c82:4097:59a0:a80b:2841:b8c8"})
Expand Down
60 changes: 24 additions & 36 deletions tollbooth_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,9 +33,8 @@ func TestLimitByKeys(t *testing.T) {

func TestDefaultBuildKeys(t *testing.T) {
lmt := NewLimiter(1, nil).SetIPLookup(limiter.IPLookup{
Name: "X-Real-IP",
HeaderIndexFromRight: 0,
IndexFromRight: 0,
Name: "X-Real-IP",
IndexFromRight: 0,
})

request, err := http.NewRequest("GET", "/", strings.NewReader("Hello, world!"))
Expand Down Expand Up @@ -64,9 +63,8 @@ func TestDefaultBuildKeys(t *testing.T) {
func TestIgnoreURLBuildKeys(t *testing.T) {
lmt := NewLimiter(1, nil).
SetIPLookup(limiter.IPLookup{
Name: "X-Real-IP",
HeaderIndexFromRight: 0,
IndexFromRight: 0,
Name: "X-Real-IP",
IndexFromRight: 0,
}).
SetIgnoreURL(true)

Expand All @@ -89,9 +87,8 @@ func TestIgnoreURLBuildKeys(t *testing.T) {
func TestBasicAuthBuildKeys(t *testing.T) {
lmt := NewLimiter(1, nil).
SetIPLookup(limiter.IPLookup{
Name: "X-Real-IP",
HeaderIndexFromRight: 0,
IndexFromRight: 0,
Name: "X-Real-IP",
IndexFromRight: 0,
}).
SetBasicAuthUsers([]string{"bro"})

Expand Down Expand Up @@ -124,9 +121,8 @@ func TestBasicAuthBuildKeys(t *testing.T) {
func TestCustomHeadersBuildKeys(t *testing.T) {
lmt := NewLimiter(1, nil).
SetIPLookup(limiter.IPLookup{
Name: "X-Real-IP",
HeaderIndexFromRight: 0,
IndexFromRight: 0,
Name: "X-Real-IP",
IndexFromRight: 0,
}).
SetHeader("X-Auth-Token", []string{"totally-top-secret", "another-secret"})

Expand Down Expand Up @@ -158,9 +154,8 @@ func TestCustomHeadersBuildKeys(t *testing.T) {
func TestRequestMethodBuildKeys(t *testing.T) {
lmt := NewLimiter(1, nil).
SetIPLookup(limiter.IPLookup{
Name: "X-Real-IP",
HeaderIndexFromRight: 0,
IndexFromRight: 0,
Name: "X-Real-IP",
IndexFromRight: 0,
}).
SetMethods([]string{"GET"})

Expand Down Expand Up @@ -191,9 +186,8 @@ func TestRequestMethodBuildKeys(t *testing.T) {
func TestContextValueBuildKeys(t *testing.T) {
lmt := NewLimiter(1, nil).
SetIPLookup(limiter.IPLookup{
Name: "X-Real-IP",
HeaderIndexFromRight: 0,
IndexFromRight: 0,
Name: "X-Real-IP",
IndexFromRight: 0,
}).
SetContextValue("API-access-level", []string{"basic"})

Expand Down Expand Up @@ -226,9 +220,8 @@ func TestContextValueBuildKeys(t *testing.T) {
func TestRequestMethodAndCustomHeadersBuildKeys(t *testing.T) {
lmt := NewLimiter(1, nil).
SetIPLookup(limiter.IPLookup{
Name: "X-Real-IP",
HeaderIndexFromRight: 0,
IndexFromRight: 0,
Name: "X-Real-IP",
IndexFromRight: 0,
}).
SetMethods([]string{"GET"}).
SetHeader("X-Auth-Token", []string{"totally-top-secret", "another-secret"})
Expand Down Expand Up @@ -263,9 +256,8 @@ func TestRequestMethodAndCustomHeadersBuildKeys(t *testing.T) {
func TestRequestMethodAndBasicAuthUsersBuildKeys(t *testing.T) {
lmt := NewLimiter(1, nil).
SetIPLookup(limiter.IPLookup{
Name: "X-Real-IP",
HeaderIndexFromRight: 0,
IndexFromRight: 0,
Name: "X-Real-IP",
IndexFromRight: 0,
}).
SetMethods([]string{"GET"}).
SetBasicAuthUsers([]string{"bro"})
Expand Down Expand Up @@ -298,9 +290,8 @@ func TestRequestMethodAndBasicAuthUsersBuildKeys(t *testing.T) {
func TestRequestMethodCustomHeadersAndBasicAuthUsersBuildKeys(t *testing.T) {
lmt := NewLimiter(1, nil).
SetIPLookup(limiter.IPLookup{
Name: "X-Real-IP",
HeaderIndexFromRight: 0,
IndexFromRight: 0,
Name: "X-Real-IP",
IndexFromRight: 0,
}).
SetMethods([]string{"GET"}).
SetHeader("X-Auth-Token", []string{"totally-top-secret", "another-secret"}).
Expand Down Expand Up @@ -338,9 +329,8 @@ func TestRequestMethodCustomHeadersAndBasicAuthUsersBuildKeys(t *testing.T) {
func TestRequestMethodCustomHeadersAndBasicAuthUsersAndContextValuesBuildKeys(t *testing.T) {
lmt := NewLimiter(1, nil).
SetIPLookup(limiter.IPLookup{
Name: "X-Real-IP",
HeaderIndexFromRight: 0,
IndexFromRight: 0,
Name: "X-Real-IP",
IndexFromRight: 0,
}).
SetMethods([]string{"GET"}).
SetHeader("X-Auth-Token", []string{"totally-top-secret", "another-secret"}).
Expand Down Expand Up @@ -382,9 +372,8 @@ func TestRequestMethodCustomHeadersAndBasicAuthUsersAndContextValuesBuildKeys(t
func TestLimitHandler(t *testing.T) {
lmt := limiter.New(nil).SetMax(1).SetBurst(1).
SetIPLookup(limiter.IPLookup{
Name: "X-Real-IP",
HeaderIndexFromRight: 0,
IndexFromRight: 0,
Name: "X-Real-IP",
IndexFromRight: 0,
}).
SetMethods([]string{"POST"})

Expand Down Expand Up @@ -432,9 +421,8 @@ func TestLimitHandler(t *testing.T) {
func TestOverrideForResponseWriter(t *testing.T) {
lmt := limiter.New(nil).SetMax(1).SetBurst(1).
SetIPLookup(limiter.IPLookup{
Name: "X-Real-IP",
HeaderIndexFromRight: 0,
IndexFromRight: 0,
Name: "X-Real-IP",
IndexFromRight: 0,
}).
SetMethods([]string{"POST"}).
SetOverrideDefaultResponseWriter(true)
Expand Down

0 comments on commit 4a4cde1

Please sign in to comment.