fixed rate limiter rules matching to acount for the Audience field

This commit is contained in:
Gani Georgiev
2024-11-18 14:46:06 +02:00
parent 487e83c84e
commit 70df03ffbb
7 changed files with 82 additions and 34 deletions
+9 -3
View File
@@ -7,6 +7,7 @@ import (
"fmt"
"os"
"regexp"
"slices"
"strconv"
"strings"
"sync"
@@ -560,13 +561,17 @@ type RateLimitsConfig struct {
}
// FindRateLimitRule returns the first matching rule based on the provided labels.
func (c *RateLimitsConfig) FindRateLimitRule(searchLabels []string) (RateLimitRule, bool) {
//
// Optionally you can further specify a list of valid RateLimitRule.Audience values to further filter the matching rule
// (aka. the rule Audience will have to exist in one of the specified options).
func (c *RateLimitsConfig) FindRateLimitRule(searchLabels []string, optOnlyAudience ...string) (RateLimitRule, bool) {
var prefixRules []int
for i, label := range searchLabels {
// check for direct match
for j := range c.Rules {
if label == c.Rules[j].Label {
if label == c.Rules[j].Label &&
(len(optOnlyAudience) == 0 || slices.Contains(optOnlyAudience, c.Rules[j].Audience)) {
return c.Rules[j], true
}
@@ -578,7 +583,8 @@ func (c *RateLimitsConfig) FindRateLimitRule(searchLabels []string) (RateLimitRu
// check for prefix match
if len(prefixRules) > 0 {
for j := range prefixRules {
if strings.HasPrefix(label+"/", c.Rules[prefixRules[j]].Label) {
if strings.HasPrefix(label+"/", c.Rules[prefixRules[j]].Label) &&
(len(optOnlyAudience) == 0 || slices.Contains(optOnlyAudience, c.Rules[prefixRules[j]].Audience)) {
return c.Rules[prefixRules[j]], true
}
}
+26 -16
View File
@@ -634,33 +634,43 @@ func TestRateLimitsFindRateLimitRule(t *testing.T) {
limits := core.RateLimitsConfig{
Rules: []core.RateLimitRule{
{Label: "abc"},
{Label: "POST /test/a/"},
{Label: "/test/a/"},
{Label: "def", Audience: core.RateLimitRuleAudienceGuest},
{Label: "/test/a", Audience: core.RateLimitRuleAudienceGuest},
{Label: "POST /test/a"},
{Label: "/test/a"},
{Label: "/test/a/", Audience: core.RateLimitRuleAudienceAuth},
{Label: "POST /test/a/"},
},
}
scenarios := []struct {
labels []string
audience []string
expected string
}{
{[]string{}, ""},
{[]string{"missing"}, ""},
{[]string{"abc"}, "abc"},
{[]string{"/test"}, ""},
{[]string{"/test/a"}, "/test/a"},
{[]string{"GET /test/a"}, ""},
{[]string{"POST /test/a"}, "POST /test/a"},
{[]string{"/test/a/b/c"}, "/test/a/"},
{[]string{"GET /test/a/b/c"}, ""},
{[]string{"POST /test/a/b/c"}, "POST /test/a/"},
{[]string{"/test/a", "abc"}, "/test/a"}, // priority checks
{[]string{}, []string{}, ""},
{[]string{"missing"}, []string{}, ""},
{[]string{"abc"}, []string{}, "abc"},
{[]string{"abc"}, []string{core.RateLimitRuleAudienceGuest}, ""},
{[]string{"abc"}, []string{core.RateLimitRuleAudienceAuth}, ""},
{[]string{"def"}, []string{core.RateLimitRuleAudienceGuest}, "def"},
{[]string{"def"}, []string{core.RateLimitRuleAudienceAuth}, ""},
{[]string{"/test"}, []string{}, ""},
{[]string{"/test/a"}, []string{}, "/test/a"},
{[]string{"/test/a"}, []string{core.RateLimitRuleAudienceAuth}, "/test/a/"},
{[]string{"/test/a"}, []string{core.RateLimitRuleAudienceGuest}, "/test/a"},
{[]string{"GET /test/a"}, []string{}, ""},
{[]string{"POST /test/a"}, []string{}, "POST /test/a"},
{[]string{"/test/a/b/c"}, []string{}, "/test/a/"},
{[]string{"/test/a/b/c"}, []string{core.RateLimitRuleAudienceAuth}, "/test/a/"},
{[]string{"/test/a/b/c"}, []string{core.RateLimitRuleAudienceGuest}, ""},
{[]string{"GET /test/a/b/c"}, []string{}, ""},
{[]string{"POST /test/a/b/c"}, []string{}, "POST /test/a/"},
{[]string{"/test/a", "abc"}, []string{}, "/test/a"}, // priority checks
}
for _, s := range scenarios {
t.Run(strings.Join(s.labels, ""), func(t *testing.T) {
rule, ok := limits.FindRateLimitRule(s.labels)
t.Run(strings.Join(s.labels, "_")+":"+strings.Join(s.audience, "_"), func(t *testing.T) {
rule, ok := limits.FindRateLimitRule(s.labels, s.audience...)
hasLabel := rule.Label != ""
if hasLabel != ok {