adding FIS register handlers

agent/handlers/task_server_setup.go

@@ -31,6 +31,8 @@ import (
 	auditinterface "github.com/aws/amazon-ecs-agent/ecs-agent/logger/audit"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/metrics"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/tmds"
+	fault "github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/fault/v1/handlers"
+	faulttype "github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/fault/v1/types"
 	tp "github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/taskprotection/v1/handlers"
 	tmdsv1 "github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/v1"
 	tmdsv2 "github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/v2"
@@ -90,6 +92,9 @@ func taskServerSetup(
 	agentAPIV1HandlersSetup(muxRouter, state, credentialsManager, cluster, tmdsAgentState,
 		taskProtectionClientFactory, metricsFactory)
 
+	// TODO: Future PR to pass in TMDS server router once all of the handlers have been implemented.
+	registerFaultHandlers(nil, tmdsAgentState, metricsFactory)
+
 	return tmds.NewServer(auditLogger,
 		tmds.WithHandler(muxRouter),
 		tmds.WithListenAddress(tmds.AddressIPv4()),
@@ -184,6 +189,67 @@ func agentAPIV1HandlersSetup(
 		Methods("GET")
 }
 
+// registerFaultHandlers adds handlers for fault endpoints
+// TODO: Will need to be called in taskServerSetup once all of the handlers have been implemented
+func registerFaultHandlers(
+	muxRouter *mux.Router,
+	agentState *v4.TMDSAgentState,
+	metricsFactory metrics.EntryFactory,
+) {
+	handler := fault.FaultHandler{
+		AgentState:     agentState,
+		MetricsFactory: metricsFactory,
+	}
+
+	if muxRouter == nil {
+		return
+	}
+
+	// Setting up handler endpoints for network blackhole port fault injections
+	muxRouter.HandleFunc(
+		fault.FaultNetworkFaultPath(faulttype.BlackHolePortFaultType),
+		handler.StartNetworkBlackholePort(),
+	).Methods("PUT")
+	muxRouter.HandleFunc(
+		fault.FaultNetworkFaultPath(faulttype.BlackHolePortFaultType),
+		handler.StopBlackHolePort(),
+	).Methods("DELETE")
+	muxRouter.HandleFunc(
+		fault.FaultNetworkFaultPath(faulttype.BlackHolePortFaultType),
+		handler.CheckBlackHolePortStatus(),
+	).Methods("GET")
+
+	// Setting up handler endpoints for network latency fault injections
+	muxRouter.HandleFunc(
+		fault.FaultNetworkFaultPath(faulttype.LatencyFaultType),
+		handler.StartLatency(),
+	).Methods("PUT")
+	muxRouter.HandleFunc(
+		fault.FaultNetworkFaultPath(faulttype.LatencyFaultType),
+		handler.StopLatency(),
+	).Methods("DELETE")
+	muxRouter.HandleFunc(
+		fault.FaultNetworkFaultPath(faulttype.LatencyFaultType),
+		handler.CheckLatencyStatus(),
+	).Methods("GET")
+
+	// Setting up handler endpoints for network packet loss fault injections
+	muxRouter.HandleFunc(
+		fault.FaultNetworkFaultPath(faulttype.PacketLossFaultType),
+		handler.StartPacketLoss(),
+	).Methods("PUT")
+	muxRouter.HandleFunc(
+		fault.FaultNetworkFaultPath(faulttype.PacketLossFaultType),
+		handler.StopPacketLoss(),
+	).Methods("DELETE")
+	muxRouter.HandleFunc(
+		fault.FaultNetworkFaultPath(faulttype.PacketLossFaultType),
+		handler.CheckPacketLossStatus(),
+	).Methods("GET")
+
+	seelog.Debug("Successfully set up Fault TMDS handlers")
+}
+
 // ServeTaskHTTPEndpoint serves task/container metadata, task/container stats, IAM Role Credentials, and Agent APIs
 // for tasks being managed by the agent.
 func ServeTaskHTTPEndpoint(

agent/handlers/task_server_setup_test.go

@@ -34,6 +34,7 @@ import (
 	"github.com/aws/amazon-ecs-agent/agent/config"
 	mock_dockerstate "github.com/aws/amazon-ecs-agent/agent/engine/dockerstate/mocks"
 	v3 "github.com/aws/amazon-ecs-agent/agent/handlers/v3"
+	agentV4 "github.com/aws/amazon-ecs-agent/agent/handlers/v4"
 	mock_stats "github.com/aws/amazon-ecs-agent/agent/stats/mock"
 	apicontainerstatus "github.com/aws/amazon-ecs-agent/ecs-agent/api/container/status"
 	mock_ecs "github.com/aws/amazon-ecs-agent/ecs-agent/api/ecs/mocks"
@@ -42,15 +43,18 @@ import (
 	"github.com/aws/amazon-ecs-agent/ecs-agent/credentials"
 	mock_credentials "github.com/aws/amazon-ecs-agent/ecs-agent/credentials/mocks"
 	mock_audit "github.com/aws/amazon-ecs-agent/ecs-agent/logger/audit/mocks"
+	mock_metrics "github.com/aws/amazon-ecs-agent/ecs-agent/metrics/mocks"
 	ni "github.com/aws/amazon-ecs-agent/ecs-agent/netlib/model/networkinterface"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/stats"
+	faulttype "github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/fault/v1/types"
 	tmdsresponse "github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/response"
 	tp "github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/taskprotection/v1/handlers"
 	tptypes "github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/taskprotection/v1/types"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/utils"
 	tmdsv1 "github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/v1"
 	v2 "github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/v2"
 	v4 "github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/v4/state"
+	"github.com/gorilla/mux"
 
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/awserr"
@@ -107,6 +111,7 @@ const (
 	privateDNSName             = "ip-172-31-47-69.us-west-2.compute.internal"
 	subnetGatewayIpv4Address   = "172.31.32.1/20"
 	taskCredentialsID          = "taskCredentialsId"
+	endpointId                 = "endpointId"
 )
 
 var (
@@ -3561,3 +3566,114 @@ func TestUpdateTaskProtection(t *testing.T) {
 		},
 	}))
 }
+
+func TestRegisterHandler(t *testing.T) {
+	tcs := []struct {
+		name               string
+		expectedStatusCode int
+		fault              string
+		method             string
+	}{
+		{
+			name:               "PUT blackholeport",
+			expectedStatusCode: http.StatusOK,
+			fault:              faulttype.BlackHolePortFaultType,
+			method:             "PUT",
+		},
+		{
+			name:               "DELETE blackholeport",
+			expectedStatusCode: http.StatusOK,
+			fault:              faulttype.BlackHolePortFaultType,
+			method:             "DELETE",
+		},
+		{
+			name:               "GET blackholeport",
+			expectedStatusCode: http.StatusOK,
+			fault:              faulttype.BlackHolePortFaultType,
+			method:             "GET",
+		},
+		{
+			name:               "PUT latency",
+			expectedStatusCode: http.StatusOK,
+			fault:              faulttype.LatencyFaultType,
+			method:             "PUT",
+		},
+		{
+			name:               "DELETE latency",
+			expectedStatusCode: http.StatusOK,
+			fault:              faulttype.LatencyFaultType,
+			method:             "DELETE",
+		},
+		{
+			name:               "GET latency",
+			expectedStatusCode: http.StatusOK,
+			fault:              faulttype.BlackHolePortFaultType,
+			method:             "GET",
+		},
+		{
+			name:               "PUT packet loss",
+			expectedStatusCode: http.StatusOK,
+			fault:              faulttype.PacketLossFaultType,
+			method:             "PUT",
+		},
+		{
+			name:               "DELETE packet loss",
+			expectedStatusCode: http.StatusOK,
+			fault:              faulttype.PacketLossFaultType,
+			method:             "DELETE",
+		},
+		{
+			name:               "GET packet loss",
+			expectedStatusCode: http.StatusOK,
+			fault:              faulttype.PacketLossFaultType,
+			method:             "GET",
+		},
+		{
+			name:               "PUT unknown",
+			expectedStatusCode: http.StatusNotFound,
+			fault:              "unknown",
+			method:             "PUT",
+		},
+		{
+			name:               "DELETE unknown",
+			expectedStatusCode: http.StatusNotFound,
+			fault:              "unknown",
+			method:             "DELETE",
+		},
+		{
+			name:               "GET unknown",
+			expectedStatusCode: http.StatusNotFound,
+			fault:              "unknown",
+			method:             "GET",
+		},
+	}
+
+	for _, tc := range tcs {
+		t.Run(tc.name, func(t *testing.T) {
+			// Mocks
+			ctrl := gomock.NewController(t)
+			defer ctrl.Finish()
+
+			state := mock_dockerstate.NewMockTaskEngineState(ctrl)
+			statsEngine := mock_stats.NewMockEngine(ctrl)
+			ecsClient := mock_ecs.NewMockECSClient(ctrl)
+
+			agentState := agentV4.NewTMDSAgentState(state, statsEngine, ecsClient, clusterName, availabilityzone, vpcID, containerInstanceArn)
+			metricsFactory := mock_metrics.NewMockEntryFactory(ctrl)
+
+			router := mux.NewRouter()
+
+			registerFaultHandlers(router, agentState, metricsFactory)
+			var requestBody io.Reader
+			req, err := http.NewRequest(tc.method, fmt.Sprintf("/api/%s/fault/v1/%s", endpointId, tc.fault),
+				requestBody)
+			require.NoError(t, err)
+
+			// Send the request and record the response
+			recorder := httptest.NewRecorder()
+			router.ServeHTTP(recorder, req)
+
+			assert.Equal(t, recorder.Code, tc.expectedStatusCode)
+		})
+	}
+}