From 2e85018c261afd115569fc8b27e4a19531bbf8aa Mon Sep 17 00:00:00 2001
From: Martynas Pumputis <martynasp@gmail.com>
Date: Tue, 22 May 2018 12:37:46 +0200
Subject: [PATCH] npc: Add ruleHost interface

The idea for the interface is taken from the kitt1987/ipblock work.

The interface is going to be implemented by selectorSpec and
ipBlockSpec.
---
 npc/rule.go     | 24 +++++++++++++++---------
 npc/selector.go | 19 +++++++++++++++++++
 2 files changed, 34 insertions(+), 9 deletions(-)

diff --git a/npc/rule.go b/npc/rule.go
index f875af5e..24ac5b13 100644
--- a/npc/rule.go
+++ b/npc/rule.go
@@ -10,30 +10,36 @@ import (
 	"github.com/weaveworks/weave/npc/iptables"
 )
 
+type ruleHost interface {
+	// getRuleSpec returns source or destination specification and comment which
+	// are used in an iptables rule.
+	//
+	// If src=true then the rulespec is for source, otherwise - for destination.
+	getRuleSpec(src bool) ([]string, string)
+}
+
 type ruleSpec struct {
 	key        string
 	args       []string
 	policyType policyType
 }
 
-func newRuleSpec(policyType policyType, proto *string, srcHost *selectorSpec, dstHost *selectorSpec, dstPort *string) *ruleSpec {
+func newRuleSpec(policyType policyType, proto *string, srcHost, dstHost ruleHost, dstPort *string) *ruleSpec {
 	args := []string{}
 	if proto != nil {
 		args = append(args, "-p", *proto)
 	}
 	srcComment := "anywhere"
 	if srcHost != nil {
-		args = append(args, "-m", "set", "--match-set", string(srcHost.ipsetName), "src")
-		if srcHost.nsName != "" {
-			srcComment = fmt.Sprintf("pods: namespace: %s, selector: %s", srcHost.nsName, srcHost.key)
-		} else {
-			srcComment = fmt.Sprintf("namespaces: selector: %s", srcHost.key)
-		}
+		rule, comment := srcHost.getRuleSpec(true)
+		args = append(args, rule...)
+		srcComment = comment
 	}
 	dstComment := "anywhere"
 	if dstHost != nil {
-		args = append(args, "-m", "set", "--match-set", string(dstHost.ipsetName), "dst")
-		dstComment = fmt.Sprintf("pods: namespace: %s, selector: %s", dstHost.nsName, dstHost.key)
+		rule, comment := dstHost.getRuleSpec(false)
+		args = append(args, rule...)
+		dstComment = comment
 	}
 	if dstPort != nil {
 		args = append(args, "--dport", *dstPort)
diff --git a/npc/selector.go b/npc/selector.go
index 4affe0ca..33f839a0 100644
--- a/npc/selector.go
+++ b/npc/selector.go
@@ -1,6 +1,8 @@
 package npc
 
 import (
+	"fmt"
+
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/labels"
 	"k8s.io/apimachinery/pkg/types"
@@ -37,6 +39,23 @@ func newSelectorSpec(json *metav1.LabelSelector, policyType []policyType, nsName
 		nsName:    nsName}, nil
 }
 
+func (spec *selectorSpec) getRuleSpec(src bool) ([]string, string) {
+	dir := "dst"
+	if src {
+		dir = "src"
+	}
+	rule := []string{"-m", "set", "--match-set", string(spec.ipsetName), dir}
+
+	comment := "anywhere"
+	if spec.nsName != "" {
+		comment = fmt.Sprintf("pods: namespace: %s, selector: %s", spec.nsName, spec.key)
+	} else {
+		comment = fmt.Sprintf("namespaces: selector: %s", spec.key)
+	}
+
+	return rule, comment
+}
+
 type selector struct {
 	ips  ipset.Interface
 	spec *selectorSpec