[middleware/proxy] add buffer size configuration (#1292)

* add buffer size

* add tests

middleware/proxy/config.go

@@ -28,6 +28,15 @@ type Config struct {
 	//
 	// Optional. Default: nil
 	ModifyResponse fiber.Handler
+
+	// Per-connection buffer size for requests' reading.
+	// This also limits the maximum header size.
+	// Increase this buffer if your clients send multi-KB RequestURIs
+	// and/or multi-KB headers (for example, BIG cookies).
+	ReadBufferSize int
+
+	// Per-connection buffer size for responses' writing.
+	WriteBufferSize int
 }
 
 // ConfigDefault is the default config

middleware/proxy/proxy.go

@@ -40,6 +40,9 @@ func Balancer(config Config) fiber.Handler {
 			NoDefaultUserAgentHeader: true,
 			DisablePathNormalizing:   true,
 			Addr:                     u.Host,
+
+			ReadBufferSize:  config.ReadBufferSize,
+			WriteBufferSize: config.WriteBufferSize,
 		}
 
 		lbc.Clients = append(lbc.Clients, client)

middleware/proxy/proxy_test.go

@@ -4,6 +4,7 @@ import (
 	"io/ioutil"
 	"net"
 	"net/http/httptest"
+	"strings"
 	"testing"
 	"time"
 
@@ -178,3 +179,41 @@ func Test_Proxy_Modify_Request(t *testing.T) {
 	utils.AssertEqual(t, nil, err)
 	utils.AssertEqual(t, "modified request", string(b))
 }
+
+func Test_Proxy_Buffer_Size_Response(t *testing.T) {
+	t.Parallel()
+
+	target := fiber.New(fiber.Config{DisableStartupMessage: true})
+	target.Get("/", func(c *fiber.Ctx) error {
+		long := strings.Join(make([]string, 5000), "-")
+		c.Response().Header.Set("Very-Long-Header", long)
+		return c.SendString("ok")
+	})
+
+	ln, err := net.Listen(fiber.NetworkTCP4, "127.0.0.1:0")
+	utils.AssertEqual(t, nil, err)
+
+	go func() {
+		utils.AssertEqual(t, nil, target.Listener(ln))
+	}()
+
+	time.Sleep(2 * time.Second)
+	addr := ln.Addr().String()
+
+	app := fiber.New()
+	app.Use(Balancer(Config{Servers: []string{addr}}))
+
+	resp, err := app.Test(httptest.NewRequest("GET", "/", nil))
+	utils.AssertEqual(t, nil, err)
+	utils.AssertEqual(t, fiber.StatusInternalServerError, resp.StatusCode)
+
+	app = fiber.New()
+	app.Use(Balancer(Config{
+		Servers:        []string{addr},
+		ReadBufferSize: 1024 * 8,
+	}))
+
+	resp, err = app.Test(httptest.NewRequest("GET", "/", nil))
+	utils.AssertEqual(t, nil, err)
+	utils.AssertEqual(t, fiber.StatusOK, resp.StatusCode)
+}