Track2 sdk: override default UserAgent header in sdk

pkg/azclient/arm_conf.go

@@ -22,6 +22,7 @@ import (
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
 
+	"sigs.k8s.io/cloud-provider-azure/pkg/azclient/policy/useragent"
 	"sigs.k8s.io/cloud-provider-azure/pkg/azclient/utils"
 )
 
@@ -52,7 +53,10 @@ func GetAzCoreClientOption(armConfig *ARMClientConfig) (*policy.ClientOptions, e
 	azCoreClientConfig := utils.GetDefaultAzCoreClientOption()
 	if armConfig != nil {
 		//update user agent header
-		azCoreClientConfig.Telemetry.ApplicationID = strings.TrimSpace(armConfig.UserAgent)
+		if userAgent := strings.TrimSpace(armConfig.UserAgent); userAgent != "" {
+			azCoreClientConfig.Telemetry.Disabled = true
+			azCoreClientConfig.PerCallPolicies = append(azCoreClientConfig.PerCallPolicies, useragent.NewCustomUserAgentPolicy(userAgent))
+		}
 		//set cloud
 		cloudConfig, err := GetAzureCloudConfig(armConfig)
 		if err != nil {

pkg/azclient/policy/useragent/user_agent.go

@@ -0,0 +1,46 @@
+/*
+Copyright 2023 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package useragent
+
+import (
+	"net/http"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
+)
+
+type CustomUserAgentPolicy struct {
+	CustomUserAgent string
+}
+
+const HeaderUserAgent = "User-Agent"
+
+func NewCustomUserAgentPolicy(customUserAgent string) policy.Policy {
+	return &CustomUserAgentPolicy{
+		CustomUserAgent: customUserAgent,
+	}
+}
+
+func (p CustomUserAgentPolicy) Do(req *policy.Request) (*http.Response, error) {
+	if p.CustomUserAgent == "" {
+		return req.Next()
+	}
+	// preserve the existing User-Agent string
+	if ua := req.Raw().Header.Get(HeaderUserAgent); ua == "" {
+		req.Raw().Header.Set(HeaderUserAgent, p.CustomUserAgent)
+	}
+	return req.Next()
+}

pkg/azclient/policy/useragent/user_agent_test.go

@@ -0,0 +1,93 @@
+/*
+Copyright 2023 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package useragent_test
+
+import (
+	"context"
+	"net/http"
+	"strings"
+	"sync"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime"
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/streaming"
+	"github.com/onsi/ginkgo/v2"
+	"github.com/onsi/gomega"
+
+	"sigs.k8s.io/cloud-provider-azure/pkg/azclient/policy/useragent"
+	"sigs.k8s.io/cloud-provider-azure/pkg/azclient/utils"
+)
+
+var _ = ginkgo.Describe("useragent", func() {
+	ginkgo.Describe("useragent", func() {
+		ginkgo.It("should respect useragent", func() {
+			once := sync.Once{}
+			userAgentPolicy := &useragent.CustomUserAgentPolicy{}
+			pipeline := runtime.NewPipeline("testmodule", "v0.1.0", runtime.PipelineOptions{}, &policy.ClientOptions{
+				Telemetry: policy.TelemetryOptions{
+					Disabled: true,
+				},
+				PerCallPolicies: []policy.Policy{
+					userAgentPolicy,
+					utils.FuncPolicyWrapper(
+						func(*policy.Request) (*http.Response, error) {
+							resp := &http.Response{
+								StatusCode: http.StatusOK,
+								Body:       http.NoBody,
+							}
+							once.Do(func() {
+								resp = &http.Response{
+									StatusCode: http.StatusTooManyRequests,
+									Body:       http.NoBody,
+									Header: http.Header{
+										"Retry-After": []string{"10"},
+									},
+								}
+							})
+							return resp, nil
+						},
+					),
+				},
+			})
+			userAgentPolicy.CustomUserAgent = "test"
+			req, err := runtime.NewRequest(context.Background(), http.MethodPut, "http://localhost:8080")
+			gomega.Expect(err).NotTo(gomega.HaveOccurred())
+			err = req.SetBody(streaming.NopCloser(strings.NewReader(`{"etag":"etag"}`)), "application/json")
+			gomega.Expect(err).NotTo(gomega.HaveOccurred())
+			_, err = pipeline.Do(req)
+			gomega.Expect(err).NotTo(gomega.HaveOccurred())
+			gomega.Expect(req.Raw().Header.Get(useragent.HeaderUserAgent)).To(gomega.Equal("test"))
+			userAgentPolicy.CustomUserAgent = ""
+			req, err = runtime.NewRequest(context.Background(), http.MethodPut, "http://localhost:8080")
+			req.Raw().Header.Set(useragent.HeaderUserAgent, "test-override")
+			gomega.Expect(err).NotTo(gomega.HaveOccurred())
+			err = req.SetBody(streaming.NopCloser(strings.NewReader(`{"etag":"etag"}`)), "application/json")
+			gomega.Expect(err).NotTo(gomega.HaveOccurred())
+			_, err = pipeline.Do(req)
+			gomega.Expect(err).NotTo(gomega.HaveOccurred())
+			gomega.Expect(req.Raw().Header.Get(useragent.HeaderUserAgent)).To(gomega.Equal("test-override"))
+			userAgentPolicy.CustomUserAgent = ""
+			req, err = runtime.NewRequest(context.Background(), http.MethodPut, "http://localhost:8080")
+			gomega.Expect(err).NotTo(gomega.HaveOccurred())
+			err = req.SetBody(streaming.NopCloser(strings.NewReader(`{"etag":"etag"}`)), "application/json")
+			gomega.Expect(err).NotTo(gomega.HaveOccurred())
+			_, err = pipeline.Do(req)
+			gomega.Expect(err).NotTo(gomega.HaveOccurred())
+			gomega.Expect(req.Raw().Header.Get(useragent.HeaderUserAgent)).To(gomega.BeEmpty())
+		})
+	})
+})

pkg/azclient/policy/useragent/useragent_suite_test.go

@@ -0,0 +1,29 @@
+/*
+Copyright 2023 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package useragent_test
+
+import (
+	"testing"
+
+	"github.com/onsi/ginkgo/v2"
+	"github.com/onsi/gomega"
+)
+
+func TestUserAgent(t *testing.T) {
+	gomega.RegisterFailHandler(ginkgo.Fail)
+	ginkgo.RunSpecs(t, "UserAgent Suite")
+}