From 31d01fce07c41724c0c4cfdeb71c78f3a2edda2a Mon Sep 17 00:00:00 2001
From: Felipe Cardozo <fraifelipe@gmail.com>
Date: Wed, 20 Dec 2023 16:42:48 -0300
Subject: [PATCH] fix validation of retry

---
 main.go              |  2 +-
 middlewares/retry.go |  7 ++++++-
 utils/balancer.go    | 18 ++++++++----------
 3 files changed, 15 insertions(+), 12 deletions(-)

diff --git a/main.go b/main.go
index bb4792b..5d7789f 100644
--- a/main.go
+++ b/main.go
@@ -32,8 +32,8 @@ func main() {
 	e.Use(middlewares.RateLimit(config))
 	e.Use(middlewares.DenyRoutes(config))
 	e.Use(middlewares.Cache(config))
-	e.Use(middlewares.Gzip(config))
 	e.Use(middlewares.Retry(config))
+	e.Use(middlewares.Gzip(config))
 	e.Use(middleware.ProxyWithConfig(*config.ProxyConfig))
 
 	// Start metrics server
diff --git a/middlewares/retry.go b/middlewares/retry.go
index ade6336..ac3a26b 100644
--- a/middlewares/retry.go
+++ b/middlewares/retry.go
@@ -1,6 +1,8 @@
 package middlewares
 
 import (
+	"strings"
+
 	"github.com/labstack/echo/v4"
 	"github.com/marigold-dev/tzproxy/utils"
 )
@@ -14,9 +16,12 @@ func Retry(config *utils.Config) echo.MiddlewareFunc {
 			err = next(c)
 
 			status := c.Response().Status
-			if status == 404 || status == 410 {
+			if status == 404 || status == 410 &&
+				strings.Contains(c.Request().URL.Path, "blocks") {
 				c.Set("retry", true)
 				return next(c)
+			} else {
+				c.Set("retry", false)
 			}
 
 			return err
diff --git a/utils/balancer.go b/utils/balancer.go
index 93a8478..a57c957 100644
--- a/utils/balancer.go
+++ b/utils/balancer.go
@@ -57,7 +57,7 @@ func (b *sameNodeBalancer) Next(c echo.Context) *middleware.ProxyTarget {
 	b.mutex.Lock()
 	defer b.mutex.Unlock()
 
-	if b.retryTarget != nil && c.Get("retry") != nil {
+	if b.retryTarget != nil && c.Get("retry") != nil && c.Get("retry").(bool) {
 		return b.retryTarget
 	}
 
@@ -67,16 +67,14 @@ func (b *sameNodeBalancer) Next(c echo.Context) *middleware.ProxyTarget {
 		return b.targets[0]
 	}
 
-	var i int
-	ipKey := []byte(c.RealIP())
-
-	got, err := b.store.Get(c.Request().Context(), ipKey)
+	ctx := c.Request().Context()
+	ip := []byte(c.RealIP())
+	got, err := b.store.Get(ctx, ip)
 	if err != nil {
-		i = b.random.Intn(len(b.targets))
-		b.store.Set(c.Request().Context(), ipKey, []byte{byte(i)}, b.TTL)
-	} else {
-		i = int(got[0])
+		i := b.random.Intn(len(b.targets))
+		b.store.Set(ctx, ip, []byte{byte(i)}, b.TTL)
+		return b.targets[i]
 	}
 
-	return b.targets[i]
+	return b.targets[int(got[0])]
 }