From 5f730436d4a2170ba408c27b33b922d7081a7eaa Mon Sep 17 00:00:00 2001
From: Guillaume Tardif <guillaume.tardif@gmail.com>
Date: Fri, 11 Dec 2020 16:32:21 +0100
Subject: [PATCH] ACI: Allow setting protocol when publishing ports

Signed-off-by: Guillaume Tardif <guillaume.tardif@gmail.com>
---
 aci/convert/container.go  |  1 +
 aci/convert/ports.go      | 17 +++++++-
 aci/convert/ports_test.go | 82 +++++++++++++++++++++++++++++++++++++++
 3 files changed, 98 insertions(+), 2 deletions(-)

diff --git a/aci/convert/container.go b/aci/convert/container.go
index c0bb4cd67..ab168e9ef 100644
--- a/aci/convert/container.go
+++ b/aci/convert/container.go
@@ -32,6 +32,7 @@ func ContainerToComposeProject(r containers.ContainerConfig) (types.Project, err
 		ports = append(ports, types.ServicePortConfig{
 			Target:    p.ContainerPort,
 			Published: p.HostPort,
+			Protocol:  p.Protocol,
 		})
 	}
 
diff --git a/aci/convert/ports.go b/aci/convert/ports.go
index 4bb3693b0..df55b5466 100644
--- a/aci/convert/ports.go
+++ b/aci/convert/ports.go
@@ -37,12 +37,25 @@ func convertPortsToAci(service serviceConfigAciHelper) ([]containerinstance.Cont
 			return nil, nil, nil, errors.New(msg)
 		}
 		portNumber := int32(portConfig.Target)
+		var groupProtocol containerinstance.ContainerGroupNetworkProtocol
+		var containerProtocol containerinstance.ContainerNetworkProtocol
+		switch portConfig.Protocol {
+		case "tcp", "":
+			groupProtocol = containerinstance.TCP
+			containerProtocol = containerinstance.ContainerNetworkProtocolTCP
+		case "udp":
+			groupProtocol = containerinstance.UDP
+			containerProtocol = containerinstance.ContainerNetworkProtocolUDP
+		default:
+			return nil, nil, nil, fmt.Errorf("unknown protocol %q in exposed port for service %q", portConfig.Protocol, service.Name)
+		}
 		containerPorts = append(containerPorts, containerinstance.ContainerPort{
-			Port: to.Int32Ptr(portNumber),
+			Port:     to.Int32Ptr(portNumber),
+			Protocol: containerProtocol,
 		})
 		groupPorts = append(groupPorts, containerinstance.Port{
 			Port:     to.Int32Ptr(portNumber),
-			Protocol: containerinstance.TCP,
+			Protocol: groupProtocol,
 		})
 	}
 	var dnsLabelName *string = nil
diff --git a/aci/convert/ports_test.go b/aci/convert/ports_test.go
index a5c1acc8a..7452f7111 100644
--- a/aci/convert/ports_test.go
+++ b/aci/convert/ports_test.go
@@ -173,3 +173,85 @@ func TestPortConvert(t *testing.T) {
 		})
 	}
 }
+
+func TestConvertTCPPortsToAci(t *testing.T) {
+	service := types.ServiceConfig{
+		Name: "myService",
+		Ports: []types.ServicePortConfig{
+			{
+				Protocol:  "",
+				Target:    80,
+				Published: 80,
+			},
+			{
+				Protocol:  "tcp",
+				Target:    90,
+				Published: 90,
+			},
+		},
+	}
+	containerPorts, groupPports, _, err := convertPortsToAci(serviceConfigAciHelper(service))
+	assert.NilError(t, err)
+	assert.DeepEqual(t, containerPorts, []containerinstance.ContainerPort{
+		{
+			Port:     to.Int32Ptr(80),
+			Protocol: containerinstance.ContainerNetworkProtocolTCP,
+		},
+		{
+			Port:     to.Int32Ptr(90),
+			Protocol: containerinstance.ContainerNetworkProtocolTCP,
+		},
+	})
+	assert.DeepEqual(t, groupPports, []containerinstance.Port{
+		{
+			Port:     to.Int32Ptr(80),
+			Protocol: containerinstance.TCP,
+		},
+		{
+			Port:     to.Int32Ptr(90),
+			Protocol: containerinstance.TCP,
+		},
+	})
+}
+
+func TestConvertUDPPortsToAci(t *testing.T) {
+	service := types.ServiceConfig{
+		Name: "myService",
+		Ports: []types.ServicePortConfig{
+			{
+				Protocol:  "udp",
+				Target:    80,
+				Published: 80,
+			},
+		},
+	}
+	containerPorts, groupPports, _, err := convertPortsToAci(serviceConfigAciHelper(service))
+	assert.NilError(t, err)
+	assert.DeepEqual(t, containerPorts, []containerinstance.ContainerPort{
+		{
+			Port:     to.Int32Ptr(80),
+			Protocol: containerinstance.ContainerNetworkProtocolUDP,
+		},
+	})
+	assert.DeepEqual(t, groupPports, []containerinstance.Port{
+		{
+			Port:     to.Int32Ptr(80),
+			Protocol: containerinstance.UDP,
+		},
+	})
+}
+
+func TestConvertErrorOnMappingPorts(t *testing.T) {
+	service := types.ServiceConfig{
+		Name: "myService",
+		Ports: []types.ServicePortConfig{
+			{
+				Protocol:  "",
+				Target:    80,
+				Published: 90,
+			},
+		},
+	}
+	_, _, _, err := convertPortsToAci(serviceConfigAciHelper(service))
+	assert.Error(t, err, "Port mapping is not supported with ACI, cannot map port 90 to 80 for container myService")
+}