From f27d4d2117ea3b2e3a9b1cfa4042a111eca2b7c5 Mon Sep 17 00:00:00 2001
From: Martynas Pumputis <martynasp@gmail.com>
Date: Sat, 19 May 2018 12:33:25 +0200
Subject: [PATCH] npc: Add defaultAllowIPSetName helper

---
 npc/namespace.go | 25 +++++++++++--------------
 1 file changed, 11 insertions(+), 14 deletions(-)

diff --git a/npc/namespace.go b/npc/namespace.go
index 5b5607cd..41fb6af3 100644
--- a/npc/namespace.go
+++ b/npc/namespace.go
@@ -118,10 +118,7 @@ func (ns *ns) onNewDstPodSelector(selector *selector, policyType policyType) err
 	for _, pod := range ns.pods {
 		if hasIP(pod) {
 			// Remove the pod from default-allow if dst podselector matches the pod
-			ipset := ns.ingressDefaultAllowIPSet
-			if policyType == egressPolicy {
-				ipset = ns.egressDefaultAllowIPSet
-			}
+			ipset := ns.defaultAllowIPSetName(policyType)
 			if selector.matches(pod.ObjectMeta.Labels) {
 				if err := ns.ips.DelEntry(pod.ObjectMeta.UID, ipset, pod.Status.PodIP); err != nil {
 					return err
@@ -162,10 +159,7 @@ func (ns *ns) addToDefaultAllowIfNoMatching(pod *coreapi.Pod, policyType policyT
 		}
 	}
 	if !found {
-		ipset := ns.ingressDefaultAllowIPSet
-		if policyType == egressPolicy {
-			ipset = ns.egressDefaultAllowIPSet
-		}
+		ipset := ns.defaultAllowIPSetName(policyType)
 		if err := ns.ips.AddEntry(pod.ObjectMeta.UID, ipset, pod.Status.PodIP, podComment(pod)); err != nil {
 			return err
 		}
@@ -306,11 +300,7 @@ func (ns *ns) updatePod(oldObj, newObj *coreapi.Pod) error {
 }
 
 func (ns *ns) addOrRemoveToDefaultAllowIPSet(ps *selector, oldObj, newObj *coreapi.Pod, oldMatch, newMatch bool, policyType policyType) error {
-	ipset := ns.ingressDefaultAllowIPSet
-	if policyType == egressPolicy {
-		ipset = ns.egressDefaultAllowIPSet
-	}
-
+	ipset := ns.defaultAllowIPSetName(policyType)
 	if ns.podSelectors.dstSelectorExist(ps, policyType) {
 		switch {
 		case !oldMatch && newMatch:
@@ -323,7 +313,6 @@ func (ns *ns) addOrRemoveToDefaultAllowIPSet(ps *selector, oldObj, newObj *corea
 			}
 		}
 	}
-
 	return nil
 }
 
@@ -601,6 +590,14 @@ func podComment(pod *coreapi.Pod) string {
 	return fmt.Sprintf("namespace: %s, pod: %s", pod.ObjectMeta.Namespace, pod.ObjectMeta.Name)
 }
 
+func (ns *ns) defaultAllowIPSetName(pt policyType) ipset.Name {
+	ipset := ns.ingressDefaultAllowIPSet
+	if pt == egressPolicy {
+		ipset = ns.egressDefaultAllowIPSet
+	}
+	return ipset
+}
+
 func (ns *ns) analyse(obj interface{}) (
 	uid types.UID,
 	rules map[string]*ruleSpec,