Add tests for ip masquerade iptables rules

Makefile

@@ -8,7 +8,7 @@ TAG?=$(shell git describe --tags --dirty)
 ARCH?=amd64
 
 # These variables can be overridden by setting an environment variable.
-TEST_PACKAGES?=pkg/ip subnet subnet/etcdv2
+TEST_PACKAGES?=pkg/ip subnet subnet/etcdv2 network
 TEST_PACKAGES_EXPANDED=$(TEST_PACKAGES:%=github.com/coreos/flannel/%)
 PACKAGES?=$(TEST_PACKAGES) network
 PACKAGES_EXPANDED=$(PACKAGES:%=github.com/coreos/flannel/%)

main.go

@@ -28,6 +28,7 @@ import (
 	"strings"
 	"syscall"
 
+	"github.com/coreos/go-iptables/iptables"
 	"github.com/coreos/pkg/flagutil"
 	log "github.com/golang/glog"
 	"golang.org/x/net/context"
@@ -284,17 +285,7 @@ func main() {
 
 	// Set up ipMasq if needed
 	if opts.ipMasq {
-		err = network.SetupIPMasq(config.Network, bn.Lease())
-		if err != nil {
-			// Continue, even though it failed.
-			log.Errorf("Failed to set up IP Masquerade: %v", err)
-		}
-
-		defer func() {
-			if err := network.TeardownIPMasq(config.Network, bn.Lease()); err != nil {
-				log.Errorf("Failed to tear down IP Masquerade: %v", err)
-			}
-		}()
+		setupIPMasq(config, bn)
 	}
 
 	if err := WriteSubnetFile(opts.subnetFile, config.Network, opts.ipMasq, bn); err != nil {
@@ -562,6 +553,25 @@ func mustRunHealthz() {
 	}
 }
 
+func setupIPMasq(config *subnet.Config, bn backend.Network) {
+	ipt, err := iptables.New()
+	if err != nil {
+		log.Errorf("Failed to set up IP Masquerade. iptables was not found: %v", err)
+		return
+	}
+	err = network.SetupIPMasq(ipt, config.Network, bn.Lease())
+	if err != nil {
+		// Continue, even though it failed.
+		log.Errorf("Failed to set up IP Masquerade: %v", err)
+	}
+
+	defer func() {
+		if err := network.TeardownIPMasq(ipt, config.Network, bn.Lease()); err != nil {
+			log.Errorf("Failed to tear down IP Masquerade: %v", err)
+		}
+	}()
+}
+
 func ReadSubnetFromSubnetFile(path string) ip.IP4Net {
 	var prevSubnet ip.IP4Net
 	if _, err := os.Stat(path); !os.IsNotExist(err) {

network/ipmasq.go

@@ -18,13 +18,17 @@ import (
 	"fmt"
 	"strings"
 
-	"github.com/coreos/go-iptables/iptables"
 	log "github.com/golang/glog"
 
 	"github.com/coreos/flannel/pkg/ip"
 	"github.com/coreos/flannel/subnet"
 )
 
+type IPTablesRules interface {
+	AppendUnique(table string, chain string, rulespec ...string) error
+	Delete(table string, chain string, rulespec ...string) error
+}
+
 func rules(ipn ip.IP4Net, lease *subnet.Lease) [][]string {
 	n := ipn.String()
 	sn := lease.Subnet.String()
@@ -41,15 +45,10 @@ func rules(ipn ip.IP4Net, lease *subnet.Lease) [][]string {
 	}
 }
 
-func SetupIPMasq(ipn ip.IP4Net, lease *subnet.Lease) error {
-	ipt, err := iptables.New()
-	if err != nil {
-		return fmt.Errorf("failed to set up IP Masquerade. iptables was not found")
-	}
-
+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...)
+		err := ipt.AppendUnique("nat", "POSTROUTING", rule...)
 		if err != nil {
 			return fmt.Errorf("failed to insert IP masquerade rule: %v", err)
 		}
@@ -58,15 +57,10 @@ func SetupIPMasq(ipn ip.IP4Net, lease *subnet.Lease) error {
 	return nil
 }
 
-func TeardownIPMasq(ipn ip.IP4Net, lease *subnet.Lease) error {
-	ipt, err := iptables.New()
-	if err != nil {
-		return fmt.Errorf("failed to teardown IP Masquerade. iptables was not found")
-	}
-
+func TeardownIPMasq(ipt IPTablesRules, ipn ip.IP4Net, lease *subnet.Lease) error {
 	for _, rule := range rules(ipn, lease) {
 		log.Info("Deleting iptables rule: ", strings.Join(rule, " "))
-		err = ipt.Delete("nat", "POSTROUTING", rule...)
+		err := ipt.Delete("nat", "POSTROUTING", rule...)
 		if err != nil {
 			return fmt.Errorf("failed to delete IP masquerade rule: %v", err)
 		}

network/ipmasq_test.go

@@ -0,0 +1,77 @@
+// Copyright 2015 flannel authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package network
+
+import (
+	"github.com/coreos/flannel/pkg/ip"
+	"github.com/coreos/flannel/subnet"
+	"net"
+	"reflect"
+	"testing"
+)
+
+func lease() *subnet.Lease {
+	_, net, _ := net.ParseCIDR("192.168.0.0/16")
+	return &subnet.Lease{
+		Subnet: ip.FromIPNet(net),
+	}
+}
+
+type MockIPTablesRule struct {
+	table    string
+	chain    string
+	rulespec []string
+}
+
+type MockIPTables struct {
+	rules []MockIPTablesRule
+}
+
+func (mock *MockIPTables) ruleIndex(table string, chain string, rulespec []string) int {
+	for i, rule := range mock.rules {
+		if rule.table == table && rule.chain == chain && reflect.DeepEqual(rule.rulespec, rulespec) {
+			return i
+		}
+	}
+	return -1
+}
+
+func (mock *MockIPTables) Delete(table string, chain string, rulespec ...string) error {
+	var ruleIndex = mock.ruleIndex(table, chain, rulespec)
+	if ruleIndex != -1 {
+		mock.rules = append(mock.rules[:ruleIndex], mock.rules[ruleIndex+1:]...)
+	}
+	return nil
+}
+
+func (mock *MockIPTables) AppendUnique(table string, chain string, rulespec ...string) error {
+	var ruleIndex = mock.ruleIndex(table, chain, rulespec)
+	if ruleIndex == -1 {
+		mock.rules = append(mock.rules, MockIPTablesRule{table: table, chain: chain, rulespec: rulespec})
+	}
+	return nil
+}
+
+func TestDeleteRules(t *testing.T) {
+	ipt := &MockIPTables{}
+	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())
+	if len(ipt.rules) != 0 {
+		t.Errorf("Should be 0 rules, there are actually %d: %#v", len(ipt.rules), ipt.rules)
+	}
+}