From e91ddc49ed2d1e139d718fdb7b477286818633e3 Mon Sep 17 00:00:00 2001
From: mye956 <michaelye956@gmail.com>
Date: Thu, 23 May 2024 00:36:10 +0000
Subject: [PATCH] Defining cloudwatch endpoint for all regions for tasks using
 awslogs as log driver

---
 agent/engine/docker_task_engine.go      | 43 ++++++++++++-------------
 agent/engine/docker_task_engine_test.go | 18 +++++++++--
 2 files changed, 36 insertions(+), 25 deletions(-)

diff --git a/agent/engine/docker_task_engine.go b/agent/engine/docker_task_engine.go
index 13aaee0e316..4046fcce8dc 100644
--- a/agent/engine/docker_task_engine.go
+++ b/agent/engine/docker_task_engine.go
@@ -1854,32 +1854,31 @@ func (engine *DockerTaskEngine) createContainer(task *apitask.Task, container *a
 		}
 	}
 
-	// This is a short term solution only for specific regions until AWS SDK Go is upgraded to V2
+	// We're currently relying on AWS SDK Go V1 to obtain/format the correct cloudwatch endpoint.
+	// This is will need to be updated once we upgrade to AWS SDK Go V2.
 	if hostConfig.LogConfig.Type == logDriverTypeAwslogs {
-		region := engine.cfg.AWSRegion
-		if region == "eu-isoe-west-1" || region == "us-isof-south-1" || region == "us-isof-east-1" {
-			endpoint := ""
-			dnsSuffix := ""
-			partition, ok := ep.PartitionForRegion(ep.DefaultPartitions(), region)
-			if !ok {
-				logger.Warn("No partition resolved for region. Using AWS default", logger.Fields{
-					"region":           region,
-					"defaultDNSSuffix": ep.AwsPartition().DNSSuffix(),
-				})
-				dnsSuffix = ep.AwsPartition().DNSSuffix()
+		region := hostConfig.LogConfig.Config["awslogs-region"]
+		endpoint := ""
+		dnsSuffix := ""
+		partition, ok := ep.PartitionForRegion(ep.DefaultPartitions(), region)
+		if !ok {
+			logger.Warn("No partition resolved for region. Using AWS default", logger.Fields{
+				"region":           region,
+				"defaultDNSSuffix": ep.AwsPartition().DNSSuffix(),
+			})
+			dnsSuffix = ep.AwsPartition().DNSSuffix()
+		} else {
+			resolvedEndpoint, err := partition.EndpointFor("logs", region)
+			if err == nil {
+				endpoint = resolvedEndpoint.URL
 			} else {
-				resolvedEndpoint, err := partition.EndpointFor("logs", region)
-				if err == nil {
-					endpoint = resolvedEndpoint.URL
-				} else {
-					dnsSuffix = partition.DNSSuffix()
-				}
+				dnsSuffix = partition.DNSSuffix()
 			}
-			if endpoint == "" {
-				endpoint = fmt.Sprintf("https://logs.%s.%s", region, dnsSuffix)
-			}
-			hostConfig.LogConfig.Config["awslogs-endpoint"] = endpoint
 		}
+		if endpoint == "" {
+			endpoint = fmt.Sprintf("https://logs.%s.%s", region, dnsSuffix)
+		}
+		hostConfig.LogConfig.Config["awslogs-endpoint"] = endpoint
 	}
 
 	//Apply the log driver secret into container's LogConfig and Env secrets to container.Environment
diff --git a/agent/engine/docker_task_engine_test.go b/agent/engine/docker_task_engine_test.go
index faea238fe0f..d2208816382 100644
--- a/agent/engine/docker_task_engine_test.go
+++ b/agent/engine/docker_task_engine_test.go
@@ -2954,7 +2954,17 @@ func TestCreateContainerAwslogsLogDriver(t *testing.T) {
 		{
 			name:                      "test container that uses awslogs log driver in IAD",
 			region:                    "us-east-1",
-			expectedLogConfigEndpoint: "",
+			expectedLogConfigEndpoint: "https://logs.us-east-1.amazonaws.com",
+		},
+		{
+			name:                      "test container that uses awslogs log driver in BJS",
+			region:                    "cn-north-1",
+			expectedLogConfigEndpoint: "https://logs.cn-north-1.amazonaws.com.cn",
+		},
+		{
+			name:                      "test container that uses awslogs log driver in OSU",
+			region:                    "us-gov-east-1",
+			expectedLogConfigEndpoint: "https://logs.us-gov-east-1.amazonaws.com",
 		},
 		{
 			name:                      "test container that uses awslogs log driver in NCL",
@@ -2983,8 +2993,10 @@ func TestCreateContainerAwslogsLogDriver(t *testing.T) {
 
 			rawHostConfigInput := dockercontainer.HostConfig{
 				LogConfig: dockercontainer.LogConfig{
-					Type:   "awslogs",
-					Config: map[string]string{},
+					Type: "awslogs",
+					Config: map[string]string{
+						"awslogs-region": tc.region,
+					},
 				},
 			}
 			rawHostConfig, err := json.Marshal(&rawHostConfigInput)