From 7ed719b46af0be975d920254ceed7f7514dbf78e Mon Sep 17 00:00:00 2001
From: Aliaksandr Valialkin <valyala@victoriametrics.com>
Date: Wed, 17 Jul 2024 11:30:15 +0200
Subject: [PATCH] app/vmauth: properly handle the case when zero backend hosts
 are resolved at SRV DNS

When zero backend hosts are resolved, then vmauth must return 'no backend hosts' error instead of crashing with panic

This is a follow-up for 590aeccd7d769cf48a6afbdecae03fbcabd130ac and 3a45bbb4e0f21b8d343d41bcd3cb8dd87956ac05

Updates https://github.com/VictoriaMetrics/VictoriaMetrics/pull/6401
---
 app/vmauth/auth_config.go     |  6 +++++
 app/vmauth/main.go            |  3 +++
 app/vmauth/target_url_test.go | 49 +++++++++++++++++++++++++++++++++++
 3 files changed, 58 insertions(+)

diff --git a/app/vmauth/auth_config.go b/app/vmauth/auth_config.go
index 5b73e4fdb7..151dd2986e 100644
--- a/app/vmauth/auth_config.go
+++ b/app/vmauth/auth_config.go
@@ -312,12 +312,18 @@ func (up *URLPrefix) getBackendsCount() int {
 
 // getBackendURL returns the backendURL depending on the load balance policy.
 //
+// It can return nil if there are no backend urls available at the moment.
+//
 // backendURL.put() must be called on the returned backendURL after the request is complete.
 func (up *URLPrefix) getBackendURL() *backendURL {
 	up.discoverBackendAddrsIfNeeded()
 
 	pbus := up.bus.Load()
 	bus := *pbus
+	if len(bus) == 0 {
+		return nil
+	}
+
 	if up.loadBalancingPolicy == "first_available" {
 		return getFirstAvailableBackendURL(bus)
 	}
diff --git a/app/vmauth/main.go b/app/vmauth/main.go
index 79010ce2bb..2c266af50c 100644
--- a/app/vmauth/main.go
+++ b/app/vmauth/main.go
@@ -208,6 +208,9 @@ func processRequest(w http.ResponseWriter, r *http.Request, ui *UserInfo) {
 	maxAttempts := up.getBackendsCount()
 	for i := 0; i < maxAttempts; i++ {
 		bu := up.getBackendURL()
+		if bu == nil {
+			break
+		}
 		targetURL := bu.url
 		// Don't change path and add request_path query param for default route.
 		if isDefault {
diff --git a/app/vmauth/target_url_test.go b/app/vmauth/target_url_test.go
index 3d72653f36..86ee43490f 100644
--- a/app/vmauth/target_url_test.go
+++ b/app/vmauth/target_url_test.go
@@ -342,6 +342,55 @@ func TestUserInfoGetBackendURL_SRV(t *testing.T) {
 	f(ui, `/test`, "http://non-exist-dns-addr/test")
 }
 
+func TestUserInfoGetBackendURL_SRVZeroBackends(t *testing.T) {
+	f := func(ui *UserInfo, requestURI string) {
+		t.Helper()
+
+		u, err := url.Parse(requestURI)
+		if err != nil {
+			t.Fatalf("cannot parse %q: %s", requestURI, err)
+		}
+		u = normalizeURL(u)
+		up, _ := ui.getURLPrefixAndHeaders(u, nil)
+		if up == nil {
+			t.Fatalf("cannot match available backend: %s", err)
+		}
+		bu := up.getBackendURL()
+		if bu != nil {
+			t.Fatalf("expecting nil backendURL; got %v", bu)
+		}
+	}
+
+	customResolver := &fakeResolver{
+		Resolver: &net.Resolver{},
+		lookupSRVResults: map[string][]*net.SRV{
+			"vmselect": {},
+		},
+	}
+	origResolver := netutil.Resolver
+	netutil.Resolver = customResolver
+	defer func() {
+		netutil.Resolver = origResolver
+	}()
+
+	allowed := true
+	ui := &UserInfo{
+		URLMaps: []URLMap{
+			{
+				SrcPaths:  getRegexs([]string{"/select/.+"}),
+				URLPrefix: mustParseURL("http://srv+vmselect"),
+			},
+		},
+		DiscoverBackendIPs: &allowed,
+		URLPrefix:          mustParseURL("http://non-exist-dns-addr"),
+	}
+	if err := ui.initURLs(); err != nil {
+		t.Fatalf("cannot initialize urls inside UserInfo: %s", err)
+	}
+
+	f(ui, `/select/0/prometheus/api/v1/query?query=up`)
+}
+
 func TestCreateTargetURLFailure(t *testing.T) {
 	f := func(ui *UserInfo, requestURI string) {
 		t.Helper()