Skip to content

Commit

Permalink
refactor rate limit whitelist
Browse files Browse the repository at this point in the history
  • Loading branch information
sethpollack committed Aug 23, 2017
1 parent 5293425 commit f045fa6
Show file tree
Hide file tree
Showing 4 changed files with 47 additions and 31 deletions.
44 changes: 23 additions & 21 deletions controllers/nginx/pkg/template/template.go
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ import (
"k8s.io/apimachinery/pkg/util/sets"
"k8s.io/ingress/controllers/nginx/pkg/config"
"k8s.io/ingress/core/pkg/ingress"
"k8s.io/ingress/core/pkg/ingress/annotations/ratelimit"
ing_net "k8s.io/ingress/core/pkg/net"
"k8s.io/ingress/core/pkg/watch"
)
Expand Down Expand Up @@ -132,8 +133,7 @@ var (
"buildAuthLocation": buildAuthLocation,
"buildAuthResponseHeaders": buildAuthResponseHeaders,
"buildProxyPass": buildProxyPass,
"buildWhitelistVariable": buildWhitelistVariable,
"whitelistExists": whitelistExists,
"filterRateLimits": filterRateLimits,
"buildRateLimitZones": buildRateLimitZones,
"buildRateLimit": buildRateLimit,
"buildResolvers": buildResolvers,
Expand Down Expand Up @@ -337,25 +337,29 @@ func buildProxyPass(host string, b interface{}, loc interface{}) string {
return defProxyPass
}

var (
whitelistVarMap = map[string]string{}
)

func whitelistExists(s string) bool {
_, ok := whitelistVarMap[s]
return ok
}
func filterRateLimits(input interface{}) []ratelimit.RateLimit {
ratelimits := []ratelimit.RateLimit{}
found := map[string]bool{}

func buildWhitelistVariable(s string) string {
if _, ok := whitelistVarMap[s]; !ok {
whitelistVarMap[s] = buildRandomUUID()
servers, ok := input.([]*ingress.Server)
if !ok {
return ratelimits
}
return whitelistVarMap[s]
for _, server := range servers {
for _, loc := range server.Locations {
if loc.RateLimit.ID != "" && !found[loc.RateLimit.ID] {
found[loc.RateLimit.ID] = true
ratelimits = append(ratelimits, loc.RateLimit)
}
}
}
return ratelimits
}

// buildRateLimitZones produces an array of limit_conn_zone in order to allow
// rate limiting of request. Each Ingress rule could have up to two zones, one
// for connection limit by IP address and other for limiting request per second
// rate limiting of request. Each Ingress rule could have up to three zones, one
// for connection limit by IP address, one for limiting requests per minute, and
// one for limiting requests per second.
func buildRateLimitZones(input interface{}) []string {
zones := sets.String{}

Expand All @@ -366,11 +370,9 @@ func buildRateLimitZones(input interface{}) []string {

for _, server := range servers {
for _, loc := range server.Locations {
whitelistVar := buildWhitelistVariable(loc.RateLimit.Name)

if loc.RateLimit.Connections.Limit > 0 {
zone := fmt.Sprintf("limit_conn_zone $limit_%s zone=%v:%vm;",
whitelistVar,
loc.RateLimit.ID,
loc.RateLimit.Connections.Name,
loc.RateLimit.Connections.SharedSize)
if !zones.Has(zone) {
Expand All @@ -380,7 +382,7 @@ func buildRateLimitZones(input interface{}) []string {

if loc.RateLimit.RPM.Limit > 0 {
zone := fmt.Sprintf("limit_req_zone $limit_%s zone=%v:%vm rate=%vr/m;",
whitelistVar,
loc.RateLimit.ID,
loc.RateLimit.RPM.Name,
loc.RateLimit.RPM.SharedSize,
loc.RateLimit.RPM.Limit)
Expand All @@ -391,7 +393,7 @@ func buildRateLimitZones(input interface{}) []string {

if loc.RateLimit.RPS.Limit > 0 {
zone := fmt.Sprintf("limit_req_zone $limit_%s zone=%v:%vm rate=%vr/s;",
whitelistVar,
loc.RateLimit.ID,
loc.RateLimit.RPS.Name,
loc.RateLimit.RPS.SharedSize,
loc.RateLimit.RPS.Limit)
Expand Down
18 changes: 8 additions & 10 deletions controllers/nginx/rootfs/etc/nginx/template/nginx.tmpl
Original file line number Diff line number Diff line change
Expand Up @@ -290,25 +290,23 @@ http {
}
{{ end }}
{{ end }}
{{ end }}
{{ end }}

{{ if ne $location.RateLimit.Name "" }}
{{ if ne (whitelistExists $location.RateLimit.Name) true }}
# Ratelimit {{ $location.RateLimit.Name }}
geo $whitelist_{{ buildWhitelistVariable $location.RateLimit.Name }} {
{{ range $rl := (filterRateLimits $servers ) }}
# Ratelimit {{ $rl.Name }}
geo $whitelist_{{ $rl.ID }} {
default 0;
{{ range $ip := $location.RateLimit.Whitelist }}
{{ range $ip := $rl.Whitelist }}
{{ $ip }} 1;{{ end }}
}

# Ratelimit {{ $location.RateLimit.Name }}
map $whitelist_{{ buildWhitelistVariable $location.RateLimit.Name }} $limit_{{ buildWhitelistVariable $location.RateLimit.Name }} {
# Ratelimit {{ $rl.Name }}
map $whitelist_{{ $rl.ID }} $limit_{{ $rl.ID }} {
0 {{ $cfg.LimitConnZoneVariable }};
1 "";
}
{{ end }}
{{ end }}
{{ end }}
{{ end }}

{{/* build all the required rate limit zones. Each annotation requires a dedicated zone */}}
{{/* 1MB -> 16 thousand 64-byte states or about 8 thousand 128-byte states */}}
Expand Down
12 changes: 12 additions & 0 deletions core/pkg/base64/base64.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
package base64

import (
"encoding/base64"
"strings"
)

// Base64Encode
func Base64Encode(s string) string {
str := base64.URLEncoding.EncodeToString([]byte(s))
return strings.Replace(str, "=", "", -1)
}
4 changes: 4 additions & 0 deletions core/pkg/ingress/annotations/ratelimit/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ import (

extensions "k8s.io/api/extensions/v1beta1"

"k8s.io/ingress/core/pkg/base64"
"k8s.io/ingress/core/pkg/ingress/annotations/parser"
"k8s.io/ingress/core/pkg/ingress/resolver"
"k8s.io/ingress/core/pkg/net"
Expand Down Expand Up @@ -62,6 +63,8 @@ type RateLimit struct {

Name string `json:"name"`

ID string `json:"id"`

Whitelist []string `json:"whitelist"`
}

Expand Down Expand Up @@ -209,6 +212,7 @@ func (a ratelimit) Parse(ing *extensions.Ingress) (interface{}, error) {
LimitRate: lr,
LimitRateAfter: lra,
Name: zoneName,
ID: base64.Base64Encode(zoneName),
Whitelist: cidrs,
}, nil
}
Expand Down

0 comments on commit f045fa6

Please sign in to comment.