Initial commit

.gitignore

@@ -0,0 +1,3 @@
+.idea/
+.vscode/
+bin/

LICENSE.md

@@ -0,0 +1,19 @@
+Copyright (c) 2019 coinpaprika
+
+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,115 @@
+# ratelimiter
+
+Simple rate limiter for any resources inspired by Cloudflare's approach: [How we built rate limiting capable of scaling to millions of domains.](https://blog.cloudflare.com/counting-things-a-lot-of-different-things/)
+
+## Usage
+
+### Getting started
+
+```go
+package main
+
+import (
+	"fmt"
+	"log"
+	"time"
+
+	"github.com/coinpaprika/ratelimiter"
+)
+
+func main() {
+	limitedKey := "key"
+	windowSize := 1 * time.Minute
+
+	dataStore := ratelimiter.NewMapLimitStore(2*windowSize, 10*time.Second) // create map data store for rate limiter and set each element's expiration time to 2*windowSize and old data flush interval to 10*time.Second
+
+	var maxLimit int64 = 5
+	rateLimiter := ratelimiter.New(dataStore, maxLimit, windowSize) // allow 5 requests per windowSize (1 minute)
+
+	for i := 0; i < 10; i++ {
+		limitStatus, err := rateLimiter.Check(limitedKey)
+		if err != nil {
+			log.Fatal(err)
+		}
+		if limitStatus.IsLimited {
+			fmt.Printf("too high rate for key: %s: rate: %f, limit: %d\nsleep: %s", limitedKey, limitStatus.CurrentRate, maxLimit, *limitStatus.LimitDuration)
+			time.Sleep(*limitStatus.LimitDuration)
+		} else {
+			err := rateLimiter.Inc(limitedKey)
+			if err != nil {
+				log.Fatal(err)
+			}
+		}
+	}
+}
+```
+
+### Rate-limit IP requests in http middleware
+
+```go
+func rateLimitMiddleware(rateLimiter *ratelimiter.RateLimiter) func(http.Handler) http.Handler {
+	return func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			remoteIP := GetRemoteIP([]string{"X-Forwarded-For", "RemoteAddr", "X-Real-IP"}, 0, r)
+			key := fmt.Sprintf("%s_%s_%s", remoteIP, r.URL.String(), r.Method)
+
+			limitStatus, err := rateLimiter.Check(key)
+			if err != nil {
+				// if rate limit error then pass the request
+				next.ServeHTTP(w, r)
+			}
+			if limitStatus.IsLimited {
+				w.WriteHeader(http.StatusTooManyRequests)
+				return
+			} else {
+				rateLimiter.Inc(key)
+			}
+
+			next.ServeHTTP(w, r)
+		})
+	}
+}
+
+func hello(w http.ResponseWriter, r *http.Request) {
+	fmt.Fprintf(w, "Hello, %q", html.EscapeString(r.URL.Path))
+}
+
+func main() {
+	windowSize := 1 * time.Minute
+	dataStore := ratelimiter.NewMapLimitStore(2*windowSize, 10*time.Second) // create map data store for rate limiter and set each element's expiration time to 2*windowSize and old data flush interval to 10*time.Second
+	rateLimiter := ratelimiter.New(dataStore, 5, windowSize)                // allow 5 requests per windowSize (1 minute)
+
+	rateLimiterHandler := rateLimitMiddleware(rateLimiter)
+	helloHandler := http.HandlerFunc(hello)
+	http.Handle("/", rateLimiterHandler(helloHandler))
+
+	log.Fatal(http.ListenAndServe(":8080", nil))
+
+}
+```
+See full [example](./examples/http_middleware/http_middleware.go)
+
+### Implement your own limit data store
+To use custom data store (memcached, Redis, MySQL etc.) you just need to implement [LimitStore](./limit_store.go) interface:
+```go
+type FakeDataStore struct{}
+
+func (f FakeDataStore) Inc(key string, window time.Time) error {
+	return nil
+}
+
+func (f FakeDataStore) Get(key string, previousWindow, currentWindow time.Time) (prevValue int64, currValue int64, err error) {
+	return 0, 0, nil
+}
+// ...
+rateLimiter := ratelimiter.New(FakeDataStore{}, maxLimit, windowSize)
+```
+
+## Examples 
+
+Check out the [examples](./examples) directory.
+
+
+## License
+
+ratelimiter is available under the MIT license. See the [LICENSE file](./LICENSE.md) for more info.

examples/http_middleware/http_middleware.go

@@ -0,0 +1,90 @@
+package main
+
+import (
+	"fmt"
+	"html"
+	"log"
+	"net"
+	"net/http"
+	"strings"
+	"time"
+
+	"github.com/coinpaprika/ratelimiter"
+)
+
+// copied from https://github.com/didip/tollbooth/blob/master/libstring/libstring.go#L21
+func GetRemoteIP(ipLookups []string, forwardedForIndexFromBehind int, r *http.Request) string {
+	realIP := r.Header.Get("X-Real-IP")
+	forwardedFor := r.Header.Get("X-Forwarded-For")
+
+	for _, lookup := range ipLookups {
+		if lookup == "RemoteAddr" {
+			// 1. Cover the basic use cases for both ipv4 and ipv6
+			ip, _, err := net.SplitHostPort(r.RemoteAddr)
+			if err != nil {
+				// 2. Upon error, just return the remote addr.
+				return r.RemoteAddr
+			}
+			return ip
+		}
+		if lookup == "X-Forwarded-For" && forwardedFor != "" {
+			// X-Forwarded-For is potentially a list of addresses separated with ","
+			parts := strings.Split(forwardedFor, ",")
+			for i, p := range parts {
+				parts[i] = strings.TrimSpace(p)
+			}
+
+			partIndex := len(parts) - 1 - forwardedForIndexFromBehind
+			if partIndex < 0 {
+				partIndex = 0
+			}
+
+			return parts[partIndex]
+		}
+		if lookup == "X-Real-IP" && realIP != "" {
+			return realIP
+		}
+	}
+
+	return ""
+}
+
+func rateLimitMiddleware(rateLimiter *ratelimiter.RateLimiter) func(http.Handler) http.Handler {
+	return func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			remoteIP := GetRemoteIP([]string{"X-Forwarded-For", "RemoteAddr", "X-Real-IP"}, 0, r)
+			key := fmt.Sprintf("%s_%s_%s", remoteIP, r.URL.String(), r.Method)
+
+			limitStatus, err := rateLimiter.Check(key)
+			if err != nil {
+				// if rate limit error then pass the request
+				next.ServeHTTP(w, r)
+			}
+			if limitStatus.IsLimited {
+				w.WriteHeader(http.StatusTooManyRequests)
+				return
+			} else {
+				rateLimiter.Inc(key)
+			}
+
+			next.ServeHTTP(w, r)
+		})
+	}
+}
+
+func hello(w http.ResponseWriter, r *http.Request) {
+	fmt.Fprintf(w, "Hello, %q", html.EscapeString(r.URL.Path))
+}
+
+func main() {
+	windowSize := 1 * time.Minute
+	dataStore := ratelimiter.NewMapLimitStore(2*windowSize, 10*time.Second) // create map data store for rate limiter and set each element's expiration time to 2*windowSize and old data flush interval to 10*time.Second
+	rateLimiter := ratelimiter.New(dataStore, 5, windowSize)                // allow 5 requests per windowSize (1 minute)
+
+	rateLimiterHandler := rateLimitMiddleware(rateLimiter)
+	helloHandler := http.HandlerFunc(hello)
+	http.Handle("/", rateLimiterHandler(helloHandler))
+
+	log.Fatal(http.ListenAndServe(":8080", nil))
+
+}

examples/simple/simple.go

@@ -0,0 +1,35 @@
+package main
+
+import (
+	"fmt"
+	"log"
+	"time"
+
+	"github.com/coinpaprika/ratelimiter"
+)
+
+func main() {
+	limitedKey := "key"
+	windowSize := 1 * time.Minute
+
+	dataStore := ratelimiter.NewMapLimitStore(2*windowSize, 10*time.Second) // create map data store for rate limiter and set each element's expiration time to 2*windowSize and old data flush interval to 10*time.Second
+
+	var maxLimit int64 = 5
+	rateLimiter := ratelimiter.New(dataStore, maxLimit, windowSize) // allow 5 requests per windowSize (1 minute)
+
+	for i := 0; i < 10; i++ {
+		limitStatus, err := rateLimiter.Check(limitedKey)
+		if err != nil {
+			log.Fatal(err)
+		}
+		if limitStatus.IsLimited {
+			fmt.Printf("too high rate for key: %s: rate: %f, limit: %d\nsleep: %s", limitedKey, limitStatus.CurrentRate, maxLimit, *limitStatus.LimitDuration)
+			time.Sleep(*limitStatus.LimitDuration)
+		} else {
+			err := rateLimiter.Inc(limitedKey)
+			if err != nil {
+				log.Fatal(err)
+			}
+		}
+	}
+}

go.mod

@@ -0,0 +1,3 @@
+module github.com/coinpaprika/ratelimiter
+
+require github.com/stretchr/testify v1.3.0

go.sum

@@ -0,0 +1,7 @@
+github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8=
+github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+github.com/stretchr/testify v1.3.0 h1:TivCn/peBQ7UY8ooIcPgZFpTNSz0Q2U6UrFlUfqbe0Q=
+github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=

limit_store.go

@@ -0,0 +1,11 @@
+package ratelimiter
+
+import "time"
+
+// LimitStore is the interface that represents limiter internal data store. Any database struct that implements LimitStore should have functions for incrementing counter of a given key and getting counter values of a given key for previous and current window
+type LimitStore interface {
+	// Inc increments current window limit counter for key
+	Inc(key string, window time.Time) error
+	// Get gets value of previous window counter and current window counter for key
+	Get(key string, previousWindow, currentWindow time.Time) (prevValue int64, currValue int64, err error)
+}

map_limit_store.go

@@ -0,0 +1,67 @@
+package ratelimiter
+
+import (
+	"fmt"
+	"sync"
+	"time"
+)
+
+type limitValue struct {
+	val        int64
+	lastUpdate time.Time
+}
+
+// MapLimitStore represents internal limiter data database where data are stored in golang maps
+type MapLimitStore struct {
+	data           map[string]limitValue
+	mutex          sync.RWMutex
+	expirationTime time.Duration
+}
+
+// NewMapLimitStore creates new in-memory data store for internal limiter data. Each element of MapLimitStore is set as expired after expirationTime from its last counter increment. Expired elements are removed with a period specified by the flushInterval argument
+func NewMapLimitStore(expirationTime time.Duration, flushInterval time.Duration) (m *MapLimitStore) {
+	m = &MapLimitStore{
+		data:           make(map[string]limitValue),
+		expirationTime: expirationTime,
+	}
+	go func() {
+		ticker := time.NewTicker(flushInterval)
+		for range ticker.C {
+			var deletedKeys []string
+			for key, val := range m.data {
+				if val.lastUpdate.Before(time.Now().UTC().Add(-m.expirationTime)) {
+					m.mutex.Lock()
+					delete(m.data, key)
+					deletedKeys = append(deletedKeys, key)
+					m.mutex.Unlock()
+				}
+			}
+		}
+	}()
+	return m
+}
+
+// Inc increments current window limit counter for key
+func (m *MapLimitStore) Inc(key string, window time.Time) error {
+	m.mutex.Lock()
+	defer m.mutex.Unlock()
+
+	data := m.data[mapKey(key, window)]
+	data.val++
+	data.lastUpdate = time.Now().UTC()
+	m.data[mapKey(key, window)] = data
+	return nil
+}
+
+// Get gets value of previous window counter and current window counter for key
+func (m *MapLimitStore) Get(key string, previousWindow, currentWindow time.Time) (prevValue int64, currValue int64, err error) {
+	m.mutex.RLock()
+	defer m.mutex.RUnlock()
+	prevValue = m.data[mapKey(key, previousWindow)].val
+	currValue = m.data[mapKey(key, currentWindow)].val
+	return prevValue, currValue, nil
+}
+
+func mapKey(key string, window time.Time) string {
+	return fmt.Sprintf("%s_%s", key, window.Format(time.RFC3339))
+}