Merge pull request #1224 from sethpollack/whitelist_refactor

refactor rate limit whitelist
kubernetes · Aug 23, 2017 · f0144a1 · f0144a1
2 parents 5293425 + f045fa6
commit f0144a1
Show file tree

Hide file tree

Showing 4 changed files with 47 additions and 31 deletions.
diff --git a/controllers/nginx/pkg/template/template.go b/controllers/nginx/pkg/template/template.go
@@ -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"
 )
@@ -132,8 +133,7 @@ var (
 		"buildAuthLocation":        buildAuthLocation,
 		"buildAuthResponseHeaders": buildAuthResponseHeaders,
 		"buildProxyPass":           buildProxyPass,
-		"buildWhitelistVariable":   buildWhitelistVariable,
-		"whitelistExists":          whitelistExists,
+		"filterRateLimits":         filterRateLimits,
 		"buildRateLimitZones":      buildRateLimitZones,
 		"buildRateLimit":           buildRateLimit,
 		"buildResolvers":           buildResolvers,
@@ -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{}
 
@@ -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) {
@@ -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)
@@ -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)

diff --git a/controllers/nginx/rootfs/etc/nginx/template/nginx.tmpl b/controllers/nginx/rootfs/etc/nginx/template/nginx.tmpl
@@ -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 */}}

diff --git a/core/pkg/base64/base64.go b/core/pkg/base64/base64.go
@@ -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)
+}
diff --git a/core/pkg/ingress/annotations/ratelimit/main.go b/core/pkg/ingress/annotations/ratelimit/main.go
@@ -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"
@@ -62,6 +63,8 @@ type RateLimit struct {
 
 	Name string `json:"name"`
 
+	ID string `json:"id"`
+
 	Whitelist []string `json:"whitelist"`
 }
 
@@ -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
 }