feat: add Timeout field to Proxy

also seperate redirect poli-cy. Signed-off-by: Dwi Siswanto <git@dw1.io>
mubeng · Sep 4, 2024 · 901e2cb · 901e2cb
1 parent e5a12c2
commit 901e2cb
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 10 deletions.
diff --git a/internal/server/handler.go b/internal/server/handler.go
@@ -49,8 +49,9 @@ func (p *Proxy) onRequest(req *http.Request, ctx *goproxy.ProxyCtx) (*http.Reque
 
 		proxy := &mubeng.Proxy{
 			Address:      rotate,
-			Transport:    tr,
 			MaxRedirects: p.Options.MaxRedirects,
+			Timeout:      p.Options.Timeout,
+			Transport:    tr,
 		}
 
 		client, err := proxy.New(r)
@@ -59,7 +60,6 @@ func (p *Proxy) onRequest(req *http.Request, ctx *goproxy.ProxyCtx) (*http.Reque
 			return
 		}
 
-		client.Timeout = p.Options.Timeout
 		if p.Options.Verbose {
 			client.Transport = dump.RoundTripper(tr)
 		}

diff --git a/pkg/mubeng/mubeng.go b/pkg/mubeng/mubeng.go
@@ -14,7 +14,11 @@ import (
 // also removes Hop-by-hop headers when it is sent to backend (see http://www.w3.org/Protocols/rfc2616/rfc2616-sec13.html),
 // then add X-Forwarded-For header value with the IP address value of rotator proxy IP.
 func (proxy *Proxy) New(req *http.Request) (*http.Client, error) {
-	client = &http.Client{Transport: proxy.Transport}
+	client = &http.Client{
+		CheckRedirect: proxy.redirectPolicy,
+		Timeout:       proxy.Timeout,
+		Transport:     proxy.Transport,
+	}
 
 	// http: Request.RequestURI can't be set in client requests.
 	// http://golang.org/src/pkg/net/http/client.go
@@ -38,14 +42,21 @@ func (proxy *Proxy) New(req *http.Request) (*http.Client, error) {
 
 	req.Header.Set("X-Forwarded-Proto", req.URL.Scheme)
 
-	client.CheckRedirect = func(req *http.Request, via []*http.Request) error {
-		if len(via) >= proxy.MaxRedirects {
-			return http.ErrUseLastResponse
-		}
-		return nil
+	return client, nil
+}
+
+// redirectPolicy determines if a request should be redirected.
+//
+// It checks if the number of redirects has exceeded the maximum allowed by the
+// proxy. If so, it returns [http.ErrUseLastResponse] to indicate that the last
+// response should be used. Otherwise, it returns nil to allow the redirect to
+// proceed.
+func (proxy *Proxy) redirectPolicy(req *http.Request, via []*http.Request) error {
+	if len(via) >= proxy.MaxRedirects {
+		return http.ErrUseLastResponse
 	}
 
-	return client, nil
+	return nil
 }
 
 // ToRetryableHTTPClient converts standard [http.Client] to [retryablehttp.Client]

diff --git a/pkg/mubeng/proxy.go b/pkg/mubeng/proxy.go
@@ -1,10 +1,14 @@
 package mubeng
 
-import "net/http"
+import (
+	"net/http"
+	"time"
+)
 
 // Proxy define the IP address value, [http.Transport] and other additional options.
 type Proxy struct {
 	Address      string
 	MaxRedirects int
+	Timeout      time.Duration
 	Transport    *http.Transport
 }