From 01216f86a41af8a2a4a241648def2bdd30255004 Mon Sep 17 00:00:00 2001
From: Pengfei Ni <feiskyer@gmail.com>
Date: Mon, 28 Aug 2023 14:43:52 +0800
Subject: [PATCH] feat: add support for disable tcp reset

---
 pkg/consts/consts.go                    |  3 +++
 pkg/consts/helpers.go                   |  7 ++++++-
 pkg/provider/azure_loadbalancer.go      |  5 +++--
 pkg/provider/azure_loadbalancer_test.go | 23 ++++++++++++++++++++++-
 4 files changed, 34 insertions(+), 4 deletions(-)

diff --git a/pkg/consts/consts.go b/pkg/consts/consts.go
index c9e075d2fb..d9ddea625e 100644
--- a/pkg/consts/consts.go
+++ b/pkg/consts/consts.go
@@ -333,6 +333,9 @@ const (
 	// The list is separated by comma. It will be omitted if multi-slb is not used.
 	ServiceAnnotationLoadBalancerConfigurations = "service.beta.kubernetes.io/azure-load-balancer-configurations"
 
+	// ServiceAnnotationDisableTCPReset is the annotation used on the service to disable TCP reset on the load balancer.
+	ServiceAnnotationDisableTCPReset = "service.beta.kubernetes.io/azure-load-balancer-disable-tcp-reset"
+
 	// ServiceTagKey is the service key applied for public IP tags.
 	ServiceTagKey       = "k8s-azure-service"
 	LegacyServiceTagKey = "service"
diff --git a/pkg/consts/helpers.go b/pkg/consts/helpers.go
index 9706b55aa9..c6636b0d20 100644
--- a/pkg/consts/helpers.go
+++ b/pkg/consts/helpers.go
@@ -65,6 +65,11 @@ func IsPLSEnabled(annotations map[string]string) bool {
 	return expectAttributeInSvcAnnotationBeEqualTo(annotations, ServiceAnnotationPLSCreation, TrueAnnotationValue)
 }
 
+// IsTCPResetDisabled return true if ServiceAnnotationDisableTCPReset is true
+func IsTCPResetDisabled(annotations map[string]string) bool {
+	return expectAttributeInSvcAnnotationBeEqualTo(annotations, ServiceAnnotationDisableTCPReset, TrueAnnotationValue)
+}
+
 // Getint32ValueFromK8sSvcAnnotation get health probe configuration for port
 func Getint32ValueFromK8sSvcAnnotation(annotations map[string]string, key string, validators ...Int32BusinessValidator) (*int32, error) {
 	val, err := GetAttributeValueInSvcAnnotation(annotations, key)
@@ -74,7 +79,7 @@ func Getint32ValueFromK8sSvcAnnotation(annotations map[string]string, key string
 	return nil, err
 }
 
-// BuildHealthProbeAnnotationKeyForPort get health probe configuration key for port
+// BuildAnnotationKeyForPort get health probe configuration key for port
 func BuildAnnotationKeyForPort(port int32, key PortParams) string {
 	return fmt.Sprintf(PortAnnotationPrefixPattern, port, string(key))
 }
diff --git a/pkg/provider/azure_loadbalancer.go b/pkg/provider/azure_loadbalancer.go
index 21d9b404ab..825a602561 100644
--- a/pkg/provider/azure_loadbalancer.go
+++ b/pkg/provider/azure_loadbalancer.go
@@ -2717,7 +2717,7 @@ func (az *Cloud) getExpectedLoadBalancingRulePropertiesForPort(
 		IdleTimeoutInMinutes: lbIdleTimeout,
 	}
 	if strings.EqualFold(string(transportProto), string(network.TransportProtocolTCP)) && az.useStandardLoadBalancer() {
-		props.EnableTCPReset = pointer.Bool(true)
+		props.EnableTCPReset = pointer.Bool(!consts.IsTCPResetDisabled(service.Annotations))
 	}
 
 	// Azure ILB does not support secondary IPs as floating IPs on the LB. Therefore, floating IP needs to be turned
@@ -2738,7 +2738,8 @@ func (az *Cloud) getExpectedHAModeLoadBalancingRuleProperties(
 	if err != nil {
 		return nil, fmt.Errorf("error generate lb rule for ha mod loadbalancer. err: %w", err)
 	}
-	props.EnableTCPReset = pointer.Bool(true)
+	props.EnableTCPReset = pointer.Bool(!consts.IsTCPResetDisabled(service.Annotations))
+
 	return props, nil
 }
 
diff --git a/pkg/provider/azure_loadbalancer_test.go b/pkg/provider/azure_loadbalancer_test.go
index f9a2f27972..c5e3b9d842 100644
--- a/pkg/provider/azure_loadbalancer_test.go
+++ b/pkg/provider/azure_loadbalancer_test.go
@@ -2664,6 +2664,15 @@ func TestReconcileLoadBalancerRuleCommon(t *testing.T) {
 			expectedRules:  getDefaultTestRules(false),
 			expectedProbes: getDefaultTestProbes("Http", "/"),
 		},
+		{
+			desc: "getExpectedLBRules should disable tcp reset when annotation is set",
+			service: getTestServiceDualStack("test1", v1.ProtocolTCP, map[string]string{
+				"service.beta.kubernetes.io/azure-load-balancer-disable-tcp-reset": "true",
+			}, 80),
+			loadBalancerSku: "standard",
+			expectedRules:   getTCPResetTestRules(false),
+			expectedProbes:  getDefaultTestProbes("Tcp", ""),
+		},
 		{
 			desc: "getExpectedLBRules should prioritize port specific probe protocol over appProtocol",
 			service: getTestServiceDualStack("test1", v1.ProtocolTCP, map[string]string{
@@ -2989,7 +2998,7 @@ func getDefaultTestRules(enableTCPReset bool) map[bool][]network.LoadBalancingRu
 
 // getDefaultInternalIPv6Rules returns a rule for IPv6 single stack.
 func getDefaultInternalIPv6Rules(enableTCPReset bool) map[bool][]network.LoadBalancingRule {
-	rule := getTestRule(true, 80, false)
+	rule := getTestRule(enableTCPReset, 80, false)
 	rule.EnableFloatingIP = pointer.Bool(false)
 	rule.BackendPort = pointer.Int32(getBackendPort(*rule.FrontendPort))
 	rule.BackendAddressPool.ID = pointer.String("backendPoolID-IPv6")
@@ -2998,6 +3007,18 @@ func getDefaultInternalIPv6Rules(enableTCPReset bool) map[bool][]network.LoadBal
 	}
 }
 
+// getTCPResetTestRules returns rules with TCPReset always set.
+func getTCPResetTestRules(enableTCPReset bool) map[bool][]network.LoadBalancingRule {
+	IPv4Rule := getTestRule(enableTCPReset, 80, consts.IPVersionIPv4)
+	IPv6Rule := getTestRule(enableTCPReset, 80, consts.IPVersionIPv6)
+	IPv4Rule.EnableTCPReset = pointer.Bool(enableTCPReset)
+	IPv6Rule.EnableTCPReset = pointer.Bool(enableTCPReset)
+	return map[bool][]network.LoadBalancingRule{
+		consts.IPVersionIPv4: {IPv4Rule},
+		consts.IPVersionIPv6: {IPv6Rule},
+	}
+}
+
 // getTestRule returns a rule for dualStack.
 func getTestRule(enableTCPReset bool, port int32, isIPv6 bool) network.LoadBalancingRule {
 	suffix := ""