feat: add disable switch for imds default timeout (#2572)

.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"
+    ]
+}

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

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

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)

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.