From 92ea032f0888d4ab59c36e1c87e03ced69b6cf7f Mon Sep 17 00:00:00 2001
From: Adrian Moreno <amorenoz@redhat.com>
Date: Fri, 11 Sep 2020 19:50:11 +0200
Subject: [PATCH] portforward: Fix UDP-only ports calculation

If a service has both TCP and UDP but the TCP port appears before in the
range loop, it will be considered a UDP-only port and the forwarding
will fail.

Fix that by calculating the difference between UDP ports and TCP ports.

Signed-off-by: Adrian Moreno <amorenoz@redhat.com>

Kubernetes-commit: c194bc85319298e7c2fd002031d0ff23ded54591
---
 pkg/cmd/portforward/portforward.go      | 18 +++++++-------
 pkg/cmd/portforward/portforward_test.go | 32 +++++++++++++++++++++++--
 2 files changed, 40 insertions(+), 10 deletions(-)

diff --git a/pkg/cmd/portforward/portforward.go b/pkg/cmd/portforward/portforward.go
index 8800a9e7..5755055e 100644
--- a/pkg/cmd/portforward/portforward.go
+++ b/pkg/cmd/portforward/portforward.go
@@ -260,35 +260,37 @@ func checkUDPPorts(udpOnlyPorts sets.Int, ports []string, obj metav1.Object) err
 // checkUDPPortInService returns an error if remote port in Service is a UDP port
 // TODO: remove this check after #47862 is solved
 func checkUDPPortInService(ports []string, svc *corev1.Service) error {
-	udpOnlyPorts := sets.NewInt()
+	udpPorts := sets.NewInt()
+	tcpPorts := sets.NewInt()
 	for _, port := range svc.Spec.Ports {
 		portNum := int(port.Port)
 		switch port.Protocol {
 		case corev1.ProtocolUDP:
-			udpOnlyPorts.Insert(portNum)
+			udpPorts.Insert(portNum)
 		case corev1.ProtocolTCP:
-			udpOnlyPorts.Delete(portNum)
+			tcpPorts.Insert(portNum)
 		}
 	}
-	return checkUDPPorts(udpOnlyPorts, ports, svc)
+	return checkUDPPorts(udpPorts.Difference(tcpPorts), ports, svc)
 }
 
 // checkUDPPortInPod returns an error if remote port in Pod is a UDP port
 // TODO: remove this check after #47862 is solved
 func checkUDPPortInPod(ports []string, pod *corev1.Pod) error {
-	udpOnlyPorts := sets.NewInt()
+	udpPorts := sets.NewInt()
+	tcpPorts := sets.NewInt()
 	for _, ct := range pod.Spec.Containers {
 		for _, ctPort := range ct.Ports {
 			portNum := int(ctPort.ContainerPort)
 			switch ctPort.Protocol {
 			case corev1.ProtocolUDP:
-				udpOnlyPorts.Insert(portNum)
+				udpPorts.Insert(portNum)
 			case corev1.ProtocolTCP:
-				udpOnlyPorts.Delete(portNum)
+				tcpPorts.Insert(portNum)
 			}
 		}
 	}
-	return checkUDPPorts(udpOnlyPorts, ports, pod)
+	return checkUDPPorts(udpPorts.Difference(tcpPorts), ports, pod)
 }
 
 // Complete completes all the required options for port-forward cmd.
diff --git a/pkg/cmd/portforward/portforward_test.go b/pkg/cmd/portforward/portforward_test.go
index 6a04de55..06b99081 100644
--- a/pkg/cmd/portforward/portforward_test.go
+++ b/pkg/cmd/portforward/portforward_test.go
@@ -880,7 +880,7 @@ func TestCheckUDPPort(t *testing.T) {
 			expectError: true,
 		},
 		{
-			name: "Pod has ports with both TCP and UDP protocol",
+			name: "Pod has ports with both TCP and UDP protocol (UDP first)",
 			pod: &corev1.Pod{
 				Spec: corev1.PodSpec{
 					Containers: []corev1.Container{
@@ -895,6 +895,22 @@ func TestCheckUDPPort(t *testing.T) {
 			},
 			ports: []string{":53"},
 		},
+		{
+			name: "Pod has ports with both TCP and UDP protocol (TCP first)",
+			pod: &corev1.Pod{
+				Spec: corev1.PodSpec{
+					Containers: []corev1.Container{
+						{
+							Ports: []corev1.ContainerPort{
+								{Protocol: corev1.ProtocolTCP, ContainerPort: 53},
+								{Protocol: corev1.ProtocolUDP, ContainerPort: 53},
+							},
+						},
+					},
+				},
+			},
+			ports: []string{":53"},
+		},
 
 		{
 			name: "forward to a UDP port in a Service",
@@ -921,7 +937,7 @@ func TestCheckUDPPort(t *testing.T) {
 			expectError: true,
 		},
 		{
-			name: "Service has ports with both TCP and UDP protocol",
+			name: "Service has ports with both TCP and UDP protocol (UDP first)",
 			service: &corev1.Service{
 				Spec: corev1.ServiceSpec{
 					Ports: []corev1.ServicePort{
@@ -932,6 +948,18 @@ func TestCheckUDPPort(t *testing.T) {
 			},
 			ports: []string{"53"},
 		},
+		{
+			name: "Service has ports with both TCP and UDP protocol (TCP first)",
+			service: &corev1.Service{
+				Spec: corev1.ServiceSpec{
+					Ports: []corev1.ServicePort{
+						{Protocol: corev1.ProtocolTCP, Port: 53},
+						{Protocol: corev1.ProtocolUDP, Port: 53},
+					},
+				},
+			},
+			ports: []string{"53"},
+		},
 	}
 	for _, tc := range tests {
 		var err error