First commit: Generic rate limiter for Go request handlers.

LICENSE

@@ -0,0 +1,21 @@
+The MIT License (MIT)
+
+Copyright (c) 2015 Didip Kerabat
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.

README.md

@@ -0,0 +1,32 @@
+[![GoDoc](https://godoc.org/github.com/didip/tollbooth?status.svg)](http://godoc.org/github.com/didip/tollbooth)
+[![license](http://img.shields.io/badge/license-MIT-red.svg?style=flat)](https://raw.githubusercontent.com/didip/tollbooth/master/LICENSE)
+
+## Tollbooth
+
+This is a generic middleware to rate limit HTTP requests.
+
+
+## Usage
+```
+package main
+
+import (
+    "github.com/didip/tollbooth"
+    "github.com/didip/tollbooth/storages"
+    "net/http"
+    "time"
+)
+
+func HelloHandler(w http.ResponseWriter, req *http.Request) {
+    w.Write([]byte("Hello, World!"))
+}
+
+func main() {
+    // 1. Create a request limiter storage.
+    storage := storages.NewInMemory()
+
+    // 2. Create a request limiter per handler.
+    http.Handle("/", tollbooth.RemoteIPLimiterFuncHandler(storage, tollbooth.NewRequestLimit(1, time.Second), HelloHandler))
+    http.ListenAndServe(":12345", nil)
+}
+```

storages/base.go

@@ -0,0 +1,10 @@
+package storages
+
+import (
+	"time"
+)
+
+type ICounterStorage interface {
+	IncrBy(string, int64, time.Duration)
+	Get(string) (int64, bool)
+}

storages/inmemory.go

@@ -0,0 +1,126 @@
+// Package storages store rate counters.
+package storages
+
+import (
+	"sync"
+	"time"
+)
+
+// NewInMemory is a constructor to InMemory struct.
+func NewInMemory() *InMemory {
+	inmem := &InMemory{
+		items: map[string]*InMemoryItem{},
+	}
+	inmem.startCleanupTimer()
+	return inmem
+}
+
+// InMemory is a synchronised map of items that auto-expires.
+type InMemory struct {
+	mutex sync.RWMutex
+	items map[string]*InMemoryItem
+}
+
+// IncrBy creates a new item on map or increment existing onr by num.
+func (inmem *InMemory) IncrBy(key string, num int64, ttl time.Duration) {
+	existing, found := inmem.GetItem(key)
+	if found {
+		inmem.mutex.Lock()
+		existing.IncrBy(num)
+		inmem.mutex.Unlock()
+
+	} else {
+		inmem.mutex.Lock()
+		inmem.items[key] = NewInMemoryItem(num, ttl)
+		inmem.mutex.Unlock()
+	}
+}
+
+// Get a count from map.
+func (inmem *InMemory) Get(key string) (count int64, found bool) {
+	item, found := inmem.GetItem(key)
+	if found {
+		return item.count, found
+	} else {
+		return int64(-1), found
+	}
+}
+
+// GetItem InMemoryItem struct from map.
+func (inmem *InMemory) GetItem(key string) (item *InMemoryItem, found bool) {
+	var exists bool
+
+	inmem.mutex.Lock()
+	item, exists = inmem.items[key]
+	if !exists || item.expired() {
+		found = false
+	} else {
+		item.touch()
+		found = true
+	}
+	inmem.mutex.Unlock()
+	return
+}
+
+func (inmem *InMemory) cleanup() {
+	inmem.mutex.Lock()
+	for key, item := range inmem.items {
+		if item.expired() {
+			delete(inmem.items, key)
+		}
+	}
+	inmem.mutex.Unlock()
+}
+
+func (inmem *InMemory) startCleanupTimer() {
+	duration := time.Second
+	ticker := time.Tick(duration)
+	go (func() {
+		for {
+			select {
+			case <-ticker:
+				inmem.cleanup()
+			}
+		}
+	})()
+}
+
+// NewInMemoryItem is a constructor to InMemoryItem struct.
+func NewInMemoryItem(num int64, ttl time.Duration) *InMemoryItem {
+	item := &InMemoryItem{count: num, ttl: ttl}
+	item.touch()
+	return item
+}
+
+// InMemoryItem represents a single record.
+type InMemoryItem struct {
+	sync.RWMutex
+	count   int64
+	ttl     time.Duration
+	expires *time.Time
+}
+
+func (item *InMemoryItem) IncrBy(num int64) {
+	item.Lock()
+	item.count = item.count + num
+	item.Unlock()
+}
+
+func (item *InMemoryItem) touch() {
+	item.Lock()
+	expiration := time.Now().Add(item.ttl)
+	item.expires = &expiration
+	item.Unlock()
+}
+
+func (item *InMemoryItem) expired() bool {
+	var value bool
+	item.RLock()
+	if item.expires == nil {
+		value = true
+	} else {
+		value = item.expires.Before(time.Now())
+	}
+	item.RUnlock()
+	return value
+}

storages/inmemory_test.go

@@ -0,0 +1,32 @@
+package storages
+
+import (
+	"testing"
+	"time"
+)
+
+func TestCRUD(t *testing.T) {
+	inmem := NewInMemory()
+
+	key := "/|127.0.0.1"
+
+	count, exists := inmem.Get(key)
+	if exists || count > 0 {
+		t.Errorf("Expected empty inmem to return no count")
+	}
+
+	inmem.IncrBy("/|127.0.0.1", int64(1), time.Second)
+	count, exists = inmem.Get(key)
+	if !exists {
+		t.Errorf("Expected inmem to return count for key: %v", key)
+	}
+	if count != 1 {
+		t.Errorf("Expected inmem to return 1 for key: %v", key)
+	}
+
+	<-time.After(2 * time.Second)
+	_, exists = inmem.Get(key)
+	if exists {
+		t.Errorf("Expected key: %v to have expired.", key)
+	}
+}

tollbooth.go

@@ -0,0 +1,62 @@
+// Package tollbooth provides core rate-limit logic.
+package tollbooth
+
+import (
+	"fmt"
+	"github.com/didip/tollbooth/storages"
+	"net/http"
+	"time"
+)
+
+func NewRequestLimit(max int64, ttl time.Duration) *RequestLimit {
+	return &RequestLimit{Max: max, TTL: ttl}
+}
+
+type RequestLimit struct {
+	Max int64
+	TTL time.Duration
+}
+
+type HTTPError struct {
+	Message    string
+	StatusCode int
+}
+
+func (httperror *HTTPError) Error() string {
+	return fmt.Sprintf("%v: %v", httperror.StatusCode, httperror.Message)
+}
+
+func LimitByRemoteIP(storage storages.ICounterStorage, reqLimit *RequestLimit, r *http.Request) *HTTPError {
+	remoteIP := r.Header.Get("REMOTE_ADDR")
+	path := r.URL.Path
+	key := path + "|" + remoteIP
+
+	storage.IncrBy(key, int64(1), reqLimit.TTL)
+	currentCount, _ := storage.Get(key)
+
+	// Check if the returned counter exceeds our limit
+	if currentCount > reqLimit.Max {
+		return &HTTPError{Message: "You have reached maximum request limit.", StatusCode: 429}
+	}
+	return nil
+}
+
+// RemoteIPLimiterHandler is a middleware that limits by RemoteIP.
+func RemoteIPLimiterHandler(storage storages.ICounterStorage, reqLimit *RequestLimit, next http.Handler) http.Handler {
+	middle := func(w http.ResponseWriter, r *http.Request) {
+		httpError := LimitByRemoteIP(storage, reqLimit, r)
+
+		if httpError != nil {
+			http.Error(w, httpError.Message, httpError.StatusCode)
+			return
+		} else {
+			next.ServeHTTP(w, r)
+		}
+	}
+	return http.HandlerFunc(middle)
+}
+
+func RemoteIPLimiterFuncHandler(storage storages.ICounterStorage, reqLimit *RequestLimit, nextFunc func(http.ResponseWriter, *http.Request)) http.Handler {
+	next := http.HandlerFunc(nextFunc)
+	return RemoteIPLimiterHandler(storage, reqLimit, next)
+}