From 05b42601c3ce870161878419bd055db09d2c0b87 Mon Sep 17 00:00:00 2001
From: Zakhar Bessarab <z.bessarab@victoriametrics.com>
Date: Fri, 9 Dec 2022 06:29:10 +0400
Subject: [PATCH] lib/promscrape/discovery/azure: remove API server from URL
 returned by azure (#3403)

* lib/promscrape/discovery/azure: remove API server from URL returned by azure

* lib/promscrape/discovery/azure: validate nextLink contains same URL as apiServer
---
 lib/promscrape/discovery/azure/machine.go | 27 ++++++++++++++++++-----
 lib/promscrape/discoveryutils/client.go   |  5 +++++
 2 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/lib/promscrape/discovery/azure/machine.go b/lib/promscrape/discovery/azure/machine.go
index b2fe24f81..5b11cda77 100644
--- a/lib/promscrape/discovery/azure/machine.go
+++ b/lib/promscrape/discovery/azure/machine.go
@@ -3,6 +3,7 @@ package azure
 import (
 	"encoding/json"
 	"fmt"
+	"net/url"
 	"sync"
 
 	"github.com/VictoriaMetrics/VictoriaMetrics/lib/cgroup"
@@ -61,24 +62,38 @@ type listAPIResponse struct {
 
 // visitAllAPIObjects iterates over list API with pagination and applies cb for each response object
 func visitAllAPIObjects(ac *apiConfig, apiURL string, cb func(data json.RawMessage) error) error {
-	nextLink := apiURL
-	for nextLink != "" {
-		resp, err := ac.c.GetAPIResponseWithReqParams(nextLink, func(request *fasthttp.Request) {
+	nextLinkURI := apiURL
+	for {
+		resp, err := ac.c.GetAPIResponseWithReqParams(nextLinkURI, func(request *fasthttp.Request) {
 			request.Header.Set("Authorization", "Bearer "+ac.mustGetAuthToken())
 		})
 		if err != nil {
-			return fmt.Errorf("cannot execute azure api request at %s: %w", nextLink, err)
+			return fmt.Errorf("cannot execute azure api request at %s: %w", nextLinkURI, err)
 		}
 		var lar listAPIResponse
 		if err := json.Unmarshal(resp, &lar); err != nil {
-			return fmt.Errorf("cannot parse azure api response %q obtained from %s: %w", resp, nextLink, err)
+			return fmt.Errorf("cannot parse azure api response %q obtained from %s: %w", resp, nextLinkURI, err)
 		}
 		for i := range lar.Value {
 			if err := cb(lar.Value[i]); err != nil {
 				return err
 			}
 		}
-		nextLink = lar.NextLink
+
+		// Azure API returns NextLink with apiServer in it, so we need to remove it
+		if lar.NextLink == "" {
+			break
+		}
+		nextURL, err := url.Parse(lar.NextLink)
+		if err != nil {
+			return fmt.Errorf("cannot parse nextLink from response %q: %w", lar.NextLink, err)
+		}
+
+		if nextURL.Host != "" && nextURL.Host != ac.c.APIServer() {
+			return fmt.Errorf("unexpected nextLink host %q, expecting %q", nextURL.Host, ac.c.APIServer())
+		}
+
+		nextLinkURI = nextURL.RequestURI()
 	}
 	return nil
 }
diff --git a/lib/promscrape/discoveryutils/client.go b/lib/promscrape/discoveryutils/client.go
index 842bda1bc..0ac62056d 100644
--- a/lib/promscrape/discoveryutils/client.go
+++ b/lib/promscrape/discoveryutils/client.go
@@ -240,6 +240,11 @@ func (c *Client) getAPIResponseWithParamsAndClient(client *fasthttp.HostClient,
 	return data, nil
 }
 
+// APIServer returns the API server address
+func (c *Client) APIServer() string {
+	return c.apiServer
+}
+
 // DoRequestWithPossibleRetry performs the given req at hc and stores the response at resp.
 func DoRequestWithPossibleRetry(hc *fasthttp.HostClient, req *fasthttp.Request, resp *fasthttp.Response, deadline time.Time, requestCounter, retryCounter *metrics.Counter) error {
 	sleepTime := time.Second