Merge pull request #2 from ecix/rate-limit

Added Rate Limiting

bird/bird.go

@@ -9,6 +9,10 @@ import (
 
 var ClientConf BirdConfig
 var StatusConf StatusConfig
+var RateLimitConf struct {
+	sync.RWMutex
+	Conf RateLimitConfig
+}
 
 var Cache = struct {
 	sync.RWMutex
@@ -35,11 +39,49 @@ func Run(args string) ([]byte, error) {
 	return exec.Command(ClientConf.BirdCmd, argsList...).Output()
 }
 
+func InstallRateLimitReset() {
+	go func() {
+		c := time.Tick(time.Second)
+
+		for _ = range c {
+			RateLimitConf.Lock()
+			RateLimitConf.Conf.Reqs = 0
+			RateLimitConf.Unlock()
+		}
+	}()
+}
+
+func checkRateLimit() bool {
+	RateLimitConf.RLock()
+	check := !RateLimitConf.Conf.Enabled
+	RateLimitConf.RUnlock()
+	if check {
+		return true
+	}
+
+	RateLimitConf.RLock()
+	check = RateLimitConf.Conf.Reqs > RateLimitConf.Conf.Max
+	RateLimitConf.RUnlock()
+	if check {
+		return false
+	}
+
+	RateLimitConf.Lock()
+	RateLimitConf.Conf.Reqs += 1
+	RateLimitConf.Unlock()
+
+	return true
+}
+
 func RunAndParse(cmd string, parser func([]byte) Parsed) (Parsed, bool) {
 	if val, ok := fromCache(cmd); ok {
 		return val, true
 	}
 
+	if !checkRateLimit() {
+		return nil, false
+	}
+
 	out, err := Run(cmd)
 
 	if err != nil {

bird/config.go

@@ -14,3 +14,9 @@ type BirdConfig struct {
 	ConfigFilename string `toml:"config"`
 	BirdCmd        string `toml:"birdc"`
 }
+
+type RateLimitConfig struct {
+	Reqs    int `toml:"requests_per_minute"`
+	Max     int `toml:"requests_per_minute"`
+	Enabled bool
+}

birdwatcher.go

@@ -93,6 +93,7 @@ func main() {
 	bird6 := flag.Bool("6", false, "Use bird6 instead of bird")
 	flag.Parse()
 
+	bird.InstallRateLimitReset()
 	// Load configurations
 	conf, err := LoadConfigs([]string{
 		"./etc/ecix/birdwatcher.conf",
@@ -115,6 +116,7 @@ func main() {
 	// Configuration
 	bird.ClientConf = birdConf
 	bird.StatusConf = conf.Status
+	bird.RateLimitConf.Conf = conf.Ratelimit
 	endpoints.Conf = conf.Server
 
 	// Make server

config.go

@@ -15,9 +15,10 @@ import (
 type Config struct {
 	Server endpoints.ServerConfig
 
-	Status bird.StatusConfig
-	Bird   bird.BirdConfig
-	Bird6  bird.BirdConfig
+	Ratelimit bird.RateLimitConfig
+	Status    bird.StatusConfig
+	Bird      bird.BirdConfig
+	Bird6     bird.BirdConfig
 }
 
 // Try to load configfiles as specified in the files

endpoints/endpoint.go

@@ -52,6 +52,10 @@ func Endpoint(wrapped func(httprouter.Params) (bird.Parsed, bool)) httprouter.Ha
 		res := make(map[string]interface{})
 
 		ret, from_cache := wrapped(ps)
+		if ret == nil {
+			w.WriteHeader(http.StatusTooManyRequests)
+			return
+		}
 		res["api"] = GetApiInfo(from_cache)
 
 		for k, v := range ret {

etc/ecix/birdwatcher.conf

@@ -1,5 +1,3 @@
-
-
 #
 # Birdwatcher Configuration
 #
@@ -36,6 +34,10 @@ reconfig_timestamp_match = "# Created: (.*)"
 # Remove fields e.g. last_reboot
 filter_fields = []
 
+[ratelimit]
+enabled = true
+requests_per_minute = 10
+
 
 [bird]
 listen = "0.0.0.0:29188"