Add support for ipv4 and ipv6 dual stack management

.travis.yml

@@ -9,6 +9,6 @@ install:
 
 script:
 - go vet ./...
-- go test ./... -cover=1 -coverprofile=_c.cov
+- go test ./... -v -cover=1 -coverprofile=_c.cov
 - go test ./... -race
 - $GOPATH/bin/goveralls -service=travis-ci -coverprofile=_c.cov

address/granter.go

@@ -23,7 +23,8 @@ type Manager interface {
 // IPManager supports granting IP subnet access using iptables or ip6tables.
 type IPManager struct {
 	*semaphore.Weighted
-	origRules []byte
+	origRules4 []byte
+	origRules6 []byte
 }
 
 // ErrMaxConcurrent is returned when the max concurrent grants has already been reached.
@@ -71,17 +72,17 @@ func (r *IPManager) Revoke(ip net.IP) error {
 	return err
 }
 
-func ipTable(command string, ip net.IP) pipe.Pipe {
+func ipTable(action string, ip net.IP) pipe.Pipe {
 	// Parameters are the same for IPv4 and IPv6 addresses, but the command is not.
 	cmd, subnet := cmdForIP(ip)
-	return pipe.Exec(cmd, "--"+command+"=INPUT", "--source="+ip.String()+subnet, "--jump=ACCEPT")
+	return pipe.Exec(cmd, "--"+action+"=INPUT", "--source="+ip.String()+subnet, "--jump=ACCEPT")
 }
 
 func cmdForIP(ip net.IP) (string, string) {
 	if ip.To4() != nil {
-		return iptables, "/24"
+		return ip4tables, "/24"
 	}
-	return iptables, "/64"
+	return ip6tables, "/64"
 }
 
 // NullManager implements the address.Manager interface while doing nothing.

address/granter_test.go

@@ -59,11 +59,12 @@ func TestIPManager_Grant(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			// NOTE: the fake iptables commands read and exit using the EXIT environment.
 			defer osx.MustSetenv("IPTABLES_EXIT", tt.grantExit)()
+			defer osx.MustSetenv("IP6TABLES_EXIT", tt.grantExit)()
 
 			if tt.ip.To4() != nil {
-				iptables = "./testdata/iptables"
+				ip4tables = "./testdata/iptables"
 			} else {
-				iptables = "./testdata/ip6tables"
+				ip6tables = "./testdata/ip6tables"
 			}
 
 			r := NewIPManager(tt.max)
@@ -84,7 +85,7 @@ func TestIPManager_Grant(t *testing.T) {
 }
 
 func TestIPManager(t *testing.T) {
-	iptables = "./testdata/iptables"
+	ip4tables = "./testdata/iptables"
 	wg := sync.WaitGroup{}
 	mgr := NewIPManager(10)
 	ip := net.ParseIP("127.0.0.2")

address/setup.go

@@ -11,20 +11,29 @@ import (
 )
 
 var (
-	iptables        string
-	iptablesSave    string
-	iptablesRestore string
+	ip4tables        string
+	ip4tablesSave    string
+	ip4tablesRestore string
+
+	ip6tables        string
+	ip6tablesSave    string
+	ip6tablesRestore string
 )
 
 func init() {
-	// NOTE: because ip6tables is flag-compatible with iptables, these flags
-	// support either ipv4 or ipv6 exclusively.
-	// TODO: support both ipv4 and ip6tables.
-	flag.StringVar(&iptables, "address.iptables", "/sbin/iptables",
+	// NOTE: because ip6tables is flag-compatible with iptables.
+	flag.StringVar(&ip4tables, "address.iptables", "/sbin/iptables",
+		"The absolute path to the iptables command")
+	flag.StringVar(&ip4tablesSave, "address.iptables-save", "/sbin/iptables-save",
+		"The absolute path to the iptables-save command")
+	flag.StringVar(&ip4tablesRestore, "address.iptables-restore", "/sbin/iptables-restore",
+		"The absolute path to the iptables-restore command")
+
+	flag.StringVar(&ip6tables, "address.ip6tables", "/sbin/ip6tables",
 		"The absolute path to the iptables command")
-	flag.StringVar(&iptablesSave, "address.iptables-save", "/sbin/iptables-save",
+	flag.StringVar(&ip6tablesSave, "address.ip6tables-save", "/sbin/ip6tables-save",
 		"The absolute path to the iptables-save command")
-	flag.StringVar(&iptablesRestore, "address.iptables-restore", "/sbin/iptables-restore",
+	flag.StringVar(&ip6tablesRestore, "address.ip6tables-restore", "/sbin/ip6tables-restore",
 		"The absolute path to the iptables-restore command")
 }
 
@@ -36,19 +45,31 @@ func init() {
 // during shutdown.
 func (r *IPManager) Start(port, device string) error {
 	// Save original rules.
-	origRules, err := pipe.Output(pipe.Exec(iptablesSave))
+	var err error
+	r.origRules4, err = start(ip4tablesSave, ip4tables, port, device)
+	if err != nil {
+		return err
+	}
+	r.origRules6, err = start(ip6tablesSave, ip6tables, port, device)
 	if err != nil {
 		return err
 	}
-	r.origRules = origRules
+	return nil
+}
+
+func start(iptablesSave, iptables, port, device string) ([]byte, error) {
+	origRules, err := pipe.Output(pipe.Exec(iptablesSave))
+	if err != nil {
+		return nil, err
+	}
 
 	// Collect commands to allow traffic from allowed (i.e. unmanaged) interfaces.
 	//
 	// NOTE: this finds and allows connections to all unmanaged devices. This
 	// allows traffic to private and local networks. This is necessary for
 	// intra-container communications on loopback and for monitoring traffic
 	// over the private network.
-	allowed := allowedInterfaces(device)
+	allowed := allowedInterfaces(iptables, device)
 
 	startCommands := []pipe.Pipe{
 		// Flushing existing rules does not change default policy.
@@ -82,23 +103,32 @@ func (r *IPManager) Start(port, device string) error {
 	err = pipe.Run(
 		pipe.Script("Setup iptables for managing access: "+device, commands...),
 	)
-	return err
+	return origRules, err
 }
 
 // Stop restores the iptables rules originally found before running Start().
 func (r *IPManager) Stop() ([]byte, error) {
-	if r.origRules == nil {
+	b4, err := stop(ip4tablesRestore, r.origRules4)
+	if err != nil {
+		return b4, err
+	}
+	b6, err := stop(ip6tablesRestore, r.origRules6)
+	return append(b4, b6...), err
+}
+
+func stop(iptablesRestore string, rules []byte) ([]byte, error) {
+	if rules == nil {
 		return nil, fmt.Errorf("cannot restore uninitialized rules")
 	}
-	b := bytes.NewBuffer(r.origRules)
+	b := bytes.NewBuffer(rules)
 	restore := pipe.Script("Restoring original iptables rules",
 		pipe.Read(b),
 		pipe.Exec(iptablesRestore),
 	)
 	return pipe.Output(restore)
 }
 
-func allowedInterfaces(name string) []pipe.Pipe {
+func allowedInterfaces(iptables, name string) []pipe.Pipe {
 	ifaces, err := net.Interfaces()
 	rtx.Must(err, "failed to list interfaces")
 	pipes := []pipe.Pipe{}

address/setup_test.go

@@ -7,25 +7,41 @@ import (
 )
 
 func TestIPManager_Start(t *testing.T) {
-	iptables = "./testdata/iptables"
-	iptablesSave = "./testdata/iptables-save"
+	ip4tables = "./testdata/iptables"
+	ip4tablesSave = "./testdata/iptables-save"
+	ip6tables = "./testdata/ip6tables"
+	ip6tablesSave = "./testdata/ip6tables-save"
 
 	tests := []struct {
-		name             string
-		iptablesSaveCode string
-		iptablesCode     string
-		wantErr          bool
+		name              string
+		iptablesSaveCode  string
+		iptablesCode      string
+		iptablesSaveCode6 string
+		iptablesCode6     string
+		wantErr           bool
 	}{
 		{
-			name:             "success",
-			iptablesSaveCode: "0",
-			iptablesCode:     "0",
+			name:              "success",
+			iptablesSaveCode:  "0",
+			iptablesCode:      "0",
+			iptablesSaveCode6: "0",
+			iptablesCode6:     "0",
 		},
 		{
-			name:             "success",
-			iptablesSaveCode: "1",
-			iptablesCode:     "0",
-			wantErr:          true,
+			name:              "error-save-ipv4-failure",
+			iptablesSaveCode:  "1",
+			iptablesCode:      "0",
+			iptablesSaveCode6: "0",
+			iptablesCode6:     "0",
+			wantErr:           true,
+		},
+		{
+			name:              "error-save-ipv6-failure",
+			iptablesSaveCode:  "0",
+			iptablesCode:      "0",
+			iptablesSaveCode6: "1",
+			iptablesCode6:     "0",
+			wantErr:           true,
 		},
 	}
 	for _, tt := range tests {
@@ -34,6 +50,9 @@ func TestIPManager_Start(t *testing.T) {
 			defer osx.MustSetenv("IPTABLES_SAVE_EXIT", tt.iptablesSaveCode)()
 			defer osx.MustSetenv("IPTABLES_EXIT", tt.iptablesCode)()
 
+			defer osx.MustSetenv("IP6TABLES_SAVE_EXIT", tt.iptablesSaveCode6)()
+			defer osx.MustSetenv("IP6TABLES_EXIT", tt.iptablesCode6)()
+
 			if err := r.Start("1234", "eth0"); (err != nil) != tt.wantErr {
 				t.Errorf("IPManager.Start() error = %v, wantErr %v", err, tt.wantErr)
 			}
@@ -42,28 +61,30 @@ func TestIPManager_Start(t *testing.T) {
 }
 
 func TestIPManager_Stop(t *testing.T) {
-	iptablesRestore = "./testdata/iptables-restore"
+	ip4tablesRestore = "./testdata/iptables-restore"
+	ip6tablesRestore = "./testdata/iptables-restore"
 
 	tests := []struct {
 		name        string
-		origRules   []byte
+		origRules4  []byte
+		origRules6  []byte
 		restoreExit string
 		wantErr     bool
 	}{
 		{
 			name:        "success",
 			restoreExit: "0",
-			origRules:   []byte("sample-input-message"),
+			origRules4:  []byte("sample-input-message"),
+			origRules6:  []byte(""),
 		},
 		{
-			name:      "failure-to-restore-empty-rules",
-			origRules: nil,
-			wantErr:   true,
+			name:    "failure-to-restore-empty-rules",
+			wantErr: true,
 		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			r := &IPManager{origRules: tt.origRules}
+			r := &IPManager{origRules4: tt.origRules4, origRules6: tt.origRules6}
 			defer osx.MustSetenv("IPTABLES_RESTORE_EXIT", tt.restoreExit)()
 
 			b, err := r.Stop()

address/testdata/ip6tables

@@ -1,3 +1,3 @@
 #!/bin/bash
 
-exit ${IPTABLES_EXIT:-1}
+exit ${IP6TABLES_EXIT:-1}

address/testdata/ip6tables-save

@@ -0,0 +1,3 @@
+#!/bin/bash
+
+exit ${IP6TABLES_SAVE_EXIT:-1}

cmd/envelope/main_test.go

@@ -39,6 +39,9 @@ func Test_main(t *testing.T) {
 	flag.Set("address.iptables", "../../address/testdata/iptables")
 	flag.Set("address.iptables-save", "../../address/testdata/iptables-save")
 	flag.Set("address.iptables-restore", "../../address/testdata/iptables-restore")
+	flag.Set("address.ip6tables", "../../address/testdata/ip6tables")
+	flag.Set("address.ip6tables-save", "../../address/testdata/ip6tables-save")
+	flag.Set("address.ip6tables-restore", "../../address/testdata/iptables-restore")
 
 	// Load fake public verify key.
 	insecurePublicTestKey := []byte(`{"use":"sig","kty":"EC","kid":"112","crv":"P-256","alg":"ES256",` +
@@ -51,7 +54,9 @@ func Test_main(t *testing.T) {
 	verifyKeys = flagx.FileBytesArray{}
 	verifyKeys.Set(f.Name())
 	defer osx.MustSetenv("IPTABLES_EXIT", "0")()
+	defer osx.MustSetenv("IP6TABLES_EXIT", "0")()
 	defer osx.MustSetenv("IPTABLES_SAVE_EXIT", "0")()
+	defer osx.MustSetenv("IP6TABLES_SAVE_EXIT", "0")()
 
 	// Simulate unencrypted server.
 	listenAddr = ":0"