Skip to content

Commit

Permalink
Updated UT
Browse files Browse the repository at this point in the history
  • Loading branch information
lukipro committed Nov 15, 2023
1 parent 1ff9529 commit 57551b1
Show file tree
Hide file tree
Showing 3 changed files with 31 additions and 14 deletions.
22 changes: 11 additions & 11 deletions src/config/config_impl.go
Original file line number Diff line number Diff line change
Expand Up @@ -40,9 +40,9 @@ type YamlRoot struct {
}

type rateLimitDescriptor struct {
descriptors map[string]*rateLimitDescriptor
limit *RateLimit
wildcardKeys []string
descriptors map[string]*rateLimitDescriptor
limit *RateLimit
wildcardValues []string
}

type rateLimitDomain struct {
Expand Down Expand Up @@ -190,8 +190,8 @@ func (this *rateLimitDescriptor) loadDescriptors(config RateLimitConfigToLoad, p
this.descriptors[finalKey] = newDescriptor

// Preload keys starting or ending with "*" symbol.
if strings.HasPrefix(finalKey, "*") || strings.HasSuffix(finalKey, "*") {
this.wildcardKeys = append(this.wildcardKeys, finalKey)
if strings.HasPrefix(descriptorConfig.Value, "*") || strings.HasSuffix(descriptorConfig.Value, "*") {
this.wildcardValues = append(this.wildcardValues, descriptorConfig.Value)
}
}
}
Expand Down Expand Up @@ -313,14 +313,14 @@ func (this *rateLimitConfigImpl) GetLimit(
logger.Debugf("looking up key: %s", finalKey)
nextDescriptor := descriptorsMap[finalKey]

if nextDescriptor == nil && len(prevDescriptor.wildcardKeys) > 0 {
for _, wildcardKey := range prevDescriptor.wildcardKeys {
if strings.HasSuffix(finalKey, strings.TrimPrefix(wildcardKey, "*")) {
nextDescriptor = descriptorsMap[wildcardKey]
if nextDescriptor == nil && len(prevDescriptor.wildcardValues) > 0 {
for _, wildcardValue := range prevDescriptor.wildcardValues {
if strings.HasSuffix(entry.Value, strings.TrimPrefix(wildcardValue, "*")) {
nextDescriptor = descriptorsMap[entry.Key+"_"+wildcardValue]
break
}
if strings.HasPrefix(finalKey, strings.TrimSuffix(wildcardKey, "*")) {
nextDescriptor = descriptorsMap[wildcardKey]
if strings.HasPrefix(entry.Value, strings.TrimSuffix(wildcardValue, "*")) {
nextDescriptor = descriptorsMap[entry.Key+"_"+wildcardValue]
break
}
}
Expand Down
16 changes: 14 additions & 2 deletions test/config/config_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -591,21 +591,33 @@ func TestWildcardConfig(t *testing.T) {
wildcard1 := rlConfig.GetLimit(
context.TODO(), "test-domain",
&pb_struct.RateLimitDescriptor{
Entries: []*pb_struct.RateLimitDescriptor_Entry{{Key: "wild", Value: "foo1"}},
Entries: []*pb_struct.RateLimitDescriptor_Entry{{Key: "rightWild", Value: "foo1"}},
})
wildcard2 := rlConfig.GetLimit(
context.TODO(), "test-domain",
&pb_struct.RateLimitDescriptor{
Entries: []*pb_struct.RateLimitDescriptor_Entry{{Key: "wild", Value: "foo2"}},
Entries: []*pb_struct.RateLimitDescriptor_Entry{{Key: "rightWild", Value: "foo2"}},
})
wildcard3 := rlConfig.GetLimit(
context.TODO(), "test-domain",
&pb_struct.RateLimitDescriptor{
Entries: []*pb_struct.RateLimitDescriptor_Entry{{Key: "leftWild", Value: "afoo"}},
})
wildcard4 := rlConfig.GetLimit(
context.TODO(), "test-domain",
&pb_struct.RateLimitDescriptor{
Entries: []*pb_struct.RateLimitDescriptor_Entry{{Key: "leftWild", Value: "bfoo"}},
})
wildcard5 := rlConfig.GetLimit(
context.TODO(), "test-domain",
&pb_struct.RateLimitDescriptor{
Entries: []*pb_struct.RateLimitDescriptor_Entry{{Key: "nestedWild", Value: "val1"}, {Key: "wild", Value: "goo2"}},
})
assert.NotNil(wildcard1)
assert.Equal(wildcard1, wildcard2)
assert.NotNil(wildcard3)
assert.Equal(wildcard3, wildcard4)
assert.NotNil(wildcard5)

// Doesn't match non-matching values
noMatch := rlConfig.GetLimit(
Expand Down
7 changes: 6 additions & 1 deletion test/config/wildcard.yaml
Original file line number Diff line number Diff line change
@@ -1,11 +1,16 @@
# Basic configuration for testing.
domain: test-domain
descriptors:
- key: wild
- key: rightWild
value: foo*
rate_limit:
unit: minute
requests_per_unit: 20
- key: leftWild
value: "*foo"
rate_limit:
unit: minute
requests_per_unit: 20
- key: noWild
value: foo
rate_limit:
Expand Down

0 comments on commit 57551b1

Please sign in to comment.