refactor: address review concerns

Author: cloud-j-luna

internal/proxy/balancer.go

@@ -108,6 +108,11 @@ func (rr *LatencyBased) Next() *Server {
 	rr.mu.Lock()
 	defer rr.mu.Unlock()
 
+	// Return nil if no servers are available
+	if len(rr.servers) == 0 {
+		return nil
+	}
+
 	r := rr.randomizer.Float64()
 	cumulative := 0.0
 

internal/proxy/proxy.go

@@ -133,9 +133,9 @@ func (p *Proxy) Start(ctx context.Context, update Updater) {
 
 // newRedirectFollowingReverseProxy creates a configured httputil.ReverseProxy that automatically follows 301 redirects.
 // It handles 301 redirects by following them automatically and returning the final response status and content.
-func newRedirectFollowingReverseProxy(srv *Server, log *slog.Logger) *httputil.ReverseProxy {
+func newRedirectFollowingReverseProxy(srv *Server, log *slog.Logger, proxyType string) *httputil.ReverseProxy {
 	// Create a custom HTTP client that doesn't follow redirects automatically
-	client := &http.Client{
+	redirectClient := &http.Client{
 		CheckRedirect: func(req *http.Request, via []*http.Request) error {
 			// Stop automatic redirect following
 			return http.ErrUseLastResponse
@@ -149,13 +149,13 @@ func newRedirectFollowingReverseProxy(srv *Server, log *slog.Logger) *httputil.R
 			request.URL.Path = srv.Url.Path + request.URL.Path
 			request.Host = srv.Url.Host
 		},
-		Transport: client.Transport,
+		Transport: http.DefaultTransport,
 		ModifyResponse: func(response *http.Response) error {
 			cors.DeleteCorsHeaders(response)
 
-			// Handle 301 redirects by following them and returning content with final status
-			if response.StatusCode == http.StatusMovedPermanently {
-				metrics.IncrementRequestStatusCount("rest", srv.Url.String(), http.StatusMovedPermanently)
+			// Handle redirect responses by following them and returning content with final status
+			if isRedirect(response.StatusCode) {
+				metrics.IncrementRequestStatusCount(proxyType, srv.Url.String(), response.StatusCode)
 
 				location := response.Header.Get("Location")
 				if location != "" {
@@ -164,20 +164,27 @@ func newRedirectFollowingReverseProxy(srv *Server, log *slog.Logger) *httputil.R
 						return fmt.Errorf("failed to parse redirect location %q: %w", location, err)
 					}
 
-					log.Info("following 301 redirect", "location", location, "resolved_url", redirectURL.String())
+					log.Info("following redirect", "original_status", response.StatusCode, "location", location, "resolved_url", redirectURL.String())
 
-					redirectReq, err := http.NewRequest(response.Request.Method, redirectURL.String(), response.Request.Body)
+					// Only follow redirects for safe/idempotent methods to avoid body consumption issues
+					if !isIdempotentMethod(response.Request.Method) {
+						log.Warn("skipping redirect for non-idempotent method", "method", response.Request.Method)
+						metrics.IncrementRequestStatusCount(proxyType, srv.Url.String(), response.StatusCode)
+						return nil
+					}
+
+					redirectReq, err := http.NewRequestWithContext(response.Request.Context(), response.Request.Method, redirectURL.String(), nil)
 					if err != nil {
 						return fmt.Errorf("failed to create redirect request to %q: %w", redirectURL.String(), err)
 					}
 
-					for name, headers := range response.Request.Header {
-						for _, h := range headers {
-							redirectReq.Header.Add(name, h)
-						}
+					copyHeaders(response.Request.Header, redirectReq.Header)
+
+					if response.Body != nil {
+						response.Body.Close()
 					}
 
-					redirectResp, err := client.Do(redirectReq)
+					redirectResp, err := redirectClient.Do(redirectReq)
 					if err != nil {
 						return fmt.Errorf("failed to follow redirect to %q: %w", redirectURL.String(), err)
 					}
@@ -188,13 +195,16 @@ func newRedirectFollowingReverseProxy(srv *Server, log *slog.Logger) *httputil.R
 					response.ContentLength = redirectResp.ContentLength
 					response.Header = redirectResp.Header.Clone()
 
-					log.Info("successfully handled 301 redirect", "original_status", "301", "final_status", redirectResp.StatusCode)
-					metrics.IncrementRequestStatusCount("rest", srv.Url.String(), response.StatusCode)
+					cors.DeleteCorsHeaders(response)
+
+					log.Info("successfully handled redirect", "original_status", "redirect", "final_status", redirectResp.StatusCode)
+					metrics.IncrementRequestStatusCount(proxyType, srv.Url.String(), response.StatusCode)
 				} else {
-					log.Warn("301 redirect without Location header, serving as-is")
+					log.Warn("redirect without Location header, serving as-is", "status", response.StatusCode)
+					metrics.IncrementRequestStatusCount(proxyType, srv.Url.String(), response.StatusCode)
 				}
 			} else {
-				metrics.IncrementRequestStatusCount("rest", srv.Url.String(), response.StatusCode)
+				metrics.IncrementRequestStatusCount(proxyType, srv.Url.String(), response.StatusCode)
 			}
 
 			return nil
@@ -205,3 +215,49 @@ func newRedirectFollowingReverseProxy(srv *Server, log *slog.Logger) *httputil.R
 		},
 	}
 }
+
+// isRedirect checks if the status code represents a redirect
+func isRedirect(statusCode int) bool {
+	switch statusCode {
+	case http.StatusMovedPermanently,
+		http.StatusFound,
+		http.StatusTemporaryRedirect,
+		http.StatusPermanentRedirect:
+		return true
+	default:
+		return false
+	}
+}
+
+// isIdempotentMethod checks if the HTTP method is safe to replay without side effects
+func isIdempotentMethod(method string) bool {
+	switch method {
+	case http.MethodGet, http.MethodHead, http.MethodOptions:
+		return true
+	default:
+		return false
+	}
+}
+
+// copyHeaders copies headers from src to dst, excluding hop-by-hop headers
+func copyHeaders(src, dst http.Header) {
+	// Hop-by-hop headers that should not be forwarded
+	hopByHopHeaders := map[string]bool{
+		"Connection":          true,
+		"Keep-Alive":          true,
+		"Proxy-Authenticate":  true,
+		"Proxy-Authorization": true,
+		"Te":                  true,
+		"Trailers":            true,
+		"Transfer-Encoding":   true,
+		"Upgrade":             true,
+	}
+
+	for name, headers := range src {
+		if !hopByHopHeaders[name] {
+			for _, h := range headers {
+				dst.Add(name, h)
+			}
+		}
+	}
+}

internal/proxy/proxy_test.go

@@ -191,7 +191,7 @@ func TestNewReverseProxy(t *testing.T) {
 			}
 
 			logger := slog.New(slog.NewTextHandler(io.Discard, nil))
-			proxy := newRedirectFollowingReverseProxy(srv, logger)
+			proxy := newRedirectFollowingReverseProxy(srv, logger, "rest")
 
 			// Create test request
 			req := httptest.NewRequest("GET", tt.reqPath, nil)
@@ -210,7 +210,7 @@ func TestReverseProxy_ModifyResponse(t *testing.T) {
 	targetURL, _ := url.Parse("http://node.com")
 	srv := &Server{Url: targetURL}
 	logger := slog.New(slog.NewTextHandler(io.Discard, nil))
-	proxy := newRedirectFollowingReverseProxy(srv, logger)
+	proxy := newRedirectFollowingReverseProxy(srv, logger, "rest")
 
 	resp := &http.Response{Header: make(http.Header)}
 	resp.Header.Set("Access-Control-Allow-Origin", "*")
@@ -229,7 +229,7 @@ func TestReverseProxy_ErrorHandler(t *testing.T) {
 	targetURL, _ := url.Parse("http://node.com")
 	srv := &Server{Url: targetURL}
 	logger := slog.New(slog.NewTextHandler(io.Discard, nil))
-	proxy := newRedirectFollowingReverseProxy(srv, logger)
+	proxy := newRedirectFollowingReverseProxy(srv, logger, "rest")
 
 	w := httptest.NewRecorder()
 	req := httptest.NewRequest("GET", "/test", nil)

internal/proxy/rest.go

@@ -48,7 +48,7 @@ func (p *RestProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
 	r.URL.Path = strings.TrimPrefix(r.URL.Path, "/rest")
 	if srv := p.lb.Next(); srv != nil {
-		proxy := newRedirectFollowingReverseProxy(srv, p.log)
+		proxy := newRedirectFollowingReverseProxy(srv, p.log, "rest")
 		proxy.ServeHTTP(w, r)
 		metrics.IncrementRequestCount("rest", srv.Url.String())
 		return

internal/proxy/rpc.go

@@ -48,7 +48,7 @@ func (p *RPCProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
 	r.URL.Path = strings.TrimPrefix(r.URL.Path, "/rpc")
 	if srv := p.lb.Next(); srv != nil {
-		proxy := newRedirectFollowingReverseProxy(srv, p.log)
+		proxy := newRedirectFollowingReverseProxy(srv, p.log, "rpc")
 		proxy.ServeHTTP(w, r)
 		metrics.IncrementRequestCount("rpc", srv.Url.String())
 		return