├── README ├── varyby.go └── ratelimit.go /README: -------------------------------------------------------------------------------- 1 | webrate 2 | ------- 3 | 4 | 5 | Package webrate contains a modified RateLimit function for 6 | http://github.com/PuerkitoBio/throttled, which rate limits only specified 7 | HTTP methods and doesn't add headers to response. 8 | 9 | 10 | EXAMPLE 11 | 12 | import ( 13 | "net/http" 14 | 15 | "github.com/PuerkitoBio/throttled" 16 | "github.com/PuerkitoBio/throttled/store" 17 | "github.com/dchest/throttled-webrate" 18 | ) 19 | 20 | func main() { 21 | 22 | // ... 23 | 24 | th := webrate.RateLimit( 25 | throttled.PerMin(30), 26 | []string{"POST"}, 27 | webrate.VaryByPathAndIP("X-Real-IP"), 28 | store.NewMemStore(1000)) 29 | 30 | http.Handle("/login", th.Throttle(loginHandler)) 31 | 32 | // ... 33 | } 34 | -------------------------------------------------------------------------------- /varyby.go: -------------------------------------------------------------------------------- 1 | package webrate 2 | 3 | import ( 4 | "net" 5 | "net/http" 6 | 7 | "github.com/PuerkitoBio/throttled" 8 | ) 9 | 10 | // VaryByIP returns a *throttled.VaryBy, which varies request based on client's 11 | // IP address from the given header. If header is empty, extracts IP address 12 | // from r.RemoteAddr. 13 | func VaryByIP(headerName string) *throttled.VaryBy { 14 | return &throttled.VaryBy{ 15 | Custom: func(r *http.Request) string { 16 | return getRequestIP(r, headerName) 17 | }, 18 | } 19 | } 20 | 21 | // VaryByPathAndIP is like VaryByIP but also adds request path. 22 | func VaryByPathAndIP(headerName string) *throttled.VaryBy { 23 | return &throttled.VaryBy{ 24 | Custom: func(r *http.Request) string { 25 | return r.URL.Path + "\n" + getRequestIP(r, headerName) 26 | }, 27 | } 28 | } 29 | 30 | // getRequestIP returns a remote IP address of the client that made the 31 | // request. The address is take from the given header name, or from RemoteAddr 32 | // if the header name is an empty string. 33 | func getRequestIP(r *http.Request, headerName string) string { 34 | if headerName == "" { 35 | return extractIP(r.RemoteAddr) 36 | } 37 | return extractIP(r.Header.Get(headerName)) 38 | } 39 | 40 | // extractIP extracts IP address (or host) from the given string, 41 | // which may have host and port in it. 42 | func extractIP(addr string) string { 43 | ip, _, err := net.SplitHostPort(addr) 44 | if err != nil { 45 | return addr 46 | } 47 | return ip 48 | } 49 | -------------------------------------------------------------------------------- /ratelimit.go: -------------------------------------------------------------------------------- 1 | // Based https://github.com/PuerkitoBio/throttled/blob/master/rate.go: 2 | // 3 | // Copyright (c) 2014, Martin Angers 4 | // All rights reserved. 5 | // 6 | // Redistribution and use in source and binary forms, with or without 7 | // modification, are permitted provided that the following conditions are met: 8 | // 9 | // * Redistributions of source code must retain the above copyright notice, this 10 | // * list of conditions and the following disclaimer. 11 | // 12 | // * Redistributions in binary form must reproduce the above copyright notice, 13 | // this list of conditions and the following disclaimer in the documentation 14 | // and/or other materials provided with the distribution. 15 | // 16 | // * Neither the name of the author nor the names of its contributors may be used 17 | // to endorse or promote products derived from this software without specific 18 | // prior written permission. 19 | // 20 | // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND 21 | // ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED 22 | // WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 23 | // DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 24 | // FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 25 | // DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 26 | // SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 27 | // CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 28 | // OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 29 | // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 30 | // 31 | // -- 32 | // Modifications by Dmitry Chestnykh are in public domain. 33 | 34 | // Package webrate contains a modified RateLimit function for 35 | // http://github.com/PuerkitoBio/throttled, which rate limits only specified 36 | // HTTP methods and doesn't add headers to response. 37 | package webrate 38 | 39 | import ( 40 | "net/http" 41 | "time" 42 | 43 | "github.com/PuerkitoBio/throttled" 44 | ) 45 | 46 | // RateLimit creates a throttler that limits the number of requests allowed 47 | // in a certain time window defined by the Quota q. The q parameter specifies 48 | // the requests per time window, and it is silently set to at least 1 request 49 | // and at least a 1 second window if it is less than that. The time window 50 | // starts when the first request is made outside an existing window. Fractions 51 | // of seconds are not supported, they are truncated. 52 | // 53 | // The vary parameter indicates what criteria should be used to group requests 54 | // for which the limit must be applied (ex.: rate limit based on the remote address). 55 | // See varyby.go for the various options. 56 | // 57 | // The specified store is used to keep track of the request count and the 58 | // time remaining in the window. The throttled package comes with some stores 59 | // in the throttled/store package. Custom stores can be created too, by implementing 60 | // the Store interface. 61 | // 62 | // Requests that bust the rate limit are denied access and go through the denied handler, 63 | // which may be specified on the Throttler and that defaults to the package-global 64 | // variable DefaultDeniedHandler. 65 | func RateLimit(q throttled.Quota, methods []string, vary *throttled.VaryBy, store throttled.Store) *throttled.Throttler { 66 | // Extract requests and window 67 | reqs, win := q.Quota() 68 | 69 | // Create and return the throttler 70 | return throttled.Custom(&rateLimiter{ 71 | reqs: reqs, 72 | window: win, 73 | vary: vary, 74 | store: store, 75 | methods: methods, 76 | }) 77 | } 78 | 79 | // The rate limiter implements limiting the request to a certain quota 80 | // based on the vary-by criteria. State is saved in the store. 81 | type rateLimiter struct { 82 | reqs int 83 | window time.Duration 84 | methods []string 85 | vary *throttled.VaryBy 86 | store throttled.Store 87 | } 88 | 89 | func (r *rateLimiter) hasMethod(method string) bool { 90 | for _, s := range r.methods { 91 | if s == method { 92 | return true 93 | } 94 | } 95 | return false 96 | } 97 | 98 | // Start initializes the limiter for execution. 99 | func (r *rateLimiter) Start() { 100 | if r.reqs < 1 { 101 | r.reqs = 1 102 | } 103 | if r.window < time.Second { 104 | r.window = time.Second 105 | } 106 | } 107 | 108 | // Limit is called for each request to the throttled handler. It checks if 109 | // the request can go through and signals it via the returned channel. 110 | // It returns an error if the operation fails. 111 | func (r *rateLimiter) Limit(w http.ResponseWriter, req *http.Request) (<-chan bool, error) { 112 | // Create return channel and initialize 113 | ch := make(chan bool, 1) 114 | ok := true 115 | key := r.vary.Key(req) 116 | 117 | if r.hasMethod(req.Method) { 118 | // Get the current count and remaining seconds 119 | cnt, secs, err := r.store.Incr(key, r.window) 120 | // Handle the possible situations: error, begin new window, or increment current window. 121 | switch { 122 | case err != nil && err != throttled.ErrNoSuchKey: 123 | // An unexpected error occurred 124 | return nil, err 125 | case err == throttled.ErrNoSuchKey || secs <= 0: 126 | // Reset counter 127 | if err := r.store.Reset(key, r.window); err != nil { 128 | return nil, err 129 | } 130 | cnt = 1 131 | secs = int(r.window.Seconds()) 132 | default: 133 | // If the limit is reached, deny access 134 | if cnt > r.reqs { 135 | ok = false 136 | } 137 | } 138 | } 139 | // Send response via the return channel 140 | ch <- ok 141 | return ch, nil 142 | } 143 | --------------------------------------------------------------------------------