From 44096b8f71032f0c29a3322f4a74ecb454b22647 Mon Sep 17 00:00:00 2001
From: elweb9858 <elweb@microsoft.com>
Date: Wed, 3 Jun 2020 13:41:59 -0700
Subject: [PATCH] Adding windows implementation for sessionaffinity

---
 pkg/proxy/winkernel/hnsV2.go   | 41 +++++++++++++++++++++++++---------
 pkg/proxy/winkernel/proxier.go | 37 ++++++++++++++++++------------
 2 files changed, 54 insertions(+), 24 deletions(-)

diff --git a/pkg/proxy/winkernel/hnsV2.go b/pkg/proxy/winkernel/hnsV2.go
index e86e5ce431e..62053f122bf 100644
--- a/pkg/proxy/winkernel/hnsV2.go
+++ b/pkg/proxy/winkernel/hnsV2.go
@@ -232,16 +232,37 @@ func (hns hnsV2) getLoadBalancer(endpoints []endpointsInfo, flags loadBalancerFl
 		lbFlags |= hcn.LoadBalancerFlagsDSR
 	}
 
-	lb, err := hcn.AddLoadBalancer(
-		hnsEndpoints,
-		lbFlags,
-		lbPortMappingFlags,
-		sourceVip,
-		vips,
-		protocol,
-		internalPort,
-		externalPort,
-	)
+	lbDistributionType := hcn.LoadBalancerDistributionNone
+
+	if flags.sessionAffinity {
+		lbDistributionType = hcn.LoadBalancerDistributionSourceIP
+	}
+
+	loadBalancer := &hcn.HostComputeLoadBalancer{
+		SourceVIP: sourceVip,
+		PortMappings: []hcn.LoadBalancerPortMapping{
+			{
+				Protocol:         uint32(protocol),
+				InternalPort:     internalPort,
+				ExternalPort:     externalPort,
+				DistributionType: lbDistributionType,
+				Flags:            lbPortMappingFlags,
+			},
+		},
+		FrontendVIPs: vips,
+		SchemaVersion: hcn.SchemaVersion{
+			Major: 2,
+			Minor: 0,
+		},
+		Flags: lbFlags,
+	}
+
+	for _, endpoint := range hnsEndpoints {
+		loadBalancer.HostComputeEndpoints = append(loadBalancer.HostComputeEndpoints, endpoint.Id)
+	}
+
+	lb, err := loadBalancer.Create()
+
 	if err != nil {
 		return nil, err
 	}
diff --git a/pkg/proxy/winkernel/proxier.go b/pkg/proxy/winkernel/proxier.go
index 3b577e96c01..27bbef14639 100644
--- a/pkg/proxy/winkernel/proxier.go
+++ b/pkg/proxy/winkernel/proxier.go
@@ -95,11 +95,12 @@ type loadBalancerInfo struct {
 }
 
 type loadBalancerFlags struct {
-	isILB          bool
-	isDSR          bool
-	localRoutedVIP bool
-	useMUX         bool
-	preserveDIP    bool
+	isILB           bool
+	isDSR           bool
+	localRoutedVIP  bool
+	useMUX          bool
+	preserveDIP     bool
+	sessionAffinity bool
 }
 
 // internal struct for string service information
@@ -488,11 +489,12 @@ type Proxier struct {
 	// precomputing some number of those and cache for future reuse.
 	precomputedProbabilities []string
 
-	hns       HostNetworkService
-	network   hnsNetworkInfo
-	sourceVip string
-	hostMac   string
-	isDSR     bool
+	hns               HostNetworkService
+	network           hnsNetworkInfo
+	sourceVip         string
+	hostMac           string
+	isDSR             bool
+	supportedFeatures hcn.SupportedFeatures
 }
 
 type localPort struct {
@@ -649,6 +651,7 @@ func NewProxier(
 		sourceVip:           sourceVip,
 		hostMac:             hostMac,
 		isDSR:               isDSR,
+		supportedFeatures:   supportedFeatures,
 	}
 
 	burstSyncs := 2
@@ -1226,9 +1229,15 @@ func (proxier *Proxier) syncProxyRules() {
 		if containsPublicIP || containsNodeIP {
 			sourceVip = proxier.nodeIP.String()
 		}
+
+		sessionAffinityClientIP := svcInfo.sessionAffinityType == v1.ServiceAffinityClientIP
+		if sessionAffinityClientIP && !proxier.supportedFeatures.SessionAffinity {
+			klog.Warningf("Session Affinity is not supported on this version of Windows.")
+		}
+
 		hnsLoadBalancer, err := hns.getLoadBalancer(
 			hnsEndpoints,
-			loadBalancerFlags{isDSR: proxier.isDSR},
+			loadBalancerFlags{isDSR: proxier.isDSR, sessionAffinity: sessionAffinityClientIP},
 			sourceVip,
 			svcInfo.clusterIP.String(),
 			Enum(svcInfo.protocol),
@@ -1253,7 +1262,7 @@ func (proxier *Proxier) syncProxyRules() {
 			}
 			hnsLoadBalancer, err := hns.getLoadBalancer(
 				nodePortEndpoints,
-				loadBalancerFlags{localRoutedVIP: true},
+				loadBalancerFlags{localRoutedVIP: true, sessionAffinity: sessionAffinityClientIP},
 				sourceVip,
 				"",
 				Enum(svcInfo.protocol),
@@ -1274,7 +1283,7 @@ func (proxier *Proxier) syncProxyRules() {
 			// Try loading existing policies, if already available
 			hnsLoadBalancer, err = hns.getLoadBalancer(
 				hnsEndpoints,
-				loadBalancerFlags{},
+				loadBalancerFlags{sessionAffinity: sessionAffinityClientIP},
 				sourceVip,
 				externalIP.ip,
 				Enum(svcInfo.protocol),
@@ -1297,7 +1306,7 @@ func (proxier *Proxier) syncProxyRules() {
 			}
 			hnsLoadBalancer, err := hns.getLoadBalancer(
 				lbIngressEndpoints,
-				loadBalancerFlags{isDSR: svcInfo.preserveDIP || proxier.isDSR, useMUX: svcInfo.preserveDIP, preserveDIP: svcInfo.preserveDIP},
+				loadBalancerFlags{isDSR: svcInfo.preserveDIP || proxier.isDSR, useMUX: svcInfo.preserveDIP, preserveDIP: svcInfo.preserveDIP, sessionAffinity: sessionAffinityClientIP},
 				sourceVip,
 				lbIngressIP.ip,
 				Enum(svcInfo.protocol),