// 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)
}
// Make sure rates from different ips are controlled separatedly
func (s *LimiterSuite) TestIsolation(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)
// The request from other source can proceed
re, _, err = testutils.Get(srv.URL, testutils.Header("Source", "b"))
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)
}
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 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)
}
// If the rate set from the HTTP header has more then one rate for the same
// time period defined, then the one mentioned in the list last is used.
func (s *RateLimitSuite) TestRequestProcessingAmbiguousConfig(c *C) {
// Given
rl, _ := FromOther(
RateLimit{
PeriodSeconds: 1,
Requests: 1,
Burst: 1,
Variable: "client.ip",
RateVar: "request.header.X-Rates",
clock: s.clock,
})
handler := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
w.Write([]byte("hello"))
})
rli, err := rl.NewHandler(handler)
c.Assert(err, IsNil)
srv := httptest.NewServer(rli)
defer srv.Close()
// When/Then: The last of configured rates with the same period is applied,
// which 2 request/second, note that the default rate is 1 request/second.
hdr := testutils.Header("X-Rates", `[{"PeriodSeconds": 1, "Requests": 10},
{"PeriodSeconds": 1, "Requests": 2}]`)
re, _, err := testutils.Get(srv.URL, hdr)
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, http.StatusOK)
re, _, err = testutils.Get(srv.URL, hdr)
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, http.StatusOK)
re, _, err = testutils.Get(srv.URL, hdr)
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, 429)
s.clock.Sleep(time.Second)
re, _, err = testutils.Get(srv.URL, hdr)
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, http.StatusOK)
}
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")
}
// If the rate set from the HTTP header has more then one rate for the same
// time period defined, then the one mentioned in the list last is used.
func (s *RateLimitSuite) TestRequestInvalidConfig(c *C) {
// Given
rl, _ := FromOther(
RateLimit{
PeriodSeconds: 1,
Requests: 1,
Burst: 1,
Variable: "client.ip",
RateVar: "request.header.X-Rates",
clock: s.clock,
})
handler := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
w.Write([]byte("hello"))
})
rli, err := rl.NewHandler(handler)
c.Assert(err, IsNil)
srv := httptest.NewServer(rli)
defer srv.Close()
// When/Then: The default rate of 1 request/second is used.
hdr := testutils.Header("X-Rates", `[{"PeriodSeconds": -1, "Requests": 10}]`)
re, _, err := testutils.Get(srv.URL, hdr)
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, http.StatusOK)
re, _, err = testutils.Get(srv.URL, hdr)
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, 429)
s.clock.Sleep(time.Second)
re, _, err = testutils.Get(srv.URL, hdr)
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, http.StatusOK)
}
func (s *TraceSuite) TestHandler(c *C) {
os.Remove("/tmp/vulcand_trace_test.sock")
unixAddr, err := net.ResolveUnixAddr("unixgram", "/tmp/vulcand_trace_test.sock")
c.Assert(err, IsNil)
conn, err := net.ListenUnixgram("unixgram", unixAddr)
c.Assert(err, IsNil)
defer conn.Close()
outC := make(chan []byte, 1000)
closeC := make(chan bool)
defer close(closeC)
go func() {
for {
buf := make([]byte, 65536)
bytes, err := conn.Read(buf)
if err != nil {
return
}
outbuf := make([]byte, bytes)
copy(outbuf, buf)
select {
case <-closeC:
return
case outC <- outbuf:
continue
}
}
}()
responder := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
w.Header().Add("X-Resp-A", "h2")
w.Write([]byte("hello"))
})
h, err := New("syslog:///tmp/vulcand_trace_test.sock", []string{"X-Req-A"}, []string{"X-Resp-A"})
c.Assert(err, IsNil)
handler, err := h.NewHandler(responder)
c.Assert(err, IsNil)
srv := httptest.NewServer(handler)
defer srv.Close()
re, _, err := testutils.Get(srv.URL+"/hello", testutils.Header("X-Req-A", "yo"))
c.Assert(err, IsNil)
c.Assert(re.StatusCode, Equals, http.StatusOK)
var buf []byte
select {
case buf = <-outC:
case <-time.After(10 * time.Millisecond):
c.Fatalf("timeout")
}
vals := strings.Split(string(buf), SyslogPrefix)
var r *oxytrace.Record
c.Assert(json.Unmarshal([]byte(vals[1]), &r), IsNil)
c.Assert(r.Request.URL, Equals, "/hello")
c.Assert(r.Request.Headers, DeepEquals, http.Header{"X-Req-A": []string{"yo"}})
c.Assert(r.Response.Headers, DeepEquals, http.Header{"X-Resp-A": []string{"h2"}})
}