From a00b2a5c2120e9bcfb13beb4edf36837fa3a689b Mon Sep 17 00:00:00 2001
From: Fan Shang Xiang <fanshangxiang@gmail.com>
Date: Thu, 27 Jul 2023 16:08:33 +0800
Subject: [PATCH] Update azure_loadbalancer_healthprobe.go

---
 .../azure_loadbalancer_healthprobe.go         | 32 +++++++++----------
 pkg/provider/azure_loadbalancer_test.go       | 26 +++++++++++++++
 2 files changed, 42 insertions(+), 16 deletions(-)

diff --git a/pkg/provider/azure_loadbalancer_healthprobe.go b/pkg/provider/azure_loadbalancer_healthprobe.go
index 6d3eb81466..8664f907a7 100644
--- a/pkg/provider/azure_loadbalancer_healthprobe.go
+++ b/pkg/provider/azure_loadbalancer_healthprobe.go
@@ -99,29 +99,22 @@ func (az *Cloud) buildHealthProbeRulesForPort(serviceManifest *v1.Service, port
 		if s == nil {
 			return nil
 		}
+		//not a integer
+		for _, item := range serviceManifest.Spec.Ports {
+			if strings.EqualFold(item.Name, *s) {
+				//found the port
+				return nil
+			}
+		}
 		//nolint:gosec
 		port, err := strconv.Atoi(*s)
 		if err != nil {
-			//not a integer
-			for _, item := range serviceManifest.Spec.Ports {
-				if strings.EqualFold(item.Name, *s) {
-					//found the port
-					return nil
-				}
-			}
 			return fmt.Errorf("port %s not found in service", *s)
 		}
 		if port < 0 || port > 65535 {
 			return fmt.Errorf("port %d is out of range", port)
 		}
-		for _, item := range serviceManifest.Spec.Ports {
-			//nolint:gosec
-			if item.Port == int32(port) {
-				//found the port
-				return nil
-			}
-		}
-		return fmt.Errorf("port %s not found in service", *s)
+		return nil
 	})
 	if err != nil {
 		return nil, fmt.Errorf("failed to parse annotation %s: %w", consts.BuildHealthProbeAnnotationKeyForPort(port.Port, consts.HealthProbeParamsPort), err)
@@ -129,7 +122,7 @@ func (az *Cloud) buildHealthProbeRulesForPort(serviceManifest *v1.Service, port
 
 	if probePort != nil {
 		//nolint:gosec
-		port, err := strconv.Atoi(*probePort)
+		port, err := strconv.ParseInt(*probePort, 10, 32)
 		if err != nil {
 			//not a integer
 			for _, item := range serviceManifest.Spec.Ports {
@@ -140,13 +133,20 @@ func (az *Cloud) buildHealthProbeRulesForPort(serviceManifest *v1.Service, port
 			}
 		} else {
 			// Not need to verify probePort is in correct range again.
+			var found bool
 			for _, item := range serviceManifest.Spec.Ports {
 				//nolint:gosec
 				if item.Port == int32(port) {
 					//found the port
 					properties.Port = pointer.Int32(item.NodePort)
+					found = true
+					break
 				}
 			}
+			if !found {
+				//nolint:gosec
+				properties.Port = pointer.Int32(int32(port))
+			}
 		}
 	}
 
diff --git a/pkg/provider/azure_loadbalancer_test.go b/pkg/provider/azure_loadbalancer_test.go
index 6a84f7bc98..4c12f1d2ab 100644
--- a/pkg/provider/azure_loadbalancer_test.go
+++ b/pkg/provider/azure_loadbalancer_test.go
@@ -2746,6 +2746,32 @@ func TestReconcileLoadBalancerRule(t *testing.T) {
 				},
 			},
 		},
+		{
+			desc: "getExpectedLBRules should support customize health probe port ",
+			service: getTestServiceDualStack("test1", v1.ProtocolTCP, map[string]string{
+				"service.beta.kubernetes.io/port_8000_health-probe_port": "5080",
+			}, 80, 8000),
+			expectedRules: map[bool][]network.LoadBalancingRule{
+				consts.IPVersionIPv4: {
+					getTestRule(false, 80, consts.IPVersionIPv4),
+					getTestRule(false, 8000, consts.IPVersionIPv4),
+				},
+				consts.IPVersionIPv6: {
+					getTestRule(false, 80, consts.IPVersionIPv6),
+					getTestRule(false, 8000, consts.IPVersionIPv6),
+				},
+			},
+			expectedProbes: map[bool][]network.Probe{
+				consts.IPVersionIPv4: {
+					getTestProbe("Tcp", "/", pointer.Int32(5), pointer.Int32(80), pointer.Int32(10080), pointer.Int32(2), consts.IPVersionIPv4),
+					getTestProbe("Tcp", "/", pointer.Int32(5), pointer.Int32(8000), pointer.Int32(5080), pointer.Int32(2), consts.IPVersionIPv4),
+				},
+				consts.IPVersionIPv6: {
+					getTestProbe("Tcp", "/", pointer.Int32(5), pointer.Int32(80), pointer.Int32(10080), pointer.Int32(2), consts.IPVersionIPv6),
+					getTestProbe("Tcp", "/", pointer.Int32(5), pointer.Int32(8000), pointer.Int32(5080), pointer.Int32(2), consts.IPVersionIPv6),
+				},
+			},
+		},
 	}
 	rulesDualStack := getDefaultTestRules(true)
 	for _, rules := range rulesDualStack {