From 35f1ba17e1441e0ba050154314803775601795ec Mon Sep 17 00:00:00 2001
From: Martynas Pumputis <martynasp@gmail.com>
Date: Sat, 26 May 2018 16:51:17 +0200
Subject: [PATCH] npc: Provision egress rules

---
 npc/rule.go | 54 +++++++++++++++++++++++++++++++++++++----------------
 1 file changed, 38 insertions(+), 16 deletions(-)

diff --git a/npc/rule.go b/npc/rule.go
index 03de4170..37238c72 100644
--- a/npc/rule.go
+++ b/npc/rule.go
@@ -46,7 +46,6 @@ func newRuleSpec(policyType policyType, proto *string, srcHost, dstHost ruleHost
 	if dstPort != nil {
 		args = append(args, "--dport", *dstPort)
 	}
-	args = append(args, "-j", "ACCEPT")
 	// NOTE: if you remove the comment bellow, then embed `policyType` into `key`.
 	// Otherwise, the rule won't be provisioned if it exists for other policy type.
 	args = append(args, "-m", "comment", "--comment", fmt.Sprintf("%s -> %s (%s)", srcComment, dstComment, policyTypeStr(policyType)))
@@ -55,6 +54,31 @@ func newRuleSpec(policyType policyType, proto *string, srcHost, dstHost ruleHost
 	return &ruleSpec{key, args, policyType}
 }
 
+func (spec *ruleSpec) iptChain() string {
+	if spec.policyType == egressPolicy {
+		return EgressCustomChain
+	}
+	return IngressChain
+}
+
+func (spec *ruleSpec) iptRuleSpecs() [][]string {
+	if spec.policyType == ingressPolicy {
+		rule := make([]string, len(spec.args))
+		copy(rule, spec.args)
+		rule = append(rule, "-j", "ACCEPT")
+		return [][]string{rule}
+	}
+
+	// egressPolicy
+	ruleMark := make([]string, len(spec.args))
+	copy(ruleMark, spec.args)
+	ruleMark = append(ruleMark, "-j", "MARK", "--set-xmark", EgressMark)
+	ruleReturn := make([]string, len(spec.args))
+	copy(ruleReturn, spec.args)
+	ruleReturn = append(ruleReturn, "-j", "RETURN")
+	return [][]string{ruleMark, ruleReturn}
+}
+
 type ruleSet struct {
 	ipt   iptables.Interface
 	users map[string]map[types.UID]struct{}
@@ -69,11 +93,14 @@ func (rs *ruleSet) deprovision(user types.UID, current, desired map[string]*rule
 		if _, found := desired[key]; !found {
 			delete(rs.users[key], user)
 			if len(rs.users[key]) == 0 {
-				chain := rulesChainByPolicyType(spec.policyType)
-				common.Log.Infof("deleting rule %v from %q chain", spec.args, chain)
-				if err := rs.ipt.Delete(TableFilter, chain, spec.args...); err != nil {
-					return err
+				chain := spec.iptChain()
+				for _, rule := range spec.iptRuleSpecs() {
+					common.Log.Infof("deleting rule %v from %q chain", rule, chain)
+					if err := rs.ipt.Delete(TableFilter, chain, rule...); err != nil {
+						return err
+					}
 				}
+
 				delete(rs.users, key)
 			}
 		}
@@ -86,10 +113,12 @@ func (rs *ruleSet) provision(user types.UID, current, desired map[string]*ruleSp
 	for key, spec := range desired {
 		if _, found := current[key]; !found {
 			if _, found := rs.users[key]; !found {
-				chain := rulesChainByPolicyType(spec.policyType)
-				common.Log.Infof("adding rule %v to %q chain", spec.args, chain)
-				if err := rs.ipt.Append(TableFilter, chain, spec.args...); err != nil {
-					return err
+				chain := spec.iptChain()
+				for _, rule := range spec.iptRuleSpecs() {
+					common.Log.Infof("adding rule %v to %q chain", rule, chain)
+					if err := rs.ipt.Append(TableFilter, chain, rule...); err != nil {
+						return err
+					}
 				}
 				rs.users[key] = make(map[types.UID]struct{})
 			}
@@ -99,10 +128,3 @@ func (rs *ruleSet) provision(user types.UID, current, desired map[string]*ruleSp
 
 	return nil
 }
-
-func rulesChainByPolicyType(policyType policyType) string {
-	if policyType == egressPolicy {
-		return EgressCustomChain
-	}
-	return IngressChain
-}