Add option to append extra string to user agent

Signed-off-by: Eddie Torres <[email protected]>

Makefile

@@ -248,7 +248,7 @@ generate-kustomize: bin/helm
 	cd charts/aws-ebs-csi-driver && ../../bin/helm template kustomize . -s templates/clusterrolebinding-provisioner.yaml > ../../deploy/kubernetes/base/clusterrolebinding-provisioner.yaml
 	cd charts/aws-ebs-csi-driver && ../../bin/helm template kustomize . -s templates/clusterrolebinding-resizer.yaml > ../../deploy/kubernetes/base/clusterrolebinding-resizer.yaml
 	cd charts/aws-ebs-csi-driver && ../../bin/helm template kustomize . -s templates/clusterrolebinding-snapshotter.yaml > ../../deploy/kubernetes/base/clusterrolebinding-snapshotter.yaml
-	cd charts/aws-ebs-csi-driver && ../../bin/helm template kustomize . -s templates/controller.yaml --api-versions 'snapshot.storage.k8s.io/v1' | sed -e "/namespace: /d" > ../../deploy/kubernetes/base/controller.yaml
+	cd charts/aws-ebs-csi-driver && ../../bin/helm template kustomize . -s templates/controller.yaml --api-versions 'snapshot.storage.k8s.io/v1' --set 'controller.userAgentExtra=kustomize' | sed -e "/namespace: /d" > ../../deploy/kubernetes/base/controller.yaml
 	cd charts/aws-ebs-csi-driver && ../../bin/helm template kustomize . -s templates/csidriver.yaml > ../../deploy/kubernetes/base/csidriver.yaml
 	cd charts/aws-ebs-csi-driver && ../../bin/helm template kustomize . -s templates/node.yaml | sed -e "/namespace: /d" > ../../deploy/kubernetes/base/node.yaml
 	cd charts/aws-ebs-csi-driver && ../../bin/helm template kustomize . -s templates/poddisruptionbudget-controller.yaml --api-versions 'policy/v1/PodDisruptionBudget' | sed -e "/namespace: /d" > ../../deploy/kubernetes/base/poddisruptionbudget-controller.yaml

charts/aws-ebs-csi-driver/templates/controller.yaml

@@ -90,6 +90,9 @@ spec:
             {{- with .Values.controller.loggingFormat }}
             - --logging-format={{ . }}
             {{- end }}
+            {{- with .Values.controller.userAgentExtra }}
+            - --user-agent-extra={{ . }}
+            {{- end }}
             - --v={{ .Values.controller.logLevel }}
             {{- range .Values.controller.additionalArgs }}
             - {{ . }}

charts/aws-ebs-csi-driver/values.yaml

@@ -166,6 +166,7 @@ controller:
   # ID of the Kubernetes cluster used for tagging provisioned EBS volumes (optional).
   k8sTagClusterId:
   logLevel: 2
+  userAgentExtra: "helm"
   nodeSelector: {}
   podAnnotations: {}
   podLabels: {}

cmd/main.go

@@ -61,6 +61,7 @@ func main() {
 		driver.WithKubernetesClusterID(options.ControllerOptions.KubernetesClusterID),
 		driver.WithAwsSdkDebugLog(options.ControllerOptions.AwsSdkDebugLog),
 		driver.WithWarnOnInvalidTag(options.ControllerOptions.WarnOnInvalidTag),
+		driver.WithUserAgentExtra(options.ControllerOptions.UserAgentExtra),
 	)
 	if err != nil {
 		klog.ErrorS(err, "failed to create driver")

cmd/options/controller_options.go

@@ -37,6 +37,8 @@ type ControllerOptions struct {
 	AwsSdkDebugLog bool
 	// flag to warn on invalid tag, instead of returning an error
 	WarnOnInvalidTag bool
+	// flag to set user agent
+	UserAgentExtra string
 }
 
 func (s *ControllerOptions) AddFlags(fs *flag.FlagSet) {
@@ -45,4 +47,5 @@ func (s *ControllerOptions) AddFlags(fs *flag.FlagSet) {
 	fs.StringVar(&s.KubernetesClusterID, "k8s-tag-cluster-id", "", "ID of the Kubernetes cluster used for tagging provisioned EBS volumes (optional).")
 	fs.BoolVar(&s.AwsSdkDebugLog, "aws-sdk-debug-log", false, "To enable the aws sdk debug log level (default to false).")
 	fs.BoolVar(&s.WarnOnInvalidTag, "warn-on-invalid-tag", false, "To warn on invalid tags, instead of returning an error")
+	fs.StringVar(&s.UserAgentExtra, "user-agent-extra", "", "Set user agent.")
 }

cmd/options/controller_options_test.go

@@ -43,6 +43,11 @@ func TestControllerOptions(t *testing.T) {
 			flag:  "aws-sdk-debug-log",
 			found: true,
 		},
+		{
+			name:  "lookup user-agent-extra",
+			flag:  "user-agent-extra",
+			found: true,
+		},
 		{
 			name:  "fail for non-desired flag",
 			flag:  "some-other-flag",

cmd/options_test.go

@@ -51,6 +51,8 @@ func TestGetOptions(t *testing.T) {
 		awsSdkDebugFlagValue := true
 		VolumeAttachLimitFlagName := "volume-attach-limit"
 		var VolumeAttachLimit int64 = 42
+		userAgentExtraFlag := "user-agent-extra"
+		userAgentExtraFlagValue := "test"
 
 		args := append([]string{
 			"aws-ebs-csi-driver",
@@ -62,6 +64,7 @@ func TestGetOptions(t *testing.T) {
 		if withControllerOptions {
 			args = append(args, "--"+extraTagsFlagName+"="+extraTagKey+"="+extraTagValue)
 			args = append(args, "--"+awsSdkDebugFlagName+"="+strconv.FormatBool(awsSdkDebugFlagValue))
+			args = append(args, "--"+userAgentExtraFlag+"="+userAgentExtraFlagValue)
 		}
 		if withNodeOptions {
 			args = append(args, "--"+VolumeAttachLimitFlagName+"="+strconv.FormatInt(VolumeAttachLimit, 10))
@@ -97,6 +100,9 @@ func TestGetOptions(t *testing.T) {
 			if options.ControllerOptions.AwsSdkDebugLog != awsSdkDebugFlagValue {
 				t.Fatalf("expected sdk debug flag to be %v but it is %v", awsSdkDebugFlagValue, options.ControllerOptions.AwsSdkDebugLog)
 			}
+			if options.ControllerOptions.UserAgentExtra != userAgentExtraFlagValue {
+				t.Fatalf("expected user agent string to be %q but it is %q", userAgentExtraFlagValue, options.ControllerOptions.UserAgentExtra)
+			}
 		}
 
 		if withNodeOptions {

deploy/kubernetes/base/controller.yaml

@@ -67,6 +67,7 @@ spec:
             # - {all,controller,node} # specify the driver mode
             - --endpoint=$(CSI_ENDPOINT)
             - --logging-format=text
+            - --user-agent-extra=kustomize
             - --v=2
           env:
             - name: CSI_ENDPOINT

docs/options.md

@@ -9,3 +9,4 @@ There are a couple of driver options that can be passed as arguments when starti
 | k8s-tag-cluster-id          | aws-cluster-id-1                                  |                                                     | ID of the Kubernetes cluster used for tagging provisioned EBS volumes|
 | aws-sdk-debug-log           | true                                              | false                                               | If set to true, the driver will enable the aws sdk debug log level|
 | logging-format              | json                                              | text                                                | Sets the log format. Permitted formats: text, json|
+| user-agent-extra            | csi-ebs                                           | helm                                                | Extra string appended to user agent|

pkg/cloud/cloud.go

@@ -236,12 +236,12 @@ var _ Cloud = &cloud{}
 
 // NewCloud returns a new instance of AWS cloud
 // It panics if session is invalid
-func NewCloud(region string, awsSdkDebugLog bool) (Cloud, error) {
+func NewCloud(region string, awsSdkDebugLog bool, userAgentExtra string) (Cloud, error) {
 	RegisterMetrics()
-	return newEC2Cloud(region, awsSdkDebugLog)
+	return newEC2Cloud(region, awsSdkDebugLog, userAgentExtra)
 }
 
-func newEC2Cloud(region string, awsSdkDebugLog bool) (Cloud, error) {
+func newEC2Cloud(region string, awsSdkDebugLog bool, userAgentExtra string) (Cloud, error) {
 	awsConfig := &aws.Config{
 		Region:                        aws.String(region),
 		CredentialsChainVerboseErrors: aws.Bool(true),
@@ -268,7 +268,11 @@ func newEC2Cloud(region string, awsSdkDebugLog bool) (Cloud, error) {
 	}
 
 	// Set the env var so that the session appends custom user agent string
-	os.Setenv("AWS_EXECUTION_ENV", "aws-ebs-csi-driver-"+driverVersion)
+	if userAgentExtra != "" {
+		os.Setenv("AWS_EXECUTION_ENV", "aws-ebs-csi-driver-"+driverVersion+"-"+userAgentExtra)
+	} else {
+		os.Setenv("AWS_EXECUTION_ENV", "aws-ebs-csi-driver-"+driverVersion)
+	}
 
 	svc := ec2.New(session.Must(session.NewSession(awsConfig)))
 	svc.Handlers.AfterRetry.PushFrontNamed(request.NamedHandler{

pkg/driver/controller.go

@@ -88,7 +88,7 @@ func newControllerService(driverOptions *DriverOptions) controllerService {
 		region = metadata.GetRegion()
 	}
 
-	cloudSrv, err := NewCloudFunc(region, driverOptions.awsSdkDebugLog)
+	cloudSrv, err := NewCloudFunc(region, driverOptions.awsSdkDebugLog, driverOptions.userAgentExtra)
 	if err != nil {
 		panic(err)
 	}

pkg/driver/controller_test.go

@@ -51,8 +51,8 @@ func TestNewControllerService(t *testing.T) {
 		testErr    = errors.New("test error")
 		testRegion = "test-region"
 
-		getNewCloudFunc = func(expectedRegion string, _ bool) func(region string, awsSdkDebugLog bool) (cloud.Cloud, error) {
-			return func(region string, awsSdkDebugLog bool) (cloud.Cloud, error) {
+		getNewCloudFunc = func(expectedRegion string, _ bool) func(region string, awsSdkDebugLog bool, userAgentExtra string) (cloud.Cloud, error) {
+			return func(region string, awsSdkDebugLog bool, userAgentExtra string) (cloud.Cloud, error) {
 				if region != expectedRegion {
 					t.Fatalf("expected region %q but got %q", expectedRegion, region)
 				}
@@ -64,7 +64,7 @@ func TestNewControllerService(t *testing.T) {
 	testCases := []struct {
 		name                  string
 		region                string
-		newCloudFunc          func(string, bool) (cloud.Cloud, error)
+		newCloudFunc          func(string, bool, string) (cloud.Cloud, error)
 		newMetadataFuncErrors bool
 		expectPanic           bool
 	}{
@@ -76,7 +76,7 @@ func TestNewControllerService(t *testing.T) {
 		{
 			name:   "AWS_REGION variable set, newCloud errors",
 			region: "foo",
-			newCloudFunc: func(region string, awsSdkDebugLog bool) (cloud.Cloud, error) {
+			newCloudFunc: func(region string, awsSdkDebugLog bool, userAgentExtra string) (cloud.Cloud, error) {
 				return nil, testErr
 			},
 			expectPanic: true,

pkg/driver/driver.go

@@ -67,6 +67,7 @@ type DriverOptions struct {
 	kubernetesClusterID string
 	awsSdkDebugLog      bool
 	warnOnInvalidTag    bool
+	userAgentExtra      string
 }
 
 func NewDriver(options ...func(*DriverOptions)) (*Driver, error) {
@@ -198,3 +199,9 @@ func WithWarnOnInvalidTag(warnOnInvalidTag bool) func(*DriverOptions) {
 		o.warnOnInvalidTag = warnOnInvalidTag
 	}
 }
+
+func WithUserAgentExtra(userAgentExtra string) func(*DriverOptions) {
+	return func(o *DriverOptions) {
+		o.userAgentExtra = userAgentExtra
+	}
+}

pkg/driver/driver_test.go

@@ -94,3 +94,12 @@ func TestWithAwsSdkDebugLog(t *testing.T) {
 		t.Fatalf("expected awsSdkDebugLog option got set to %v but is set to %v", enableSdkDebugLog, options.awsSdkDebugLog)
 	}
 }
+
+func TestWithUserAgentExtra(t *testing.T) {
+	var userAgentExtra string = "test-user-agent"
+	options := &DriverOptions{}
+	WithUserAgentExtra(userAgentExtra)(options)
+	if options.userAgentExtra != userAgentExtra {
+		t.Fatalf("expected userAgentExtra option got set to %s but is set to %s", userAgentExtra, options.userAgentExtra)
+	}
+}

tests/e2e/dynamic_provisioning.go

@@ -374,7 +374,7 @@ var _ = Describe("[ebs-csi-e2e] [single-az] Dynamic Provisioning", func() {
 		availabilityZones := strings.Split(os.Getenv(awsAvailabilityZonesEnv), ",")
 		availabilityZone := availabilityZones[rand.Intn(len(availabilityZones))]
 		region := availabilityZone[0 : len(availabilityZone)-1]
-		cloud, err := awscloud.NewCloud(region, false)
+		cloud, err := awscloud.NewCloud(region, false, "")
 		if err != nil {
 			Fail(fmt.Sprintf("could not get NewCloud: %v", err))
 		}

tests/e2e/pre_provsioning.go

@@ -88,7 +88,7 @@ var _ = Describe("[ebs-csi-e2e] [single-az] Pre-Provisioned", func() {
 			Tags:             map[string]string{awscloud.VolumeNameTagKey: dummyVolumeName, awscloud.AwsEbsDriverTagKey: "true"},
 		}
 		var err error
-		cloud, err = awscloud.NewCloud(region, false)
+		cloud, err = awscloud.NewCloud(region, false, "")
 		if err != nil {
 			Fail(fmt.Sprintf("could not get NewCloud: %v", err))
 		}