🔥 Add DialDualStack option to proxy middleware for upstream IPv6 support

negrel · negrel · commit 472bd2166f36 · 2024-03-05T14:18:23.000+01:00
diff --git a/docs/api/middleware/proxy.md b/docs/api/middleware/proxy.md
@@ -9,7 +9,7 @@ Proxy middleware for [Fiber](https://github.com/gofiber/fiber) that allows you t
 ## Signatures
 
 ```go
-// Balancer create a load balancer among multiple upstrem servers.
+// Balancer create a load balancer among multiple upstream servers.
 func Balancer(config Config) fiber.Handler
 // Forward performs the given http request and fills the given http response.
 func Forward(addr string, clients ...*fasthttp.Client) fiber.Handler
@@ -21,9 +21,9 @@ func DoRedirects(c fiber.Ctx, addr string, maxRedirectsCount int, clients ...*fa
 func DoDeadline(c fiber.Ctx, addr string, deadline time.Time, clients ...*fasthttp.Client) error
 // DoTimeout performs the given request and waits for response during the given timeout duration.
 func DoTimeout(c fiber.Ctx, addr string, timeout time.Duration, clients ...*fasthttp.Client) error
-// DomainForward the given http request based on the given domain and fills the given http response
+// DomainForward the given http request based on the given domain and fills the given http response.
 func DomainForward(hostname string, addr string, clients ...*fasthttp.Client) fiber.Handler
-// BalancerForward performs the given http request based round robin balancer and fills the given http response
+// BalancerForward performs the given http request based round robin balancer and fills the given http response.
 func BalancerForward(servers []string, clients ...*fasthttp.Client) fiber.Handler
 ```
 
@@ -137,6 +137,17 @@ app.Use(proxy.BalancerForward([]string{
     "http://localhost:3002",
     "http://localhost:3003",
 }))
+
+
+// Make round robin balancer with IPv6 support.
+app.Use(proxy.Balancer(proxy.Config{
+    Servers: []string{
+        "http://[::1]:3001",
+        "http://127.0.0.1:3002",
+        "http://localhost:3003",
+    },
+    DialDualStack: true,
+}))
 ```
 
 ## Config
@@ -151,6 +162,7 @@ app.Use(proxy.BalancerForward([]string{
 | ReadBufferSize  | `int`                                          | 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). | (Not specified) |
 | WriteBufferSize | `int`                                          | Per-connection buffer size for responses' writing.                                                                                                                                                             | (Not specified) |
 | TlsConfig       | `*tls.Config` (or `*fasthttp.TLSConfig` in v3) | TLS config for the HTTP client.                                                                                                                                                                                | `nil`           |
+| DialDualStack   | `bool`                                         | Client will attempt to connect to both IPv4 and IPv6 host addresses if set to true.                                                                                                                            | `false`         |
 | Client          | `*fasthttp.LBClient`                           | Client is a custom client when client config is complex.                                                                                                                                                       | `nil`           |
 
 ## Default Config
diff --git a/middleware/proxy/config.go b/middleware/proxy/config.go
@@ -51,9 +51,17 @@ type Config struct {
 	TlsConfig *tls.Config //nolint:stylecheck,revive // TODO: Rename to "TLSConfig" in v3
 
 	// Client is custom client when client config is complex.
-	// Note that Servers, Timeout, WriteBufferSize, ReadBufferSize and TlsConfig
-	// will not be used if the client are set.
+	// Note that Servers, Timeout, WriteBufferSize, ReadBufferSize, TlsConfig
+	// and DialDualStack will not be used if the client are set.
 	Client *fasthttp.LBClient
+
+	// Attempt to connect to both ipv4 and ipv6 host addresses if set to true.
+	//
+	// By default client connects only to ipv4 addresses, since unfortunately ipv6
+	// remains broken in many networks worldwide :)
+	//
+	// Optional. Default: false
+	DialDualStack bool
 }
 
 // ConfigDefault is the default config
diff --git a/middleware/proxy/proxy.go b/middleware/proxy/proxy.go
@@ -45,6 +45,8 @@ func Balancer(config Config) fiber.Handler {
 				WriteBufferSize: config.WriteBufferSize,
 
 				TLSConfig: config.TlsConfig,
+
+				DialDualStack: config.DialDualStack,
 			}
 
 			lbc.Clients = append(lbc.Clients, client)
diff --git a/middleware/proxy/proxy_test.go b/middleware/proxy/proxy_test.go
@@ -18,13 +18,13 @@ import (
 	"github.com/valyala/fasthttp"
 )
 
-func createProxyTestServer(t *testing.T, handler fiber.Handler) (*fiber.App, string) {
+func createProxyTestServer(t *testing.T, handler fiber.Handler, network, address string) (*fiber.App, string) {
 	t.Helper()
 
 	target := fiber.New()
 	target.Get("/", handler)
 
-	ln, err := net.Listen(fiber.NetworkTCP4, "127.0.0.1:0")
+	ln, err := net.Listen(network, address)
 	require.NoError(t, err)
 
 	go func() {
@@ -39,6 +39,16 @@ func createProxyTestServer(t *testing.T, handler fiber.Handler) (*fiber.App, str
 	return target, addr
 }
 
+func createProxyTestServerIPv4(t *testing.T, handler fiber.Handler) (*fiber.App, string) {
+	t.Helper()
+	return createProxyTestServer(t, handler, fiber.NetworkTCP4, "127.0.0.1:0")
+}
+
+func createProxyTestServerIPv6(t *testing.T, handler fiber.Handler) (*fiber.App, string) {
+	t.Helper()
+	return createProxyTestServer(t, handler, fiber.NetworkTCP6, "[::1]:0")
+}
+
 // go test -run Test_Proxy_Empty_Host
 func Test_Proxy_Empty_Upstream_Servers(t *testing.T) {
 	t.Parallel()
@@ -86,7 +96,7 @@ func Test_Proxy_Next(t *testing.T) {
 func Test_Proxy(t *testing.T) {
 	t.Parallel()
 
-	target, addr := createProxyTestServer(t, func(c fiber.Ctx) error {
+	target, addr := createProxyTestServerIPv4(t, func(c fiber.Ctx) error {
 		return c.SendStatus(fiber.StatusTeapot)
 	})
 
@@ -148,11 +158,86 @@ func Test_Proxy_Balancer_WithTlsConfig(t *testing.T) {
 	resp.Close()
 }
 
+// go test -run Test_Proxy_Balancer_IPv6_Upstream
+func Test_Proxy_Balancer_IPv6_Upstream(t *testing.T) {
+	t.Parallel()
+
+	target, addr := createProxyTestServerIPv6(t, func(c fiber.Ctx) error {
+		return c.SendStatus(fiber.StatusTeapot)
+	})
+
+	resp, err := target.Test(httptest.NewRequest(fiber.MethodGet, "/", nil), 2*time.Second)
+	require.NoError(t, err)
+	require.Equal(t, fiber.StatusTeapot, resp.StatusCode)
+
+	app := fiber.New()
+
+	app.Use(Balancer(Config{Servers: []string{addr}}))
+
+	req := httptest.NewRequest(fiber.MethodGet, "/", nil)
+	req.Host = addr
+	resp, err = app.Test(req)
+	require.NoError(t, err)
+	require.Equal(t, fiber.StatusInternalServerError, resp.StatusCode)
+}
+
+// go test -run Test_Proxy_Balancer_IPv6_Upstream
+func Test_Proxy_Balancer_IPv6_Upstream_With_DialDualStack(t *testing.T) {
+	t.Parallel()
+
+	target, addr := createProxyTestServerIPv6(t, func(c fiber.Ctx) error {
+		return c.SendStatus(fiber.StatusTeapot)
+	})
+
+	resp, err := target.Test(httptest.NewRequest(fiber.MethodGet, "/", nil), 2*time.Second)
+	require.NoError(t, err)
+	require.Equal(t, fiber.StatusTeapot, resp.StatusCode)
+
+	app := fiber.New()
+
+	app.Use(Balancer(Config{
+		Servers:       []string{addr},
+		DialDualStack: true,
+	}))
+
+	req := httptest.NewRequest(fiber.MethodGet, "/", nil)
+	req.Host = addr
+	resp, err = app.Test(req)
+	require.NoError(t, err)
+	require.Equal(t, fiber.StatusTeapot, resp.StatusCode)
+}
+
+// go test -run Test_Proxy_Balancer_IPv6_Upstream
+func Test_Proxy_Balancer_IPv4_Upstream_With_DialDualStack(t *testing.T) {
+	t.Parallel()
+
+	target, addr := createProxyTestServerIPv4(t, func(c fiber.Ctx) error {
+		return c.SendStatus(fiber.StatusTeapot)
+	})
+
+	resp, err := target.Test(httptest.NewRequest(fiber.MethodGet, "/", nil), 2*time.Second)
+	require.NoError(t, err)
+	require.Equal(t, fiber.StatusTeapot, resp.StatusCode)
+
+	app := fiber.New()
+
+	app.Use(Balancer(Config{
+		Servers:       []string{addr},
+		DialDualStack: true,
+	}))
+
+	req := httptest.NewRequest(fiber.MethodGet, "/", nil)
+	req.Host = addr
+	resp, err = app.Test(req)
+	require.NoError(t, err)
+	require.Equal(t, fiber.StatusTeapot, resp.StatusCode)
+}
+
 // go test -run Test_Proxy_Forward_WithTlsConfig_To_Http
 func Test_Proxy_Forward_WithTlsConfig_To_Http(t *testing.T) {
 	t.Parallel()
 
-	_, targetAddr := createProxyTestServer(t, func(c fiber.Ctx) error {
+	_, targetAddr := createProxyTestServerIPv4(t, func(c fiber.Ctx) error {
 		return c.SendString("hello from target")
 	})
 
@@ -193,7 +278,7 @@ func Test_Proxy_Forward(t *testing.T) {
 
 	app := fiber.New()
 
-	_, addr := createProxyTestServer(t, func(c fiber.Ctx) error {
+	_, addr := createProxyTestServerIPv4(t, func(c fiber.Ctx) error {
 		return c.SendString("forwarded")
 	})
 
@@ -255,7 +340,7 @@ func Test_Proxy_Forward_WithClient_TLSConfig(t *testing.T) {
 func Test_Proxy_Modify_Response(t *testing.T) {
 	t.Parallel()
 
-	_, addr := createProxyTestServer(t, func(c fiber.Ctx) error {
+	_, addr := createProxyTestServerIPv4(t, func(c fiber.Ctx) error {
 		return c.Status(500).SendString("not modified")
 	})
 
@@ -281,7 +366,7 @@ func Test_Proxy_Modify_Response(t *testing.T) {
 func Test_Proxy_Modify_Request(t *testing.T) {
 	t.Parallel()
 
-	_, addr := createProxyTestServer(t, func(c fiber.Ctx) error {
+	_, addr := createProxyTestServerIPv4(t, func(c fiber.Ctx) error {
 		b := c.Request().Body()
 		return c.SendString(string(b))
 	})
@@ -308,7 +393,7 @@ func Test_Proxy_Modify_Request(t *testing.T) {
 func Test_Proxy_Timeout_Slow_Server(t *testing.T) {
 	t.Parallel()
 
-	_, addr := createProxyTestServer(t, func(c fiber.Ctx) error {
+	_, addr := createProxyTestServerIPv4(t, func(c fiber.Ctx) error {
 		time.Sleep(300 * time.Millisecond)
 		return c.SendString("fiber is awesome")
 	})
@@ -332,7 +417,7 @@ func Test_Proxy_Timeout_Slow_Server(t *testing.T) {
 func Test_Proxy_With_Timeout(t *testing.T) {
 	t.Parallel()
 
-	_, addr := createProxyTestServer(t, func(c fiber.Ctx) error {
+	_, addr := createProxyTestServerIPv4(t, func(c fiber.Ctx) error {
 		time.Sleep(1 * time.Second)
 		return c.SendString("fiber is awesome")
 	})
@@ -356,7 +441,7 @@ func Test_Proxy_With_Timeout(t *testing.T) {
 func Test_Proxy_Buffer_Size_Response(t *testing.T) {
 	t.Parallel()
 
-	_, addr := createProxyTestServer(t, func(c fiber.Ctx) error {
+	_, addr := createProxyTestServerIPv4(t, func(c fiber.Ctx) error {
 		long := strings.Join(make([]string, 5000), "-")
 		c.Set("Very-Long-Header", long)
 		return c.SendString("ok")
@@ -383,7 +468,7 @@ func Test_Proxy_Buffer_Size_Response(t *testing.T) {
 // go test -race -run Test_Proxy_Do_RestoreOriginalURL
 func Test_Proxy_Do_RestoreOriginalURL(t *testing.T) {
 	t.Parallel()
-	_, addr := createProxyTestServer(t, func(c fiber.Ctx) error {
+	_, addr := createProxyTestServerIPv4(t, func(c fiber.Ctx) error {
 		return c.SendString("proxied")
 	})
 
@@ -470,7 +555,7 @@ func Test_Proxy_DoRedirects_TooManyRedirects(t *testing.T) {
 func Test_Proxy_DoTimeout_RestoreOriginalURL(t *testing.T) {
 	t.Parallel()
 
-	_, addr := createProxyTestServer(t, func(c fiber.Ctx) error {
+	_, addr := createProxyTestServerIPv4(t, func(c fiber.Ctx) error {
 		return c.SendString("proxied")
 	})
 
@@ -492,7 +577,7 @@ func Test_Proxy_DoTimeout_RestoreOriginalURL(t *testing.T) {
 func Test_Proxy_DoTimeout_Timeout(t *testing.T) {
 	t.Parallel()
 
-	_, addr := createProxyTestServer(t, func(c fiber.Ctx) error {
+	_, addr := createProxyTestServerIPv4(t, func(c fiber.Ctx) error {
 		time.Sleep(time.Second * 5)
 		return c.SendString("proxied")
 	})
@@ -510,7 +595,7 @@ func Test_Proxy_DoTimeout_Timeout(t *testing.T) {
 func Test_Proxy_DoDeadline_RestoreOriginalURL(t *testing.T) {
 	t.Parallel()
 
-	_, addr := createProxyTestServer(t, func(c fiber.Ctx) error {
+	_, addr := createProxyTestServerIPv4(t, func(c fiber.Ctx) error {
 		return c.SendString("proxied")
 	})
 
@@ -532,7 +617,7 @@ func Test_Proxy_DoDeadline_RestoreOriginalURL(t *testing.T) {
 func Test_Proxy_DoDeadline_PastDeadline(t *testing.T) {
 	t.Parallel()
 
-	_, addr := createProxyTestServer(t, func(c fiber.Ctx) error {
+	_, addr := createProxyTestServerIPv4(t, func(c fiber.Ctx) error {
 		time.Sleep(time.Second * 5)
 		return c.SendString("proxied")
 	})
@@ -550,7 +635,7 @@ func Test_Proxy_DoDeadline_PastDeadline(t *testing.T) {
 func Test_Proxy_Do_HTTP_Prefix_URL(t *testing.T) {
 	t.Parallel()
 
-	_, addr := createProxyTestServer(t, func(c fiber.Ctx) error {
+	_, addr := createProxyTestServerIPv4(t, func(c fiber.Ctx) error {
 		return c.SendString("hello world")
 	})
 
@@ -641,7 +726,7 @@ func Test_Proxy_Forward_Local_Client(t *testing.T) {
 func Test_ProxyBalancer_Custom_Client(t *testing.T) {
 	t.Parallel()
 
-	target, addr := createProxyTestServer(t, func(c fiber.Ctx) error {
+	target, addr := createProxyTestServerIPv4(t, func(c fiber.Ctx) error {
 		return c.SendStatus(fiber.StatusTeapot)
 	})
 
@@ -721,7 +806,7 @@ func Test_Proxy_Balancer_Forward_Local(t *testing.T) {
 
 	app := fiber.New()
 
-	_, addr := createProxyTestServer(t, func(c fiber.Ctx) error {
+	_, addr := createProxyTestServerIPv4(t, func(c fiber.Ctx) error {
 		return c.SendString("forwarded")
 	})
 

Original file line number	Diff line number	Diff line change
`@@ -45,6 +45,8 @@ func Balancer(config Config) fiber.Handler {`
`45`	`45`	`WriteBufferSize: config.WriteBufferSize,`
`46`	`46`
`47`	`47`	`TLSConfig: config.TlsConfig,`
	`48`	`+`
	`49`	`+ DialDualStack: config.DialDualStack,`
`48`	`50`	`}`
`49`	`51`
`50`	`52`	`lbc.Clients = append(lbc.Clients, client)`