Detect when running agent in AKS and set hostname accordingly (#26860)

pkg/util/cloudproviders/azure/azure.go

@@ -29,6 +29,36 @@ var (
 )
 
 const hostnameStyleSetting = "azure_hostname_style"
+const aksManagedOrchestratorTag = "aks-managed-orchestrator"
+const kubernetesTagValue = "Kubernetes"
+
+type metadata struct {
+	VMID              string
+	Name              string
+	ResourceGroupName string
+	SubscriptionID    string
+	TagsList          []map[string]string
+	OsProfile         struct {
+		ComputerName string
+	}
+}
+
+func (metadata metadata) GetFromStyle(style string) (string, error) {
+	switch style {
+	case "vmid":
+		return metadata.VMID, nil
+	case "name":
+		return metadata.Name, nil
+	case "name_and_resource_group":
+		return fmt.Sprintf("%s.%s", metadata.Name, metadata.ResourceGroupName), nil
+	case "full":
+		return fmt.Sprintf("%s.%s.%s", metadata.Name, metadata.ResourceGroupName, metadata.SubscriptionID), nil
+	case "os_computer_name":
+		return strings.ToLower(metadata.OsProfile.ComputerName), nil
+	default:
+		return "", fmt.Errorf("invalid azure_hostname_style value: %s", style)
+	}
+}
 
 // IsRunningOn returns true if the agent is running on Azure
 func IsRunningOn(ctx context.Context) bool {
@@ -45,15 +75,29 @@ var vmIDFetcher = cachedfetch.Fetcher{
 			metadataURL+"/metadata/instance/compute/vmId?api-version=2017-04-02&format=text",
 			config.Datadog().GetInt("metadata_endpoints_max_hostname_size"))
 		if err != nil {
-			return nil, fmt.Errorf("Azure HostAliases: unable to query metadata endpoint: %s", err)
+			return nil, fmt.Errorf("Azure HostAliases: unable to query metadata VM ID endpoint: %s", err)
 		}
 		return []string{res}, nil
 	},
 }
 
 // GetHostAliases returns the VM ID from the Azure Metadata api
 func GetHostAliases(ctx context.Context) ([]string, error) {
-	return vmIDFetcher.FetchStringSlice(ctx)
+	aliases := []string{}
+	vm, err := vmIDFetcher.FetchStringSlice(ctx)
+	if err == nil {
+		aliases = append(aliases, vm...)
+	}
+
+	metadata, err := getMetadata(ctx)
+	if err != nil {
+		return aliases, fmt.Errorf("Azure GetHostAliases: unable to query metadata endpoint: %s", err)
+	}
+
+	if isKubernetesTag(metadata.TagsList) {
+		aliases = append(aliases, metadata.OsProfile.ComputerName)
+	}
+	return aliases, err
 }
 
 var resourceGroupNameFetcher = cachedfetch.Fetcher{
@@ -122,48 +166,57 @@ var instanceMetaFetcher = cachedfetch.Fetcher{
 	Name: "Azure Instance Metadata",
 	Attempt: func(ctx context.Context) (interface{}, error) {
 		metadataJSON, err := getResponse(ctx,
-			metadataURL+"/metadata/instance/compute?api-version=2017-08-01")
+			metadataURL+"/metadata/instance/compute?api-version=2021-02-01")
 		if err != nil {
 			return "", fmt.Errorf("failed to get Azure instance metadata: %s", err)
 		}
 		return metadataJSON, nil
 	},
 }
 
-func getHostnameWithConfig(ctx context.Context, config config.Config) (string, error) {
-	style := config.GetString(hostnameStyleSetting)
-
-	if style == "os" {
-		return "", fmt.Errorf("azure_hostname_style is set to 'os'")
+func isKubernetesTag(tagsList []map[string]string) bool {
+	for _, tag := range tagsList {
+		if tag["name"] == aksManagedOrchestratorTag && strings.Contains(tag["value"], kubernetesTagValue) {
+			return true
+		}
 	}
+	return false
+}
 
+func getMetadata(ctx context.Context) (metadata, error) {
+	metadataInfo := metadata{}
 	metadataJSON, err := instanceMetaFetcher.FetchString(ctx)
 	if err != nil {
-		return "", err
+		return metadataInfo, err
 	}
 
-	var metadata struct {
-		VMID              string
-		Name              string
-		ResourceGroupName string
-		SubscriptionID    string
+	if err := json.Unmarshal([]byte(metadataJSON), &metadataInfo); err != nil {
+		return metadataInfo, fmt.Errorf("failed to parse Azure instance metadata: %s", err)
 	}
-	if err := json.Unmarshal([]byte(metadataJSON), &metadata); err != nil {
-		return "", fmt.Errorf("failed to parse Azure instance metadata: %s", err)
+	return metadataInfo, nil
+}
+
+func getHostnameWithConfig(ctx context.Context, config config.Config) (string, error) {
+	style := config.GetString(hostnameStyleSetting)
+	metadata, err := getMetadata(ctx)
+	if err != nil {
+		return "", err
 	}
 
-	var name string
-	switch style {
-	case "vmid":
-		name = metadata.VMID
-	case "name":
-		name = metadata.Name
-	case "name_and_resource_group":
-		name = fmt.Sprintf("%s.%s", metadata.Name, metadata.ResourceGroupName)
-	case "full":
-		name = fmt.Sprintf("%s.%s.%s", metadata.Name, metadata.ResourceGroupName, metadata.SubscriptionID)
-	default:
-		return "", fmt.Errorf("invalid azure_hostname_style value: %s", style)
+	isKubernetes := isKubernetesTag(metadata.TagsList)
+
+	if style == "os" {
+		if isKubernetes {
+			// If running in AKS, use the node name as the hostname
+			style = "os_computer_name"
+		} else {
+			return "", fmt.Errorf("azure_hostname_style is set to 'os'")
+		}
+	}
+
+	name, err := metadata.GetFromStyle(style)
+	if err != nil {
+		return "", err
 	}
 
 	if err := validate.ValidHostname(name); err != nil {

pkg/util/cloudproviders/azure/azure_test.go

@@ -22,22 +22,43 @@ import (
 
 func TestGetAlias(t *testing.T) {
 	ctx := context.Background()
-	expected := "5d33a910-a7a0-4443-9f01-6a807801b29b"
+	expectedNodeName := "node-name-A"
+	expectedVM := "5d33a910-a7a0-4443-9f01-6a807801b29b"
+	responseIdx := 0
+	responses := []func(w http.ResponseWriter, r *http.Request){
+		func(w http.ResponseWriter, r *http.Request) {
+			w.Header().Set("Content-Type", "text/plain")
+			io.WriteString(w, expectedVM)
+		},
+		func(w http.ResponseWriter, r *http.Request) {
+			w.Header().Set("Content-Type", "application/json")
+			io.WriteString(w, fmt.Sprintf(`{
+				"name": "vm-name",
+				"resourceGroupName": "my-resource-group",
+				"subscriptionId": "2370ac56-5683-45f8-a2d4-d1054292facb",
+				"vmId": "b33fa46-6aff-4dfa-be0a-9e922ca3ac6d",
+				"osProfile": {"computerName":"%s"},
+				"tagsList": [{"name":"aks-managed-orchestrator","value":"Kubernetes"}]
+			}`, expectedNodeName))
+		},
+	}
 	var lastRequest *http.Request
 	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		w.Header().Set("Content-Type", "text/plain")
-		io.WriteString(w, expected)
+		responses[responseIdx](w, r)
+		responseIdx++
 		lastRequest = r
 	}))
+
 	defer ts.Close()
 	metadataURL = ts.URL
 
 	aliases, err := GetHostAliases(ctx)
 	assert.NoError(t, err)
-	require.Len(t, aliases, 1)
-	assert.Equal(t, expected, aliases[0])
-	assert.Equal(t, lastRequest.URL.Path, "/metadata/instance/compute/vmId")
-	assert.Equal(t, lastRequest.URL.RawQuery, "api-version=2017-04-02&format=text")
+	require.Len(t, aliases, 2)
+	assert.Equal(t, expectedVM, aliases[0])
+	assert.Equal(t, expectedNodeName, aliases[1])
+	assert.Equal(t, lastRequest.URL.Path, "/metadata/instance/compute")
+	assert.Equal(t, lastRequest.URL.RawQuery, "api-version=2021-02-01")
 }
 
 func TestGetClusterName(t *testing.T) {
@@ -90,9 +111,27 @@ func TestGetNTPHosts(t *testing.T) {
 	ctx := context.Background()
 	expectedHosts := []string{"time.windows.com"}
 
+	responseIdx := 0
+	responses := []func(w http.ResponseWriter, r *http.Request){
+		func(w http.ResponseWriter, r *http.Request) {
+			w.Header().Set("Content-Type", "text/plain")
+			io.WriteString(w, "test")
+		},
+		func(w http.ResponseWriter, r *http.Request) {
+			w.Header().Set("Content-Type", "application/json")
+			io.WriteString(w, fmt.Sprintf(`{
+				"name": "vm-name",
+				"resourceGroupName": "my-resource-group",
+				"subscriptionId": "2370ac56-5683-45f8-a2d4-d1054292facb",
+				"vmId": "b33fa46-6aff-4dfa-be0a-9e922ca3ac6d",
+				"osProfile": {"computerName":"%s"},
+				"tagsList": [{"name":"aks-managed-orchestrator","value":"Kubernetes"}]
+			}`, "node-name-a"))
+		},
+	}
 	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		w.Header().Set("Content-Type", "text/plain")
-		io.WriteString(w, "test")
+		responses[responseIdx](w, r)
+		responseIdx++
 	}))
 	defer ts.Close()
 
@@ -139,6 +178,44 @@ func TestGetHostname(t *testing.T) {
 	}
 }
 
+func TestGetHostnameKubernetesTag(t *testing.T) {
+	ctx := context.Background()
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		io.WriteString(w, `{
+			"name": "vm-name",
+			"resourceGroupName": "my-resource-group",
+			"subscriptionId": "2370ac56-5683-45f8-a2d4-d1054292facb",
+			"vmId": "b33fa46-6aff-4dfa-be0a-9e922ca3ac6d",
+			"osProfile": {"computerName":"node-name-A"},
+			"tagsList": [{"name":"aks-managed-orchestrator","value":"Kubernetes"}]
+		}`)
+	}))
+	defer ts.Close()
+	metadataURL = ts.URL
+
+	cases := []struct {
+		style, value string
+		err          bool
+	}{
+		{"os", "node-name-a", false}, // use osProfile.computerName when running in AKS
+		{"vmid", "b33fa46-6aff-4dfa-be0a-9e922ca3ac6d", false},
+		{"name", "vm-name", false},
+		{"name_and_resource_group", "vm-name.my-resource-group", false},
+		{"full", "vm-name.my-resource-group.2370ac56-5683-45f8-a2d4-d1054292facb", false},
+		{"invalid", "", true},
+	}
+
+	mockConfig := config.Mock(t)
+
+	for _, tt := range cases {
+		mockConfig.SetWithoutSource(hostnameStyleSetting, tt.style)
+		hostname, err := getHostnameWithConfig(ctx, mockConfig)
+		assert.Equal(t, tt.value, hostname)
+		assert.Equal(t, tt.err, (err != nil))
+	}
+}
+
 func TestGetHostnameWithInvalidMetadata(t *testing.T) {
 	ctx := context.Background()
 	mockConfig := config.Mock(t)

releasenotes/notes/use-cloud-provided-hostname-aks-929bfb01e0e336e5.yaml

@@ -0,0 +1,12 @@
+# Each section from every release note are combined when the
+# CHANGELOG.rst is rendered. So the text needs to be worded so that
+# it does not depend on any information only available in another
+# section. This may mean repeating some details, but each section
+# must be readable independently of the other.
+#
+# Each section note must be formatted as reStructuredText.
+---
+enhancements:
+  - |
+    Use cloud-provided hostname as default when running the Agent
+    in AKS.