Skip to content
Snippets Groups Projects
Commit fe97baec authored by ale's avatar ale
Browse files

fix setting memcache counters for the first time

parent 3ea4d874
No related branches found
No related tags found
No related merge requests found
...@@ -13,6 +13,7 @@ var ( ...@@ -13,6 +13,7 @@ var (
dosKeyNamespace = "dos" dosKeyNamespace = "dos"
memcacheAddr = flag.String("memcache", "127.0.0.1:11211", "Memcache address") memcacheAddr = flag.String("memcache", "127.0.0.1:11211", "Memcache address")
whitelistedAddrs = flag.String("whitelist-addrs", "127.0.0.1", "Whitelisted addresses (no rate limiting), comma separated")
) )
type Limiter struct { type Limiter struct {
...@@ -31,7 +32,7 @@ func NewLimiter(param string, count int, period int, extractfn func(*http.Reques ...@@ -31,7 +32,7 @@ func NewLimiter(param string, count int, period int, extractfn func(*http.Reques
Count: count, Count: count,
Period: period, Period: period,
extractValue: extractfn, extractValue: extractfn,
whitelistedValues: make([]string, 0), whitelistedValues: strings.Split(*whitelistedAddrs, ","),
mc: memcache.New(*memcacheAddr), mc: memcache.New(*memcacheAddr),
} }
} }
...@@ -58,7 +59,7 @@ func (l *Limiter) Limit(handler http.Handler) http.Handler { ...@@ -58,7 +59,7 @@ func (l *Limiter) Limit(handler http.Handler) http.Handler {
result, err := l.mc.Increment(key, 1) result, err := l.mc.Increment(key, 1)
if err != nil { if err != nil {
result = 1 result = 1
if err = l.mc.Add(&memcache.Item{Key: key, Value: []byte{0, 0, 0, 0}, Expiration: int32(l.Period)}); err != nil { if err = l.mc.Add(&memcache.Item{Key: key, Value: []byte("1"), Expiration: int32(l.Period)}); err != nil {
// Possible race for who adds to the cache first. // Possible race for who adds to the cache first.
if result, err = l.mc.Increment(key, 1); err != nil { if result, err = l.mc.Increment(key, 1); err != nil {
log.Printf("Memcache failed for '%s': %s", key, err) log.Printf("Memcache failed for '%s': %s", key, err)
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment