// 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 *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")
}
// 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) 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)
}
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) 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) 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 *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)
}
// 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)
}
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 *TraceSuite) TestTraceSimple(c *C) {
handler := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
w.Header().Set("Content-Length", "5")
w.Write([]byte("hello"))
})
buf := &bytes.Buffer{}
l := utils.NewFileLogger(buf, utils.INFO)
trace := &bytes.Buffer{}
t, err := New(handler, trace, Logger(l))
c.Assert(err, IsNil)
srv := httptest.NewServer(t)
defer srv.Close()
re, _, err := testutils.MakeRequest(srv.URL+"/hello", testutils.Method("POST"), testutils.Body("123456"))
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.Method, Equals, "POST")
c.Assert(r.Request.URL, Equals, "/hello")
c.Assert(r.Response.Code, Equals, http.StatusOK)
c.Assert(r.Request.BodyBytes, Equals, int64(6))
c.Assert(r.Response.Roundtrip, Not(Equals), float64(0))
c.Assert(r.Response.BodyBytes, Equals, int64(5))
}
func new(c *C, p string) (*roundrobin.RoundRobin, *Streamer) {
logger := utils.NewFileLogger(os.Stdout, utils.INFO)
// forwarder will proxy the request to whatever destination
fwd, err := forward.New(forward.Logger(logger))
c.Assert(err, IsNil)
// load balancer will round robin request
lb, err := roundrobin.New(fwd)
c.Assert(err, IsNil)
// stream handler will forward requests to redirect, make sure it uses files
st, err := New(lb, Logger(logger), Retry(p), MemRequestBodyBytes(1))
c.Assert(err, IsNil)
return lb, st
}
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) 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)
}
func (s *TraceSuite) TestTraceTLS(c *C) {
handler := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
w.Write([]byte("hello"))
})
buf := &bytes.Buffer{}
l := utils.NewFileLogger(buf, utils.INFO)
trace := &bytes.Buffer{}
t, err := New(handler, trace, Logger(l))
c.Assert(err, IsNil)
srv := httptest.NewUnstartedServer(t)
srv.StartTLS()
defer srv.Close()
config := &tls.Config{
InsecureSkipVerify: true,
}
u, err := url.Parse(srv.URL)
c.Assert(err, IsNil)
conn, err := tls.Dial("tcp", u.Host, config)
c.Assert(err, IsNil)
fmt.Fprintf(conn, "GET / HTTP/1.0\r\n\r\n")
status, err := bufio.NewReader(conn).ReadString('\n')
c.Assert(status, Equals, "HTTP/1.0 200 OK\r\n")
state := conn.ConnectionState()
conn.Close()
var r *Record
c.Assert(json.Unmarshal(trace.Bytes(), &r), IsNil)
c.Assert(r.Request.TLS.Version, Equals, versionToString(state.Version))
}
func (s *RBSuite) SetUpSuite(c *C) {
s.log = utils.NewFileLogger(os.Stdout, utils.INFO)
}