Feature: Add DoRedirects, DoTimeout and DoDeadline to Proxy middleware

middleware/proxy/README.md

@@ -18,6 +18,8 @@ func Balancer(config Config) fiber.Handler
 func Forward(addr string, clients ...*fasthttp.Client) fiber.Handler
 // Do performs the given http request and fills the given http response.
 func Do(c *fiber.Ctx, addr string, clients ...*fasthttp.Client) error
+// DoRedirects performs the given http request and fills the given http response while following up to maxRedirectsCount redirects.
+func DoRedirects(c *fiber.Ctx, addr string, maxRedirectsCount int, clients ...*fasthttp.Client) error
 // 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

middleware/proxy/proxy.go

@@ -171,6 +171,39 @@ func Do(c *fiber.Ctx, addr string, clients ...*fasthttp.Client) error {
 	return nil
 }
 
+func DoRedirects(c *fiber.Ctx, addr string, maxRedirectsCount int, clients ...*fasthttp.Client) error {
+	var cli *fasthttp.Client
+	if len(clients) != 0 {
+		// Set local client
+		cli = clients[0]
+	} else {
+		// Set global client
+		lock.RLock()
+		cli = client
+		lock.RUnlock()
+	}
+	req := c.Request()
+	res := c.Response()
+	originalURL := utils.CopyString(c.OriginalURL())
+	defer req.SetRequestURI(originalURL)
+
+	copiedURL := utils.CopyString(addr)
+	req.SetRequestURI(copiedURL)
+	// NOTE: if req.isTLS is true, SetRequestURI keeps the scheme as https.
+	// issue reference:
+	// https://github.com/gofiber/fiber/issues/1762
+	if scheme := getScheme(utils.UnsafeBytes(copiedURL)); len(scheme) > 0 {
+		req.URI().SetSchemeBytes(scheme)
+	}
+
+	req.Header.Del(fiber.HeaderConnection)
+	if err := cli.DoRedirects(req, res, maxRedirectsCount); err != nil {
+		return err
+	}
+	res.Header.Del(fiber.HeaderConnection)
+	return nil
+}
+
 func getScheme(uri []byte) []byte {
 	i := bytes.IndexByte(uri, '/')
 	if i < 1 || uri[i-1] != ':' || i == len(uri)-1 || uri[i+1] != '/' {

middleware/proxy/proxy_test.go

@@ -350,19 +350,89 @@ func Test_Proxy_Do_RestoreOriginalURL(t *testing.T) {
 		return c.SendString("ok")
 	})
 	app.Get("/test", func(c *fiber.Ctx) error {
-		originalURL := utils.CopyString(c.OriginalURL())
 		if err := Do(c, "/proxy"); err != nil {
 			return err
 		}
-		utils.AssertEqual(t, originalURL, c.OriginalURL())
 		return c.SendString("ok")
 	})
-	_, err1 := app.Test(httptest.NewRequest(fiber.MethodGet, "/test", nil))
-	// This test requires multiple requests due to zero allocation used in fiber
-	_, err2 := app.Test(httptest.NewRequest(fiber.MethodGet, "/test", nil))
+	resp, err1 := app.Test(httptest.NewRequest(fiber.MethodGet, "/test", nil))
+	utils.AssertEqual(t, nil, err1)
+	utils.AssertEqual(t, "/test", resp.Request.URL.String())
+}
 
+// go test -race -run Test_Proxy_Do_WithRealURL
+func Test_Proxy_Do_WithRealURL(t *testing.T) {
+	t.Parallel()
+	app := fiber.New()
+	app.Get("/test", func(c *fiber.Ctx) error {
+		if err := Do(c, "https://www.google.com"); err != nil {
+			return err
+		}
+		return c.SendString("ok")
+	})
+
+	resp, err1 := app.Test(httptest.NewRequest(fiber.MethodGet, "/test", nil))
 	utils.AssertEqual(t, nil, err1)
-	utils.AssertEqual(t, nil, err2)
+	utils.AssertEqual(t, fiber.StatusOK, resp.StatusCode)
+	utils.AssertEqual(t, "/test", resp.Request.URL.String())
+}
+
+// go test -race -run Test_Proxy_Do_WithRedirect
+func Test_Proxy_Do_WithRedirect(t *testing.T) {
+	t.Parallel()
+	app := fiber.New()
+	app.Get("/test", func(c *fiber.Ctx) error {
+		if err := Do(c, "https://google.com"); err != nil {
+			return err
+		}
+		return c.SendString("ok")
+	})
+
+	resp, err1 := app.Test(httptest.NewRequest(fiber.MethodGet, "/test", nil))
+	utils.AssertEqual(t, nil, err1)
+	body, err := io.ReadAll(resp.Body)
+	utils.AssertEqual(t, nil, err)
+	utils.AssertEqual(t, "ok", string(body))
+	utils.AssertEqual(t, 301, resp.StatusCode)
+}
+
+// go test -race -run Test_Proxy_DoRedirects_RestoreOriginalURL
+func Test_Proxy_DoRedirects_RestoreOriginalURL(t *testing.T) {
+	t.Parallel()
+	app := fiber.New()
+	app.Get("/test", func(c *fiber.Ctx) error {
+		if err := DoRedirects(c, "http://google.com", 1); err != nil {
+			return err
+		}
+		return c.SendString("ok")
+	})
+
+	resp, err1 := app.Test(httptest.NewRequest(fiber.MethodGet, "/test", nil))
+	utils.AssertEqual(t, nil, err1)
+	_, err := io.ReadAll(resp.Body)
+	utils.AssertEqual(t, nil, err)
+	utils.AssertEqual(t, fiber.StatusOK, resp.StatusCode)
+	utils.AssertEqual(t, "/test", resp.Request.URL.String())
+}
+
+// go test -race -run Test_Proxy_DoRedirects_TooManyRedirects
+func Test_Proxy_DoRedirects_TooManyRedirects(t *testing.T) {
+	t.Parallel()
+	app := fiber.New()
+	app.Get("/test", func(c *fiber.Ctx) error {
+		if err := DoRedirects(c, "http://google.com", 0); err != nil {
+			return err
+		}
+		return c.SendString("ok")
+	})
+
+	resp, err1 := app.Test(httptest.NewRequest(fiber.MethodGet, "/test", nil))
+	utils.AssertEqual(t, nil, err1)
+	body, err := io.ReadAll(resp.Body)
+	utils.AssertEqual(t, nil, err)
+	utils.AssertEqual(t, string(body), "too many redirects detected when doing the request")
+	utils.AssertEqual(t, fiber.StatusInternalServerError, resp.StatusCode)
+	utils.AssertEqual(t, "/test", resp.Request.URL.String())
 }
 
 // go test -race -run Test_Proxy_Do_HTTP_Prefix_URL