network: refactor ipmasq management into it's own file

main.go

@@ -28,7 +28,6 @@ import (
 	"strings"
 	"syscall"
 
-	"github.com/coreos/go-iptables/iptables"
 	"github.com/coreos/pkg/flagutil"
 	log "github.com/golang/glog"
 	"golang.org/x/net/context"
@@ -285,7 +284,7 @@ func main() {
 
 	// Set up ipMasq if needed
 	if opts.ipMasq {
-		go setupIPMasq(config, bn)
+		go network.SetupAndEnsureIPMasq(config.Network, bn.Lease())
 	}
 
 	if err := WriteSubnetFile(opts.subnetFile, config.Network, opts.ipMasq, bn); err != nil {
@@ -553,26 +552,6 @@ func mustRunHealthz() {
 	}
 }
 
-func setupIPMasq(config *subnet.Config, bn backend.Network) {
-	ipt, err := iptables.New()
-	if err != nil {
-		// if we can't find iptables, give up and return
-		log.Errorf("Failed to set up IP Masquerade. iptables was not found: %v", err)
-		return
-	}
-	defer func() {
-		network.TeardownIPMasq(ipt, config.Network, bn.Lease())
-	}()
-	for {
-		// Ensure that all the rules exist every 5 seconds
-		if err := network.EnsureIPMasq(ipt, config.Network, bn.Lease()); err != nil {
-			log.Errorf("Failed to ensure IP Masquerade: %v", err)
-		}
-		time.Sleep(5 * time.Second)
-	}
-
-}
-
 func ReadSubnetFromSubnetFile(path string) ip.IP4Net {
 	var prevSubnet ip.IP4Net
 	if _, err := os.Stat(path); !os.IsNotExist(err) {

network/ipmasq.go

@@ -22,6 +22,8 @@ import (
 
 	"github.com/coreos/flannel/pkg/ip"
 	"github.com/coreos/flannel/subnet"
+	"github.com/coreos/go-iptables/iptables"
+	"time"
 )
 
 type IPTablesRules interface {
@@ -61,7 +63,29 @@ func ipMasqRulesExist(ipt IPTablesRules, ipn ip.IP4Net, lease *subnet.Lease) (bo
 	return true, nil
 }
 
-func EnsureIPMasq(ipt IPTablesRules, ipn ip.IP4Net, lease *subnet.Lease) error {
+func SetupAndEnsureIPMasq(network ip.IP4Net, lease *subnet.Lease) {
+	ipt, err := iptables.New()
+	if err != nil {
+		// if we can't find iptables, give up and return
+		log.Errorf("Failed to setup IP Masquerade.  IPTables was not found: %v", err)
+		return
+	}
+
+	defer func() {
+		teardownIPMasq(ipt, network, lease)
+	}()
+
+	for {
+		// Ensure that all the rules exist every 5 seconds
+		if err := ensureIPMasq(ipt, network, lease); err != nil {
+			log.Errorf("Failed to ensure IP Masquerade: %v", err)
+		}
+
+		time.Sleep(5 * time.Second)
+	}
+}
+
+func ensureIPMasq(ipt IPTablesRules, ipn ip.IP4Net, lease *subnet.Lease) error {
 	exists, err := ipMasqRulesExist(ipt, ipn, lease)
 	if err != nil {
 		return fmt.Errorf("Error checking rule existence: %v", err)
@@ -73,14 +97,14 @@ func EnsureIPMasq(ipt IPTablesRules, ipn ip.IP4Net, lease *subnet.Lease) error {
 	// Otherwise, teardown all the rules and set them up again
 	// We do this because the order of the rules is important
 	log.Info("Some iptables rules are missing; deleting and recreating rules")
-	TeardownIPMasq(ipt, ipn, lease)
-	if err = SetupIPMasq(ipt, ipn, lease); err != nil {
+	teardownIPMasq(ipt, ipn, lease)
+	if err = setupIPMasq(ipt, ipn, lease); err != nil {
 		return fmt.Errorf("Error setting up rules: %v", err)
 	}
 	return nil
 }
 
-func SetupIPMasq(ipt IPTablesRules, ipn ip.IP4Net, lease *subnet.Lease) error {
+func setupIPMasq(ipt IPTablesRules, ipn ip.IP4Net, lease *subnet.Lease) error {
 	for _, rule := range rules(ipn, lease) {
 		log.Info("Adding iptables rule: ", strings.Join(rule, " "))
 		err := ipt.AppendUnique("nat", "POSTROUTING", rule...)
@@ -92,7 +116,7 @@ func SetupIPMasq(ipt IPTablesRules, ipn ip.IP4Net, lease *subnet.Lease) error {
 	return nil
 }
 
-func TeardownIPMasq(ipt IPTablesRules, ipn ip.IP4Net, lease *subnet.Lease) {
+func teardownIPMasq(ipt IPTablesRules, ipn ip.IP4Net, lease *subnet.Lease) {
 	for _, rule := range rules(ipn, lease) {
 		log.Info("Deleting iptables rule: ", strings.Join(rule, " "))
 		// We ignore errors here because if there's an error it's almost certainly because the rule

network/ipmasq_test.go

@@ -74,11 +74,11 @@ func (mock *MockIPTables) AppendUnique(table string, chain string, rulespec ...s
 
 func TestDeleteRules(t *testing.T) {
 	ipt := &MockIPTables{}
-	SetupIPMasq(ipt, ip.IP4Net{}, lease())
+	setupIPMasq(ipt, ip.IP4Net{}, lease())
 	if len(ipt.rules) != 4 {
 		t.Errorf("Should be 4 rules, there are actually %d: %#v", len(ipt.rules), ipt.rules)
 	}
-	TeardownIPMasq(ipt, ip.IP4Net{}, lease())
+	teardownIPMasq(ipt, ip.IP4Net{}, lease())
 	if len(ipt.rules) != 0 {
 		t.Errorf("Should be 0 rules, there are actually %d: %#v", len(ipt.rules), ipt.rules)
 	}
@@ -87,13 +87,13 @@ func TestDeleteRules(t *testing.T) {
 func TestEnsureRules(t *testing.T) {
 	// If any rules are missing, they should be all deleted and recreated in the correct order
 	ipt_correct := &MockIPTables{}
-	SetupIPMasq(ipt_correct, ip.IP4Net{}, lease())
-	// setup a mock instance where we delete some rules and run `EnsureIPMasq`
+	setupIPMasq(ipt_correct, ip.IP4Net{}, lease())
+	// setup a mock instance where we delete some rules and run `ensureIPMasq`
 	ipt_recreate := &MockIPTables{}
-	SetupIPMasq(ipt_recreate, ip.IP4Net{}, lease())
+	setupIPMasq(ipt_recreate, ip.IP4Net{}, lease())
 	ipt_recreate.rules = ipt_recreate.rules[0:2]
-	EnsureIPMasq(ipt_recreate, ip.IP4Net{}, lease())
+	ensureIPMasq(ipt_recreate, ip.IP4Net{}, lease())
 	if !reflect.DeepEqual(ipt_recreate.rules, ipt_correct.rules) {
-		t.Errorf("iptables rules after EnsureIPMasq are incorrected. Expected: %#v, Actual: %#v", ipt_recreate.rules, ipt_correct.rules)
+		t.Errorf("iptables rules after ensureIPMasq are incorrected. Expected: %#v, Actual: %#v", ipt_recreate.rules, ipt_correct.rules)
 	}
 }