Move TMDS task protection types to ecs-agent

agent/handlers/agentapi/taskprotection/factory.go

@@ -0,0 +1,50 @@
+// Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"). You may
+// not use this file except in compliance with the License. A copy of the
+// License is located at
+//
+//	http://aws.amazon.com/apache2.0/
+//
+// or in the "license" file accompanying this file. This file is distributed
+// on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+// express or implied. See the License for the specific language governing
+// permissions and limitations under the License.
+package taskprotection
+
+import (
+	"github.com/aws/amazon-ecs-agent/agent/api/ecsclient"
+	"github.com/aws/amazon-ecs-agent/agent/httpclient"
+
+	"github.com/aws/amazon-ecs-agent/ecs-agent/api"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/credentials"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/ecs_client/model/ecs"
+
+	"github.com/aws/aws-sdk-go/aws"
+	awscreds "github.com/aws/aws-sdk-go/aws/credentials"
+	"github.com/aws/aws-sdk-go/aws/session"
+)
+
+// TaskProtectionClientFactory implements TaskProtectionClientFactoryInterface
+type TaskProtectionClientFactory struct {
+	Region             string
+	Endpoint           string
+	AcceptInsecureCert bool
+}
+
+// Helper function for retrieving credential from credentials manager and create ecs client
+func (factory TaskProtectionClientFactory) NewTaskProtectionClient(
+	taskRoleCredential credentials.TaskIAMRoleCredentials,
+) api.ECSTaskProtectionSDK {
+	taskCredential := taskRoleCredential.GetIAMRoleCredentials()
+	cfg := aws.NewConfig().
+		WithCredentials(awscreds.NewStaticCredentials(taskCredential.AccessKeyID,
+			taskCredential.SecretAccessKey,
+			taskCredential.SessionToken)).
+		WithRegion(factory.Region).
+		WithHTTPClient(httpclient.New(ecsclient.RoundtripTimeout, factory.AcceptInsecureCert)).
+		WithEndpoint(factory.Endpoint)
+
+	ecsClient := ecs.New(session.Must(session.NewSession()), cfg)
+	return ecsClient
+}

agent/handlers/agentapi/taskprotection/factory_test.go

@@ -0,0 +1,56 @@
+// Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"). You may
+// not use this file except in compliance with the License. A copy of the
+// License is located at
+//
+//	http://aws.amazon.com/apache2.0/
+//
+// or in the "license" file accompanying this file. This file is distributed
+// on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+// express or implied. See the License for the specific language governing
+// permissions and limitations under the License.
+package taskprotection
+
+import (
+	"testing"
+
+	"github.com/aws/amazon-ecs-agent/ecs-agent/api"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/credentials"
+	"github.com/golang/mock/gomock"
+	"github.com/stretchr/testify/assert"
+)
+
+const (
+	testAccessKey          = "accessKey"
+	testSecretKey          = "secretKey"
+	testSessionToken       = "sessionToken"
+	testRegion             = "region"
+	testECSEndpoint        = "endpoint"
+	testAcceptInsecureCert = false
+)
+
+// TestGetECSClientHappyCase tests newTaskProtectionClient uses credential in credentials manager and
+// returns an ECS client with correct status code and error
+func TestGetECSClientHappyCase(t *testing.T) {
+	testIAMRoleCredentials := credentials.TaskIAMRoleCredentials{
+		IAMRoleCredentials: credentials.IAMRoleCredentials{
+			AccessKeyID:     testAccessKey,
+			SecretAccessKey: testSecretKey,
+			SessionToken:    testSessionToken,
+		},
+	}
+
+	factory := TaskProtectionClientFactory{
+		Region: testRegion, Endpoint: testECSEndpoint, AcceptInsecureCert: testAcceptInsecureCert,
+	}
+
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	ret := factory.NewTaskProtectionClient(testIAMRoleCredentials)
+	_, ok := ret.(api.ECSTaskProtectionSDK)
+
+	// Assert response
+	assert.True(t, ok)
+}

agent/handlers/agentapi/taskprotection/v1/handlers/handlers.go

@@ -21,23 +21,19 @@ import (
 	"net/http"
 	"time"
 
-	"github.com/aws/amazon-ecs-agent/agent/api/ecsclient"
 	apitask "github.com/aws/amazon-ecs-agent/agent/api/task"
 	"github.com/aws/amazon-ecs-agent/agent/engine/dockerstate"
-	"github.com/aws/amazon-ecs-agent/agent/handlers/agentapi/taskprotection/v1/types"
 	v3 "github.com/aws/amazon-ecs-agent/agent/handlers/v3"
-	"github.com/aws/amazon-ecs-agent/agent/httpclient"
-	"github.com/aws/amazon-ecs-agent/ecs-agent/api"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/credentials"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/ecs_client/model/ecs"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/logger"
 	loggerfield "github.com/aws/amazon-ecs-agent/ecs-agent/logger/field"
+	tpinterface "github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/taskprotection/v1/handlers"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/taskprotection/v1/types"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/utils"
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/awserr"
-	awscreds "github.com/aws/aws-sdk-go/aws/credentials"
 	"github.com/aws/aws-sdk-go/aws/request"
-	"github.com/aws/aws-sdk-go/aws/session"
 )
 
 const (
@@ -63,17 +59,10 @@ type TaskProtectionRequest struct {
 	ExpiresInMinutes  *int64
 }
 
-// TaskProtectionClientFactory implements TaskProtectionClientFactoryInterface
-type TaskProtectionClientFactory struct {
-	Region             string
-	Endpoint           string
-	AcceptInsecureCert bool
-}
-
 // UpdateTaskProtectionHandler returns an HTTP request handler function for
 // UpdateTaskProtection API
 func UpdateTaskProtectionHandler(state dockerstate.TaskEngineState, credentialsManager credentials.Manager,
-	factory TaskProtectionClientFactoryInterface, cluster string) func(http.ResponseWriter, *http.Request) {
+	factory tpinterface.TaskProtectionClientFactoryInterface, cluster string) func(http.ResponseWriter, *http.Request) {
 	return func(w http.ResponseWriter, r *http.Request) {
 		updateTaskProtectionRequestType := "api/UpdateTaskProtection/v1"
 
@@ -193,7 +182,7 @@ func UpdateTaskProtectionHandler(state dockerstate.TaskEngineState, credentialsM
 
 // GetTaskProtectionHandler returns a handler function for GetTaskProtection API
 func GetTaskProtectionHandler(state dockerstate.TaskEngineState, credentialsManager credentials.Manager,
-	factory TaskProtectionClientFactoryInterface, cluster string) func(http.ResponseWriter, *http.Request) {
+	factory tpinterface.TaskProtectionClientFactoryInterface, cluster string) func(http.ResponseWriter, *http.Request) {
 	return func(w http.ResponseWriter, r *http.Request) {
 		getTaskProtectionRequestType := "api/GetTaskProtection/v1"
 
@@ -286,21 +275,6 @@ func GetTaskProtectionHandler(state dockerstate.TaskEngineState, credentialsMana
 	}
 }
 
-// Helper function for retrieving credential from credentials manager and create ecs client
-func (factory TaskProtectionClientFactory) NewTaskProtectionClient(taskRoleCredential credentials.TaskIAMRoleCredentials) api.ECSTaskProtectionSDK {
-	taskCredential := taskRoleCredential.GetIAMRoleCredentials()
-	cfg := aws.NewConfig().
-		WithCredentials(awscreds.NewStaticCredentials(taskCredential.AccessKeyID,
-			taskCredential.SecretAccessKey,
-			taskCredential.SessionToken)).
-		WithRegion(factory.Region).
-		WithHTTPClient(httpclient.New(ecsclient.RoundtripTimeout, factory.AcceptInsecureCert)).
-		WithEndpoint(factory.Endpoint)
-
-	ecsClient := ecs.New(session.Must(session.NewSession()), cfg)
-	return ecsClient
-}
-
 // Helper function to parse error to get ErrorCode, ExceptionMessage, HttpStatusCode, RequestID.
 // RequestID will be empty if the request is not able to reach AWS
 func getErrorCodeAndStatusCode(err error) (string, string, int, *string) {

agent/handlers/agentapi/taskprotection/v1/handlers/handlers_test.go

@@ -26,14 +26,15 @@ import (
 	"github.com/aws/amazon-ecs-agent/agent/api/task"
 	"github.com/aws/amazon-ecs-agent/agent/engine/dockerstate"
 	mock_dockerstate "github.com/aws/amazon-ecs-agent/agent/engine/dockerstate/mocks"
-	"github.com/aws/amazon-ecs-agent/agent/handlers/agentapi/taskprotection/v1/types"
+	tpfactory "github.com/aws/amazon-ecs-agent/agent/handlers/agentapi/taskprotection"
 	v3 "github.com/aws/amazon-ecs-agent/agent/handlers/v3"
 	"github.com/aws/amazon-ecs-agent/agent/utils"
-	"github.com/aws/amazon-ecs-agent/ecs-agent/api"
 	mock_api "github.com/aws/amazon-ecs-agent/ecs-agent/api/mocks"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/credentials"
 	mock_credentials "github.com/aws/amazon-ecs-agent/ecs-agent/credentials/mocks"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/ecs_client/model/ecs"
+	tpinterface "github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/taskprotection/v1/handlers"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/taskprotection/v1/types"
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/awserr"
 	"github.com/aws/aws-sdk-go/aws/request"
@@ -67,32 +68,6 @@ func TestTaskProtectionPath(t *testing.T) {
 	assert.Equal(t, "/api/{v3EndpointIDMuxName:[^/]*}/task-protection/v1/state", TaskProtectionPath())
 }
 
-// TestGetECSClientHappyCase tests newTaskProtectionClient uses credential in credentials manager and
-// returns an ECS client with correct status code and error
-func TestGetECSClientHappyCase(t *testing.T) {
-
-	testIAMRoleCredentials := credentials.TaskIAMRoleCredentials{
-		IAMRoleCredentials: credentials.IAMRoleCredentials{
-			AccessKeyID:     testAccessKey,
-			SecretAccessKey: testSecretKey,
-			SessionToken:    testSessionToken,
-		},
-	}
-
-	factory := TaskProtectionClientFactory{
-		Region: testRegion, Endpoint: testECSEndpoint, AcceptInsecureCert: testAcceptInsecureCert,
-	}
-
-	ctrl := gomock.NewController(t)
-	defer ctrl.Finish()
-
-	ret := factory.NewTaskProtectionClient(testIAMRoleCredentials)
-	_, ok := ret.(api.ECSTaskProtectionSDK)
-
-	// Assert response
-	assert.True(t, ok)
-}
-
 func getRequestWithUnknownFields(t *testing.T) map[string]interface{} {
 	request := TaskProtectionRequest{ProtectionEnabled: utils.BoolPtr(false)}
 	requestJSON, err := json.Marshal(request)
@@ -107,7 +82,8 @@ func getRequestWithUnknownFields(t *testing.T) map[string]interface{} {
 
 // Helper function for running tests for UpdateTaskProtection handler
 func testUpdateTaskProtectionHandler(t *testing.T, state dockerstate.TaskEngineState,
-	v3EndpointID string, credentialsManager credentials.Manager, factory TaskProtectionClientFactoryInterface,
+	v3EndpointID string, credentialsManager credentials.Manager,
+	factory tpinterface.TaskProtectionClientFactoryInterface,
 	request interface{}, expectedResponse interface{}, expectedResponseCode int) {
 	// Prepare request
 	requestBytes, err := json.Marshal(request)
@@ -249,7 +225,7 @@ func TestUpdateTaskProtectionHandlerTaskRoleCredentialsNotFound(t *testing.T) {
 	}
 	testTask.SetCredentialsID(testTaskCredentialsId)
 
-	factory := TaskProtectionClientFactory{
+	factory := tpfactory.TaskProtectionClientFactory{
 		Region: testRegion, Endpoint: testECSEndpoint, AcceptInsecureCert: testAcceptInsecureCert,
 	}
 
@@ -399,7 +375,7 @@ func TestUpdateTaskProtectionHandler_PostCall(t *testing.T) {
 
 			mockState := mock_dockerstate.NewMockTaskEngineState(ctrl)
 			mockManager := mock_credentials.NewMockManager(ctrl)
-			mockFactory := NewMockTaskProtectionClientFactoryInterface(ctrl)
+			mockFactory := tpinterface.NewMockTaskProtectionClientFactoryInterface(ctrl)
 			mockECSClient := mock_api.NewMockECSTaskProtectionSDK(ctrl)
 
 			mockState.EXPECT().TaskARNByV3EndpointID(gomock.Eq(testV3EndpointId)).Return(testTaskArn, true)
@@ -423,7 +399,9 @@ func TestUpdateTaskProtectionHandler_PostCall(t *testing.T) {
 }
 
 func testGetTaskProtectionHandler(t *testing.T, state dockerstate.TaskEngineState,
-	v3EndpointID string, credentialsManager credentials.Manager, factory TaskProtectionClientFactoryInterface, expectedResponse interface{}, expectedResponseCode int) {
+	v3EndpointID string, credentialsManager credentials.Manager,
+	factory tpinterface.TaskProtectionClientFactoryInterface,
+	expectedResponse interface{}, expectedResponseCode int) {
 	// Prepare request
 	bodyReader := bytes.NewReader([]byte{})
 	req, err := http.NewRequest("GET", "", bodyReader)
@@ -492,7 +470,7 @@ func TestGetTaskProtectionHandlerTaskRoleCredentialsNotFound(t *testing.T) {
 	}
 	testTask.SetCredentialsID(testTaskCredentialsId)
 
-	factory := TaskProtectionClientFactory{
+	factory := tpfactory.TaskProtectionClientFactory{
 		Region: testRegion, Endpoint: testECSEndpoint, AcceptInsecureCert: testAcceptInsecureCert,
 	}
 
@@ -638,7 +616,7 @@ func TestGetTaskProtectionHandler_PostCall(t *testing.T) {
 
 			mockState := mock_dockerstate.NewMockTaskEngineState(ctrl)
 			mockManager := mock_credentials.NewMockManager(ctrl)
-			mockFactory := NewMockTaskProtectionClientFactoryInterface(ctrl)
+			mockFactory := tpinterface.NewMockTaskProtectionClientFactoryInterface(ctrl)
 			mockECSClient := mock_api.NewMockECSTaskProtectionSDK(ctrl)
 
 			mockState.EXPECT().TaskARNByV3EndpointID(gomock.Eq(testV3EndpointId)).Return(testTaskArn, true)

agent/handlers/task_server_setup.go

@@ -21,7 +21,8 @@ import (
 	"github.com/aws/amazon-ecs-agent/agent/api"
 	"github.com/aws/amazon-ecs-agent/agent/config"
 	"github.com/aws/amazon-ecs-agent/agent/engine/dockerstate"
-	agentAPITaskProtectionV1 "github.com/aws/amazon-ecs-agent/agent/handlers/agentapi/taskprotection/v1/handlers"
+	tpfactory "github.com/aws/amazon-ecs-agent/agent/handlers/agentapi/taskprotection"
+	tphandlers "github.com/aws/amazon-ecs-agent/agent/handlers/agentapi/taskprotection/v1/handlers"
 	v2 "github.com/aws/amazon-ecs-agent/agent/handlers/v2"
 	v3 "github.com/aws/amazon-ecs-agent/agent/handlers/v3"
 	v4 "github.com/aws/amazon-ecs-agent/agent/handlers/v4"
@@ -31,6 +32,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"
+	tpinterface "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"
 	tmdsv4 "github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/v4"
@@ -60,7 +63,7 @@ func taskServerSetup(credentialsManager credentials.Manager,
 	availabilityZone string,
 	vpcID string,
 	containerInstanceArn string,
-	taskProtectionClientFactory agentAPITaskProtectionV1.TaskProtectionClientFactoryInterface,
+	taskProtectionClientFactory tpinterface.TaskProtectionClientFactoryInterface,
 ) (*http.Server, error) {
 
 	muxRouter := mux.NewRouter()
@@ -156,17 +159,17 @@ func agentAPIV1HandlersSetup(
 	state dockerstate.TaskEngineState,
 	credentialsManager credentials.Manager,
 	cluster string,
-	factory agentAPITaskProtectionV1.TaskProtectionClientFactoryInterface,
+	factory tpinterface.TaskProtectionClientFactoryInterface,
 ) {
 	muxRouter.
 		HandleFunc(
-			agentAPITaskProtectionV1.TaskProtectionPath(),
-			agentAPITaskProtectionV1.UpdateTaskProtectionHandler(state, credentialsManager, factory, cluster)).
+			tphandlers.TaskProtectionPath(),
+			tphandlers.UpdateTaskProtectionHandler(state, credentialsManager, factory, cluster)).
 		Methods("PUT")
 	muxRouter.
 		HandleFunc(
-			agentAPITaskProtectionV1.TaskProtectionPath(),
-			agentAPITaskProtectionV1.GetTaskProtectionHandler(state, credentialsManager, factory, cluster)).
+			tphandlers.TaskProtectionPath(),
+			tphandlers.GetTaskProtectionHandler(state, credentialsManager, factory, cluster)).
 		Methods("GET")
 }
 
@@ -192,7 +195,7 @@ func ServeTaskHTTPEndpoint(
 
 	auditLogger := audit.NewAuditLog(containerInstanceArn, cfg, logger)
 
-	taskProtectionClientFactory := agentAPITaskProtectionV1.TaskProtectionClientFactory{
+	taskProtectionClientFactory := tpfactory.TaskProtectionClientFactory{
 		Region: cfg.AWSRegion, Endpoint: cfg.APIEndpoint, AcceptInsecureCert: cfg.AcceptInsecureCert,
 	}
 	server, err := taskServerSetup(credentialsManager, auditLogger, state, ecsClient, cfg.Cluster,