More tests for forwarder.

couple of things remained to test for forwarder - check that url is not touched
This commit is contained in:
Alexander Klizhentas
2014-12-10 10:20:42 -08:00
parent ee978f2a31
commit 1567438717
11 changed files with 385 additions and 33 deletions
+30 -5
View File
@@ -7,8 +7,9 @@ import (
"net/http"
"net/url"
"os"
"time"
"github.com/mailgun/oxy/netutils"
"github.com/mailgun/oxy/utils"
)
// ReqRewriter can alter request headers and body
@@ -40,10 +41,18 @@ func ErrorHandler(h http.Handler) optSetter {
}
}
func Logger(l utils.Logger) optSetter {
return func(f *Forwarder) error {
f.log = l
return nil
}
}
type Forwarder struct {
errHandler http.Handler
roundTripper http.RoundTripper
rewriter ReqRewriter
log utils.Logger
}
func New(setters ...optSetter) (*Forwarder, error) {
@@ -62,28 +71,40 @@ func New(setters ...optSetter) (*Forwarder, error) {
h = "localhost"
}
f.rewriter = &HeaderRewriter{TrustForwardHeader: true, Hostname: h}
}
if f.errHandler == nil {
f.errHandler = &utils.BadGatewayHandler{}
}
if f.log == nil {
f.log = &utils.NOPLogger{}
}
return f, nil
}
func (f *Forwarder) ServeHTTP(w http.ResponseWriter, req *http.Request) {
start := time.Now().UTC()
response, err := f.roundTripper.RoundTrip(f.copyRequest(req, req.URL))
if err != nil {
f.log.Errorf("Error forwarding to %v, err: %v", req.URL, err)
f.errHandler.ServeHTTP(w, req)
return
}
netutils.CopyHeaders(w.Header(), response.Header)
f.log.Infof("Forward to %v, code: %v, duration: %v",
req.URL, response.StatusCode, time.Now().UTC().Sub(start))
utils.CopyHeaders(w.Header(), response.Header)
w.WriteHeader(response.StatusCode)
io.Copy(w, response.Body)
response.Body.Close()
}
func (l *Forwarder) copyRequest(req *http.Request, u *url.URL) *http.Request {
func (f *Forwarder) copyRequest(req *http.Request, u *url.URL) *http.Request {
outReq := new(http.Request)
*outReq = *req // includes shallow copies of maps, but we handle this below
outReq.URL = utils.CopyURL(req.URL)
outReq.URL.Scheme = u.Scheme
outReq.URL.Host = u.Host
outReq.URL.Opaque = req.RequestURI
@@ -98,6 +119,10 @@ func (l *Forwarder) copyRequest(req *http.Request, u *url.URL) *http.Request {
outReq.Close = false
outReq.Header = make(http.Header)
netutils.CopyHeaders(outReq.Header, req.Header)
utils.CopyHeaders(outReq.Header, req.Header)
if f.rewriter != nil {
f.rewriter.Rewrite(outReq)
}
return outReq
}
+157 -2
View File
@@ -1,10 +1,14 @@
package forward
import (
"bytes"
"net/http"
"strings"
"testing"
"time"
"github.com/mailgun/oxy/testutils"
"github.com/mailgun/oxy/utils"
. "gopkg.in/check.v1"
)
@@ -15,10 +19,13 @@ type FwdSuite struct{}
var _ = Suite(&FwdSuite{})
func (s *FwdSuite) TestForwardSimple(c *C) {
// Makes sure hop-by-hop headers are removed
func (s *FwdSuite) TestForwardHopHeaders(c *C) {
called := false
var outHeaders http.Header
srv := testutils.NewTestServer(func(w http.ResponseWriter, req *http.Request) {
called = true
outHeaders = req.Header
w.Write([]byte("hello"))
})
defer srv.Close()
@@ -32,9 +39,157 @@ func (s *FwdSuite) TestForwardSimple(c *C) {
})
defer proxy.Close()
re, body, err := testutils.Get(proxy.URL, testutils.Opts{})
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, "")
}
func (s *FwdSuite) TestDefaultErrHandler(c *C) {
f, err := New()
c.Assert(err, IsNil)
proxy := testutils.NewTestServer(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(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
w.WriteHeader(http.StatusTeapot)
w.Write([]byte(http.StatusText(http.StatusTeapot)))
})))
c.Assert(err, IsNil)
proxy := testutils.NewTestServer(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))
}
// Makes sure hop-by-hop headers are removed
func (s *FwdSuite) TestForwardedHeaders(c *C) {
var outHeaders http.Header
srv := testutils.NewTestServer(func(w http.ResponseWriter, req *http.Request) {
outHeaders = req.Header
w.Write([]byte("hello"))
})
defer srv.Close()
f, err := New()
c.Assert(err, IsNil)
proxy := testutils.NewTestServer(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"},
}
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)
}
func (s *FwdSuite) TestCustomRewriter(c *C) {
var outHeaders http.Header
srv := testutils.NewTestServer(func(w http.ResponseWriter, req *http.Request) {
outHeaders = req.Header
w.Write([]byte("hello"))
})
defer srv.Close()
f, err := New(Rewriter(&HeaderRewriter{TrustForwardHeader: false, Hostname: "hello"}))
c.Assert(err, IsNil)
proxy := testutils.NewTestServer(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"},
}
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, "http")
c.Assert(strings.Contains(outHeaders.Get(XForwardedFor), "192.168.1.1"), Equals, false)
}
func (s *FwdSuite) TestCustomTransport(c *C) {
srv := testutils.NewTestServer(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.NewTestServer(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.StatusBadGateway)
}
func (s *FwdSuite) TestCustomLogger(c *C) {
srv := testutils.NewTestServer(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.NewTestServer(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)
}
+2 -2
View File
@@ -5,7 +5,7 @@ import (
"net/http"
"strings"
"github.com/mailgun/oxy/netutils"
"github.com/mailgun/oxy/utils"
)
// Rewriter is responsible for removing hop-by-hop headers and setting forwarding headers
@@ -39,5 +39,5 @@ func (rw *HeaderRewriter) Rewrite(req *http.Request) {
// Remove hop-by-hop headers to the backend. Especially important is "Connection" because we want a persistent
// connection, regardless of what the client sent to us.
netutils.RemoveHeaders(req.Header, HopHeaders...)
utils.RemoveHeaders(req.Header, HopHeaders...)
}
+2 -2
View File
@@ -4,7 +4,7 @@ import (
"io"
"net/http"
"github.com/mailgun/oxy/netutils"
"github.com/mailgun/oxy/utils"
)
type Retry struct {
@@ -38,7 +38,7 @@ func New(next http.Handler, predicate string, settings ...optSetter) (*Retry, er
func (r *Retry) ServeHTTP(w http.ResponseWriter, req *http.Request) {
c := &context{r: req}
for i := 0; i < r.maxAttempts; i++ {
pw := &netutils.ProxyWriter{W: w}
pw := &utils.ProxyWriter{W: w}
// this is not the first attempt, we need to rewind the body
if i != 0 {
+3 -3
View File
@@ -8,7 +8,7 @@ import (
"time"
"github.com/mailgun/oxy/memmetrics"
"github.com/mailgun/oxy/netutils"
"github.com/mailgun/oxy/utils"
"github.com/mailgun/timetools"
)
@@ -98,7 +98,7 @@ func NewRebalancer(handler balancerHandler, opts ...rbOptSetter) (*Rebalancer, e
}
func (rb *Rebalancer) ServeHTTP(w http.ResponseWriter, req *http.Request) {
pw := &netutils.ProxyWriter{W: w}
pw := &utils.ProxyWriter{W: w}
start := rb.clock.UtcNow()
rb.next.ServeHTTP(pw, req)
rb.recordMetrics(req.URL, pw.Code, rb.clock.UtcNow().Sub(start))
@@ -163,7 +163,7 @@ func (rb *Rebalancer) upsertServer(u *url.URL, weight int) {
if s, i := rb.findServer(u); i != -1 {
s.origWeight = weight
}
rb.servers = append(rb.servers, &rbServer{url: netutils.CopyURL(u), origWeight: weight, curWeight: weight})
rb.servers = append(rb.servers, &rbServer{url: utils.CopyURL(u), origWeight: weight, curWeight: weight})
}
func (r *Rebalancer) findServer(u *url.URL) (*rbServer, int) {
+3 -3
View File
@@ -7,7 +7,7 @@ import (
"net/url"
"sync"
"github.com/mailgun/oxy/netutils"
"github.com/mailgun/oxy/utils"
)
// Weight is an optional functional argument that sets weight of the server
@@ -60,7 +60,7 @@ func (r *RoundRobin) ServeHTTP(w http.ResponseWriter, req *http.Request) {
r.errHandler.ServeHTTP(w, req)
return
}
req.URL = netutils.CopyURL(srv.url)
req.URL = utils.CopyURL(srv.url)
r.next.ServeHTTP(w, req)
}
@@ -138,7 +138,7 @@ func (rr *RoundRobin) UpsertServer(u *url.URL, options ...serverSetter) error {
return fmt.Errorf("Server %v already exists", u)
}
srv := &server{url: netutils.CopyURL(u)}
srv := &server{url: utils.CopyURL(u)}
for _, o := range options {
if err := o(srv); err != nil {
return err
+64 -14
View File
@@ -9,7 +9,7 @@ import (
"net/url"
"strings"
"github.com/mailgun/oxy/netutils"
"github.com/mailgun/oxy/utils"
)
func NewTestServer(handler http.HandlerFunc) *httptest.Server {
@@ -31,25 +31,75 @@ func ParseURI(uri string) *url.URL {
return out
}
type Opts struct {
type reqOpts struct {
Host string
Method string
Body string
Headers http.Header
}
func MakeRequest(url string, opts Opts) (*http.Response, []byte, error) {
method := "GET"
if opts.Method != "" {
method = opts.Method
type reqOptSetter func(o *reqOpts) error
func Method(m string) reqOptSetter {
return func(o *reqOpts) error {
o.Method = m
return nil
}
request, _ := http.NewRequest(method, url, strings.NewReader(opts.Body))
if opts.Headers != nil {
netutils.CopyHeaders(request.Header, opts.Headers)
}
func Host(h string) reqOptSetter {
return func(o *reqOpts) error {
o.Host = h
return nil
}
}
func Body(b string) reqOptSetter {
return func(o *reqOpts) error {
o.Body = b
return nil
}
}
func Header(name, val string) reqOptSetter {
return func(o *reqOpts) error {
if o.Headers == nil {
o.Headers = make(http.Header)
}
o.Headers.Add(name, val)
return nil
}
}
func Headers(h http.Header) reqOptSetter {
return func(o *reqOpts) error {
if o.Headers == nil {
o.Headers = make(http.Header)
}
utils.CopyHeaders(o.Headers, h)
return nil
}
}
func MakeRequest(url string, opts ...reqOptSetter) (*http.Response, []byte, error) {
o := &reqOpts{}
for _, s := range opts {
if err := s(o); err != nil {
return nil, nil, err
}
}
if len(opts.Host) != 0 {
request.Host = opts.Host
method := "GET"
if o.Method == "" {
o.Method = "GET"
}
request, _ := http.NewRequest(method, url, strings.NewReader(o.Body))
if o.Headers != nil {
utils.CopyHeaders(request.Header, o.Headers)
}
if len(o.Host) != 0 {
request.Host = o.Host
}
var tr *http.Transport
@@ -78,7 +128,7 @@ func MakeRequest(url string, opts Opts) (*http.Response, []byte, error) {
return response, nil, err
}
func Get(url string, o Opts) (*http.Response, []byte, error) {
o.Method = "GET"
return MakeRequest(url, o)
func Get(url string, opts ...reqOptSetter) (*http.Response, []byte, error) {
opts = append(opts, Method("GET"))
return MakeRequest(url, opts...)
}
+13
View File
@@ -0,0 +1,13 @@
package utils
import (
"net/http"
)
type BadGatewayHandler struct {
}
func (e *BadGatewayHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
w.WriteHeader(http.StatusBadGateway)
w.Write([]byte(http.StatusText(http.StatusBadGateway)))
}
+109
View File
@@ -0,0 +1,109 @@
package utils
import (
"io"
"log"
)
// Logger defines a simple logging interface
type Logger interface {
Infof(format string, args ...interface{})
Warningf(format string, args ...interface{})
Errorf(format string, args ...interface{})
Info(message string)
Error(message string)
Warning(message string)
}
type FileLogger struct {
info *log.Logger
warn *log.Logger
error *log.Logger
}
func NewFileLogger(w io.Writer, lvl LogLevel) *FileLogger {
l := &FileLogger{}
flag := log.Ldate | log.Ltime | log.Lmicroseconds
if lvl >= INFO {
l.info = log.New(w, "INFO: ", flag)
}
if lvl >= WARN {
l.warn = log.New(w, "WARN: ", flag)
}
if lvl >= ERROR {
l.warn = log.New(w, "ERR: ", flag)
}
return l
}
func (f *FileLogger) Infof(format string, args ...interface{}) {
if f.info == nil {
return
}
f.info.Printf(format, args...)
}
func (f *FileLogger) Warningf(format string, args ...interface{}) {
if f.warn == nil {
return
}
f.warn.Printf(format, args...)
}
func (f *FileLogger) Errorf(format string, args ...interface{}) {
if f.error == nil {
return
}
f.error.Printf(format, args...)
}
func (f *FileLogger) Info(message string) {
if f.info == nil {
return
}
f.info.Print(message)
}
func (f *FileLogger) Warning(message string) {
if f.warn == nil {
return
}
f.warn.Print(message)
}
func (f *FileLogger) Error(message string) {
if f.error == nil {
return
}
f.error.Print(message)
}
type NOPLogger struct {
}
func (*NOPLogger) Infof(format string, args ...interface{}) {
}
func (*NOPLogger) Warningf(format string, args ...interface{}) {
}
func (*NOPLogger) Errorf(format string, args ...interface{}) {
}
func (*NOPLogger) Info(string) {
}
func (*NOPLogger) Warning(string) {
}
func (*NOPLogger) Error(string) {
}
type LogLevel int
const (
INFO = iota
WARN
ERROR
)
+1 -1
View File
@@ -1,4 +1,4 @@
package netutils
package utils
import (
"net/http"
@@ -1,4 +1,4 @@
package netutils
package utils
import (
"net/http"