func (tl *TokenLimiter) ProcessRequest(r Request) (*http.Response, error) {
tl.mutex.Lock()
defer tl.mutex.Unlock()
token, amount, err := tl.mapper(r)
if err != nil {
return nil, err
}
bucketI, exists := tl.buckets.Get(token)
if !exists {
bucketI, err = NewTokenBucket(tl.rate, tl.options.Burst+1, tl.options.TimeProvider)
if err != nil {
return nil, err
}
// We set ttl as 10 times rate period. E.g. if rate is 100 requests/second per client ip
// the counters for this ip will expire after 10 seconds of inactivity
tl.buckets.Set(token, bucketI, int(tl.rate.Period/time.Second)*10+1)
}
bucket := bucketI.(*TokenBucket)
delay, err := bucket.Consume(amount)
if err != nil {
return nil, err
}
if delay > 0 {
return netutils.NewTextResponse(r.GetHttpRequest(), errors.StatusTooManyRequests, "Too many requests"), nil
}
return nil, nil
}
// Test scenario when middleware intercepts the request
func (s *LocSuite) TestMiddlewareInterceptsRequest(c *C) {
server := NewTestServer(func(w http.ResponseWriter, r *http.Request) {
w.Write([]byte("Hi, I'm endpoint"))
})
defer server.Close()
location, proxy := s.newProxy(s.newRoundRobin(server.URL))
defer proxy.Close()
calls := make(map[string]int)
auth := &MiddlewareWrapper{
OnRequest: func(r Request) (*http.Response, error) {
calls["authReq"] += 1
return netutils.NewTextResponse(
r.GetHttpRequest(),
http.StatusForbidden,
"Intercepted Request"), nil
},
OnResponse: func(r Request, a Attempt) {
calls["authRe"] += 1
},
}
location.GetMiddlewareChain().Add("auth", 0, auth)
response, bodyBytes, err := MakeRequest(proxy.URL, Opts{})
c.Assert(err, IsNil)
c.Assert(response.StatusCode, Equals, http.StatusForbidden)
c.Assert(string(bodyBytes), Equals, "Intercepted Request")
// Auth middleware has been called on response as well
c.Assert(calls["authReq"], Equals, 1)
c.Assert(calls["authRe"], Equals, 1)
}
func (tl *TokenLimiter) ProcessRequest(r request.Request) (*http.Response, error) {
tl.mutex.Lock()
defer tl.mutex.Unlock()
token, amount, err := tl.mapper(r)
if err != nil {
return nil, err
}
effectiveRates := tl.effectiveRates(r)
bucketSetI, exists := tl.bucketSets.Get(token)
var bucketSet *tokenBucketSet
if exists {
bucketSet = bucketSetI.(*tokenBucketSet)
bucketSet.update(effectiveRates)
} else {
bucketSet = newTokenBucketSet(effectiveRates, tl.clock)
// We set ttl as 10 times rate period. E.g. if rate is 100 requests/second per client ip
// the counters for this ip will expire after 10 seconds of inactivity
tl.bucketSets.Set(token, bucketSet, int(bucketSet.maxPeriod/time.Second)*10+1)
}
delay, err := bucketSet.consume(amount)
if err != nil {
return nil, err
}
if delay > 0 {
return netutils.NewTextResponse(r.GetHttpRequest(), errors.StatusTooManyRequests, "Too many requests"), nil
}
return nil, nil
}
func (cl *ConnectionLimiter) ProcessRequest(r Request) (*http.Response, error) {
cl.mutex.Lock()
defer cl.mutex.Unlock()
token, amount, err := cl.mapper(r)
if err != nil {
return nil, err
}
connections := cl.connections[token]
if connections >= cl.maxConnections {
return netutils.NewTextResponse(
r.GetHttpRequest(),
errors.StatusTooManyRequests,
fmt.Sprintf("Connection limit reached. Max is: %d, yours: %d", cl.maxConnections, connections)), nil
}
cl.connections[token] += amount
cl.totalConnections += int64(amount)
return nil, nil
}
// Test scenario when middleware intercepts the request
func (s *LocSuite) TestMultipleMiddlewaresRequestIntercepted(c *C) {
server := NewTestServer(func(w http.ResponseWriter, r *http.Request) {
w.Write([]byte("Hi, I'm endpoint"))
})
defer server.Close()
location, proxy := s.newProxy(s.newRoundRobin(server.URL))
defer proxy.Close()
calls := make(map[string]int)
auth := &MiddlewareWrapper{
OnRequest: func(r Request) (*http.Response, error) {
calls["authReq"] += 1
return netutils.NewTextResponse(
r.GetHttpRequest(),
http.StatusForbidden,
"Intercepted Request"), nil
},
OnResponse: func(r Request, a Attempt) {
calls["authRe"] += 1
},
}
cb := &MiddlewareWrapper{
OnRequest: func(r Request) (*http.Response, error) {
calls["cbReq"] += 1
return nil, nil
},
OnResponse: func(r Request, a Attempt) {
calls["cbRe"] += 1
},
}
observer := &ObserverWrapper{
OnRequest: func(r Request) {
calls["oReq"] += 1
},
OnResponse: func(r Request, a Attempt) {
calls["oRe"] += 1
},
}
location.GetMiddlewareChain().Add("auth", 0, auth)
location.GetMiddlewareChain().Add("cb", 1, cb)
location.GetObserverChain().Add("ob", observer)
response, bodyBytes, err := MakeRequest(proxy.URL, Opts{})
c.Assert(err, IsNil)
c.Assert(response.StatusCode, Equals, http.StatusForbidden)
c.Assert(string(bodyBytes), Equals, "Intercepted Request")
// Auth middleware has been called on response as well
c.Assert(calls["authReq"], Equals, 1)
c.Assert(calls["authRe"], Equals, 1)
// Callback has never got to a request, because it was intercepted
c.Assert(calls["cbReq"], Equals, 0)
c.Assert(calls["cbRe"], Equals, 0)
// Observer was called regardless
c.Assert(calls["oReq"], Equals, 1)
c.Assert(calls["oRe"], Equals, 1)
}