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
}
Example #2
0
// 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)
}
Example #3
0
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
}
Example #5
0
// 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)
}