From 541dd963ba05efab3521f0ad20d36d9f6a1b96f3 Mon Sep 17 00:00:00 2001
From: Mythri Garaga Manjunatha <mythr@amazon.com>
Date: Tue, 21 Dec 2021 11:44:40 -0800
Subject: [PATCH] Add Host IP to port response for v4 container response

---
 agent/handlers/v1/response.go      |  1 +
 agent/handlers/v2/response.go      |  3 ++
 agent/handlers/v2/response_test.go | 44 ++++++++++++++++++++++--------
 3 files changed, 37 insertions(+), 11 deletions(-)

diff --git a/agent/handlers/v1/response.go b/agent/handlers/v1/response.go
index bba57fcc1d9..069cf400159 100644
--- a/agent/handlers/v1/response.go
+++ b/agent/handlers/v1/response.go
@@ -67,6 +67,7 @@ type PortResponse struct {
 	ContainerPort uint16 `json:"ContainerPort,omitempty"`
 	Protocol      string `json:"Protocol,omitempty"`
 	HostPort      uint16 `json:"HostPort,omitempty"`
+	HostIp        string `json:"HostIp,omitempty"`
 }
 
 // NewTaskResponse creates a TaskResponse for a task.
diff --git a/agent/handlers/v2/response.go b/agent/handlers/v2/response.go
index 812031cd0d2..d97c77d88f3 100644
--- a/agent/handlers/v2/response.go
+++ b/agent/handlers/v2/response.go
@@ -281,6 +281,9 @@ func NewContainerResponse(
 		} else {
 			port.HostPort = port.ContainerPort
 		}
+		if includeV4Metadata {
+			port.HostIp = binding.BindIP
+		}
 
 		resp.Ports = append(resp.Ports, port)
 	}
diff --git a/agent/handlers/v2/response_test.go b/agent/handlers/v2/response_test.go
index 718dcd74e5c..798e51356f5 100644
--- a/agent/handlers/v2/response_test.go
+++ b/agent/handlers/v2/response_test.go
@@ -54,6 +54,7 @@ const (
 	volDestination       = "/volume"
 	availabilityZone     = "us-west-2b"
 	containerInstanceArn = "containerInstance-test"
+	hostIp               = "0.0.0.0"
 )
 
 func TestTaskResponse(t *testing.T) {
@@ -462,6 +463,19 @@ func TestTaskResponseMarshal(t *testing.T) {
 }
 
 func TestContainerResponseMarshal(t *testing.T) {
+	testCases := []struct {
+		description       string
+		includeV4Metadata bool
+	}{
+		{
+			"task container response without v4 metadata",
+			false,
+		},
+		{
+			"task container response with v4 metadata",
+			true,
+		},
+	}
 	timeRFC3339, _ := time.Parse(time.RFC3339, "2014-11-12T11:45:26Z")
 
 	expectedContainerResponseMap := map[string]interface{}{
@@ -549,20 +563,28 @@ func TestContainerResponseMarshal(t *testing.T) {
 			},
 		},
 	}
-	gomock.InOrder(
-		state.EXPECT().ContainerByID(containerID).Return(dockerContainer, true),
-		state.EXPECT().TaskByID(containerID).Return(task, true),
-	)
 
-	containerResponse, err := NewContainerResponseFromState(containerID, state, false)
-	assert.NoError(t, err)
+	for _, tc := range testCases {
+		t.Run(tc.description, func(t *testing.T) {
+			gomock.InOrder(
+				state.EXPECT().ContainerByID(containerID).Return(dockerContainer, true),
+				state.EXPECT().TaskByID(containerID).Return(task, true),
+			)
+			if tc.includeV4Metadata {
+				container.KnownPortBindingsUnsafe[0].BindIP = hostIp
+				expectedContainerResponseMap["Ports"].([]interface{})[0].(map[string]interface{})["HostIp"] = hostIp
+			}
+			containerResponse, err := NewContainerResponseFromState(containerID, state, tc.includeV4Metadata)
+			assert.NoError(t, err)
 
-	containerResponseJSON, err := json.Marshal(containerResponse)
-	assert.NoError(t, err)
+			containerResponseJSON, err := json.Marshal(containerResponse)
+			assert.NoError(t, err)
 
-	containerResponseMap := make(map[string]interface{})
-	json.Unmarshal(containerResponseJSON, &containerResponseMap)
-	assert.Equal(t, expectedContainerResponseMap, containerResponseMap)
+			containerResponseMap := make(map[string]interface{})
+			json.Unmarshal(containerResponseJSON, &containerResponseMap)
+			assert.Equal(t, expectedContainerResponseMap, containerResponseMap)
+		})
+	}
 }
 
 func TestTaskResponseWithV4TagsError(t *testing.T) {