From be660fbdcb7dd1fb5edf3d0c3ccb8f96b783aee3 Mon Sep 17 00:00:00 2001
From: Stanko <stankoa@amazon.com>
Date: Thu, 7 Feb 2019 10:37:43 -0800
Subject: [PATCH] Fix getExecutionRolePolicyARN() in regcreds

---
 .../create_task_execution_role_test.go        | 81 +++++++++++++++++++
 .../cli/regcreds/regcreds_app_helpers.go      |  3 +-
 2 files changed, 83 insertions(+), 1 deletion(-)

diff --git a/ecs-cli/modules/cli/regcreds/create_task_execution_role_test.go b/ecs-cli/modules/cli/regcreds/create_task_execution_role_test.go
index 3ab688c46..c0513b389 100644
--- a/ecs-cli/modules/cli/regcreds/create_task_execution_role_test.go
+++ b/ecs-cli/modules/cli/regcreds/create_task_execution_role_test.go
@@ -18,6 +18,7 @@ import (
 
 	"github.com/aws/amazon-ecs-cli/ecs-cli/modules/utils/regcredio"
 	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/arn"
 	"github.com/aws/aws-sdk-go/aws/awserr"
 	"github.com/aws/aws-sdk-go/service/iam"
 	"github.com/golang/mock/gomock"
@@ -63,6 +64,86 @@ func TestCreateTaskExecutionRole(t *testing.T) {
 	assert.NotNil(t, policyCreateTime, "Expected policy create time to be non-nil")
 }
 
+func TestCreateTaskExecutionRole_CnPartition(t *testing.T) {
+	testRegistry := "myreg.test.io"
+	testRegCredARN := "arn:aws-cn:secret/some-test-arn"
+	testCreds := map[string]regcredio.CredsOutputEntry{
+		testRegistry: regcredio.BuildOutputEntry(testRegCredARN, "", []string{""}),
+	}
+	testRoleName := "myNginxProjectRole"
+
+	testPolicyArn := aws.String("arn:aws-cn:iam::policy/" + testRoleName + "-policy")
+	testRoleArn := aws.String("arn:aws-cn:iam::role/" + testRoleName)
+
+	expectedManagedPolicyARN := arn.ARN{
+		Service:   "iam",
+		Resource:  "policy/service-role/AmazonECSTaskExecutionRolePolicy",
+		AccountID: "aws",
+		Partition: "aws-cn", // Expected CN Partition
+	}
+
+	mocks := setupTestController(t)
+	gomock.InOrder(
+		mocks.MockIAM.EXPECT().CreateOrFindRole(testRoleName, roleDescriptionString, assumeRolePolicyDocString).Return(*testRoleArn, nil),
+		mocks.MockIAM.EXPECT().CreateRole(gomock.Any()).Return(&iam.CreateRoleOutput{Role: &iam.Role{Arn: testRoleArn}}, nil),
+	)
+	gomock.InOrder(
+		mocks.MockIAM.EXPECT().CreatePolicy(gomock.Any()).Return(&iam.CreatePolicyOutput{Policy: &iam.Policy{Arn: testPolicyArn}}, nil),
+		mocks.MockIAM.EXPECT().AttachRolePolicy(expectedManagedPolicyARN.String(), testRoleName).Return(nil, nil), // FAIL?
+		mocks.MockIAM.EXPECT().AttachRolePolicy(*testPolicyArn, testRoleName).Return(nil, nil),
+	)
+
+	testParams := executionRoleParams{
+		CredEntries: testCreds,
+		RoleName:    testRoleName,
+		Region:      "cn-north-1",
+	}
+
+	policyCreateTime, err := createTaskExecutionRole(testParams, mocks.MockIAM, mocks.MockKMS)
+	assert.NoError(t, err, "Unexpected error when creating task execution role")
+	assert.NotNil(t, policyCreateTime, "Expected policy create time to be non-nil")
+}
+
+func TestCreateTaskExecutionRole_UsGovPartition(t *testing.T) {
+	testRegistry := "myreg.test.io"
+	testRegCredARN := "arn:aws-us-gov:secret/some-test-arn"
+	testCreds := map[string]regcredio.CredsOutputEntry{
+		testRegistry: regcredio.BuildOutputEntry(testRegCredARN, "", []string{""}),
+	}
+	testRoleName := "myNginxProjectRole"
+
+	testPolicyArn := aws.String("arn:aws-us-gov:iam::policy/" + testRoleName + "-policy")
+	testRoleArn := aws.String("arn:aws-us-gov:iam::role/" + testRoleName)
+
+	expectedManagedPolicyARN := arn.ARN{
+		Service:   "iam",
+		Resource:  "policy/service-role/AmazonECSTaskExecutionRolePolicy",
+		AccountID: "aws",
+		Partition: "aws-us-gov", // Expected us-gov Partition
+	}
+
+	mocks := setupTestController(t)
+	gomock.InOrder(
+		mocks.MockIAM.EXPECT().CreateOrFindRole(testRoleName, roleDescriptionString, assumeRolePolicyDocString).Return(*testRoleArn, nil),
+		mocks.MockIAM.EXPECT().CreateRole(gomock.Any()).Return(&iam.CreateRoleOutput{Role: &iam.Role{Arn: testRoleArn}}, nil),
+	)
+	gomock.InOrder(
+		mocks.MockIAM.EXPECT().CreatePolicy(gomock.Any()).Return(&iam.CreatePolicyOutput{Policy: &iam.Policy{Arn: testPolicyArn}}, nil),
+		mocks.MockIAM.EXPECT().AttachRolePolicy(expectedManagedPolicyARN.String(), testRoleName).Return(nil, nil), // FAIL?
+		mocks.MockIAM.EXPECT().AttachRolePolicy(*testPolicyArn, testRoleName).Return(nil, nil),
+	)
+
+	testParams := executionRoleParams{
+		CredEntries: testCreds,
+		RoleName:    testRoleName,
+		Region:      "us-gov-west-1",
+	}
+
+	policyCreateTime, err := createTaskExecutionRole(testParams, mocks.MockIAM, mocks.MockKMS)
+	assert.NoError(t, err, "Unexpected error when creating task execution role")
+	assert.NotNil(t, policyCreateTime, "Expected policy create time to be non-nil")
+}
+
 func TestCreateTaskExecutionRole_NoKMSKey(t *testing.T) {
 	testRegistry := "myreg.test.io"
 	testRegCredARN := "arn:aws:secret/some-test-arn"
diff --git a/ecs-cli/modules/cli/regcreds/regcreds_app_helpers.go b/ecs-cli/modules/cli/regcreds/regcreds_app_helpers.go
index 7e46bb4ad..7976c0ac2 100644
--- a/ecs-cli/modules/cli/regcreds/regcreds_app_helpers.go
+++ b/ecs-cli/modules/cli/regcreds/regcreds_app_helpers.go
@@ -48,12 +48,13 @@ func getExecutionRolePolicyARN(region string) string {
 		AccountID: "aws",
 	}
 
+	// TODO: use utils.GetPartition func once merged
 	if regionToPartition[region] != "" {
 		expectedARN.Partition = regionToPartition[region]
+		return expectedARN.String()
 	}
 
 	expectedARN.Partition = "aws"
-
 	return expectedARN.String()
 }