From 1dc0e18773efa049ef2323f62da913416d9712ae Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=F0=9F=90=99PiperYxzzy?= Date: Wed, 3 Aug 2022 22:17:23 +0200 Subject: [PATCH] Upgrades to rate limiting * Added extensive behaviour tests * Added regex capabilities --- .gitignore | 4 +- controllers/ratelimit.go | 28 +++++- controllers/ratelimit_test.go | 173 ++++++++++++++++++++++++++++++++++ 3 files changed, 199 insertions(+), 6 deletions(-) create mode 100644 controllers/ratelimit_test.go diff --git a/.gitignore b/.gitignore index 2eade68..05f210a 100644 --- a/.gitignore +++ b/.gitignore @@ -16,4 +16,6 @@ # Other prepack.db -test_prepack.db \ No newline at end of file +test_prepack.db + +.vscode/ diff --git a/controllers/ratelimit.go b/controllers/ratelimit.go index d4323a2..b1907d0 100644 --- a/controllers/ratelimit.go +++ b/controllers/ratelimit.go @@ -1,7 +1,9 @@ package controllers import ( + "fmt" "net/http" + "regexp" "time" "github.com/gin-gonic/gin" @@ -14,15 +16,31 @@ type rule struct { } type bucket struct { - rules map[string]rule + rules *map[string]rule access map[string]int } func (b *bucket) take(resource string) bool { - r, ex := b.rules[resource] + r, ex := (*b.rules)[resource] if !ex { - resource = "*" - r = b.rules[resource] + // does not exist, forced to try match on regex? + regexMatched := false + for attemptMatch, attemptRes := range *b.rules { + match, _ := regexp.MatchString("^"+attemptMatch+"$", resource) + if match { + resource = attemptMatch + r = attemptRes + regexMatched = true + break + } + } + + if !regexMatched { + // Default to Global + fmt.Printf("defaulting %v to global\n", resource) + resource = "" + r = (*b.rules)[resource] + } } max := r.limit duration := r.duration @@ -57,7 +75,7 @@ func (m *megabucket) take(signature string, resource string) bool { b, ex := m.buckets[signature] if !ex { b = bucket{ - rules: m.rules, + rules: &m.rules, access: map[string]int{}, } m.buckets[signature] = b diff --git a/controllers/ratelimit_test.go b/controllers/ratelimit_test.go new file mode 100644 index 0000000..46207c8 --- /dev/null +++ b/controllers/ratelimit_test.go @@ -0,0 +1,173 @@ +package controllers + +import ( + "testing" + "time" +) + +func TestBucketBehaviour(t *testing.T) { + b := bucket{ + rules: &map[string]rule{ + "": {time.Second, 0}, // Deny + "/1sec5max": {time.Second, 5}, + "/2sec1max": {time.Second * 2, 1}, + "/wildcard/.+/1sec1max": {time.Second, 1}, + "/regex/(test|woot)/1sec2max/[A-Z]{2,3}": {time.Second, 2}, + }, + access: map[string]int{}, + } + + firstTestAutoblocked := []string{ + "willnotwork", + "/invalidresource", + "/1sec5max/butThenHasThis", + "/wildcard/1sec1max", + "/wildcard//1sec1max", + "/regex/test/1sec2max/A", + "/regex/test/1sec2max/aaa", + "/regex/incorrect /1sec2max/HD", + } + + for i, fail := range firstTestAutoblocked { + if b.take(fail) { + t.Errorf("should have auto-throttled %v: '%v'", i, fail) + } + } + + // Test exhausting the whole bucket, all of these calls should return TRUE + secondTestExhaust := []string{ + "/1sec5max", + "/1sec5max", + "/1sec5max", + "/1sec5max", + "/1sec5max", + "/2sec1max", + "/wildcard/bloop/1sec1max", + "/regex/woot/1sec2max/FF", + "/regex/woot/1sec2max/ZAF", + } + + for i, succeed := range secondTestExhaust { + if !b.take(succeed) { + t.Errorf("draining buckets: should have allowed %v: '%v'", i, succeed) + } + } + + // Immediately testing this should return false for successful ones + thirdTestHasExhausted := []string{ + "/1sec5max", + "/2sec1max", + "/wildcard/cool-stuff/1sec1max", + "/regex/woot/1sec2max/JFK", + } + + for i, exhausted := range thirdTestHasExhausted { + if b.take(exhausted) { + t.Errorf("testing exhausted buckets: should have throttled %v: '%v'", i, exhausted) + } + } + + // Wait for the smallest duration, 1 second, and test + time.Sleep(time.Second + time.Millisecond*200) + + fourthTestShouldSucceedAgain := []string{ + "/1sec5max", + "/1sec5max", + "/1sec5max", + "/1sec5max", + "/1sec5max", + "/wildcard/yeehaw/1sec1max", + "/regex/woot/1sec2max/ASD", + "/regex/test/1sec2max/ZZ", + } + + for i, succeed := range fourthTestShouldSucceedAgain { + if !b.take(succeed) { + t.Errorf("1 sec has elapsed, should allow again %v: '%v'", i, succeed) + } + } + + fourthTestShouldStillFail := []string{ + "/2sec1max", + } + + for i, fail := range fourthTestShouldStillFail { + if b.take(fail) { + t.Errorf("1 sec has elapsed, should still not allow %v: '%v'", i, fail) + } + } + + time.Sleep(time.Second + time.Millisecond*200) + + fifthTestShouldNowSucceed := fourthTestShouldStillFail + + for i, succeed := range fifthTestShouldNowSucceed { + if !b.take(succeed) { + t.Errorf("1 more sec has elapsed, should now allow again %v: '%v'", i, succeed) + } + } + +} + +func TestMegabucketBehaviour(t *testing.T) { + m := megabucket{ + rules: map[string]rule{ + "": {time.Second, 0}, // Deny + "/1sec5max": {time.Second, 5}, + "/2sec1max": {time.Second * 2, 1}, + }, + buckets: map[string]bucket{}, + } + + // User up all for user1 + firstTestUser1Succeed := []string{ + "/1sec5max", + "/1sec5max", + "/1sec5max", + "/1sec5max", + "/1sec5max", + "/2sec1max", + } + for i, succeed := range firstTestUser1Succeed { + if !m.take("user1", succeed) { + t.Errorf("user1 failed to take %v: %v", i, succeed) + } + } + + if !m.take("user2", "/1sec5max") { + t.Errorf("user2 was throttled unfairly when taking /1sec5max") + } + + if !m.take("user2", "/2sec1max") { + t.Errorf("user2 was throttled unfairly when taking /2sec1max") + } + + // Now, user1 and user2 should be getting throttled on /2sec1max + if m.take("user1", "/2sec1max") { + t.Errorf("user1 was not throttled on /2sec1max") + } + + if m.take("user2", "/2sec1max") { + t.Errorf("user2 was not throttled on /2sec1max") + } + + // Wait one second, confirm that both can do 5x 1sec5max again + time.Sleep(time.Second + time.Millisecond*200) + + thirdTestUnblockedBothSucceed := []string{ + "/1sec5max", + "/1sec5max", + "/1sec5max", + "/1sec5max", + "/1sec5max", + } + for i, succeed := range thirdTestUnblockedBothSucceed { + if !m.take("user1", succeed) { + t.Errorf("user1 failed to take %v: %v", i, succeed) + } + if !m.take("user2", succeed) { + t.Errorf("user2 failed to take %v: %v", i, succeed) + } + } + +}