func (s *RRSuite) TestWeighted(c *C) {
a := testutils.NewResponder("a")
defer a.Close()
b := testutils.NewResponder("b")
defer b.Close()
fwd, err := forward.New()
c.Assert(err, IsNil)
lb, err := New(fwd)
c.Assert(err, IsNil)
lb.UpsertServer(testutils.ParseURI(a.URL), Weight(3))
lb.UpsertServer(testutils.ParseURI(b.URL), Weight(2))
proxy := httptest.NewServer(lb)
defer proxy.Close()
c.Assert(seq(c, proxy.URL, 6), DeepEquals, []string{"a", "a", "b", "a", "b", "a"})
w, ok := lb.ServerWeight(testutils.ParseURI(a.URL))
c.Assert(w, Equals, 3)
c.Assert(ok, Equals, true)
w, ok = lb.ServerWeight(testutils.ParseURI(b.URL))
c.Assert(w, Equals, 2)
c.Assert(ok, Equals, true)
w, ok = lb.ServerWeight(testutils.ParseURI("http://caramba:4000"))
c.Assert(w, Equals, -1)
c.Assert(ok, Equals, false)
}
func (s *RBSuite) TestRebalancerRemoveServer(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)
c.Assert(err, IsNil)
rb.UpsertServer(testutils.ParseURI(a.URL))
rb.UpsertServer(testutils.ParseURI(b.URL))
proxy := httptest.NewServer(rb)
defer proxy.Close()
c.Assert(seq(c, proxy.URL, 3), DeepEquals, []string{"a", "b", "a"})
c.Assert(rb.RemoveServer(testutils.ParseURI(a.URL)), IsNil)
c.Assert(seq(c, proxy.URL, 3), DeepEquals, []string{"b", "b", "b"})
}
func (s *RRSuite) TestUpsertWeight(c *C) {
a := testutils.NewResponder("a")
defer a.Close()
b := testutils.NewResponder("b")
defer b.Close()
fwd, err := forward.New()
c.Assert(err, IsNil)
lb, err := New(fwd)
c.Assert(err, IsNil)
c.Assert(lb.UpsertServer(testutils.ParseURI(a.URL)), IsNil)
c.Assert(lb.UpsertServer(testutils.ParseURI(b.URL)), IsNil)
proxy := httptest.NewServer(lb)
defer proxy.Close()
c.Assert(seq(c, proxy.URL, 3), DeepEquals, []string{"a", "b", "a"})
c.Assert(lb.UpsertServer(testutils.ParseURI(b.URL), Weight(3)), IsNil)
c.Assert(seq(c, proxy.URL, 4), DeepEquals, []string{"b", "b", "a", "b"})
}
func (s *STSuite) TestChunkedEncodingLimitReached(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)), MemRequestBodyBytes(4), MaxRequestBodyBytes(8))
c.Assert(err, IsNil)
proxy := httptest.NewServer(st)
defer proxy.Close()
conn, err := net.Dial("tcp", testutils.ParseURI(proxy.URL).Host)
c.Assert(err, IsNil)
fmt.Fprintf(conn, "POST / HTTP/1.0\r\nTransfer-Encoding: chunked\r\n\r\n4\r\ntest\r\n5\r\ntest1\r\n5\r\ntest2\r\n0\r\n\r\n")
status, err := bufio.NewReader(conn).ReadString('\n')
c.Assert(status, Equals, "HTTP/1.0 413 Request Entity Too Large\r\n")
}
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 *FwdSuite) TestEscapedURL(c *C) {
var outURL string
srv := testutils.NewHandler(func(w http.ResponseWriter, req *http.Request) {
outURL = req.RequestURI
w.Write([]byte("hello"))
})
defer srv.Close()
f, err := New()
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()
path := "/log/http%3A%2F%2Fwww.site.com%2Fsomething?a=b"
request, err := http.NewRequest("GET", proxy.URL, nil)
parsed := testutils.ParseURI(proxy.URL)
parsed.Opaque = path
request.URL = parsed
re, err := http.DefaultClient.Do(request)
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, http.StatusOK)
c.Assert(outURL, Equals, path)
}
// 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 *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")
}
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 *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)
}
// 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)
}
// 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"},
XForwardedHost: []string{"upstream-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(strings.Contains(outHeaders.Get(XForwardedHost), "upstream-foobar"), Equals, true)
c.Assert(outHeaders.Get(XForwardedServer), Equals, "hello")
}
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)
}
// 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 (s *RTSuite) TestRetryOnError(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("http://localhost:64321"))
lb.UpsertServer(testutils.ParseURI(srv.URL))
re, body, err := testutils.Get(proxy.URL, testutils.Body("some request parameters"))
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, http.StatusOK)
c.Assert(string(body), Equals, "hello")
}
func (s *RTSuite) TestRetryExceedAttempts(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("http://localhost:64321"))
lb.UpsertServer(testutils.ParseURI("http://localhost:64322"))
lb.UpsertServer(testutils.ParseURI("http://localhost:64323"))
lb.UpsertServer(testutils.ParseURI(srv.URL))
re, _, err := testutils.Get(proxy.URL)
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, http.StatusBadGateway)
}
func (s *STSuite) TestChunkedEncodingSuccess(c *C) {
var reqBody string
var contentLength int64
srv := testutils.NewHandler(func(w http.ResponseWriter, req *http.Request) {
body, err := ioutil.ReadAll(req.Body)
c.Assert(err, IsNil)
reqBody = string(body)
contentLength = req.ContentLength
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)))
c.Assert(err, IsNil)
proxy := httptest.NewServer(st)
defer proxy.Close()
conn, err := net.Dial("tcp", testutils.ParseURI(proxy.URL).Host)
c.Assert(err, IsNil)
fmt.Fprintf(conn, "POST / HTTP/1.0\r\nTransfer-Encoding: chunked\r\n\r\n4\r\ntest\r\n5\r\ntest1\r\n5\r\ntest2\r\n0\r\n\r\n")
status, err := bufio.NewReader(conn).ReadString('\n')
c.Assert(reqBody, Equals, "testtest1test2")
c.Assert(status, Equals, "HTTP/1.0 200 OK\r\n")
c.Assert(contentLength, Equals, int64(len(reqBody)))
}
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 *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))
}
func (s *RBSuite) TestRebalancerNormalOperation(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)
c.Assert(err, IsNil)
rb.UpsertServer(testutils.ParseURI(a.URL))
c.Assert(rb.Servers()[0].String(), Equals, a.URL)
proxy := httptest.NewServer(rb)
defer proxy.Close()
c.Assert(seq(c, proxy.URL, 3), DeepEquals, []string{"a", "a", "a"})
}
func (s *FwdSuite) TestCustomTransportTimeout(c *C) {
srv := testutils.NewHandler(func(w http.ResponseWriter, req *http.Request) {
time.Sleep(20 * time.Millisecond)
w.Write([]byte("hello"))
})
defer srv.Close()
f, err := New(RoundTripper(
&http.Transport{
ResponseHeaderTimeout: 5 * time.Millisecond,
}))
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()
re, _, err := testutils.Get(proxy.URL)
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, http.StatusGatewayTimeout)
}
func (s *FwdSuite) TestCustomLogger(c *C) {
srv := testutils.NewHandler(func(w http.ResponseWriter, req *http.Request) {
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 := testutils.NewHandler(func(w http.ResponseWriter, req *http.Request) {
req.URL = testutils.ParseURI(srv.URL)
f.ServeHTTP(w, req)
})
defer proxy.Close()
re, _, err := testutils.Get(proxy.URL)
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, http.StatusOK)
c.Assert(strings.Contains(buf.String(), srv.URL), Equals, true)
}
// Makes sure hop-by-hop headers are removed
func (s *FwdSuite) TestForwardHopHeaders(c *C) {
called := false
var outHeaders http.Header
var outHost, expectedHost string
srv := testutils.NewHandler(func(w http.ResponseWriter, req *http.Request) {
called = true
outHeaders = req.Header
outHost = req.Host
w.Write([]byte("hello"))
})
defer srv.Close()
f, err := New()
c.Assert(err, IsNil)
proxy := testutils.NewHandler(func(w http.ResponseWriter, req *http.Request) {
req.URL = testutils.ParseURI(srv.URL)
expectedHost = req.URL.Host
f.ServeHTTP(w, req)
})
defer proxy.Close()
headers := http.Header{
Connection: []string{"close"},
KeepAlive: []string{"timeout=600"},
}
re, body, err := testutils.Get(proxy.URL, testutils.Headers(headers))
c.Assert(err, IsNil)
c.Assert(string(body), Equals, "hello")
c.Assert(re.StatusCode, Equals, http.StatusOK)
c.Assert(called, Equals, true)
c.Assert(outHeaders.Get(Connection), Equals, "")
c.Assert(outHeaders.Get(KeepAlive), Equals, "")
c.Assert(outHost, Equals, expectedHost)
}
func (s *FwdSuite) TestChunkedResponseConversion(c *C) {
srv := testutils.NewHandler(func(w http.ResponseWriter, req *http.Request) {
h := w.(http.Hijacker)
conn, _, _ := h.Hijack()
fmt.Fprintf(conn, "HTTP/1.0 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n4\r\ntest\r\n5\r\ntest1\r\n5\r\ntest2\r\n0\r\n\r\n")
conn.Close()
})
defer srv.Close()
f, err := New()
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()
re, body, err := testutils.Get(proxy.URL)
c.Assert(err, IsNil)
c.Assert(string(body), Equals, "testtest1test2")
c.Assert(re.StatusCode, Equals, http.StatusOK)
c.Assert(re.Header.Get("Content-Length"), Equals, fmt.Sprintf("%d", len("testtest1test2")))
}
func (s *RRSuite) TestRemoveBadServer(c *C) {
lb, err := New(nil)
c.Assert(err, IsNil)
c.Assert(lb.RemoveServer(testutils.ParseURI("http://google.com")), NotNil)
}