feat: add --rotate-on-error & --max-errors options (#249)

* feat(mubeng): rm global `client` var Signed-off-by: Dwi Siswanto <git@dw1.io> * refactor(mubeng): do not join prior proxy hosts Signed-off-by: Dwi Siswanto <git@dw1.io> * feat: add `--rotate-on-error` & `--max-errors` options Signed-off-by: Dwi Siswanto <git@dw1.io> * docs(README): add `--rotate-on-error` & `--max-errors` flags Signed-off-by: Dwi Siswanto <git@dw1.io> * fix(server): use sprint for remaining errs format instead Signed-off-by: Dwi Siswanto <git@dw1.io> --------- Signed-off-by: Dwi Siswanto <git@dw1.io>
mubeng · Sep 7, 2024 · 82ae8ac · 82ae8ac
1 parent 901e2cb
commit 82ae8ac
Show file tree

Hide file tree

Showing 8 changed files with 114 additions and 83 deletions.
diff --git a/README.md b/README.md
@@ -141,6 +141,9 @@ Here are all the options it supports.
 |     --only-cc `<AA>,<BB>`       | Only show specific country code (comma separated).            |
 | -t, --timeout                   | Max. time allowed for proxy server/check (default: 30s).      |
 | -r, --rotate `<AFTER>`          | Rotate proxy IP for every `AFTER` request (default: 1).       |
+|     --rotate-on-error `<N>`     | Rotate proxy IP and retry failed HTTP requests                |
+|     --max-errors `<N>`          | Max. errors allowed during rotation (default:3).              |
+|                                 | Use this with `--rotate-on-error`.                            |
 |     --max-redirs `<N>`          | Max. redirects allowed (default: 10).                         |
 |     --max-retries `<N>`         | Max. retries for failed HTTP requests (default: 0).           |
 | -m, --method `<METHOD>`         | Rotation method (sequent/random) (default: sequent).          |

diff --git a/common/options.go b/common/options.go
@@ -24,9 +24,11 @@ type Options struct {
 	Method       string
 	Output       string
 	Rotate       int
+	RotateOnErr  bool
 	Sync         bool
 	Verbose      bool
 	Watch        bool
-	MaxRetries   int
+	MaxErrors    int
 	MaxRedirects int
+	MaxRetries   int
 }
diff --git a/common/vars.go b/common/vars.go
@@ -39,6 +39,9 @@ Options:
     -d, --daemon                     Daemonize proxy server
     -m, --method <METHOD>            Rotation method (sequent/random) (default: sequent)
     -r, --rotate <N>                 Rotate proxy IP after N request (default: 1)
+        --rotate-on-error            Rotate proxy IP and retry failed HTTP requests
+        --max-errors <N>             Max. errors allowed during rotation (default: 3)
+                                     Use this with --rotate-on-error
         --max-redirs <N>             Max. redirects allowed (default: 10)
         --max-retries <N>            Max. retries for failed HTTP requests (default: 0)
     -s, --sync                       Syncrounus mode

diff --git a/internal/runner/options.go b/internal/runner/options.go
@@ -33,6 +33,8 @@ func Options() *common.Options {
 	flag.IntVar(&opt.Rotate, "r", 1, "")
 	flag.IntVar(&opt.Rotate, "rotate", 1, "")
 
+	flag.BoolVar(&opt.RotateOnErr, "rotate-on-error", false, "")
+
 	flag.StringVar(&opt.Method, "m", "sequent", "")
 	flag.StringVar(&opt.Method, "method", "sequent", "")
 
@@ -60,8 +62,8 @@ func Options() *common.Options {
 	flag.IntVar(&opt.Goroutine, "g", 50, "")
 	flag.IntVar(&opt.Goroutine, "goroutine", 50, "")
 
+	flag.IntVar(&opt.MaxErrors, "max-errors", 3, "")
 	flag.IntVar(&opt.MaxRedirects, "max-redirs", 10, "")
-
 	flag.IntVar(&opt.MaxRetries, "max-retries", 0, "")
 
 	flag.Usage = func() {

diff --git a/internal/server/handler.go b/internal/server/handler.go
@@ -25,77 +25,55 @@ func (p *Proxy) onRequest(req *http.Request, ctx *goproxy.ProxyCtx) (*http.Reque
 		return req, serverErr(req)
 	}
 
-	// Rotate proxy IP for every AFTER request
-	if (rotate == "") || (ok >= p.Options.Rotate) {
-		rotate = p.Options.ProxyManager.Rotate(p.Options.Method)
-
-		if ok >= p.Options.Rotate {
-			ok = 1
-		}
-	} else {
-		ok++
-	}
-
 	resChan := make(chan interface{})
 
 	go func(r *http.Request) {
 		log.Debugf("%s %s %s", r.RemoteAddr, r.Method, r.URL)
 
-		tr, err := mubeng.Transport(rotate)
-		if err != nil {
-			resChan <- err
-			return
-		}
+		for i := 0; i <= p.Options.MaxErrors; i++ {
+			retryablehttpClient, err := p.getClient(r, p.rotateProxy())
+			if err != nil {
+				resChan <- err
 
-		proxy := &mubeng.Proxy{
-			Address:      rotate,
-			MaxRedirects: p.Options.MaxRedirects,
-			Timeout:      p.Options.Timeout,
-			Transport:    tr,
-		}
+				return
+			}
 
-		client, err := proxy.New(r)
-		if err != nil {
-			resChan <- err
-			return
-		}
+			retryablehttpRequest, err := retryablehttp.FromRequest(r)
+			if err != nil {
+				resChan <- err
 
-		if p.Options.Verbose {
-			client.Transport = dump.RoundTripper(tr)
-		}
+				return
+			}
 
-		retryablehttpClient := mubeng.ToRetryableHTTPClient(client)
-		retryablehttpClient.RetryMax = p.Options.MaxRetries
-		retryablehttpClient.RetryWaitMin = client.Timeout
-		retryablehttpClient.RetryWaitMax = client.Timeout
-		retryablehttpClient.Logger = ReleveledLogo{
-			Logger:  log,
-			Request: r,
-			Verbose: p.Options.Verbose,
-		}
+			resp, err := retryablehttpClient.Do(retryablehttpRequest)
+			if err != nil {
+				if p.Options.RotateOnErr && i < p.Options.MaxErrors {
+					log.Debugf(
+						"%s Retrying (rotated) %s %s [remaining=%q]",
+						r.RemoteAddr, r.Method, r.URL, fmt.Sprint(p.Options.MaxErrors-i),
+					)
+
+					continue
+				} else {
+					resChan <- err
+
+					return
+				}
+			}
+			defer resp.Body.Close()
 
-		retryablehttpRequest, err := retryablehttp.FromRequest(r)
-		if err != nil {
-			resChan <- err
-			return
-		}
+			buf, err := io.ReadAll(resp.Body)
+			if err != nil {
+				resChan <- err
 
-		resp, err := retryablehttpClient.Do(retryablehttpRequest)
-		if err != nil {
-			resChan <- err
-			return
-		}
-		defer resp.Body.Close()
+				return
+			}
+			resp.Body = io.NopCloser(bytes.NewBuffer(buf))
+
+			resChan <- resp
 
-		buf, err := io.ReadAll(resp.Body)
-		if err != nil {
-			resChan <- err
 			return
 		}
-
-		resp.Body = io.NopCloser(bytes.NewBuffer(buf))
-
-		resChan <- resp
 	}(req)
 
 	var resp *http.Response
@@ -152,6 +130,57 @@ func (p *Proxy) onResponse(resp *http.Response, ctx *goproxy.ProxyCtx) *http.Res
 	return resp
 }
 
+func (p *Proxy) rotateProxy() string {
+	var proxyAddr string
+
+	if ok >= p.Options.Rotate {
+		proxyAddr = p.Options.ProxyManager.Rotate(p.Options.Method)
+
+		if ok >= p.Options.Rotate {
+			ok = 1
+		}
+	} else {
+		ok++
+	}
+
+	return proxyAddr
+}
+
+func (p *Proxy) getClient(req *http.Request, proxyAddr string) (*retryablehttp.Client, error) {
+	tr, err := mubeng.Transport(proxyAddr)
+	if err != nil {
+		return nil, err
+	}
+
+	proxy := &mubeng.Proxy{
+		Address:      proxyAddr,
+		MaxRedirects: p.Options.MaxRedirects,
+		Timeout:      p.Options.Timeout,
+		Transport:    tr,
+	}
+
+	client, err := proxy.New(req)
+	if err != nil {
+		return nil, err
+	}
+
+	if p.Options.Verbose {
+		client.Transport = dump.RoundTripper(tr)
+	}
+
+	retryablehttpClient := mubeng.ToRetryableHTTPClient(client)
+	retryablehttpClient.RetryMax = p.Options.MaxRetries
+	retryablehttpClient.RetryWaitMin = client.Timeout
+	retryablehttpClient.RetryWaitMax = client.Timeout
+	retryablehttpClient.Logger = ReleveledLogo{
+		Logger:  log,
+		Request: req,
+		Verbose: p.Options.Verbose,
+	}
+
+	return retryablehttpClient, nil
+}
+
 // nonProxy handles non-proxy requests
 func nonProxy(w http.ResponseWriter, req *http.Request) {
 	if common.Version != "" {

diff --git a/internal/server/vars.go b/internal/server/vars.go
@@ -9,7 +9,6 @@ import (
 )
 
 var (
-	rotate  string
 	handler *Proxy
 	server  *http.Server
 	dump    *httpretty.Logger

diff --git a/pkg/mubeng/mubeng.go b/pkg/mubeng/mubeng.go
@@ -4,7 +4,6 @@ import (
 	"net"
 	"net/http"
 	"net/url"
-	"strings"
 
 	"github.com/hashicorp/go-retryablehttp"
 )
@@ -14,7 +13,7 @@ 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{
+	client := &http.Client{
 		CheckRedirect: proxy.redirectPolicy,
 		Timeout:       proxy.Timeout,
 		Transport:     proxy.Transport,
@@ -34,9 +33,9 @@ func (proxy *Proxy) New(req *http.Request) (*http.Client, error) {
 	}
 
 	if host, _, err := net.SplitHostPort(proxyURL.Host); err == nil {
-		if prior, ok := req.Header["X-Forwarded-For"]; ok {
-			host = strings.Join(prior, ", ") + ", " + host
-		}
+		// if prior, ok := req.Header["X-Forwarded-For"]; ok {
+		// 	host = strings.Join(prior, ", ") + ", " + host
+		// }
 		req.Header.Set("X-Forwarded-For", host)
 	}
 

diff --git a/pkg/mubeng/vars.go b/pkg/mubeng/vars.go
@@ -1,20 +1,14 @@
 package mubeng
 
-import "net/http"
-
-var (
-	client *http.Client
-
-	// HopHeaders are meaningful only for a single transport-level connection, and are not stored by caches or forwarded by proxies.
-	HopHeaders = []string{
-		"Connection",
-		"Keep-Alive",
-		"Proxy-Authenticate",
-		"Proxy-Authorization",
-		"Proxy-Connection",
-		"Te", // canonicalized version of "TE"
-		"Trailers",
-		"Transfer-Encoding",
-		"Upgrade",
-	}
-)
+// HopHeaders are meaningful only for a single transport-level connection, and are not stored by caches or forwarded by proxies.
+var HopHeaders = []string{
+	"Connection",
+	"Keep-Alive",
+	"Proxy-Authenticate",
+	"Proxy-Authorization",
+	"Proxy-Connection",
+	"Te", // canonicalized version of "TE"
+	"Trailers",
+	"Transfer-Encoding",
+	"Upgrade",
+}