From 2998c91a62691a560195f22c47267bf956b98127 Mon Sep 17 00:00:00 2001
From: Luc Talatinian <102624213+lucix-aws@users.noreply.github.com>
Date: Thu, 21 Mar 2024 12:39:10 -0400
Subject: [PATCH] feat: add disable switch for imds default timeout (#2572)

---
 .../bd8974c81b064061a388acea73c74956.json     |  8 +++++++
 feature/ec2/imds/api_client.go                |  4 ++++
 feature/ec2/imds/doc.go                       |  5 ++--
 feature/ec2/imds/request_middleware.go        |  6 +++++
 feature/ec2/imds/request_middleware_test.go   | 23 +++++++++++++++++++
 5 files changed, 44 insertions(+), 2 deletions(-)
 create mode 100644 .changelog/bd8974c81b064061a388acea73c74956.json

diff --git a/.changelog/bd8974c81b064061a388acea73c74956.json b/.changelog/bd8974c81b064061a388acea73c74956.json
new file mode 100644
index 00000000000..203d7f4037d
--- /dev/null
+++ b/.changelog/bd8974c81b064061a388acea73c74956.json
@@ -0,0 +1,8 @@
+{
+    "id": "bd8974c8-1b06-4061-a388-acea73c74956",
+    "type": "feature",
+    "description": "Add config switch `DisableDefaultTimeout` that allows you to disable the default operation timeout (5 seconds) for IMDS calls.",
+    "modules": [
+        "feature/ec2/imds"
+    ]
+}
\ No newline at end of file
diff --git a/feature/ec2/imds/api_client.go b/feature/ec2/imds/api_client.go
index 46e144d9363..3f4a10e2c16 100644
--- a/feature/ec2/imds/api_client.go
+++ b/feature/ec2/imds/api_client.go
@@ -185,6 +185,10 @@ type Options struct {
 	// [configuring IMDS]: https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/configuring-instance-metadata-service.html
 	EnableFallback aws.Ternary
 
+	// By default, all IMDS client operations enforce a 5-second timeout. You
+	// can disable that behavior with this setting.
+	DisableDefaultTimeout bool
+
 	// provides the caching of API tokens used for operation calls. If unset,
 	// the API token will not be retrieved for the operation.
 	tokenProvider *tokenProvider
diff --git a/feature/ec2/imds/doc.go b/feature/ec2/imds/doc.go
index bacdb5d21f2..d5765c36b17 100644
--- a/feature/ec2/imds/doc.go
+++ b/feature/ec2/imds/doc.go
@@ -3,8 +3,9 @@
 //
 // All Client operation calls have a default timeout. If the operation is not
 // completed before this timeout expires, the operation will be canceled. This
-// timeout can be overridden by providing Context with a timeout or deadline
-// with calling the client's operations.
+// timeout can be overridden through the following:
+//   - Set the options flag DisableDefaultTimeout
+//   - Provide a Context with a timeout or deadline with calling the client's operations.
 //
 // See the EC2 IMDS user guide for more information on using the API.
 // https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/ec2-instance-metadata.html
diff --git a/feature/ec2/imds/request_middleware.go b/feature/ec2/imds/request_middleware.go
index fc948c27d89..90cf4aeb3df 100644
--- a/feature/ec2/imds/request_middleware.go
+++ b/feature/ec2/imds/request_middleware.go
@@ -56,6 +56,7 @@ func addRequestMiddleware(stack *middleware.Stack,
 
 	// Operation timeout
 	err = stack.Initialize.Add(&operationTimeout{
+		Disabled:       options.DisableDefaultTimeout,
 		DefaultTimeout: defaultOperationTimeout,
 	}, middleware.Before)
 	if err != nil {
@@ -260,6 +261,7 @@ const (
 // Otherwise the timeout cleanup will race the resource being consumed
 // upstream.
 type operationTimeout struct {
+	Disabled       bool
 	DefaultTimeout time.Duration
 }
 
@@ -270,6 +272,10 @@ func (m *operationTimeout) HandleInitialize(
 ) (
 	output middleware.InitializeOutput, metadata middleware.Metadata, err error,
 ) {
+	if m.Disabled {
+		return next.HandleInitialize(ctx, input)
+	}
+
 	if _, ok := ctx.Deadline(); !ok && m.DefaultTimeout != 0 {
 		var cancelFn func()
 		ctx, cancelFn = context.WithTimeout(ctx, m.DefaultTimeout)
diff --git a/feature/ec2/imds/request_middleware_test.go b/feature/ec2/imds/request_middleware_test.go
index ba8b070bdce..b596da89cb3 100644
--- a/feature/ec2/imds/request_middleware_test.go
+++ b/feature/ec2/imds/request_middleware_test.go
@@ -220,6 +220,29 @@ func TestOperationTimeoutMiddleware_withCustomDeadline(t *testing.T) {
 	}
 }
 
+func TestOperationTimeoutMiddleware_Disabled(t *testing.T) {
+	m := &operationTimeout{
+		Disabled:       true,
+		DefaultTimeout: time.Nanosecond,
+	}
+
+	_, _, err := m.HandleInitialize(context.Background(), middleware.InitializeInput{},
+		middleware.InitializeHandlerFunc(func(
+			ctx context.Context, input middleware.InitializeInput,
+		) (
+			out middleware.InitializeOutput, metadata middleware.Metadata, err error,
+		) {
+			if err := sdk.SleepWithContext(ctx, time.Second); err != nil {
+				return out, metadata, err
+			}
+
+			return out, metadata, nil
+		}))
+	if err != nil {
+		t.Fatalf("expect no error, got %v", err)
+	}
+}
+
 // Ensure that the response body is read in the deserialize middleware,
 // ensuring that the timeoutOperation middleware won't race canceling the
 // context with the upstream reading the response body.