From 9860d34a5e8e4bb9ce7bd9e77a412be916acc41e Mon Sep 17 00:00:00 2001 From: kinggo Date: Sun, 25 Sep 2022 02:10:11 +0800 Subject: [PATCH] optimize: add WithClient --- internal/gopsutil/net/net_darwin.go | 3 +- middleware/proxy/README.md | 23 ++++++-- middleware/proxy/config.go | 7 ++- middleware/proxy/proxy.go | 74 +++++++++++++++---------- middleware/proxy/proxy_test.go | 84 +++++++++++++++++++++++++++++ 5 files changed, 156 insertions(+), 35 deletions(-) diff --git a/internal/gopsutil/net/net_darwin.go b/internal/gopsutil/net/net_darwin.go index da294f63707..32d67c0d510 100644 --- a/internal/gopsutil/net/net_darwin.go +++ b/internal/gopsutil/net/net_darwin.go @@ -7,11 +7,12 @@ import ( "context" "errors" "fmt" - "github.com/gofiber/fiber/v2/internal/gopsutil/common" "os/exec" "regexp" "strconv" "strings" + + "github.com/gofiber/fiber/v2/internal/gopsutil/common" ) var ( diff --git a/middleware/proxy/README.md b/middleware/proxy/README.md index c3ceee9dfdc..eb77f89bc2b 100644 --- a/middleware/proxy/README.md +++ b/middleware/proxy/README.md @@ -13,8 +13,8 @@ Proxy middleware for [Fiber](https://github.com/gofiber/fiber) that allows you t ```go func Balancer(config Config) fiber.Handler -func Forward(addr string) fiber.Handler -func Do(c *fiber.Ctx, addr string) error +func Forward(addr string, clients ...*fasthttp.Client) fiber.Handler +func Do(c *fiber.Ctx, addr string, clients ...*fasthttp.Client) error ``` ### Examples @@ -37,9 +37,21 @@ proxy.WithTlsConfig(&tls.Config{ InsecureSkipVerify: true, }) +// if you need to use global self-custom client, you should use proxy.WithClient. +proxy.WithClient(&fasthttp.Client{ + NoDefaultUserAgentHeader: true, + DisablePathNormalizing: true, +}) + // Forward to url app.Get("/gif", proxy.Forward("https://i.imgur.com/IWaBepg.gif")) +// Forward to url with local custom client +app.Get("/gif", proxy.Forward("https://i.imgur.com/IWaBepg.gif", &fasthttp.Client{ + NoDefaultUserAgentHeader: true, + DisablePathNormalizing: true, +})) + // Make request within handler app.Get("/:id", func(c *fiber.Ctx) error { url := "https://i.imgur.com/"+c.Params("id")+".gif" @@ -120,8 +132,11 @@ type Config struct { // Per-connection buffer size for responses' writing. WriteBufferSize int - // tls config for the http client - TlsConfig *tls.Config + // tls config for the http client. + TlsConfig *tls.Config + + // Client is custom client when client config is complex. + Client *fasthttp.LBClient } ``` diff --git a/middleware/proxy/config.go b/middleware/proxy/config.go index 80b6472514a..ea6c474cf0c 100644 --- a/middleware/proxy/config.go +++ b/middleware/proxy/config.go @@ -47,8 +47,11 @@ type Config struct { // Per-connection buffer size for responses' writing. WriteBufferSize int - // tls config for the http client + // tls config for the http client. TlsConfig *tls.Config + + // Client is custom client when client config is complex. + Client *fasthttp.LBClient } // ConfigDefault is the default config @@ -75,7 +78,7 @@ func configDefault(config ...Config) Config { } // Set default values - if len(cfg.Servers) == 0 { + if len(cfg.Servers) == 0 && cfg.Client == nil { panic("Servers cannot be empty") } return cfg diff --git a/middleware/proxy/proxy.go b/middleware/proxy/proxy.go index bb96b03382e..2263b4b72fc 100644 --- a/middleware/proxy/proxy.go +++ b/middleware/proxy/proxy.go @@ -24,34 +24,37 @@ func Balancer(config Config) fiber.Handler { cfg := configDefault(config) // Load balanced client - var lbc fasthttp.LBClient - // Set timeout - lbc.Timeout = cfg.Timeout - - // Scheme must be provided, falls back to http - // TODO add https support - for _, server := range cfg.Servers { - if !strings.HasPrefix(server, "http") { - server = "http://" + server - } + var lbc = &fasthttp.LBClient{} + if config.Client == nil { + // Set timeout + lbc.Timeout = cfg.Timeout + // Scheme must be provided, falls back to http + for _, server := range cfg.Servers { + if !strings.HasPrefix(server, "http") { + server = "http://" + server + } - u, err := url.Parse(server) - if err != nil { - panic(err) - } + u, err := url.Parse(server) + if err != nil { + panic(err) + } - client := &fasthttp.HostClient{ - NoDefaultUserAgentHeader: true, - DisablePathNormalizing: true, - Addr: u.Host, + client := &fasthttp.HostClient{ + NoDefaultUserAgentHeader: true, + DisablePathNormalizing: true, + Addr: u.Host, - ReadBufferSize: config.ReadBufferSize, - WriteBufferSize: config.WriteBufferSize, + ReadBufferSize: config.ReadBufferSize, + WriteBufferSize: config.WriteBufferSize, - TLSConfig: config.TlsConfig, - } + TLSConfig: config.TlsConfig, + } - lbc.Clients = append(lbc.Clients, client) + lbc.Clients = append(lbc.Clients, client) + } + } else { + // Set custom client + lbc = config.Client } // Return new handler @@ -97,28 +100,43 @@ func Balancer(config Config) fiber.Handler { } } -var client = fasthttp.Client{ +var client = &fasthttp.Client{ NoDefaultUserAgentHeader: true, DisablePathNormalizing: true, } // WithTlsConfig update http client with a user specified tls.config // This function should be called before Do and Forward. +// Deprecated: use WithClient instead. func WithTlsConfig(tlsConfig *tls.Config) { client.TLSConfig = tlsConfig } +// WithClient sets the global proxy client. +// This function should be called before Do and Forward. +func WithClient(cli *fasthttp.Client) { + client = cli +} + // Forward performs the given http request and fills the given http response. // This method will return an fiber.Handler -func Forward(addr string) fiber.Handler { +func Forward(addr string, clients ...*fasthttp.Client) fiber.Handler { return func(c *fiber.Ctx) error { - return Do(c, addr) + return Do(c, addr, clients...) } } // Do performs the given http request and fills the given http response. // This method can be used within a fiber.Handler -func Do(c *fiber.Ctx, addr string) error { +func Do(c *fiber.Ctx, addr string, clients ...*fasthttp.Client) error { + var cli *fasthttp.Client + if len(clients) != 0 { + // Set local client + cli = clients[0] + } else { + // Set global client + cli = client + } req := c.Request() res := c.Response() originalURL := utils.CopyString(c.OriginalURL()) @@ -134,7 +152,7 @@ func Do(c *fiber.Ctx, addr string) error { } req.Header.Del(fiber.HeaderConnection) - if err := client.Do(req, res); err != nil { + if err := cli.Do(req, res); err != nil { return err } res.Header.Del(fiber.HeaderConnection) diff --git a/middleware/proxy/proxy_test.go b/middleware/proxy/proxy_test.go index 08160b2e868..47ec14a5fd9 100644 --- a/middleware/proxy/proxy_test.go +++ b/middleware/proxy/proxy_test.go @@ -13,6 +13,7 @@ import ( "github.com/gofiber/fiber/v2" "github.com/gofiber/fiber/v2/internal/tlstest" "github.com/gofiber/fiber/v2/utils" + "github.com/valyala/fasthttp" ) func createProxyTestServer(handler fiber.Handler, t *testing.T) (*fiber.App, string) { @@ -364,6 +365,7 @@ func Test_Proxy_Do_RestoreOriginalURL(t *testing.T) { utils.AssertEqual(t, nil, err2) } +// go test -race -run Test_Proxy_Do_HTTP_Prefix_URL func Test_Proxy_Do_HTTP_Prefix_URL(t *testing.T) { t.Parallel() @@ -390,3 +392,85 @@ func Test_Proxy_Do_HTTP_Prefix_URL(t *testing.T) { utils.AssertEqual(t, nil, err) utils.AssertEqual(t, "hello world", string(s)) } + +// go test -race -run Test_Proxy_Forward_Global_Client +func Test_Proxy_Forward_Global_Client(t *testing.T) { + t.Parallel() + ln, err := net.Listen(fiber.NetworkTCP4, "127.0.0.1:0") + utils.AssertEqual(t, nil, err) + WithClient(&fasthttp.Client{ + NoDefaultUserAgentHeader: true, + DisablePathNormalizing: true, + }) + app := fiber.New(fiber.Config{DisableStartupMessage: true}) + app.Get("/test_global_client", func(c *fiber.Ctx) error { + return c.SendString("test_global_client") + }) + + addr := ln.Addr().String() + app.Use(Forward("http://" + addr + "/test_global_client")) + go func() { utils.AssertEqual(t, nil, app.Listener(ln)) }() + + code, body, errs := fiber.Get("http://" + addr).String() + utils.AssertEqual(t, 0, len(errs)) + utils.AssertEqual(t, fiber.StatusOK, code) + utils.AssertEqual(t, "test_global_client", body) +} + +// go test -race -run Test_Proxy_Forward_Local_Client +func Test_Proxy_Forward_Local_Client(t *testing.T) { + t.Parallel() + ln, err := net.Listen(fiber.NetworkTCP4, "127.0.0.1:0") + utils.AssertEqual(t, nil, err) + app := fiber.New(fiber.Config{DisableStartupMessage: true}) + app.Get("/test_local_client", func(c *fiber.Ctx) error { + return c.SendString("test_local_client") + }) + + addr := ln.Addr().String() + app.Use(Forward("http://"+addr+"/test_local_client", &fasthttp.Client{ + NoDefaultUserAgentHeader: true, + DisablePathNormalizing: true, + Dial: func(addr string) (net.Conn, error) { + return fasthttp.Dial(addr) + }, + })) + go func() { utils.AssertEqual(t, nil, app.Listener(ln)) }() + + code, body, errs := fiber.Get("http://" + addr).String() + utils.AssertEqual(t, 0, len(errs)) + utils.AssertEqual(t, fiber.StatusOK, code) + utils.AssertEqual(t, "test_local_client", body) +} + +// go test -run Test_ProxyBalancer_Custom_Client +func Test_ProxyBalancer_Custom_Client(t *testing.T) { + t.Parallel() + + target, addr := createProxyTestServer( + func(c *fiber.Ctx) error { return c.SendStatus(fiber.StatusTeapot) }, t, + ) + + resp, err := target.Test(httptest.NewRequest("GET", "/", nil), 2000) + utils.AssertEqual(t, nil, err) + utils.AssertEqual(t, fiber.StatusTeapot, resp.StatusCode) + + app := fiber.New(fiber.Config{DisableStartupMessage: true}) + + app.Use(Balancer(Config{Client: &fasthttp.LBClient{ + Clients: []fasthttp.BalancingClient{ + &fasthttp.HostClient{ + NoDefaultUserAgentHeader: true, + DisablePathNormalizing: true, + Addr: addr, + }, + }, + Timeout: time.Second, + }})) + + req := httptest.NewRequest("GET", "/", nil) + req.Host = addr + resp, err = app.Test(req) + utils.AssertEqual(t, nil, err) + utils.AssertEqual(t, fiber.StatusTeapot, resp.StatusCode) +}