Add user-agent strings to request context

aws_config.go

@@ -168,6 +168,10 @@ func commonLoadOptions(c *Config) ([]func(*config.LoadOptions) error, error) {
 		apiOptions = append(apiOptions, awsmiddleware.AddUserAgentKey(c.UserAgent.BuildUserAgentString()))
 	}
 
+	apiOptions = append(apiOptions, func(stack *middleware.Stack) error {
+		return stack.Build.Add(userAgentFromContextMiddleware(), middleware.After)
+	})
+
 	if v := os.Getenv(constants.AppendUserAgentEnvVar); v != "" {
 		log.Printf("[DEBUG] Using additional User-Agent Info: %s", v)
 		apiOptions = append(apiOptions, awsmiddleware.AddUserAgentKey(v))

aws_config_test.go

@@ -28,6 +28,7 @@ import (
 	"github.com/hashicorp/aws-sdk-go-base/v2/internal/test"
 	"github.com/hashicorp/aws-sdk-go-base/v2/mockdata"
 	"github.com/hashicorp/aws-sdk-go-base/v2/servicemocks"
+	"github.com/hashicorp/aws-sdk-go-base/v2/useragent"
 )
 
 const (
@@ -1107,7 +1108,13 @@ func testUserAgentProducts(t *testing.T, testCase test.UserAgentTestCase) {
 
 	client := stsClient(awsConfig, testCase.Config)
 
-	_, err = client.GetCallerIdentity(context.Background(), &sts.GetCallerIdentityInput{},
+	ctx := context.Background()
+
+	if testCase.Context != nil {
+		ctx = useragent.Context(ctx, testCase.Context)
+	}
+
+	_, err = client.GetCallerIdentity(ctx, &sts.GetCallerIdentityInput{},
 		func(opts *sts.Options) {
 			opts.APIOptions = append(opts.APIOptions, func(stack *middleware.Stack) error {
 				return stack.Finalize.Add(readUserAgent, middleware.Before)

internal/test/user_agent.go

@@ -11,6 +11,7 @@ import (
 
 type UserAgentTestCase struct {
 	Config               *config.Config
+	Context              config.UserAgentProducts
 	EnvironmentVariables map[string]string
 	ExpectedUserAgent    string
 }
@@ -133,6 +134,54 @@ func TestUserAgentProducts(t *testing.T, awsSdkGoUserAgent func() string, testUs
 			},
 			ExpectedUserAgent: "APN/1.0 partner/1.0 first/1.2.3 second/1.0.2 (a comment) " + awsSdkGoUserAgent() + " third/4.5.6 fourth/2.1",
 		},
+		"context": {
+			Config: &config.Config{
+				AccessKey: servicemocks.MockStaticAccessKey,
+				Region:    "us-east-1",
+				SecretKey: servicemocks.MockStaticSecretKey,
+			},
+			Context: []config.UserAgentProduct{
+				{
+					Name:    "first",
+					Version: "1.2.3",
+				},
+				{
+					Name:    "second",
+					Version: "1.0.2",
+					Comment: "a comment",
+				},
+			},
+			ExpectedUserAgent: awsSdkGoUserAgent() + " first/1.2.3 second/1.0.2 (a comment)",
+		},
+		"User-Agent Products and context": {
+			Config: &config.Config{
+				AccessKey: servicemocks.MockStaticAccessKey,
+				Region:    "us-east-1",
+				SecretKey: servicemocks.MockStaticSecretKey,
+				UserAgent: []config.UserAgentProduct{
+					{
+						Name:    "first",
+						Version: "1.2.3",
+					},
+					{
+						Name:    "second",
+						Version: "1.0.2",
+						Comment: "a comment",
+					},
+				},
+			},
+			Context: []config.UserAgentProduct{
+				{
+					Name:    "third",
+					Version: "4.5.6",
+				},
+				{
+					Name:    "fourth",
+					Version: "2.1",
+				},
+			},
+			ExpectedUserAgent: awsSdkGoUserAgent() + " first/1.2.3 second/1.0.2 (a comment) third/4.5.6 fourth/2.1",
+		},
 	}
 
 	for name, testCase := range testCases {

user_agent.go

@@ -6,6 +6,7 @@ import (
 
 	"github.com/aws/smithy-go/middleware"
 	smithyhttp "github.com/aws/smithy-go/transport/http"
+	"github.com/hashicorp/aws-sdk-go-base/v2/useragent"
 )
 
 func apnUserAgentMiddleware(apn APNInfo) middleware.BuildMiddleware {
@@ -33,5 +34,31 @@ func prependUserAgentHeader(request *smithyhttp.Request, value string) {
 		current = value
 	}
 	request.Header["User-Agent"] = append(request.Header["User-Agent"][:0], current)
+}
+
+func userAgentFromContextMiddleware() middleware.BuildMiddleware {
+	return middleware.BuildMiddlewareFunc("CtxUserAgent",
+		func(ctx context.Context, in middleware.BuildInput, next middleware.BuildHandler) (middleware.BuildOutput, middleware.Metadata, error) {
+			request, ok := in.Request.(*smithyhttp.Request)
+			if !ok {
+				return middleware.BuildOutput{}, middleware.Metadata{}, fmt.Errorf("unknown request type %T", in.Request)
+			}
+
+			if v := useragent.BuildFromContext(ctx); v != "" {
+				appendUserAgentHeader(request, v)
+			}
+
+			return next.HandleBuild(ctx, in)
+		},
+	)
+}
 
+func appendUserAgentHeader(request *smithyhttp.Request, value string) {
+	current := request.Header.Get("User-Agent")
+	if len(current) > 0 {
+		current = current + " " + value
+	} else {
+		current = value
+	}
+	request.Header["User-Agent"] = append(request.Header["User-Agent"][:0], current)
 }

useragent/context.go

@@ -0,0 +1,26 @@
+package useragent
+
+import (
+	"context"
+
+	"github.com/hashicorp/aws-sdk-go-base/v2/internal/config"
+)
+
+type userAgentKey string
+
+const (
+	contextScopedUserAgent userAgentKey = "ContextScopedUserAgent"
+)
+
+func Context(ctx context.Context, products config.UserAgentProducts) context.Context {
+	return context.WithValue(ctx, contextScopedUserAgent, products)
+}
+
+func BuildFromContext(ctx context.Context) string {
+	ps, ok := ctx.Value(contextScopedUserAgent).(config.UserAgentProducts)
+	if !ok {
+		return ""
+	}
+
+	return ps.BuildUserAgentString()
+}

useragent/context_test.go

@@ -0,0 +1,66 @@
+package useragent
+
+import (
+	"context"
+	"testing"
+
+	"github.com/hashicorp/aws-sdk-go-base/v2/internal/config"
+)
+
+func TestFromContext(t *testing.T) {
+	testcases := map[string]struct {
+		setup    func() context.Context
+		expected string
+	}{
+		"empty": {
+			setup: func() context.Context {
+				return context.Background()
+			},
+			expected: "",
+		},
+		"UserAgentProducts": {
+			setup: func() context.Context {
+				return Context(context.Background(), config.UserAgentProducts{
+					{
+						Name:    "first",
+						Version: "1.2.3",
+					},
+					{
+						Name:    "second",
+						Version: "1.0.2",
+						Comment: "a comment",
+					},
+				})
+			},
+			expected: "first/1.2.3 second/1.0.2 (a comment)",
+		},
+		"[]UserAgentProduct": {
+			setup: func() context.Context {
+				return Context(context.Background(), []config.UserAgentProduct{
+					{
+						Name:    "first",
+						Version: "1.2.3",
+					},
+					{
+						Name:    "second",
+						Version: "1.0.2",
+						Comment: "a comment",
+					},
+				})
+			},
+			expected: "first/1.2.3 second/1.0.2 (a comment)",
+		},
+	}
+
+	for name, testcase := range testcases {
+		t.Run(name, func(t *testing.T) {
+			ctx := testcase.setup()
+
+			v := BuildFromContext(ctx)
+
+			if v != testcase.expected {
+				t.Errorf("expected %q, got %q", testcase.expected, v)
+			}
+		})
+	}
+}

v2/awsv1shim/session.go

@@ -91,6 +91,8 @@ func GetSession(awsC *awsv2.Config, c *awsbase.Config) (*session.Session, error)
 
 	SetSessionUserAgent(sess, c.APNInfo, c.UserAgent)
 
+	sess.Handlers.Build.PushBack(userAgentFromContextHandler)
+
 	// Add custom input from ENV to the User-Agent request header
 	// Reference: https://github.com/terraform-providers/terraform-provider-aws/issues/9149
 	if v := os.Getenv(constants.AppendUserAgentEnvVar); v != "" {

v2/awsv1shim/session_test.go

@@ -32,6 +32,7 @@ import (
 	"github.com/hashicorp/aws-sdk-go-base/v2/internal/constants"
 	"github.com/hashicorp/aws-sdk-go-base/v2/internal/test"
 	"github.com/hashicorp/aws-sdk-go-base/v2/servicemocks"
+	"github.com/hashicorp/aws-sdk-go-base/v2/useragent"
 )
 
 func TestGetSessionOptions(t *testing.T) {
@@ -1171,6 +1172,11 @@ func testUserAgentProducts(t *testing.T, testCase test.UserAgentTestCase) {
 
 	req := conn.NewRequest(&request.Operation{Name: "Operation"}, nil, nil)
 
+	if testCase.Context != nil {
+		ctx := useragent.Context(context.Background(), testCase.Context)
+		req.SetContext(ctx)
+	}
+
 	if err := req.Build(); err != nil {
 		t.Fatalf("expect no Request.Build() error, got %s", err)
 	}

v2/awsv1shim/user_agent.go

@@ -0,0 +1,14 @@
+package awsv1shim
+
+import (
+	"github.com/aws/aws-sdk-go/aws/request"
+	"github.com/hashicorp/aws-sdk-go-base/v2/useragent"
+)
+
+func userAgentFromContextHandler(r *request.Request) {
+	ctx := r.Context()
+
+	if v := useragent.BuildFromContext(ctx); v != "" {
+		request.AddToUserAgent(r, v)
+	}
+}