Stream package covered by tests

forward/fwd.go

@@ -73,10 +73,10 @@ func New(setters ...optSetter) (*Forwarder, error) {
 		f.rewriter = &HeaderRewriter{TrustForwardHeader: true, Hostname: h}
 	}
 	if f.log == nil {
-		f.log = &utils.NOPLogger{}
+		f.log = utils.NullLogger
 	}
 	if f.errHandler == nil {
-		f.errHandler = &utils.DefaultHandler{}
+		f.errHandler = utils.DefaultHandler
 	}
 	return f, nil
 }

roundrobin/rr.go

@@ -52,7 +52,7 @@ func New(next http.Handler, opts ...rrSetter) (*RoundRobin, error) {
 		}
 	}
 	if rr.errHandler == nil {
-		rr.errHandler = &utils.DefaultHandler{}
+		rr.errHandler = utils.DefaultHandler
 	}
 	return rr, nil
 }

stream/stream.go

@@ -1,4 +1,4 @@
-// package streamer provides a special http.Handler middleware that solves several problems when dealing with http requests:
+// package stream provides a special http.Handler middleware that solves several problems when dealing with http requests:
 //
 // * Reads the entire request and response into buffer, optionally buffering it to disk for large requests
 //
@@ -6,7 +6,7 @@
 //
 // * Changes request content-transfer-encoding from chunked and provides total size to the handlers
 //
-package streamer
+package stream
 
 import (
 	"fmt"
@@ -14,6 +14,7 @@ import (
 	"net/http"
 
 	"github.com/mailgun/multibuf"
+	"github.com/mailgun/oxy/utils"
 )
 
 const (
@@ -23,6 +24,8 @@ const (
 	DefaultMaxBodyBytes = -1
 )
 
+var errHandler utils.ErrorHandler = &SizeErrHandler{}
+
 // Streamer is responsible for streaming requests and responses
 // It buffers large reqeuests and responses to disk,
 type Streamer struct {
@@ -33,7 +36,8 @@ type Streamer struct {
 	memResponseBodyBytes int64
 
 	next       http.Handler
-	errHandler http.Handler
+	errHandler utils.ErrorHandler
+	log        utils.Logger
 }
 
 // New returns a new streamer middleware. New() function supports optional functional arguments.
@@ -56,14 +60,29 @@ func New(next http.Handler, setters ...optSetter) (*Streamer, error) {
 			return nil, err
 		}
 	}
+	if strm.errHandler == nil {
+		strm.errHandler = errHandler
+	}
+
+	if strm.log == nil {
+		strm.log = utils.NullLogger
+	}
 
 	return strm, nil
 }
 
 type optSetter func(s *Streamer) error
 
+// Logger
+func Logger(l utils.Logger) optSetter {
+	return func(s *Streamer) error {
+		s.log = l
+		return nil
+	}
+}
+
 // ErrorHandler is a functional argument that sets error handler of the server
-func ErrorHandler(h http.Handler) optSetter {
+func ErrorHandler(h utils.ErrorHandler) optSetter {
 	return func(s *Streamer) error {
 		s.errHandler = h
 		return nil
@@ -119,8 +138,9 @@ func (s *Streamer) Wrap(next http.Handler) error {
 }
 
 func (s *Streamer) ServeHTTP(w http.ResponseWriter, req *http.Request) {
-	if s.isOverLimit(req) {
-		s.errHandler.ServeHTTP(w, req)
+	if err := s.checkLimit(req); err != nil {
+		s.log.Infof("request body over limit: %v", err)
+		s.errHandler.ServeHTTP(w, req, err)
 		return
 	}
 
@@ -128,9 +148,10 @@ func (s *Streamer) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 	// to read into memory and disk. This reader returns an error if the total request size exceeds the
 	// prefefined MaxSizeBytes. This can occur if we got chunked request, in this case ContentLength would be set to -1
 	// and the reader would be unbounded bufio in the http.Server
+
 	body, err := multibuf.New(req.Body, multibuf.MaxBytes(s.maxRequestBodyBytes), multibuf.MemBytes(s.memRequestBodyBytes))
 	if err != nil || body == nil {
-		s.errHandler.ServeHTTP(w, req)
+		s.errHandler.ServeHTTP(w, req, err)
 		return
 	}
 
@@ -146,17 +167,18 @@ func (s *Streamer) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 	// set without content length or using chunked TransferEncoding
 	totalSize, err := body.Size()
 	if err != nil {
-		s.errHandler.ServeHTTP(w, req)
+		s.log.Errorf("failed to get size, err %v", err)
+		s.errHandler.ServeHTTP(w, req, err)
 		return
 	}
-	req.ContentLength = totalSize
+	outreq.ContentLength = totalSize
 	// remove TransferEncoding that could have been previously set because we have transformed the request from chunked encoding
-	req.TransferEncoding = []string{}
+	outreq.TransferEncoding = []string{}
 
 	// We create a special writer that will limit the response size, buffer it to disk if necessary
 	writer, err := multibuf.NewWriterOnce(multibuf.MaxBytes(s.maxResponseBodyBytes), multibuf.MemBytes(s.memResponseBodyBytes))
 	if err != nil {
-		s.errHandler.ServeHTTP(w, req)
+		s.errHandler.ServeHTTP(w, req, err)
 		return
 	}
 
@@ -171,7 +193,8 @@ func (s *Streamer) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 
 	reader, err := writer.Reader()
 	if err != nil {
-		s.errHandler.ServeHTTP(w, req)
+		s.log.Errorf("failed to read response, err %v", err)
+		s.errHandler.ServeHTTP(w, req, err)
 		return
 	}
 	defer reader.Close()
@@ -181,11 +204,14 @@ func (s *Streamer) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 	io.Copy(w, reader)
 }
 
-func (s *Streamer) isOverLimit(req *http.Request) bool {
+func (s *Streamer) checkLimit(req *http.Request) error {
 	if s.maxRequestBodyBytes <= 0 {
-		return false
+		return nil
 	}
-	return req.ContentLength > s.maxRequestBodyBytes
+	if req.ContentLength > s.maxRequestBodyBytes {
+		return &multibuf.MaxSizeReachedError{MaxSize: s.maxRequestBodyBytes}
+	}
+	return nil
 }
 
 type bufferWriter struct {
@@ -218,3 +244,15 @@ func copyHeaders(dst, src http.Header) {
 		}
 	}
 }
+
+type SizeErrHandler struct {
+}
+
+func (e *SizeErrHandler) ServeHTTP(w http.ResponseWriter, req *http.Request, err error) {
+	if _, ok := err.(*multibuf.MaxSizeReachedError); ok {
+		w.WriteHeader(http.StatusRequestEntityTooLarge)
+		w.Write([]byte(http.StatusText(http.StatusRequestEntityTooLarge)))
+		return
+	}
+	utils.DefaultHandler.ServeHTTP(w, req, err)
+}

stream/stream_test.go

@@ -0,0 +1,240 @@
+package stream
+
+import (
+	"bufio"
+	"fmt"
+	"io/ioutil"
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"os"
+	"testing"
+
+	"github.com/mailgun/oxy/forward"
+	"github.com/mailgun/oxy/testutils"
+	"github.com/mailgun/oxy/utils"
+
+	. "gopkg.in/check.v1"
+)
+
+func TestStream(t *testing.T) { TestingT(t) }
+
+type STSuite struct{}
+
+var _ = Suite(&STSuite{})
+
+func (s *STSuite) TestSimple(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)))
+	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")
+}
+
+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 *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 *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) TestResponseLimitReached(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
+	st, err := New(rdr, Logger(utils.NewFileLogger(os.Stdout, utils.INFO)), MaxResponseBodyBytes(4))
+	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.StatusInternalServerError)
+}
+
+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) 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)
+}

utils/handler.go

@@ -9,10 +9,12 @@ type ErrorHandler interface {
 	ServeHTTP(w http.ResponseWriter, req *http.Request, err error)
 }
 
-type DefaultHandler struct {
+var DefaultHandler ErrorHandler = &StdHandler{}
+
+type StdHandler struct {
 }
 
-func (e *DefaultHandler) ServeHTTP(w http.ResponseWriter, req *http.Request, err error) {
+func (e *StdHandler) ServeHTTP(w http.ResponseWriter, req *http.Request, err error) {
 	statusCode := http.StatusInternalServerError
 	if e, ok := err.(net.Error); ok {
 		if e.Timeout() {

utils/logging.go

@@ -5,6 +5,8 @@ import (
 	"log"
 )
 
+var NullLogger Logger = &NOPLogger{}
+
 // Logger defines a simple logging interface
 type Logger interface {
 	Infof(format string, args ...interface{})