Skip to content

Instantly share code, notes, and snippets.

@elithrar
Last active February 17, 2016 06:14

Revisions

  1. elithrar revised this gist Dec 1, 2013. 1 changed file with 1 addition and 3 deletions.
    4 changes: 1 addition & 3 deletions ratelimit.go
    Original file line number Diff line number Diff line change
    @@ -48,7 +48,7 @@ func NewRateStore(idle int, limit int64, header, net, port, password string) *Ra
    }
    }

    // RateLimit provides HTTP request limiting middleware. Requests are limited to Limit per second.
    // RateLimit provides HTTP request limiting middleware. Requests are limited to Limit per second per IP.
    // Requests that exceed the limit are served with HTTP 429 (Too Many Requests).
    func (s *RateStore) RateLimit(h http.HandlerFunc) http.HandlerFunc {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    @@ -72,8 +72,6 @@ func (s *RateStore) RateLimit(h http.HandlerFunc) http.HandlerFunc {
    return
    }

    logger.Printf("%d\n", current)

    // Set a 1s expiry on newly instantiated counters
    if current.(int64) == 1 {
    _, err := rconn.Do("EXPIRE", path+":"+remoteIP, 1)
  2. elithrar revised this gist Dec 1, 2013. 1 changed file with 14 additions and 11 deletions.
    25 changes: 14 additions & 11 deletions ratelimit.go
    Original file line number Diff line number Diff line change
    @@ -1,26 +1,29 @@
    package main

    import (
    "errors"
    "github.com/garyburd/redigo/redis"
    "net/http"
    "time"
    )

    type RateStore struct {
    Pool *redis.Pool
    Limit int64
    Interval int
    Pool *redis.Pool
    Limit int64
    Header string
    }

    // NewRateStore returns a new RateStore.
    // Depending on your setup or reverse proxy, you will need to set Header to
    // inspect either "REMOTE_ADDR" or "X-Forwarded-For".
    // Example:
    // limitStore = NewRateStore(10, 1, "tcp", ":6380", "")
    // limitStore = NewRateStore(10, 1, "REMOTE_ADDR", "tcp", ":6380", "password")
    //
    // Note: You should spin up a second Redis instance if you already have a primary for other tasks.
    func NewRateStore(limit int, interval int, net, port, password string) *RateStore {
    func NewRateStore(idle int, limit int64, header, net, port, password string) *RateStore {
    return &RateStore{
    Pool: &redis.Pool{
    MaxIdle: limit,
    MaxIdle: idle,
    IdleTimeout: 240 * time.Second,
    Dial: func() (c redis.Conn, err error) {
    c, err = redis.Dial(net, port)
    @@ -40,6 +43,8 @@ func NewRateStore(limit int, interval int, net, port, password string) *RateStor
    return err
    },
    },
    Limit: limit,
    Header: header,
    }
    }

    @@ -52,9 +57,7 @@ func (s *RateStore) RateLimit(h http.HandlerFunc) http.HandlerFunc {
    defer rconn.Close()

    path := r.URL.Path
    // TODO(matt): Allow this to be customised (boolean in struct?) to allow a
    // user to use either r.RemoteAddr, REMOTE_ADDR or other HTTP headers.
    remoteIP := r.Header.Get("REMOTE_ADDR")
    remoteIP := r.Header.Get(s.Header)
    // Invoke the next handler if the remote address is not set
    // (we cannot determine the rate without it)
    if remoteIP == "" {
    @@ -73,14 +76,14 @@ func (s *RateStore) RateLimit(h http.HandlerFunc) http.HandlerFunc {

    // Set a 1s expiry on newly instantiated counters
    if current.(int64) == 1 {
    _, err := rconn.Do("EXPIRE", path+":"+remoteIP, s.Interval)
    _, err := rconn.Do("EXPIRE", path+":"+remoteIP, 1)
    if err != nil {
    serverError(w, r, err, 500)
    return
    }
    } else if current.(int64) > s.Limit {
    // Check if the returned counter exceeds the limit
    serverError(w, r, err, 429)
    serverError(w, r, errors.New("Rate exceeded."), 429)
    return
    }

  3. elithrar revised this gist Dec 1, 2013. 1 changed file with 1 addition and 2 deletions.
    3 changes: 1 addition & 2 deletions ratelimit.go
    Original file line number Diff line number Diff line change
    @@ -54,7 +54,6 @@ func (s *RateStore) RateLimit(h http.HandlerFunc) http.HandlerFunc {
    path := r.URL.Path
    // TODO(matt): Allow this to be customised (boolean in struct?) to allow a
    // user to use either r.RemoteAddr, REMOTE_ADDR or other HTTP headers.
    between r.RemoteAddr/this.
    remoteIP := r.Header.Get("REMOTE_ADDR")
    // Invoke the next handler if the remote address is not set
    // (we cannot determine the rate without it)
    @@ -81,7 +80,7 @@ func (s *RateStore) RateLimit(h http.HandlerFunc) http.HandlerFunc {
    }
    } else if current.(int64) > s.Limit {
    // Check if the returned counter exceeds the limit
    w.WriteHeader(429)
    serverError(w, r, err, 429)
    return
    }

  4. elithrar revised this gist Dec 1, 2013. 1 changed file with 3 additions and 1 deletion.
    4 changes: 3 additions & 1 deletion ratelimit.go
    Original file line number Diff line number Diff line change
    @@ -52,7 +52,9 @@ func (s *RateStore) RateLimit(h http.HandlerFunc) http.HandlerFunc {
    defer rconn.Close()

    path := r.URL.Path
    // TODO(matt): Allow this to be customised (boolean in struct?) to alternate between r.RemoteAddr/this.
    // TODO(matt): Allow this to be customised (boolean in struct?) to allow a
    // user to use either r.RemoteAddr, REMOTE_ADDR or other HTTP headers.
    between r.RemoteAddr/this.
    remoteIP := r.Header.Get("REMOTE_ADDR")
    // Invoke the next handler if the remote address is not set
    // (we cannot determine the rate without it)
  5. elithrar revised this gist Dec 1, 2013. 1 changed file with 6 additions and 1 deletion.
    7 changes: 6 additions & 1 deletion ratelimit.go
    Original file line number Diff line number Diff line change
    @@ -12,6 +12,11 @@ type RateStore struct {
    Interval int
    }

    // NewRateStore returns a new RateStore.
    // Example:
    // limitStore = NewRateStore(10, 1, "tcp", ":6380", "")
    //
    // Note: You should spin up a second Redis instance if you already have a primary for other tasks.
    func NewRateStore(limit int, interval int, net, port, password string) *RateStore {
    return &RateStore{
    Pool: &redis.Pool{
    @@ -87,4 +92,4 @@ func (s *RateStore) RateLimit(h http.HandlerFunc) http.HandlerFunc {
    // Close closes the current connection
    func (s *RateStore) Close() {
    s.Pool.Close()
    }
    }
  6. elithrar revised this gist Dec 1, 2013. 1 changed file with 43 additions and 25 deletions.
    68 changes: 43 additions & 25 deletions ratelimit.go
    Original file line number Diff line number Diff line change
    @@ -6,18 +6,48 @@ import (
    "time"
    )

    var RedisPool *redis.Pool
    var Limit int64 = 10
    type RateStore struct {
    Pool *redis.Pool
    Limit int64
    Interval int
    }

    func NewRateStore(limit int, interval int, net, port, password string) *RateStore {
    return &RateStore{
    Pool: &redis.Pool{
    MaxIdle: limit,
    IdleTimeout: 240 * time.Second,
    Dial: func() (c redis.Conn, err error) {
    c, err = redis.Dial(net, port)
    if err != nil {
    return nil, err
    }
    if password != "" {
    if _, err := c.Do("AUTH", password); err != nil {
    c.Close()
    return nil, err
    }
    }
    return c, err
    },
    TestOnBorrow: func(c redis.Conn, t time.Time) error {
    _, err := c.Do("PING")
    return err
    },
    },
    }
    }

    // RateLimit provides HTTP request limiting middleware. Requests are limited to Limit per second.
    // Requests that exceed the limit are served with HTTP 429 (Too Many Requests).
    func RateLimit(h http.HandlerFunc) http.HandlerFunc {
    func (s *RateStore) RateLimit(h http.HandlerFunc) http.HandlerFunc {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {

    rconn := RedisPool.Get()
    rconn := s.Pool.Get()
    defer rconn.Close()

    path := r.URL.Path
    // TODO(matt): Allow this to be customised (boolean in struct?) to alternate between r.RemoteAddr/this.
    remoteIP := r.Header.Get("REMOTE_ADDR")
    // Invoke the next handler if the remote address is not set
    // (we cannot determine the rate without it)
    @@ -26,19 +56,23 @@ func RateLimit(h http.HandlerFunc) http.HandlerFunc {
    return
    }

    // INCR will increment an existing key (if any) else it creates a new one (at 1)
    current, err := rconn.Do("INCR", path+":"+remoteIP)
    if err != nil {
    serverError(w, r, err, 500)
    return
    }

    logger.Printf("%d\n", current)

    // Set a 1s expiry on newly instantiated counters
    if current.(int64) == 1 {
    _, err := rconn.Do("EXPIRE", path+":"+remoteIP, 1)
    _, err := rconn.Do("EXPIRE", path+":"+remoteIP, s.Interval)
    if err != nil {
    serverError(w, r, err, 500)
    return
    }
    } else if current.(int64) > 10 {
    } else if current.(int64) > s.Limit {
    // Check if the returned counter exceeds the limit
    w.WriteHeader(429)
    return
    @@ -50,23 +84,7 @@ func RateLimit(h http.HandlerFunc) http.HandlerFunc {
    })
    }

    func init() {
    RedisPool = &redis.Pool{
    MaxIdle: 10,
    IdleTimeout: 240 * time.Second,
    Dial: func() (c redis.Conn, err error) {
    c, err = redis.Dial("tcp", ":6380")
    if err != nil {
    return nil, err
    }
    return c, err
    },
    TestOnBorrow: func(c redis.Conn, t time.Time) error {
    _, err := c.Do("PING")
    return err
    },
    }

    c := RedisPool.Get()
    defer c.Close()
    // Close closes the current connection
    func (s *RateStore) Close() {
    s.Pool.Close()
    }
  7. elithrar created this gist Dec 1, 2013.
    72 changes: 72 additions & 0 deletions ratelimit.go
    Original file line number Diff line number Diff line change
    @@ -0,0 +1,72 @@
    package main

    import (
    "github.com/garyburd/redigo/redis"
    "net/http"
    "time"
    )

    var RedisPool *redis.Pool
    var Limit int64 = 10

    // RateLimit provides HTTP request limiting middleware. Requests are limited to Limit per second.
    // Requests that exceed the limit are served with HTTP 429 (Too Many Requests).
    func RateLimit(h http.HandlerFunc) http.HandlerFunc {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {

    rconn := RedisPool.Get()
    defer rconn.Close()

    path := r.URL.Path
    remoteIP := r.Header.Get("REMOTE_ADDR")
    // Invoke the next handler if the remote address is not set
    // (we cannot determine the rate without it)
    if remoteIP == "" {
    h.ServeHTTP(w, r)
    return
    }

    current, err := rconn.Do("INCR", path+":"+remoteIP)
    if err != nil {
    serverError(w, r, err, 500)
    return
    }

    if current.(int64) == 1 {
    _, err := rconn.Do("EXPIRE", path+":"+remoteIP, 1)
    if err != nil {
    serverError(w, r, err, 500)
    return
    }
    } else if current.(int64) > 10 {
    // Check if the returned counter exceeds the limit
    w.WriteHeader(429)
    return
    }

    // Invoke the next handler if we haven't hit the limit
    h.ServeHTTP(w, r)
    return
    })
    }

    func init() {
    RedisPool = &redis.Pool{
    MaxIdle: 10,
    IdleTimeout: 240 * time.Second,
    Dial: func() (c redis.Conn, err error) {
    c, err = redis.Dial("tcp", ":6380")
    if err != nil {
    return nil, err
    }
    return c, err
    },
    TestOnBorrow: func(c redis.Conn, t time.Time) error {
    _, err := c.Do("PING")
    return err
    },
    }

    c := RedisPool.Get()
    defer c.Close()
    }