diff --git a/bird/bird.go b/bird/bird.go index fa0abc0..bdaf6f2 100644 --- a/bird/bird.go +++ b/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 { diff --git a/bird/config.go b/bird/config.go index b5c6788..89bf0d8 100644 --- a/bird/config.go +++ b/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 +} diff --git a/birdwatcher.go b/birdwatcher.go index 3910aa4..9f53a4e 100644 --- a/birdwatcher.go +++ b/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 diff --git a/config.go b/config.go index 6127f0f..82660a9 100644 --- a/config.go +++ b/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 diff --git a/endpoints/endpoint.go b/endpoints/endpoint.go index ac95194..e364239 100644 --- a/endpoints/endpoint.go +++ b/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 { diff --git a/etc/ecix/birdwatcher.conf b/etc/ecix/birdwatcher.conf index 93cc840..5a5e20c 100644 --- a/etc/ecix/birdwatcher.conf +++ b/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"