EC2 IMDS client logging fixes (#1891)

aws · Oct 23, 2022 · 17628c4 · 17628c4
1 parent 7655449
commit 17628c4
Show file tree

Hide file tree

Showing 4 changed files with 35 additions and 2 deletions.
diff --git a/.changelog/a3c2a480e3c9401a93ab4bafab06262c.json b/.changelog/a3c2a480e3c9401a93ab4bafab06262c.json
@@ -0,0 +1,8 @@
+{
+    "id": "a3c2a480-e3c9-401a-93ab-4bafab06262c",
+    "type": "bugfix",
+    "description": "Fixes an issue that prevented logging of the API request or responses when the respective log modes were enabled.",
+    "modules": [
+        "feature/ec2/imds"
+    ]
+}
diff --git a/feature/ec2/imds/api_client.go b/feature/ec2/imds/api_client.go
@@ -106,8 +106,10 @@ func New(options Options, optFns ...func(*Options)) *Client {
 // or adding custom middleware behavior.
 func NewFromConfig(cfg aws.Config, optFns ...func(*Options)) *Client {
 	opts := Options{
-		APIOptions: append([]func(*middleware.Stack) error{}, cfg.APIOptions...),
-		HTTPClient: cfg.HTTPClient,
+		APIOptions:    append([]func(*middleware.Stack) error{}, cfg.APIOptions...),
+		HTTPClient:    cfg.HTTPClient,
+		ClientLogMode: cfg.ClientLogMode,
+		Logger:        cfg.Logger,
 	}
 
 	if cfg.Retryer != nil {

diff --git a/feature/ec2/imds/request_middleware.go b/feature/ec2/imds/request_middleware.go
@@ -86,13 +86,32 @@ func addRequestMiddleware(stack *middleware.Stack,
 		return err
 	}
 
+	err = stack.Deserialize.Add(&smithyhttp.RequestResponseLogger{
+		LogRequest:          options.ClientLogMode.IsRequest(),
+		LogRequestWithBody:  options.ClientLogMode.IsRequestWithBody(),
+		LogResponse:         options.ClientLogMode.IsResponse(),
+		LogResponseWithBody: options.ClientLogMode.IsResponseWithBody(),
+	}, middleware.After)
+	if err != nil {
+		return err
+	}
+
+	err = addSetLoggerMiddleware(stack, options)
+	if err != nil {
+		return err
+	}
+
 	// Retry support
 	return retry.AddRetryMiddlewares(stack, retry.AddRetryMiddlewaresOptions{
 		Retryer:          options.Retryer,
 		LogRetryAttempts: options.ClientLogMode.IsRetries(),
 	})
 }
 
+func addSetLoggerMiddleware(stack *middleware.Stack, o Options) error {
+	return middleware.AddSetLoggerMiddleware(stack, o.Logger)
+}
+
 type serializeRequest struct {
 	GetPath func(interface{}) (string, error)
 	Method  string

diff --git a/feature/ec2/imds/request_middleware_test.go b/feature/ec2/imds/request_middleware_test.go
@@ -42,6 +42,7 @@ func TestAddRequestMiddleware(t *testing.T) {
 			},
 			ExpectInitialize: []string{
 				(*operationTimeout)(nil).ID(),
+				"SetLogger",
 			},
 			ExpectSerialize: []string{
 				"ResolveEndpoint",
@@ -58,6 +59,7 @@ func TestAddRequestMiddleware(t *testing.T) {
 			ExpectDeserialize: []string{
 				"APITokenProvider",
 				"OperationDeserializer",
+				"RequestResponseLogger",
 			},
 		},
 
@@ -74,6 +76,7 @@ func TestAddRequestMiddleware(t *testing.T) {
 			},
 			ExpectInitialize: []string{
 				(*operationTimeout)(nil).ID(),
+				"SetLogger",
 			},
 			ExpectSerialize: []string{
 				"ResolveEndpoint",
@@ -88,6 +91,7 @@ func TestAddRequestMiddleware(t *testing.T) {
 			},
 			ExpectDeserialize: []string{
 				"OperationDeserializer",
+				"RequestResponseLogger",
 			},
 		},
 	}