app/vmauth: add missing tests for requestHandler()

app/vmauth/auth_config.go

@@ -461,7 +461,7 @@ func getLeastLoadedBackendURL(bus []*backendURL, atomicCounter *atomic.Uint32) *
 	}
 
 	// Slow path - select other backend urls.
-	n := atomicCounter.Add(1)
+	n := atomicCounter.Add(1) - 1
 
 	for i := uint32(0); i < uint32(len(bus)); i++ {
 		idx := (n + i) % uint32(len(bus))

app/vmauth/auth_config_test.go

@@ -418,7 +418,7 @@ users:
 				HeadersConf: HeadersConf{
 					RequestHeaders: []*Header{
 						mustNewHeader("'foo: bar'"),
-						mustNewHeader("'xxx: y'"),
+						mustNewHeader("'xxx:'"),
 					},
 				},
 			},
@@ -437,7 +437,7 @@ users:
     url_prefix: ["http://vminsert1/insert/0/prometheus","http://vminsert2/insert/0/prometheus"]
     headers:
     - "foo: bar"
-    - "xxx: y"
+    - "xxx:"
 `, map[string]*UserInfo{
 		getHTTPAuthBearerToken("foo"):    sharedUserInfo,
 		getHTTPAuthBasicToken("foo", ""): sharedUserInfo,
@@ -723,9 +723,9 @@ func TestGetLeastLoadedBackendURL(t *testing.T) {
 
 	fn := func(ns ...int) {
 		t.Helper()
-		bus := up.bus.Load()
-		pbus := *bus
-		for i, b := range pbus {
+		pbus := up.bus.Load()
+		bus := *pbus
+		for i, b := range bus {
 			got := int(b.concurrentRequests.Load())
 			exp := ns[i]
 			if got != exp {
@@ -735,39 +735,39 @@ func TestGetLeastLoadedBackendURL(t *testing.T) {
 	}
 
 	up.getBackendURL()
-	fn(0, 1, 0)
+	fn(1, 0, 0)
 	up.getBackendURL()
-	fn(0, 1, 1)
+	fn(1, 1, 0)
 	up.getBackendURL()
 	fn(1, 1, 1)
 
 	up.getBackendURL()
 	up.getBackendURL()
-	fn(1, 2, 2)
+	fn(2, 2, 1)
 
 	bus := up.bus.Load()
 	pbus := *bus
 	pbus[0].concurrentRequests.Add(2)
 	pbus[2].concurrentRequests.Add(5)
-	fn(3, 2, 7)
+	fn(4, 2, 6)
 
 	up.getBackendURL()
-	fn(3, 3, 7)
+	fn(4, 3, 6)
 
 	up.getBackendURL()
-	fn(3, 4, 7)
+	fn(4, 4, 6)
 
 	up.getBackendURL()
-	fn(4, 4, 7)
+	fn(4, 5, 6)
 
 	up.getBackendURL()
-	fn(5, 4, 7)
+	fn(5, 5, 6)
 
 	up.getBackendURL()
-	fn(5, 5, 7)
+	fn(6, 5, 6)
 
 	up.getBackendURL()
-	fn(6, 5, 7)
+	fn(6, 6, 6)
 
 	up.getBackendURL()
 	fn(6, 6, 7)

app/vmauth/main.go

@@ -218,7 +218,15 @@ func processRequest(w http.ResponseWriter, r *http.Request, ui *UserInfo) {
 		} else { // Update path for regular routes.
 			targetURL = mergeURLs(targetURL, u, up.dropSrcPathPrefixParts)
 		}
-		ok := tryProcessingRequest(w, r, targetURL, hc, up.retryStatusCodes, ui)
+
+		wasLocalRetry := false
+	again:
+		ok, needLocalRetry := tryProcessingRequest(w, r, targetURL, hc, up.retryStatusCodes, ui)
+		if needLocalRetry && !wasLocalRetry {
+			wasLocalRetry = true
+			goto again
+		}
+
 		bu.put()
 		if ok {
 			return
@@ -226,14 +234,14 @@ func processRequest(w http.ResponseWriter, r *http.Request, ui *UserInfo) {
 		bu.setBroken()
 	}
 	err := &httpserver.ErrorWithStatusCode{
-		Err:        fmt.Errorf("all the backends for the user %q are unavailable", ui.name()),
+		Err:        fmt.Errorf("all the %d backends for the user %q are unavailable", up.getBackendsCount(), ui.name()),
 		StatusCode: http.StatusServiceUnavailable,
 	}
 	httpserver.Errorf(w, r, "%s", err)
 	ui.backendErrors.Inc()
 }
 
-func tryProcessingRequest(w http.ResponseWriter, r *http.Request, targetURL *url.URL, hc HeadersConf, retryStatusCodes []int, ui *UserInfo) bool {
+func tryProcessingRequest(w http.ResponseWriter, r *http.Request, targetURL *url.URL, hc HeadersConf, retryStatusCodes []int, ui *UserInfo) (bool, bool) {
 	req := sanitizeRequestHeaders(r)
 
 	req.URL = targetURL
@@ -246,9 +254,7 @@ func tryProcessingRequest(w http.ResponseWriter, r *http.Request, targetURL *url
 		}
 	}
 
-	var trivialRetries int
 	rtb, rtbOK := req.Body.(*readTrackingBody)
-again:
 	res, err := ui.rt.RoundTrip(req)
 	if err != nil {
 		if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
@@ -260,7 +266,7 @@ again:
 				// Timed out request must be counted as errors, since this usually means that the backend is slow.
 				ui.backendErrors.Inc()
 			}
-			return true
+			return true, false
 		}
 		if !rtbOK || !rtb.canRetry() {
 			// Request body cannot be re-sent to another backend. Return the error to the client then.
@@ -270,19 +276,19 @@ again:
 			}
 			httpserver.Errorf(w, r, "%s", err)
 			ui.backendErrors.Inc()
-			return true
+			return true, false
 		}
-		// Retry request on trivial network errors, such as proxy idle timeout misconfiguration or socket close by OS
-		if (netutil.IsTrivialNetworkError(err) || errors.Is(err, io.EOF)) && trivialRetries < 1 {
-			trivialRetries++
-			goto again
+		if netutil.IsTrivialNetworkError(err) {
+			// Retry request at the same backend on trivial network errors, such as proxy idle timeout misconfiguration or socket close by OS
+			return false, true
 		}
+
 		// Retry the request if its body wasn't read yet. This usually means that the backend isn't reachable.
 		remoteAddr := httpserver.GetQuotedRemoteAddr(r)
 		// NOTE: do not use httpserver.GetRequestURI
 		// it explicitly reads request body, which may fail retries.
 		logger.Warnf("remoteAddr: %s; requestURI: %s; retrying the request to %s because of response error: %s", remoteAddr, req.URL, targetURL, err)
-		return false
+		return false, false
 	}
 	if slices.Contains(retryStatusCodes, res.StatusCode) {
 		_ = res.Body.Close()
@@ -296,7 +302,7 @@ again:
 			}
 			httpserver.Errorf(w, r, "%s", err)
 			ui.backendErrors.Inc()
-			return true
+			return true, false
 		}
 		// Retry requests at other backends if it matches retryStatusCodes.
 		// See https://github.com/VictoriaMetrics/VictoriaMetrics/issues/4893
@@ -305,7 +311,7 @@ again:
 		// it explicitly reads request body, which may fail retries.
 		logger.Warnf("remoteAddr: %s; requestURI: %s; retrying the request to %s because response status code=%d belongs to retry_status_codes=%d",
 			remoteAddr, req.URL, targetURL, res.StatusCode, retryStatusCodes)
-		return false
+		return false, false
 	}
 	removeHopHeaders(res.Header)
 	copyHeader(w.Header(), res.Header)
@@ -321,9 +327,9 @@ again:
 		remoteAddr := httpserver.GetQuotedRemoteAddr(r)
 		requestURI := httpserver.GetRequestURI(r)
 		logger.Warnf("remoteAddr: %s; requestURI: %s; error when proxying response body from %s: %s", remoteAddr, requestURI, targetURL, err)
-		return true
+		return true, false
 	}
-	return true
+	return true, false
 }
 
 var copyBufPool bytesutil.ByteBufferPool

app/vmauth/main_test.go

@@ -4,10 +4,14 @@ import (
 	"bytes"
 	"fmt"
 	"io"
+	"net"
 	"net/http"
 	"net/http/httptest"
 	"strings"
+	"sync/atomic"
 	"testing"
+
+	"github.com/VictoriaMetrics/VictoriaMetrics/lib/netutil"
 )
 
 func TestRequestHandler(t *testing.T) {
@@ -83,12 +87,13 @@ statusCode=200
 requested_url=http://some-host.com/foo/abc/def?bar=baz`
 	f(cfgStr, requestURL, backendHandler, responseExpected)
 
-	// override request host
+	// override request host with non-empty host
 	cfgStr = `
 unauthorized_user:
   url_prefix: "{BACKEND}/foo?bar=baz"
   headers:
-  - "Host: other-host:12345"`
+  - "Host: other-host:12345"
+  - "abc:"`
 	requestURL = "http://some-host.com/abc/def"
 	backendHandler = func(w http.ResponseWriter, r *http.Request) {
 		fmt.Fprintf(w, "requested_url=http://%s%s", r.Host, r.URL)
@@ -98,6 +103,21 @@ statusCode=200
 requested_url=http://other-host:12345/foo/abc/def?bar=baz`
 	f(cfgStr, requestURL, backendHandler, responseExpected)
 
+	// override request host with empty host
+	cfgStr = `
+unauthorized_user:
+  url_prefix: "{BACKEND}/foo?bar=baz"
+  headers:
+  - "Host:"`
+	requestURL = "http://some-host.com/abc/def"
+	backendHandler = func(w http.ResponseWriter, r *http.Request) {
+		fmt.Fprintf(w, "requested_url=http://%s%s", r.Host, r.URL)
+	}
+	responseExpected = `
+statusCode=200
+requested_url={BACKEND}/foo/abc/def?bar=baz`
+	f(cfgStr, requestURL, backendHandler, responseExpected)
+
 	// /-/reload handler failure
 	origAuthKey := reloadAuthKey.Get()
 	if err := reloadAuthKey.Set("secret"); err != nil {
@@ -315,7 +335,7 @@ unauthorized_user:
 	}
 	responseExpected = `
 statusCode=503
-remoteAddr: "42.2.3.84:6789"; requestURI: /foo/?de=fg; all the backends for the user "" are unavailable`
+remoteAddr: "42.2.3.84:6789"; requestURI: /foo/?de=fg; all the 2 backends for the user "" are unavailable`
 	f(cfgStr, requestURL, backendHandler, responseExpected)
 
 	// all the backend_urls are unavailable for authorized user
@@ -333,8 +353,72 @@ users:
 	}
 	responseExpected = `
 statusCode=503
-remoteAddr: "42.2.3.84:6789"; requestURI: /foo/?de=fg; all the backends for the user "some-user" are unavailable`
+remoteAddr: "42.2.3.84:6789"; requestURI: /foo/?de=fg; all the 2 backends for the user "some-user" are unavailable`
 	f(cfgStr, requestURL, backendHandler, responseExpected)
+
+	// zero discovered backend IPs
+	customResolver := &fakeResolver{
+		Resolver: &net.Resolver{},
+		lookupIPAddrResults: map[string][]net.IPAddr{
+			"some-addr": {},
+		},
+	}
+	origResolver := netutil.Resolver
+	netutil.Resolver = customResolver
+	cfgStr = `
+unauthorized_user:
+  url_prefix: ['http://some-addr:1234/foo/bar']
+  discover_backend_ips: true`
+	requestURL = "http://abc.com/def/?de=fg"
+	backendHandler = func(_ http.ResponseWriter, _ *http.Request) {
+		panic(fmt.Errorf("backend handler shouldn't be called"))
+	}
+	responseExpected = `
+statusCode=503
+remoteAddr: "42.2.3.84:6789"; requestURI: /def/?de=fg; all the 0 backends for the user "" are unavailable`
+	f(cfgStr, requestURL, backendHandler, responseExpected)
+	netutil.Resolver = origResolver
+
+	// retry_status_codes failure
+	var retries atomic.Int64
+	cfgStr = `
+unauthorized_user:
+  url_prefix: ['{BACKEND}/path1', '{BACKEND}/path2']
+  retry_status_codes: [500, 502]`
+	requestURL = "http://some-host.com/foo/?de=fg"
+	backendHandler = func(w http.ResponseWriter, _ *http.Request) {
+		retries.Add(1)
+		w.WriteHeader(500)
+	}
+	responseExpected = `
+statusCode=503
+remoteAddr: "42.2.3.84:6789"; requestURI: /foo/?de=fg; all the 2 backends for the user "" are unavailable`
+	f(cfgStr, requestURL, backendHandler, responseExpected)
+	if n := retries.Load(); n != 2 {
+		t.Fatalf("unexpected number of retries; got %d; want 2", n)
+	}
+
+	// retry_status_codes success
+	retries.Store(0)
+	cfgStr = `
+unauthorized_user:
+  url_prefix: ['{BACKEND}/path1', '{BACKEND}/path2']
+  retry_status_codes: [500, 502]`
+	requestURL = "http://some-host.com/foo/?de=fg"
+	backendHandler = func(w http.ResponseWriter, r *http.Request) {
+		if n := retries.Add(1); n < 2 {
+			w.WriteHeader(500)
+			return
+		}
+		fmt.Fprintf(w, "requested_url=http://%s%s", r.Host, r.URL)
+	}
+	responseExpected = `
+statusCode=200
+requested_url={BACKEND}/path2/foo/?de=fg`
+	f(cfgStr, requestURL, backendHandler, responseExpected)
+	if n := retries.Load(); n != 2 {
+		t.Fatalf("unexpected number of retries; got %d; want 2", n)
+	}
 }
 
 type fakeResponseWriter struct {

app/vmauth/target_url_test.go

@@ -335,8 +335,8 @@ func TestUserInfoGetBackendURL_SRV(t *testing.T) {
 		t.Fatalf("cannot initialize urls inside UserInfo: %s", err)
 	}
 
-	f(ui, `/select/0/prometheus/api/v1/query?query=up`, "http://10.6.142.51:8481/select/0/prometheus/api/v1/query?query=up")
 	f(ui, `/select/0/prometheus/api/v1/query?query=up`, "http://10.6.142.50:8481/select/0/prometheus/api/v1/query?query=up")
+	f(ui, `/select/0/prometheus/api/v1/query?query=up`, "http://10.6.142.51:8481/select/0/prometheus/api/v1/query?query=up")
 	f(ui, `/insert/0/prometheus/api/v1/write`, "http://10.6.142.52:8480/insert/0/prometheus/api/v1/write")
 	// unsuccessful dns resolve
 	f(ui, `/test`, "http://non-exist-dns-addr/test")