// 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)
}
func (s *ReaderSuite) TestReaderRoundRobin(c *C) {
called1 := false
srv1 := testutils.NewHandler(func(w http.ResponseWriter, req *http.Request) {
called1 = true
})
defer srv1.Close()
called2 := false
srv2 := testutils.NewHandler(func(w http.ResponseWriter, req *http.Request) {
called2 = true
})
defer srv2.Close()
target1 := srv1.URL
target2 := srv2.URL
reader := NewReader([]string{target1, target2})
proxyWrapper := testutils.NewHandler(func(w http.ResponseWriter, req *http.Request) {
reader.ServeHTTP(w, req)
})
defer proxyWrapper.Close()
re, _, err := testutils.Get(proxyWrapper.URL, testutils.Headers(http.Header{}))
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, http.StatusOK)
c.Assert(called1, Equals, true)
c.Assert(called2, Equals, false)
re, _, err = testutils.Get(proxyWrapper.URL, testutils.Headers(http.Header{}))
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, http.StatusOK)
c.Assert(called2, Equals, true)
}
// 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 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)
}
// Test scenario when increaing the weight on good endpoints made it worse
func (s *RBSuite) TestRebalancerCascading(c *C) {
a, b, d := testutils.NewResponder("a"), testutils.NewResponder("b"), testutils.NewResponder("d")
defer a.Close()
defer b.Close()
defer d.Close()
fwd, err := forward.New()
c.Assert(err, IsNil)
lb, err := New(fwd)
c.Assert(err, IsNil)
newMeter := func() (Meter, error) {
return &testMeter{}, nil
}
rb, err := NewRebalancer(lb, RebalancerMeter(newMeter), RebalancerClock(s.clock), RebalancerLogger(s.log))
c.Assert(err, IsNil)
rb.UpsertServer(testutils.ParseURI(a.URL))
rb.UpsertServer(testutils.ParseURI(b.URL))
rb.UpsertServer(testutils.ParseURI(d.URL))
rb.servers[0].meter.(*testMeter).rating = 0.3
proxy := httptest.NewServer(rb)
defer proxy.Close()
for i := 0; i < 6; i += 1 {
testutils.Get(proxy.URL)
testutils.Get(proxy.URL)
s.clock.CurrentTime = s.clock.CurrentTime.Add(rb.backoffDuration + time.Second)
}
// We have increased the load, and the situation became worse as the other servers started failing
c.Assert(rb.servers[0].curWeight, Equals, 1)
c.Assert(rb.servers[1].curWeight, Equals, FSMMaxWeight)
c.Assert(rb.servers[2].curWeight, Equals, FSMMaxWeight)
// server a is now recovering, the weights should go back to the original state
rb.servers[0].meter.(*testMeter).rating = 0.3
rb.servers[1].meter.(*testMeter).rating = 0.2
rb.servers[2].meter.(*testMeter).rating = 0.2
for i := 0; i < 6; i += 1 {
testutils.Get(proxy.URL)
testutils.Get(proxy.URL)
s.clock.CurrentTime = s.clock.CurrentTime.Add(rb.backoffDuration + time.Second)
}
// the algo reverted it back
c.Assert(rb.servers[0].curWeight, Equals, 1)
c.Assert(rb.servers[1].curWeight, Equals, 1)
c.Assert(rb.servers[2].curWeight, Equals, 1)
}
// Test scenario when one server goes down after what it recovers
func (s *RBSuite) TestRebalancerRecovery(c *C) {
a, b := testutils.NewResponder("a"), testutils.NewResponder("b")
defer a.Close()
defer b.Close()
fwd, err := forward.New()
c.Assert(err, IsNil)
lb, err := New(fwd)
c.Assert(err, IsNil)
newMeter := func() (Meter, error) {
return &testMeter{}, nil
}
rb, err := NewRebalancer(lb, RebalancerMeter(newMeter), RebalancerClock(s.clock), RebalancerLogger(s.log))
c.Assert(err, IsNil)
rb.UpsertServer(testutils.ParseURI(a.URL))
rb.UpsertServer(testutils.ParseURI(b.URL))
rb.servers[0].meter.(*testMeter).rating = 0.3
proxy := httptest.NewServer(rb)
defer proxy.Close()
for i := 0; i < 6; i += 1 {
testutils.Get(proxy.URL)
testutils.Get(proxy.URL)
s.clock.CurrentTime = s.clock.CurrentTime.Add(rb.backoffDuration + time.Second)
}
c.Assert(rb.servers[0].curWeight, Equals, 1)
c.Assert(rb.servers[1].curWeight, Equals, FSMMaxWeight)
c.Assert(lb.servers[0].weight, Equals, 1)
c.Assert(lb.servers[1].weight, Equals, FSMMaxWeight)
// server a is now recovering, the weights should go back to the original state
rb.servers[0].meter.(*testMeter).rating = 0
for i := 0; i < 6; i += 1 {
testutils.Get(proxy.URL)
testutils.Get(proxy.URL)
s.clock.CurrentTime = s.clock.CurrentTime.Add(rb.backoffDuration + time.Second)
}
c.Assert(rb.servers[0].curWeight, Equals, 1)
c.Assert(rb.servers[1].curWeight, Equals, 1)
// Make sure we have applied the weights to the inner load balancer
c.Assert(lb.servers[0].weight, Equals, 1)
c.Assert(lb.servers[1].weight, Equals, 1)
}
func (s *FwdSuite) TestForwardedProto(c *C) {
var proto string
srv := testutils.NewHandler(func(w http.ResponseWriter, req *http.Request) {
proto = req.Header.Get(XForwardedProto)
w.Write([]byte("hello"))
})
defer srv.Close()
buf := &bytes.Buffer{}
l := utils.NewFileLogger(buf, utils.INFO)
f, err := New(Logger(l))
c.Assert(err, IsNil)
proxy := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
req.URL = testutils.ParseURI(srv.URL)
f.ServeHTTP(w, req)
})
tproxy := httptest.NewUnstartedServer(proxy)
tproxy.StartTLS()
defer tproxy.Close()
re, _, err := testutils.Get(tproxy.URL)
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, http.StatusOK)
c.Assert(proto, Equals, "https")
c.Assert(strings.Contains(buf.String(), "tls"), Equals, true)
}
func (s *TraceSuite) TestTraceCaptureHeaders(c *C) {
respHeaders := http.Header{
"X-Re-1": []string{"6", "7"},
"X-Re-2": []string{"2", "3"},
}
handler := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
utils.CopyHeaders(w.Header(), respHeaders)
w.Write([]byte("hello"))
})
buf := &bytes.Buffer{}
l := utils.NewFileLogger(buf, utils.INFO)
trace := &bytes.Buffer{}
t, err := New(handler, trace, Logger(l), RequestHeaders("X-Req-B", "X-Req-A"), ResponseHeaders("X-Re-1", "X-Re-2"))
c.Assert(err, IsNil)
srv := httptest.NewServer(t)
defer srv.Close()
reqHeaders := http.Header{"X-Req-A": []string{"1", "2"}, "X-Req-B": []string{"3", "4"}}
re, _, err := testutils.Get(srv.URL+"/hello", testutils.Headers(reqHeaders))
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, http.StatusOK)
var r *Record
c.Assert(json.Unmarshal(trace.Bytes(), &r), IsNil)
c.Assert(r.Request.Headers, DeepEquals, reqHeaders)
c.Assert(r.Response.Headers, DeepEquals, respHeaders)
}
func (s *STSuite) TestFileStreamingResponse(c *C) {
srv := testutils.NewHandler(func(w http.ResponseWriter, req *http.Request) {
w.Write([]byte("hello, this response is too large to fit in memory"))
})
defer srv.Close()
// forwarder will proxy the request to whatever destination
fwd, err := forward.New()
c.Assert(err, IsNil)
// this is our redirect to server
rdr := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
req.URL = testutils.ParseURI(srv.URL)
fwd.ServeHTTP(w, req)
})
// stream handler will forward requests to redirect
st, err := New(rdr, Logger(utils.NewFileLogger(os.Stdout, utils.INFO)), MemResponseBodyBytes(4))
c.Assert(err, IsNil)
proxy := httptest.NewServer(st)
defer proxy.Close()
re, body, err := testutils.Get(proxy.URL)
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, http.StatusOK)
c.Assert(string(body), Equals, "hello, this response is too large to fit in memory")
}
// Make sure that stream handler preserves TLS settings
func (s *STSuite) TestPreservesTLS(c *C) {
srv := testutils.NewHandler(func(w http.ResponseWriter, req *http.Request) {
w.WriteHeader(http.StatusOK)
w.Write([]byte("ok"))
})
defer srv.Close()
// forwarder will proxy the request to whatever destination
fwd, err := forward.New()
c.Assert(err, IsNil)
var t *tls.ConnectionState
// this is our redirect to server
rdr := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
t = req.TLS
req.URL = testutils.ParseURI(srv.URL)
fwd.ServeHTTP(w, req)
})
// stream handler will forward requests to redirect
st, err := New(rdr, Logger(utils.NewFileLogger(os.Stdout, utils.INFO)))
c.Assert(err, IsNil)
proxy := httptest.NewUnstartedServer(st)
proxy.StartTLS()
defer proxy.Close()
re, _, err := testutils.Get(proxy.URL)
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, http.StatusOK)
c.Assert(t, NotNil)
}
func (s *STSuite) TestRequestLimitReached(c *C) {
srv := testutils.NewHandler(func(w http.ResponseWriter, req *http.Request) {
w.Write([]byte("hello"))
})
defer srv.Close()
// forwarder will proxy the request to whatever destination
fwd, err := forward.New()
c.Assert(err, IsNil)
// this is our redirect to server
rdr := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
req.URL = testutils.ParseURI(srv.URL)
fwd.ServeHTTP(w, req)
})
// stream handler will forward requests to redirect
st, err := New(rdr, Logger(utils.NewFileLogger(os.Stdout, utils.INFO)), MaxRequestBodyBytes(4))
c.Assert(err, IsNil)
proxy := httptest.NewServer(st)
defer proxy.Close()
re, _, err := testutils.Get(proxy.URL, testutils.Body("this request is too long"))
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, http.StatusRequestEntityTooLarge)
}
func (s *STSuite) TestNoBody(c *C) {
srv := testutils.NewHandler(func(w http.ResponseWriter, req *http.Request) {
w.WriteHeader(http.StatusOK)
})
defer srv.Close()
// forwarder will proxy the request to whatever destination
fwd, err := forward.New()
c.Assert(err, IsNil)
// this is our redirect to server
rdr := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
req.URL = testutils.ParseURI(srv.URL)
fwd.ServeHTTP(w, req)
})
// stream handler will forward requests to redirect
st, err := New(rdr, Logger(utils.NewFileLogger(os.Stdout, utils.INFO)))
c.Assert(err, IsNil)
proxy := httptest.NewServer(st)
defer proxy.Close()
re, _, err := testutils.Get(proxy.URL)
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, http.StatusOK)
}
func (s *CBSuite) TestRedirect(c *C) {
handler := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
w.Write([]byte("hello"))
})
cb, err := New(handler, triggerNetRatio, Clock(s.clock), Fallback(fallbackRedirect))
c.Assert(err, IsNil)
srv := httptest.NewServer(cb)
defer srv.Close()
cb.metrics = statsNetErrors(0.6)
re, _, err := testutils.Get(srv.URL)
c.Assert(err, IsNil)
client := &http.Client{
CheckRedirect: func(req *http.Request, via []*http.Request) error {
return fmt.Errorf("no redirects")
},
}
re, err = client.Get(srv.URL)
c.Assert(err, NotNil)
c.Assert(re.StatusCode, Equals, http.StatusFound)
c.Assert(re.Header.Get("Location"), Equals, "http://localhost:5000")
}
func (s *RewriteSuite) TestSameHeaderMulti(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 "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/bar",
testutils.Host("localhost"), testutils.Header("X-Header", "baz"))
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, http.StatusOK)
c.Assert(outURL, Equals, "http://localhost/foo/baz/bar/baz")
}
// 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)
}
func (s *RBSuite) TestRebalancerLive(c *C) {
a, b := testutils.NewResponder("a"), testutils.NewResponder("b")
defer a.Close()
defer b.Close()
fwd, err := forward.New()
c.Assert(err, IsNil)
lb, err := New(fwd)
c.Assert(err, IsNil)
rb, err := NewRebalancer(lb, RebalancerBackoff(time.Millisecond), RebalancerClock(s.clock))
c.Assert(err, IsNil)
rb.UpsertServer(testutils.ParseURI(a.URL))
rb.UpsertServer(testutils.ParseURI(b.URL))
rb.UpsertServer(testutils.ParseURI("http://localhost:62345"))
proxy := httptest.NewServer(rb)
defer proxy.Close()
for i := 0; i < 1000; i += 1 {
testutils.Get(proxy.URL)
if i%10 == 0 {
s.clock.CurrentTime = s.clock.CurrentTime.Add(rb.backoffDuration + time.Second)
}
}
// load balancer changed weights
c.Assert(rb.servers[0].curWeight, Equals, FSMMaxWeight)
c.Assert(rb.servers[1].curWeight, Equals, FSMMaxWeight)
c.Assert(rb.servers[2].curWeight, Equals, 1)
}
func (s *STSuite) TestCustomErrorHandler(c *C) {
srv := testutils.NewHandler(func(w http.ResponseWriter, req *http.Request) {
w.Write([]byte("hello, this response is too large"))
})
defer srv.Close()
// forwarder will proxy the request to whatever destination
fwd, err := forward.New()
c.Assert(err, IsNil)
// this is our redirect to server
rdr := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
req.URL = testutils.ParseURI(srv.URL)
fwd.ServeHTTP(w, req)
})
// stream handler will forward requests to redirect
errHandler := utils.ErrorHandlerFunc(func(w http.ResponseWriter, req *http.Request, err error) {
w.WriteHeader(http.StatusTeapot)
w.Write([]byte(http.StatusText(http.StatusTeapot)))
})
st, err := New(rdr, Logger(utils.NewFileLogger(os.Stdout, utils.INFO)), MaxResponseBodyBytes(4), ErrorHandler(errHandler))
c.Assert(err, IsNil)
proxy := httptest.NewServer(st)
defer proxy.Close()
re, _, err := testutils.Get(proxy.URL)
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, http.StatusTeapot)
}
// Makes sure hop-by-hop headers are removed
func (s *FwdSuite) TestForwardedHeaders(c *C) {
var outHeaders http.Header
srv := testutils.NewHandler(func(w http.ResponseWriter, req *http.Request) {
outHeaders = req.Header
w.Write([]byte("hello"))
})
defer srv.Close()
f, err := New(Rewriter(&HeaderRewriter{TrustForwardHeader: true, Hostname: "hello"}))
c.Assert(err, IsNil)
proxy := testutils.NewHandler(func(w http.ResponseWriter, req *http.Request) {
req.URL = testutils.ParseURI(srv.URL)
f.ServeHTTP(w, req)
})
defer proxy.Close()
headers := http.Header{
XForwardedProto: []string{"httpx"},
XForwardedFor: []string{"192.168.1.1"},
XForwardedServer: []string{"foobar"},
}
re, _, err := testutils.Get(proxy.URL, testutils.Headers(headers))
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, http.StatusOK)
c.Assert(outHeaders.Get(XForwardedProto), Equals, "httpx")
c.Assert(strings.Contains(outHeaders.Get(XForwardedFor), "192.168.1.1"), Equals, true)
c.Assert(outHeaders.Get(XForwardedServer), Equals, "hello")
}
// Removing the server resets the state
func (s *RBSuite) TestRebalancerReset(c *C) {
a, b, d := testutils.NewResponder("a"), testutils.NewResponder("b"), testutils.NewResponder("d")
defer a.Close()
defer b.Close()
defer d.Close()
fwd, err := forward.New()
c.Assert(err, IsNil)
lb, err := New(fwd)
c.Assert(err, IsNil)
newMeter := func() (Meter, error) {
return &testMeter{}, nil
}
rb, err := NewRebalancer(lb, RebalancerMeter(newMeter), RebalancerClock(s.clock), RebalancerLogger(s.log))
c.Assert(err, IsNil)
rb.UpsertServer(testutils.ParseURI(a.URL))
rb.UpsertServer(testutils.ParseURI(b.URL))
rb.UpsertServer(testutils.ParseURI(d.URL))
rb.servers[0].meter.(*testMeter).rating = 0.3
rb.servers[1].meter.(*testMeter).rating = 0
rb.servers[2].meter.(*testMeter).rating = 0
proxy := httptest.NewServer(rb)
defer proxy.Close()
for i := 0; i < 6; i += 1 {
testutils.Get(proxy.URL)
testutils.Get(proxy.URL)
s.clock.CurrentTime = s.clock.CurrentTime.Add(rb.backoffDuration + time.Second)
}
// load balancer changed weights
c.Assert(rb.servers[0].curWeight, Equals, 1)
c.Assert(rb.servers[1].curWeight, Equals, FSMMaxWeight)
c.Assert(rb.servers[2].curWeight, Equals, FSMMaxWeight)
// Removing servers has reset the state
rb.RemoveServer(testutils.ParseURI(d.URL))
c.Assert(rb.servers[0].curWeight, Equals, 1)
c.Assert(rb.servers[1].curWeight, Equals, 1)
}
func seq(c *C, url string, repeat int) []string {
out := []string{}
for i := 0; i < repeat; i++ {
_, body, err := testutils.Get(url)
c.Assert(err, IsNil)
out = append(out, string(body))
}
return out
}
// 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)
}
func (s *FwdSuite) TestDefaultErrHandler(c *C) {
f, err := New()
c.Assert(err, IsNil)
proxy := testutils.NewHandler(func(w http.ResponseWriter, req *http.Request) {
req.URL = testutils.ParseURI("http://localhost:63450")
f.ServeHTTP(w, req)
})
defer proxy.Close()
re, _, err := testutils.Get(proxy.URL)
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, http.StatusBadGateway)
}
func (s *RRSuite) TestNoServers(c *C) {
fwd, err := forward.New()
c.Assert(err, IsNil)
lb, err := New(fwd)
c.Assert(err, IsNil)
proxy := httptest.NewServer(lb)
defer proxy.Close()
re, _, err := testutils.Get(proxy.URL)
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, http.StatusInternalServerError)
}
func (s *CBSuite) TestTriggerDuringRecovery(c *C) {
handler := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
w.Write([]byte("hello"))
})
cb, err := New(handler, triggerNetRatio, Clock(s.clock), CheckPeriod(time.Microsecond))
c.Assert(err, IsNil)
srv := httptest.NewServer(cb)
defer srv.Close()
cb.metrics = statsNetErrors(0.6)
re, _, err := testutils.Get(srv.URL)
c.Assert(err, IsNil)
c.Assert(cb.state, Equals, cbState(stateTripped))
// We should be in recovering state by now
s.advanceTime(10*time.Second + time.Millisecond)
re, _, err = testutils.Get(srv.URL)
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, http.StatusServiceUnavailable)
c.Assert(cb.state, Equals, cbState(stateRecovering))
// We have matched error condition during recovery state and are going back to tripped state
s.advanceTime(5 * time.Second)
cb.metrics = statsNetErrors(0.6)
allowed := 0
for i := 0; i < 100; i++ {
re, _, err = testutils.Get(srv.URL)
if re.StatusCode == http.StatusOK && err == nil {
allowed++
}
}
c.Assert(allowed, Not(Equals), 0)
c.Assert(cb.state, Equals, cbState(stateTripped))
}
// We've hit the limit and were able to proceed once the request has completed
func (s *ConnLimiterSuite) TestFaultyExtract(c *C) {
handler := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
w.Write([]byte("hello"))
})
l, err := New(handler, faultyExtract, 1)
c.Assert(err, Equals, nil)
srv := httptest.NewServer(l)
defer srv.Close()
re, _, err := testutils.Get(srv.URL)
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, http.StatusInternalServerError)
}
func (s *RewriteSuite) TestRedirect(c *C) {
handler := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
w.Write([]byte("hello"))
})
rh, err := newRewriteHandler(handler, &Rewrite{"^http://localhost/(foo)/(bar)", "https://localhost/$2", false, true})
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"))
c.Assert(re.StatusCode, Equals, http.StatusFound)
c.Assert(re.Header.Get("Location"), Equals, "https://localhost/bar")
}
func (s *CBSuite) TestStandbyCycle(c *C) {
handler := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
w.Write([]byte("hello"))
})
cb, err := New(handler, triggerNetRatio)
c.Assert(err, IsNil)
srv := httptest.NewServer(cb)
defer srv.Close()
re, body, err := testutils.Get(srv.URL)
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, http.StatusOK)
c.Assert(string(body), Equals, "hello")
}
func (s *CBSuite) TestFullCycle(c *C) {
handler := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
w.Write([]byte("hello"))
})
cb, err := New(handler, triggerNetRatio, Clock(s.clock))
c.Assert(err, IsNil)
srv := httptest.NewServer(cb)
defer srv.Close()
re, _, err := testutils.Get(srv.URL)
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, http.StatusOK)
cb.metrics = statsNetErrors(0.6)
s.advanceTime(defaultCheckPeriod + time.Millisecond)
re, _, err = testutils.Get(srv.URL)
c.Assert(err, IsNil)
c.Assert(cb.state, Equals, cbState(stateTripped))
// Some time has passed, but we are still in trpped state.
s.advanceTime(9 * time.Second)
re, _, err = testutils.Get(srv.URL)
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, http.StatusServiceUnavailable)
c.Assert(cb.state, Equals, cbState(stateTripped))
// We should be in recovering state by now
s.advanceTime(time.Second*1 + time.Millisecond)
re, _, err = testutils.Get(srv.URL)
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, http.StatusServiceUnavailable)
c.Assert(cb.state, Equals, cbState(stateRecovering))
// 5 seconds after we should be allowing some requests to pass
s.advanceTime(5 * time.Second)
allowed := 0
for i := 0; i < 100; i++ {
re, _, err = testutils.Get(srv.URL)
if re.StatusCode == http.StatusOK && err == nil {
allowed++
}
}
c.Assert(allowed, Not(Equals), 0)
// After some time, all is good and we should be in stand by mode again
s.advanceTime(5*time.Second + time.Millisecond)
re, _, err = testutils.Get(srv.URL)
c.Assert(cb.state, Equals, cbState(stateStandby))
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, http.StatusOK)
}
func (s *RTSuite) TestSuccess(c *C) {
srv := testutils.NewHandler(func(w http.ResponseWriter, req *http.Request) {
w.Write([]byte("hello"))
})
defer srv.Close()
lb, rt := new(c, `IsNetworkError() && Attempts() <= 2`)
proxy := httptest.NewServer(rt)
defer proxy.Close()
lb.UpsertServer(testutils.ParseURI(srv.URL))
re, body, err := testutils.Get(proxy.URL)
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, http.StatusOK)
c.Assert(string(body), Equals, "hello")
}
func (s *FwdSuite) TestCustomErrHandler(c *C) {
f, err := New(ErrorHandler(utils.ErrorHandlerFunc(func(w http.ResponseWriter, req *http.Request, err error) {
w.WriteHeader(http.StatusTeapot)
w.Write([]byte(http.StatusText(http.StatusTeapot)))
})))
c.Assert(err, IsNil)
proxy := testutils.NewHandler(func(w http.ResponseWriter, req *http.Request) {
req.URL = testutils.ParseURI("http://localhost:63450")
f.ServeHTTP(w, req)
})
defer proxy.Close()
re, body, err := testutils.Get(proxy.URL)
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, http.StatusTeapot)
c.Assert(string(body), Equals, http.StatusText(http.StatusTeapot))
}