feat: add remove-on-error option

Signed-off-by: Dwi Siswanto <git@dw1.io>
mubeng · Sep 15, 2024 · 7f946f3 · 7f946f3
1 parent fd8db54
commit 7f946f3
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 4 deletions.
diff --git a/common/options.go b/common/options.go
@@ -25,6 +25,7 @@ type Options struct {
 	Output       string
 	Rotate       int
 	RotateOnErr  bool
+	RemoveOnErr  bool
 	Sync         bool
 	Verbose      bool
 	Watch        bool

diff --git a/common/vars.go b/common/vars.go
@@ -40,6 +40,7 @@ Options:
     -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
+        --remove-on-error            Remove proxy IP from proxy pool on 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)

diff --git a/internal/runner/options.go b/internal/runner/options.go
@@ -34,6 +34,7 @@ func Options() *common.Options {
 	flag.IntVar(&opt.Rotate, "rotate", 1, "")
 
 	flag.BoolVar(&opt.RotateOnErr, "rotate-on-error", false, "")
+	flag.BoolVar(&opt.RemoveOnErr, "remove-on-error", false, "")
 
 	flag.StringVar(&opt.Method, "m", "sequent", "")
 	flag.StringVar(&opt.Method, "method", "sequent", "")

diff --git a/internal/server/handler.go b/internal/server/handler.go
@@ -31,7 +31,9 @@ func (p *Proxy) onRequest(req *http.Request, ctx *goproxy.ProxyCtx) (*http.Reque
 		log.Debugf("%s %s %s", r.RemoteAddr, r.Method, r.URL)
 
 		for i := 0; i <= p.Options.MaxErrors; i++ {
-			retryablehttpClient, err := p.getClient(r, p.rotateProxy())
+			proxy := p.rotateProxy()
+
+			retryablehttpClient, err := p.getClient(r, proxy)
 			if err != nil {
 				resChan <- err
 
@@ -47,6 +49,15 @@ func (p *Proxy) onRequest(req *http.Request, ctx *goproxy.ProxyCtx) (*http.Reque
 
 			resp, err := retryablehttpClient.Do(retryablehttpRequest)
 			if err != nil {
+				if p.Options.RemoveOnErr {
+					p.removeProxy(proxy)
+
+					log.Debugf(
+						"%s Removing proxy IP from proxy pool [proxies=%q]",
+						r.RemoteAddr, fmt.Sprint(p.Options.ProxyManager.Count()),
+					)
+				}
+
 				if p.Options.RotateOnErr && i < p.Options.MaxErrors {
 					log.Debugf(
 						"%s Retrying (rotated) %s %s [remaining=%q]",
@@ -131,10 +142,14 @@ func (p *Proxy) onResponse(resp *http.Response, ctx *goproxy.ProxyCtx) *http.Res
 }
 
 func (p *Proxy) rotateProxy() string {
-	var proxyAddr string
+	var proxy string
+	var err error
 
 	if ok >= p.Options.Rotate {
-		proxyAddr = p.Options.ProxyManager.Rotate(p.Options.Method)
+		proxy, err = p.Options.ProxyManager.Rotate(p.Options.Method)
+		if err != nil {
+			log.Fatalf("Could not rotate proxy IP: %s", err)
+		}
 
 		if ok >= p.Options.Rotate {
 			ok = 1
@@ -143,7 +158,14 @@ func (p *Proxy) rotateProxy() string {
 		ok++
 	}
 
-	return proxyAddr
+	return proxy
+}
+
+func (p *Proxy) removeProxy(target string) {
+	err := p.Options.ProxyManager.RemoveProxy(target)
+	if err != nil {
+		log.Error(err)
+	}
 }
 
 func (p *Proxy) getClient(req *http.Request, proxyAddr string) (*retryablehttp.Client, error) {