func (s *RewriteSuite) TestMultipleHeaders(c *C) { var outURL string handler := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { outURL = rawURL(req) w.Write([]byte("hello")) }) rh, err := newRewriteHandler(handler, &Rewrite{ "^http://localhost/(foo)/(bar)$", `http://localhost/$1/{{.Request.Header.Get "X-Header"}}/$2/{{.Request.Header.Get "Y-Header"}}`, false, false}) c.Assert(rh, NotNil) c.Assert(err, IsNil) srv := httptest.NewServer(rh) defer srv.Close() re, _, err := testutils.Get(srv.URL+"/foo/bar", testutils.Host("localhost"), testutils.Header("X-Header", "baz"), testutils.Header("Y-Header", "bam")) c.Assert(err, IsNil) c.Assert(re.StatusCode, Equals, http.StatusOK) c.Assert(outURL, Equals, "http://localhost/foo/baz/bar/bam") }
// We've hit the limit and were able to proceed on the next time run func (s *LimiterSuite) TestHitLimit(c *C) { handler := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { w.Write([]byte("hello")) }) rates := NewRateSet() rates.Add(time.Second, 1, 1) l, err := New(handler, headerLimit, rates, Clock(s.clock)) c.Assert(err, IsNil) srv := httptest.NewServer(l) defer srv.Close() re, _, err := testutils.Get(srv.URL, testutils.Header("Source", "a")) c.Assert(err, IsNil) c.Assert(re.StatusCode, Equals, http.StatusOK) // Next request from the same source hits rate limit re, _, err = testutils.Get(srv.URL, testutils.Header("Source", "a")) c.Assert(err, IsNil) c.Assert(re.StatusCode, Equals, 429) // Second later, the request from this ip will succeed s.clock.Sleep(time.Second) re, _, err = testutils.Get(srv.URL, testutils.Header("Source", "a")) c.Assert(err, IsNil) c.Assert(re.StatusCode, Equals, http.StatusOK) }
// We've hit the limit and were able to proceed on the next time run func (s *LimiterSuite) TestOptions(c *C) { handler := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { w.Write([]byte("hello")) }) rates := NewRateSet() rates.Add(time.Second, 1, 1) errHandler := utils.ErrorHandlerFunc(func(w http.ResponseWriter, req *http.Request, err error) { w.WriteHeader(http.StatusTeapot) w.Write([]byte(http.StatusText(http.StatusTeapot))) }) buf := &bytes.Buffer{} log := utils.NewFileLogger(buf, utils.INFO) l, err := New(handler, headerLimit, rates, ErrorHandler(errHandler), Logger(log), Clock(s.clock)) c.Assert(err, IsNil) srv := httptest.NewServer(l) defer srv.Close() re, _, err := testutils.Get(srv.URL, testutils.Header("Source", "a")) c.Assert(err, IsNil) c.Assert(re.StatusCode, Equals, http.StatusOK) re, _, err = testutils.Get(srv.URL, testutils.Header("Source", "a")) c.Assert(err, IsNil) c.Assert(re.StatusCode, Equals, http.StatusTeapot) c.Assert(len(buf.String()), Not(Equals), 0) }
// If configMapper returns empty rates, then the default rate is applied. func (s *LimiterSuite) TestExtractorEmpty(c *C) { // Given extractor := func(*http.Request) (*RateSet, error) { return NewRateSet(), nil } rates := NewRateSet() rates.Add(time.Second, 1, 1) handler := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { w.Write([]byte("hello")) }) l, err := New(handler, headerLimit, rates, Clock(s.clock), ExtractRates(RateExtractorFunc(extractor))) c.Assert(err, IsNil) srv := httptest.NewServer(l) defer srv.Close() // When/Then: The default rate is applied, which 1 req/second re, _, err := testutils.Get(srv.URL, testutils.Header("Source", "a")) c.Assert(err, IsNil) c.Assert(re.StatusCode, Equals, http.StatusOK) re, _, err = testutils.Get(srv.URL, testutils.Header("Source", "a")) c.Assert(err, IsNil) c.Assert(re.StatusCode, Equals, 429) s.clock.Sleep(time.Second) re, _, err = testutils.Get(srv.URL, testutils.Header("Source", "a")) c.Assert(err, IsNil) c.Assert(re.StatusCode, Equals, http.StatusOK) }
// We've hit the limit and were able to proceed once the request has completed func (s *ConnLimiterSuite) TestCustomHandlers(c *C) { handler := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { w.Write([]byte("hello")) }) errHandler := utils.ErrorHandlerFunc(func(w http.ResponseWriter, req *http.Request, err error) { w.WriteHeader(http.StatusTeapot) w.Write([]byte(http.StatusText(http.StatusTeapot))) }) buf := &bytes.Buffer{} log := utils.NewFileLogger(buf, utils.INFO) l, err := New(handler, headerLimit, 0, ErrorHandler(errHandler), Logger(log)) c.Assert(err, Equals, nil) srv := httptest.NewServer(l) defer srv.Close() re, _, err := testutils.Get(srv.URL, testutils.Header("Limit", "a")) c.Assert(err, IsNil) c.Assert(re.StatusCode, Equals, http.StatusTeapot) c.Assert(len(buf.String()), Not(Equals), 0) }
// We've hit the limit and were able to proceed once the request has completed func (s *ConnLimiterSuite) TestHitLimitAndRelease(c *C) { wait := make(chan bool) proceed := make(chan bool) handler := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { if req.Header.Get("wait") != "" { proceed <- true <-wait } w.Write([]byte("hello")) }) l, err := New(handler, headerLimit, 1) c.Assert(err, Equals, nil) srv := httptest.NewServer(l) defer srv.Close() go func() { re, _, err := testutils.Get(srv.URL, testutils.Header("Limit", "a"), testutils.Header("wait", "yes")) c.Assert(err, IsNil) c.Assert(re.StatusCode, Equals, http.StatusOK) }() <-proceed re, _, err := testutils.Get(srv.URL, testutils.Header("Limit", "a")) c.Assert(err, IsNil) c.Assert(re.StatusCode, Equals, 429) // request from another source succeeds re, _, err = testutils.Get(srv.URL, testutils.Header("Limit", "b")) c.Assert(err, IsNil) c.Assert(re.StatusCode, Equals, http.StatusOK) // Once the first request finished, next one succeeds close(wait) re, _, err = testutils.Get(srv.URL, testutils.Header("Limit", "a")) c.Assert(err, IsNil) c.Assert(re.StatusCode, Equals, http.StatusOK) }
// We've failed to extract client ip func (s *LimiterSuite) TestFailure(c *C) { handler := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { w.Write([]byte("hello")) }) rates := NewRateSet() rates.Add(time.Second, 1, 1) l, err := New(handler, faultyExtract, rates, Clock(s.clock)) c.Assert(err, IsNil) srv := httptest.NewServer(l) defer srv.Close() re, _, err := testutils.Get(srv.URL, testutils.Header("Source", "a")) c.Assert(err, IsNil) c.Assert(re.StatusCode, Equals, http.StatusInternalServerError) }
func (s *RewriteSuite) TestDontRewriteResponseBody(c *C) { handler := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { w.Write([]byte(`{"foo": "{{.Request.Header.Get "X-Header"}}"}`)) }) rh, err := newRewriteHandler(handler, &Rewrite{"", "", false, false}) c.Assert(rh, NotNil) c.Assert(err, IsNil) srv := httptest.NewServer(rh) defer srv.Close() re, body, err := testutils.Get(srv.URL, testutils.Host("localhost"), testutils.Header("X-Header", "bar")) c.Assert(re.StatusCode, Equals, http.StatusOK) c.Assert(string(body), Equals, `{"foo": "{{.Request.Header.Get "X-Header"}}"}`) }
func (s *RewriteSuite) TestRewriteInQuery(c *C) { var outURL string handler := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { outURL = rawURL(req) w.Write([]byte("hello")) }) rh, err := newRewriteHandler(handler, &Rewrite{"^http://localhost/foo\\?(.*)=(.*)", `http://localhost/foo?$1={{.Request.Header.Get "X-Header"}}`, false, false}) c.Assert(rh, NotNil) c.Assert(err, IsNil) srv := httptest.NewServer(rh) defer srv.Close() re, _, err := testutils.Get(srv.URL+"/foo?a=b", testutils.Host("localhost"), testutils.Header("X-Header", "xxx")) c.Assert(err, IsNil) c.Assert(re.StatusCode, Equals, http.StatusOK) c.Assert(outURL, Equals, "http://localhost/foo?a=xxx") }
// TestContentLength makes sure Content-Length is re-calculated if body rewrite is enabled. func (s *RewriteSuite) TestContentLength(c *C) { handler := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { w.Header().Set("Content-Length", "45") w.WriteHeader(200) w.Write([]byte(`{"foo": "{{.Request.Header.Get "X-Header"}}"}`)) }) rh, err := newRewriteHandler(handler, &Rewrite{"", "", true, false}) c.Assert(rh, NotNil) c.Assert(err, IsNil) srv := httptest.NewServer(rh) defer srv.Close() re, _, _ := testutils.Get(srv.URL, testutils.Host("localhost"), testutils.Header("X-Header", "bar")) c.Assert(re.Header.Get("Content-Length"), Equals, "14") }