Add IP whitelisting in request throttler

Author: sosedoff

api.go

@@ -94,6 +94,12 @@ func throttleMiddleware(throttler *Throttler) gin.HandlerFunc {
 	return func(c *gin.Context) {
 		ip := strings.Split(c.Request.RemoteAddr, ":")[0]
 
+		// Bypass throttling for whitelisted IPs
+		if throttler.Whitelisted(ip) {
+			c.Next()
+			return
+		}
+
 		if err := throttler.Add(ip); err != nil {
 			errorResponse(429, err, c)
 			c.Abort()
@@ -115,6 +121,7 @@ func corsMiddleware() gin.HandlerFunc {
 
 func RunApi(config *Config, client *docker.Client) {
 	throttler := NewThrottler(config.ThrottleConcurrency, config.ThrottleQuota)
+	throttler.SetWhitelist(config.ThrottleWhitelist)
 	throttler.StartPeriodicFlush()
 
 	gin.SetMode(gin.ReleaseMode)

config.go

@@ -21,6 +21,7 @@ type Config struct {
 	RunDuration         time.Duration `json:"run_duration"`
 	ThrottleQuota       int           `json:"throttle_quota"`
 	ThrottleConcurrency int           `json:"throttle_concurrency"`
+	ThrottleWhitelist   []string      `json:"throttle_whitelist"`
 	NetworkDisabled     bool          `json:"network_disabled"`
 	MemoryLimit         int64         `json:"memory_limit"`
 	Pools               []PoolConfig  `json:"pools"`

config.json.example

@@ -6,8 +6,12 @@
   "run_duration": 10,
   "throttle_quota": 5,
   "throttle_concurrency": 1,
+  "throttle_whitelist": [
+    "127.0.0.1"
+  ],
   "network_disabled": false,
   "memory_limit": 67108864,
+  "fetch_images": true,
   "pools": [
     { "image": "bitrun/ruby:2.2", "capacity": 10 }
   ]

throttler.go

@@ -11,6 +11,7 @@ type Throttler struct {
 	Quota       int
 	Clients     map[string]int
 	Requests    map[string]int
+	Whitelist   map[string]bool
 	*sync.Mutex
 }
 
@@ -20,6 +21,7 @@ func NewThrottler(concurrency int, quota int) *Throttler {
 		Quota:       quota,
 		Clients:     make(map[string]int),
 		Requests:    make(map[string]int),
+		Whitelist:   make(map[string]bool),
 		Mutex:       &sync.Mutex{},
 	}
 }
@@ -70,3 +72,14 @@ func (t *Throttler) Flush() {
 		delete(t.Requests, k)
 	}
 }
+
+func (t *Throttler) SetWhitelist(ips []string) {
+	for _, ip := range ips {
+		t.Whitelist[ip] = true
+	}
+}
+
+func (t *Throttler) Whitelisted(ip string) bool {
+	_, ok := t.Whitelist[ip]
+	return ok
+}