From fd3693c0da0cf7e83941943690e1a369e03ac7f0 Mon Sep 17 00:00:00 2001
From: xingzhen <xingzhen@amazon.com>
Date: Fri, 16 Aug 2024 13:40:00 -0700
Subject: [PATCH] Expose FaultInjection status in TaskResponse (#4285)

* Expose FIS status in TaskResponse

* Rename the new field
---
 .../tmds/handlers/v4/state/response.go        |  1 +
 ecs-agent/tmds/handlers/v4/handlers_test.go   | 28 ++++++++++++++++++-
 ecs-agent/tmds/handlers/v4/state/response.go  |  1 +
 3 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/agent/vendor/github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/v4/state/response.go b/agent/vendor/github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/v4/state/response.go
index e7226be2a6f..15baf8e21fe 100644
--- a/agent/vendor/github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/v4/state/response.go
+++ b/agent/vendor/github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/v4/state/response.go
@@ -38,6 +38,7 @@ type TaskResponse struct {
 	EphemeralStorageMetrics *EphemeralStorageMetrics `json:"EphemeralStorageMetrics,omitempty"`
 	CredentialsID           string                   `json:"-"`
 	TaskNetworkConfig       *TaskNetworkConfig       `json:"-"`
+	FaultInjectionEnabled   bool                     `json:"-"`
 }
 
 // TaskNetworkConfig contains required network configurations for network faults injection.
diff --git a/ecs-agent/tmds/handlers/v4/handlers_test.go b/ecs-agent/tmds/handlers/v4/handlers_test.go
index 01c322bd3dd..388170f7244 100644
--- a/ecs-agent/tmds/handlers/v4/handlers_test.go
+++ b/ecs-agent/tmds/handlers/v4/handlers_test.go
@@ -128,7 +128,7 @@ var (
 	}
 )
 
-// Returns a standard agent task response
+// taskResponse returns a standard agent task response
 func taskResponse() *state.TaskResponse {
 	return &state.TaskResponse{
 		TaskResponse: &v2.TaskResponse{
@@ -175,6 +175,13 @@ func taskResponse() *state.TaskResponse {
 	}
 }
 
+// taskResponseWithFaultInjectionEnabled returns a standard agent task response with FaultInjection enabled
+func taskResponseWithFaultInjectionEnabled() *state.TaskResponse {
+	taskResponse := taskResponse()
+	taskResponse.FaultInjectionEnabled = true
+	return taskResponse
+}
+
 func TestContainerMetadata(t *testing.T) {
 	var setup = func(t *testing.T) (*mux.Router, *gomock.Controller, *mock_state.MockAgentState,
 		*mock_metrics.MockEntryFactory,
@@ -288,6 +295,25 @@ func TestTaskMetadata(t *testing.T) {
 			expectedResponseBody: *expectedTaskResponse,
 		})
 	})
+
+	t.Run("happy case with FaultInjection enabled", func(t *testing.T) {
+		metadata := taskResponseWithFaultInjectionEnabled()
+		expectedTaskResponse := taskResponseWithFaultInjectionEnabled()
+		expectedTaskResponse.CredentialsID = ""            // credentials ID not expected
+		expectedTaskResponse.TaskNetworkConfig = nil       // TaskNetworkConfig is not expected and would be used internally
+		expectedTaskResponse.FaultInjectionEnabled = false // FaultInjectionEnabled is not expected and would be used internally
+
+		handler, _, agentState, _ := setup(t)
+		agentState.EXPECT().
+			GetTaskMetadata(endpointContainerID).
+			Return(*metadata, nil)
+		testTMDSRequest(t, handler, TMDSTestCase[state.TaskResponse]{
+			path:                 path,
+			expectedStatusCode:   http.StatusOK,
+			expectedResponseBody: *expectedTaskResponse,
+		})
+	})
+
 	t.Run("task lookup failure", func(t *testing.T) {
 		handler, _, agentState, _ := setup(t)
 		agentState.EXPECT().
diff --git a/ecs-agent/tmds/handlers/v4/state/response.go b/ecs-agent/tmds/handlers/v4/state/response.go
index e7226be2a6f..15baf8e21fe 100644
--- a/ecs-agent/tmds/handlers/v4/state/response.go
+++ b/ecs-agent/tmds/handlers/v4/state/response.go
@@ -38,6 +38,7 @@ type TaskResponse struct {
 	EphemeralStorageMetrics *EphemeralStorageMetrics `json:"EphemeralStorageMetrics,omitempty"`
 	CredentialsID           string                   `json:"-"`
 	TaskNetworkConfig       *TaskNetworkConfig       `json:"-"`
+	FaultInjectionEnabled   bool                     `json:"-"`
 }
 
 // TaskNetworkConfig contains required network configurations for network faults injection.