fix: [NPM] update endpointcache after remove policy (#1804)

* update endpointcache

* update delete logic

* fix lint issue

* added check for endpoint in cache

* add unit test

* fix type lint issue

* fix naming lint error

* updated endpoint cache lock

* moved cache lock down and added check for windows

npm/pkg/dataplane/dataplane-test-cases_windows_test.go

@@ -511,6 +511,65 @@ func getAllSerialTests() []*SerialTestCase {
 				},
 			},
 		},
+		{
+			Description: "pod created to satisfy policy, then policy deleted, then pod relabeled to no longer satisfy policy, then policy re-created and pod relabeled to satisfy policy",
+			Actions: []*Action{
+				CreateEndpoint(endpoint1, ip1),
+				CreatePod("x", "a", ip1, thisNode, map[string]string{"k1": "v1"}),
+				// will apply dirty ipsets from CreatePod
+				UpdatePolicy(policyXBaseOnK1V1()),
+				DeletePolicyByObject(policyXBaseOnK1V1()),
+				UpdatePodLabels("x", "a", ip1, thisNode, map[string]string{"k1": "v1"}, map[string]string{"k2": "v2"}),
+				ApplyDP(),
+				UpdatePolicy(policyXBaseOnK1V1()),
+				ApplyDP(),
+				UpdatePodLabels("x", "a", ip1, thisNode, map[string]string{"k2": "v2"}, map[string]string{"k1": "v1"}),
+				ApplyDP(),
+			},
+			TestCaseMetadata: &TestCaseMetadata{
+				Tags: []Tag{
+					podCrudTag,
+					netpolCrudTag,
+				},
+				DpCfg:            defaultWindowsDPCfg,
+				InitialEndpoints: nil,
+				ExpectedSetPolicies: []*hcn.SetPolicySetting{
+					dptestutils.SetPolicy(emptySet),
+					dptestutils.SetPolicy(allNamespaces, emptySet.GetHashedName(), nsXSet.GetHashedName()),
+					dptestutils.SetPolicy(nsXSet, ip1),
+					dptestutils.SetPolicy(podK1Set, ip1),
+					dptestutils.SetPolicy(podK1V1Set, ip1),
+					dptestutils.SetPolicy(podK2Set),
+					dptestutils.SetPolicy(podK2V2Set),
+				},
+				ExpectedEnpdointACLs: map[string][]*hnswrapper.FakeEndpointPolicy{
+					endpoint1: {
+						{
+							ID:              "azure-acl-x-base",
+							Protocols:       "",
+							Action:          "Allow",
+							Direction:       "In",
+							LocalAddresses:  "",
+							RemoteAddresses: "",
+							LocalPorts:      "",
+							RemotePorts:     "",
+							Priority:        222,
+						},
+						{
+							ID:              "azure-acl-x-base",
+							Protocols:       "",
+							Action:          "Allow",
+							Direction:       "Out",
+							LocalAddresses:  "",
+							RemoteAddresses: "",
+							LocalPorts:      "",
+							RemotePorts:     "",
+							Priority:        222,
+						},
+					},
+				},
+			},
+		},
 	}
 }
 

npm/pkg/dataplane/dataplane.go

@@ -284,15 +284,41 @@ func (dp *DataPlane) RemovePolicy(policyKey string) error {
 	// because policy Manager will remove from policy from cache
 	// keep a local copy to remove references for ipsets
 	policy, ok := dp.policyMgr.GetPolicy(policyKey)
+	endpoints := make(map[string]string, len(policy.PodEndpoints))
+
+	for podIP, endpointID := range policy.PodEndpoints {
+		endpoints[podIP] = endpointID
+	}
+
 	if !ok {
 		klog.Infof("[DataPlane] Policy %s is not found. Might been deleted already", policyKey)
 		return nil
 	}
+
 	// Use the endpoint list saved in cache for this network policy to remove
 	err := dp.policyMgr.RemovePolicy(policy.PolicyKey)
 	if err != nil {
 		return fmt.Errorf("[DataPlane] error while removing policy: %w", err)
 	}
+
+	if dp.shouldUpdatePod() {
+
+		dp.endpointCache.Lock()
+
+		for podIP := range endpoints {
+			// if the endpoint is not in the policy's endpoint list, delete policy reference from cache
+			if _, ok := policy.PodEndpoints[podIP]; !ok {
+				// check if the endpoint is in the cache
+				if endpoint, ok := dp.endpointCache.cache[podIP]; ok {
+					delete(endpoint.netPolReference, policyKey)
+				}
+			}
+		}
+
+		dp.endpointCache.Unlock()
+
+	}
+
 	// Remove references for Rule IPSets first
 	err = dp.deleteIPSetsAndReferences(policy.RuleIPSets, policy.PolicyKey, ipsets.NetPolType)
 	if err != nil {

npm/pkg/dataplane/types_linux.go

@@ -1,4 +1,6 @@
 package dataplane
 
 // npmEndpoint holds info relevant for endpoints in windows
-type npmEndpoint struct{}
+type npmEndpoint struct {
+	netPolReference map[string]struct{}
+}