Override user agent

Author: wweiwei-li

pkg/aws/aws_config.go

@@ -2,13 +2,14 @@ package aws
 
 import (
 	"context"
+
 	"github.com/aws/aws-sdk-go-v2/aws"
-	awsmiddleware "github.com/aws/aws-sdk-go-v2/aws/middleware"
 	"github.com/aws/aws-sdk-go-v2/aws/ratelimit"
 	"github.com/aws/aws-sdk-go-v2/aws/retry"
 	"github.com/aws/aws-sdk-go-v2/config"
 	"github.com/aws/aws-sdk-go-v2/feature/ec2/imds"
 	smithymiddleware "github.com/aws/smithy-go/middleware"
+	smithyhttp "github.com/aws/smithy-go/transport/http"
 	"sigs.k8s.io/aws-load-balancer-controller/pkg/aws/throttle"
 	awsmetrics "sigs.k8s.io/aws-load-balancer-controller/pkg/metrics/aws"
 	"sigs.k8s.io/aws-load-balancer-controller/pkg/version"
@@ -50,7 +51,7 @@ func (gen *awsConfigGeneratorImpl) GenerateAWSConfig(optFns ...func(*config.Load
 		}),
 		config.WithEC2IMDSEndpointMode(gen.ec2IMDSEndpointMode),
 		config.WithAPIOptions([]func(stack *smithymiddleware.Stack) error{
-			awsmiddleware.AddUserAgentKeyValue(userAgent, version.GitVersion),
+			overrideUserAgentMiddleware(userAgent + "/" + version.GitVersion),
 		}),
 	}
 
@@ -79,3 +80,17 @@ func (gen *awsConfigGeneratorImpl) GenerateAWSConfig(optFns ...func(*config.Load
 }
 
 var _ AWSConfigGenerator = &awsConfigGeneratorImpl{}
+
+// overrideUserAgentMiddleware returns a middleware that replaces the User-Agent
+// header with the given value, stripping all SDK-generated metadata.
+func overrideUserAgentMiddleware(ua string) func(stack *smithymiddleware.Stack) error {
+	return func(stack *smithymiddleware.Stack) error {
+		return stack.Build.Add(smithymiddleware.BuildMiddlewareFunc("OverrideUserAgent",
+			func(ctx context.Context, in smithymiddleware.BuildInput, next smithymiddleware.BuildHandler) (smithymiddleware.BuildOutput, smithymiddleware.Metadata, error) {
+				if req, ok := in.Request.(*smithyhttp.Request); ok {
+					req.Header.Set("User-Agent", ua)
+				}
+				return next.HandleBuild(ctx, in)
+			}), smithymiddleware.After)
+	}
+}

pkg/aws/aws_config_test.go

@@ -0,0 +1,75 @@
+package aws
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"testing"
+
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/config"
+	"github.com/aws/aws-sdk-go-v2/feature/ec2/imds"
+	"github.com/aws/aws-sdk-go-v2/service/ec2"
+	smithyendpoints "github.com/aws/smithy-go/endpoints"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"sigs.k8s.io/aws-load-balancer-controller/pkg/version"
+)
+
+// staticEC2EndpointResolver redirects all ec2 calls to a fixed URL.
+type staticEC2EndpointResolver struct{ url string }
+
+func (r staticEC2EndpointResolver) ResolveEndpoint(_ context.Context, _ ec2.EndpointParameters) (smithyendpoints.Endpoint, error) {
+	u, err := url.Parse(r.url)
+	if err != nil {
+		return smithyendpoints.Endpoint{}, err
+	}
+	return smithyendpoints.Endpoint{URI: *u}, nil
+}
+
+func TestGenerateAWSConfig_UserAgentHeader(t *testing.T) {
+	tests := []struct {
+		name       string
+		gitVersion string
+		wantUA     string
+	}{
+		{
+			name:       "user agent is exactly elbv2.k8s.aws/<version>",
+			gitVersion: "v2.14.1",
+			wantUA:     "elbv2.k8s.aws/v2.14.1",
+		},
+		{
+			name:       "user agent reflects updated git version",
+			gitVersion: "v3.0.0",
+			wantUA:     "elbv2.k8s.aws/v3.0.0",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			version.GitVersion = tt.gitVersion
+
+			var capturedUA string
+			srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				capturedUA = r.Header.Get("User-Agent")
+				w.WriteHeader(http.StatusBadRequest)
+			}))
+			defer srv.Close()
+
+			gen := NewAWSConfigGenerator(CloudConfig{Region: "us-east-1", MaxRetries: 1}, imds.EndpointModeStateIPv4, nil)
+			awsCfg, err := gen.GenerateAWSConfig(
+				config.WithCredentialsProvider(aws.AnonymousCredentials{}),
+			)
+			require.NoError(t, err)
+
+			ec2Client := ec2.NewFromConfig(awsCfg, func(o *ec2.Options) {
+				o.EndpointResolverV2 = staticEC2EndpointResolver{url: srv.URL}
+			})
+			_, _ = ec2Client.DescribeInstances(t.Context(), &ec2.DescribeInstancesInput{})
+
+			require.NotEmpty(t, capturedUA, "no request was captured")
+			assert.Equal(t, tt.wantUA, capturedUA)
+		})
+	}
+}