func TestMemStoreLRU(t *testing.T) {
st, err := memstore.New(10)
if err != nil {
t.Fatal(err)
}
storetest.TestGCRAStore(t, st)
}
func BenchmarkMemStoreUnlimited(b *testing.B) {
st, err := memstore.New(0)
if err != nil {
b.Fatal(err)
}
storetest.BenchmarkGCRAStore(b, st)
}
// Demonstrates direct use of GCRARateLimiter's RateLimit function (and the
// more general RateLimiter interface). This should be used anywhere where
// granular control over rate limiting is required.
func ExampleGCRARateLimiter() {
store, err := memstore.New(65536)
if err != nil {
log.Fatal(err)
}
// Maximum burst of 5 which refills at 1 token per hour.
quota := throttled.RateQuota{throttled.PerHour(1), 5}
rateLimiter, err := throttled.NewGCRARateLimiter(store, quota)
if err != nil {
log.Fatal(err)
}
// Bucket according to the number i / 10 (so 1 falls into the bucket 0
// while 11 falls into the bucket 1). This has the effect of allowing a
// burst of 5 plus 1 (a single emission interval) on every ten iterations
// of the loop. See the output for better clarity here.
//
// We also refill the bucket at 1 token per hour, but that has no effect
// for the purposes of this example.
for i := 0; i < 20; i++ {
bucket := fmt.Sprintf("by-order:%v", i/10)
limited, result, err := rateLimiter.RateLimit(bucket, 1)
if err != nil {
log.Fatal(err)
}
if limited {
fmt.Printf("Iteration %2v; bucket %v: FAILED. Rate limit exceeded.\n",
i, bucket)
} else {
fmt.Printf("Iteration %2v; bucket %v: Operation successful (remaining=%v).\n",
i, bucket, result.Remaining)
}
}
// Output:
// Iteration 0; bucket by-order:0: Operation successful (remaining=5).
// Iteration 1; bucket by-order:0: Operation successful (remaining=4).
// Iteration 2; bucket by-order:0: Operation successful (remaining=3).
// Iteration 3; bucket by-order:0: Operation successful (remaining=2).
// Iteration 4; bucket by-order:0: Operation successful (remaining=1).
// Iteration 5; bucket by-order:0: Operation successful (remaining=0).
// Iteration 6; bucket by-order:0: FAILED. Rate limit exceeded.
// Iteration 7; bucket by-order:0: FAILED. Rate limit exceeded.
// Iteration 8; bucket by-order:0: FAILED. Rate limit exceeded.
// Iteration 9; bucket by-order:0: FAILED. Rate limit exceeded.
// Iteration 10; bucket by-order:1: Operation successful (remaining=5).
// Iteration 11; bucket by-order:1: Operation successful (remaining=4).
// Iteration 12; bucket by-order:1: Operation successful (remaining=3).
// Iteration 13; bucket by-order:1: Operation successful (remaining=2).
// Iteration 14; bucket by-order:1: Operation successful (remaining=1).
// Iteration 15; bucket by-order:1: Operation successful (remaining=0).
// Iteration 16; bucket by-order:1: FAILED. Rate limit exceeded.
// Iteration 17; bucket by-order:1: FAILED. Rate limit exceeded.
// Iteration 18; bucket by-order:1: FAILED. Rate limit exceeded.
// Iteration 19; bucket by-order:1: FAILED. Rate limit exceeded.
}
func throttledHandler(h http.HandlerFunc) http.Handler {
// for now just use an inmemory storage, so per server. In the future we
// can change to store the state on a remote DB if we want to distribute
// the counts
store, err := memstore.New(65536)
if err != nil {
// panics only if memstore.New() receives an integer number, so this is
// OK, this means it's a human error and needs to be fixed
log.Fatal(err)
}
// Based on datadog metrics, kloud.info is called on average 200
// req/minute.
quota := throttled.RateQuota{
MaxRate: throttled.PerMin(200),
MaxBurst: 300,
}
rateLimiter, err := throttled.NewGCRARateLimiter(store, quota)
if err != nil {
// we exit because this is code error and must be handled
log.Fatalln(err)
}
httpRateLimiter := throttled.HTTPRateLimiter{
RateLimiter: rateLimiter,
}
return httpRateLimiter.RateLimit(http.HandlerFunc(h))
}
// DEPRECATED. NewMemStore is a compatible alias for mem.New
func NewMemStore(maxKeys int) *memstore.MemStore {
st, err := memstore.New(maxKeys)
if err != nil {
// As of this writing, `lru.New` can only return an error if you pass
// maxKeys <= 0 so this should never occur.
panic(err)
}
return st
}
func TestRateLimitUpdateFailures(t *testing.T) {
rq := throttled.RateQuota{throttled.PerSec(1), 1}
mst, err := memstore.New(0)
if err != nil {
t.Fatal(err)
}
st := testStore{store: mst, failUpdates: true}
rl, err := throttled.NewGCRARateLimiter(&st, rq)
if err != nil {
t.Fatal(err)
}
if _, _, err := rl.RateLimit("foo", 1); err == nil {
t.Error("Expected limiting to fail when store updates fail")
}
}
func buildLimiter(quota throttled.RateQuota) (*throttled.HTTPRateLimiter, error) {
store, err := memstore.New(65536)
if err != nil {
return nil, err
}
rateLimiter, err := throttled.NewGCRARateLimiter(store, quota)
if err != nil {
return nil, err
}
return &throttled.HTTPRateLimiter{
RateLimiter: rateLimiter,
VaryBy: &throttled.VaryBy{RemoteAddr: true},
}, nil
}
func StartServer() {
l4g.Info(utils.T("api.server.start_server.starting.info"))
l4g.Info(utils.T("api.server.start_server.listening.info"), utils.Cfg.ServiceSettings.ListenAddress)
var handler http.Handler = &CorsWrapper{Srv.Router}
if *utils.Cfg.RateLimitSettings.Enable {
l4g.Info(utils.T("api.server.start_server.rate.info"))
store, err := memstore.New(utils.Cfg.RateLimitSettings.MemoryStoreSize)
if err != nil {
l4g.Critical(utils.T("api.server.start_server.rate_limiting_memory_store"))
return
}
quota := throttled.RateQuota{
MaxRate: throttled.PerSec(utils.Cfg.RateLimitSettings.PerSec),
MaxBurst: *utils.Cfg.RateLimitSettings.MaxBurst,
}
rateLimiter, err := throttled.NewGCRARateLimiter(store, quota)
if err != nil {
l4g.Critical(utils.T("api.server.start_server.rate_limiting_rate_limiter"))
return
}
httpRateLimiter := throttled.HTTPRateLimiter{
RateLimiter: rateLimiter,
VaryBy: &VaryBy{},
DeniedHandler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
l4g.Error("%v: Denied due to throttling settings code=429 ip=%v", r.URL.Path, GetIpAddress(r))
throttled.DefaultDeniedHandler.ServeHTTP(w, r)
}),
}
handler = httpRateLimiter.RateLimit(handler)
}
go func() {
err := manners.ListenAndServe(utils.Cfg.ServiceSettings.ListenAddress, handlers.RecoveryHandler(handlers.PrintRecoveryStack(true))(handler))
if err != nil {
l4g.Critical(utils.T("api.server.start_server.starting.critical"), err)
time.Sleep(time.Second)
}
}()
}
// WithRateLimit wraps an ctxhttp.Handler to limit incoming requests.
// Requests that are not limited will be passed to the handler
// unchanged. Limited requests will be passed to the DeniedHandler.
// X-RateLimit-Limit, X-RateLimit-Remaining, X-RateLimit-Reset and
// Retry-After headers will be written to the response based on the
// values in the RateLimitResult.
func (t *HTTPRateLimit) WithRateLimit(rlStore throttled.GCRAStore, h ctxhttp.Handler) ctxhttp.Handler {
if t.Config == nil {
t.Config = config.DefaultManager
}
if t.DeniedHandler == nil {
t.DeniedHandler = DefaultDeniedHandler
}
if t.RateLimiter == nil {
if rlStore == nil {
var err error
rlStore, err = memstore.New(65536)
if err != nil {
panic(err)
}
}
var err error
t.RateLimiter, err = throttled.NewGCRARateLimiter(rlStore, t.quota())
if err != nil {
panic(err)
}
}
return ctxhttp.HandlerFunc(func(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
var k string
if t.VaryBy != nil {
k = t.VaryBy.Key(r)
}
limited, context, err := t.RateLimiter.RateLimit(k, 1)
if err != nil {
return err
}
setRateLimitHeaders(w, context)
if !limited {
return h.ServeHTTPContext(ctx, w, r)
}
return t.DeniedHandler.ServeHTTPContext(ctx, w, r)
})
}
// ExampleHTTPRateLimiter demonstrates the usage of HTTPRateLimiter
// for rate-limiting access to an http.Handler to 20 requests per path
// per minute with a maximum burst of 5 requests.
func ExampleHTTPRateLimiter() {
store, err := memstore.New(65536)
if err != nil {
log.Fatal(err)
}
quota := throttled.RateQuota{throttled.PerMin(20), 5}
rateLimiter, err := throttled.NewGCRARateLimiter(store, quota)
if err != nil {
log.Fatal(err)
}
httpRateLimiter := throttled.HTTPRateLimiter{
RateLimiter: rateLimiter,
VaryBy: &throttled.VaryBy{Path: true},
}
http.ListenAndServe(":8080", httpRateLimiter.RateLimit(myHandler))
}
func throttle(next http.Handler, o ServerOptions) http.Handler {
store, err := memstore.New(65536)
if err != nil {
return throttleError(err)
}
quota := throttled.RateQuota{throttled.PerSec(o.Concurrency), o.Burst}
rateLimiter, err := throttled.NewGCRARateLimiter(store, quota)
if err != nil {
return throttleError(err)
}
httpRateLimiter := throttled.HTTPRateLimiter{
RateLimiter: rateLimiter,
VaryBy: &throttled.VaryBy{Method: true},
}
return httpRateLimiter.RateLimit(next)
}
// Run is the main function
func Run(cmd *cobra.Command, args []string) {
database.InitDB()
address := viper.GetString("address")
// Throttling control
store, err := memstore.New(65536)
if err != nil {
logrus.Fatal(err)
}
quota := throttled.RateQuota{throttled.PerMin(20), 5}
rateLimiter, err := throttled.NewGCRARateLimiter(store, quota)
if err != nil {
logrus.Fatal(err)
}
httpRateLimiter := throttled.HTTPRateLimiter{
RateLimiter: rateLimiter,
VaryBy: &throttled.VaryBy{Path: true},
}
mux := http.NewServeMux()
mux.HandleFunc("/authenticate/free", handleFreeCookie)
mux.HandleFunc("/authenticate/verify", handleAuthenticate)
mux.Handle("/", http.StripPrefix("/authenticate/", http.FileServer(http.Dir("./static"))))
database.InitDB()
if buildst == "none" {
buildst = "[This Dev version is not compiled using regular procedure]"
}
logrus.Info("Starting App, ", Version)
logrus.Infof("2FA HTTP layer listening on %s", address)
logrus.Infof("Domain for cookies is %s", viper.GetString("domain"))
logrus.Infof("Cookie max age is %d hour(s)", viper.GetInt("cookiemaxage"))
logrus.Info("Starting instance on ", time.Now())
if err := http.ListenAndServe(address, httpRateLimiter.RateLimit(mux)); err != nil {
logrus.Fatal("Unable to create HTTP layer", err)
}
}
func NewWallClockLimiter(burst, rate time.Duration, clock Clock) (Limiter, error) {
var l wallClockLimiter
l.estimates = make(map[string]time.Duration)
store, err := memstore.New(65536)
if err != nil {
return nil, err
}
quota := throttled.RateQuota{
throttled.PerSec(int(rate / time.Millisecond)),
int(burst / time.Millisecond),
}
l.limiter, err = throttled.NewGCRARateLimiter(&clockedMemStore{store, clock}, quota)
if err != nil {
return nil, err
}
return &l, nil
}
// NewDefaultRateLimiter creates rate limiter with sane configuration for koding
func NewDefaultRateLimiter() *throttled.HTTPRateLimiter {
memStore, err := memstore.New(65536)
if err != nil {
// errors only for non positve numbers, so no worries :)
log.Fatal(err)
}
quota := throttled.RateQuota{
MaxRate: throttled.PerSec(11),
MaxBurst: 12,
}
rateLimiter, err := throttled.NewGCRARateLimiter(memStore, quota)
if err != nil {
// we exit because this is code error and must be handled
log.Fatalln(err)
}
return &throttled.HTTPRateLimiter{
RateLimiter: rateLimiter,
VaryBy: &throttled.VaryBy{Cookies: []string{"clientId"}},
}
}
func TestRateLimit(t *testing.T) {
limit := 5
rq := throttled.RateQuota{throttled.PerSec(1), limit - 1}
start := time.Unix(0, 0)
cases := []struct {
now time.Time
volume, remaining int
reset, retry time.Duration
limited bool
}{
// You can never make a request larger than the maximum
0: {start, 6, 5, 0, -1, true},
// Rate limit normal requests appropriately
1: {start, 1, 4, time.Second, -1, false},
2: {start, 1, 3, 2 * time.Second, -1, false},
3: {start, 1, 2, 3 * time.Second, -1, false},
4: {start, 1, 1, 4 * time.Second, -1, false},
5: {start, 1, 0, 5 * time.Second, -1, false},
6: {start, 1, 0, 5 * time.Second, time.Second, true},
7: {start.Add(3000 * time.Millisecond), 1, 2, 3000 * time.Millisecond, -1, false},
8: {start.Add(3100 * time.Millisecond), 1, 1, 3900 * time.Millisecond, -1, false},
9: {start.Add(4000 * time.Millisecond), 1, 1, 4000 * time.Millisecond, -1, false},
10: {start.Add(8000 * time.Millisecond), 1, 4, 1000 * time.Millisecond, -1, false},
11: {start.Add(9500 * time.Millisecond), 1, 4, 1000 * time.Millisecond, -1, false},
// Zero-volume request just peeks at the state
12: {start.Add(9500 * time.Millisecond), 0, 4, time.Second, -1, false},
// High-volume request uses up more of the limit
13: {start.Add(9500 * time.Millisecond), 2, 2, 3 * time.Second, -1, false},
// Large requests cannot exceed limits
14: {start.Add(9500 * time.Millisecond), 5, 2, 3 * time.Second, 3 * time.Second, true},
}
mst, err := memstore.New(0)
if err != nil {
t.Fatal(err)
}
st := testStore{store: mst}
rl, err := throttled.NewGCRARateLimiter(&st, rq)
if err != nil {
t.Fatal(err)
}
// Start the server
for i, c := range cases {
st.clock = c.now
limited, context, err := rl.RateLimit("foo", c.volume)
if err != nil {
t.Fatalf("%d: %#v", i, err)
}
if limited != c.limited {
t.Errorf("%d: expected Limited to be %t but got %t", i, c.limited, limited)
}
if have, want := context.Limit, limit; have != want {
t.Errorf("%d: expected Limit to be %d but got %d", i, want, have)
}
if have, want := context.Remaining, c.remaining; have != want {
t.Errorf("%d: expected Remaining to be %d but got %d", i, want, have)
}
if have, want := context.ResetAfter, c.reset; have != want {
t.Errorf("%d: expected ResetAfter to be %s but got %s", i, want, have)
}
if have, want := context.RetryAfter, c.retry; have != want {
t.Errorf("%d: expected RetryAfter to be %d but got %d", i, want, have)
}
}
}
func main() {
maxProcs := runtime.NumCPU()
runtime.GOMAXPROCS(maxProcs)
var (
listen_host string
fs_basepath string
fs_cachepath string
http_origin string
error_imgpath string
secret string
rate_limit_by_path int
rate_limit_bursts int
concurrency int
newrelic_key string
)
flag.StringVar(&listen_host, "listen", "localhost:3000", "HTTP host:port to listen for incoming requests")
flag.StringVar(&fs_basepath, "basepath", "", "File system base path to lookup images")
flag.StringVar(&fs_cachepath, "cachepath", "", "File system base path to cache images")
flag.StringVar(&http_origin, "httporigin", "", "HTTP endpoint to lookup origin images")
flag.StringVar(&error_imgpath, "errorimgpath", "./fixtures/error.png", "Error image path in local filesystem")
flag.StringVar(&secret, "secret", "", "Secret used to check image signature")
flag.StringVar(&newrelic_key, "newrelic", "", "NewRelic application key")
flag.IntVar(&rate_limit_by_path, "ratelimit", 0, "Rate limit. Concurrent images per request path")
flag.IntVar(&rate_limit_bursts, "bursts", 5, "Rate limits bursts")
flag.IntVar(&concurrency, "concurrency", 0, "Max concurrent requests")
flag.Parse()
var store btcdn.Store
var storeErr error
var newrelic *gorelic.Agent
if fs_basepath != "" {
store, storeErr = btcdn.NewFsStore(fs_basepath)
} else {
store, storeErr = btcdn.NewHttpStore(http_origin)
}
if storeErr != nil {
log.Fatal(storeErr)
}
pr, err := processor.New()
if err != nil {
log.Fatal(err)
}
imgSrv, err := btcdn.NewImageServer(store, pr, error_imgpath)
if err != nil {
log.Fatal(err)
}
router := mux.NewRouter()
var srv http.Handler
srv = router
if secret == "" { // no security
log.Println("No secret. Using ImageServer with no DoS protection")
router.HandleFunc("/", imgSrv.ServeHTTP)
} else { // secure server
log.Println("Secret provided. Using ImageServer with DoS protection")
srv = btcdn.NewSecureServer(imgSrv, secret, router)
}
// Only cache named sizes
// to the filesystem for now
if fs_cachepath != "" {
log.Println("Caching named sizes to", fs_cachepath)
cachedEndpoint := btcdn.NewCachedEndpoint(imgSrv, fs_cachepath)
srv = btcdn.NewNamedSizesServer(cachedEndpoint, router)
} else {
srv = btcdn.NewNamedSizesServer(imgSrv, router)
}
// Setup rate-limited server
if rate_limit_by_path > 0 {
store, err := tstore.New(65536)
if err != nil {
log.Fatal(err)
}
quota := throttled.RateQuota{throttled.PerSec(rate_limit_by_path), rate_limit_bursts}
rateLimiter, err := throttled.NewGCRARateLimiter(store, quota)
if err != nil {
log.Fatal(err)
}
httpRateLimiter := throttled.HTTPRateLimiter{
RateLimiter: rateLimiter,
VaryBy: &throttled.VaryBy{Path: true},
}
log.Println("Rate limiting by path with ", rate_limit_by_path, "reqs. per second with bursts of ", rate_limit_bursts)
srv = httpRateLimiter.RateLimit(srv)
}
// http.HandleFunc("/favicon.ico", FaviconHandler)
var concurrencyPrompt string
if concurrency > 0 {
srv = httpool.Wrap(srv, concurrency)
concurrencyPrompt = fmt.Sprintf("concurrency: %d", concurrency)
}
if newrelic_key != "" {
newrelic = gorelic.NewAgent()
newrelic.NewrelicLicense = newrelic_key
newrelic.NewrelicName = "Btcdn (Go image resizer"
newrelic.CollectHTTPStat = true
newrelic.Run()
srv = newrelic.WrapHTTPHandler(srv)
log.Println("Newrelic support enabled")
}
log.Println("Serving on", listen_host, "cores:", maxProcs, "store:", concurrencyPrompt, store.String(), "error img:", error_imgpath)
log.Fatal(http.ListenAndServe(listen_host, srv))
}
func StartServer() {
l4g.Info(utils.T("api.server.start_server.starting.info"))
var handler http.Handler = &CorsWrapper{Srv.Router}
if *utils.Cfg.RateLimitSettings.Enable {
l4g.Info(utils.T("api.server.start_server.rate.info"))
store, err := memstore.New(utils.Cfg.RateLimitSettings.MemoryStoreSize)
if err != nil {
l4g.Critical(utils.T("api.server.start_server.rate_limiting_memory_store"))
return
}
quota := throttled.RateQuota{
MaxRate: throttled.PerSec(utils.Cfg.RateLimitSettings.PerSec),
MaxBurst: *utils.Cfg.RateLimitSettings.MaxBurst,
}
rateLimiter, err := throttled.NewGCRARateLimiter(store, quota)
if err != nil {
l4g.Critical(utils.T("api.server.start_server.rate_limiting_rate_limiter"))
return
}
httpRateLimiter := throttled.HTTPRateLimiter{
RateLimiter: rateLimiter,
VaryBy: &VaryBy{},
DeniedHandler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
l4g.Error("%v: Denied due to throttling settings code=429 ip=%v", r.URL.Path, GetIpAddress(r))
throttled.DefaultDeniedHandler.ServeHTTP(w, r)
}),
}
handler = httpRateLimiter.RateLimit(handler)
}
Srv.GracefulServer = &graceful.Server{
Timeout: TIME_TO_WAIT_FOR_CONNECTIONS_TO_CLOSE_ON_SERVER_SHUTDOWN,
Server: &http.Server{
Addr: utils.Cfg.ServiceSettings.ListenAddress,
Handler: handlers.RecoveryHandler(handlers.PrintRecoveryStack(true))(handler),
ReadTimeout: time.Duration(*utils.Cfg.ServiceSettings.ReadTimeout) * time.Second,
WriteTimeout: time.Duration(*utils.Cfg.ServiceSettings.WriteTimeout) * time.Second,
},
}
l4g.Info(utils.T("api.server.start_server.listening.info"), utils.Cfg.ServiceSettings.ListenAddress)
if *utils.Cfg.ServiceSettings.Forward80To443 {
go func() {
listener, err := net.Listen("tcp", ":80")
if err != nil {
l4g.Error("Unable to setup forwarding")
return
}
defer listener.Close()
http.Serve(listener, http.HandlerFunc(redirectHTTPToHTTPS))
}()
}
go func() {
var err error
if *utils.Cfg.ServiceSettings.ConnectionSecurity == model.CONN_SECURITY_TLS {
if *utils.Cfg.ServiceSettings.UseLetsEncrypt {
var m letsencrypt.Manager
m.CacheFile(*utils.Cfg.ServiceSettings.LetsEncryptCertificateCacheFile)
tlsConfig := &tls.Config{
GetCertificate: m.GetCertificate,
}
tlsConfig.NextProtos = append(tlsConfig.NextProtos, "h2")
err = Srv.GracefulServer.ListenAndServeTLSConfig(tlsConfig)
} else {
err = Srv.GracefulServer.ListenAndServeTLS(*utils.Cfg.ServiceSettings.TLSCertFile, *utils.Cfg.ServiceSettings.TLSKeyFile)
}
} else {
err = Srv.GracefulServer.ListenAndServe()
}
if err != nil {
l4g.Critical(utils.T("api.server.start_server.starting.critical"), err)
time.Sleep(time.Second)
}
}()
}
