Merge pull request #3171 from kevinGC:ipv6-kokoro

PiperOrigin-RevId: 320290162

4 files changed, 46 insertions, 5 deletions

diff --git a/test/iptables/filter_input.go b/test/iptables/filter_input.go
index 872021358..068f228bd 100644
--- a/test/iptables/filter_input.go
+++ b/test/iptables/filter_input.go
@@ -618,7 +618,7 @@ func (FilterInputDestination) Name() string {
 
 // ContainerAction implements TestCase.ContainerAction.
 func (FilterInputDestination) ContainerAction(ip net.IP) error {
-	addrs, err := localAddrs()
+	addrs, err := localAddrs(false)
 	if err != nil {
 		return err
 	}
diff --git a/test/iptables/iptables_test.go b/test/iptables/iptables_test.go
index 9dc64f655..f5ac79370 100644
--- a/test/iptables/iptables_test.go
+++ b/test/iptables/iptables_test.go
@@ -18,6 +18,7 @@ import (
 	"context"
 	"fmt"
 	"net"
+	"reflect"
 	"testing"
 
 	"gvisor.dev/gvisor/pkg/test/dockerutil"
@@ -317,3 +318,28 @@ func TestInputSource(t *testing.T) {
 func TestInputInvertSource(t *testing.T) {
 	singleTest(t, FilterInputInvertSource{})
 }
+
+func TestFilterAddrs(t *testing.T) {
+	tcs := []struct {
+		ipv6  bool
+		addrs []string
+		want  []string
+	}{
+		{
+			ipv6:  false,
+			addrs: []string{"192.168.0.1", "192.168.0.2/24", "::1", "::2/128"},
+			want:  []string{"192.168.0.1", "192.168.0.2"},
+		},
+		{
+			ipv6:  true,
+			addrs: []string{"192.168.0.1", "192.168.0.2/24", "::1", "::2/128"},
+			want:  []string{"::1", "::2"},
+		},
+	}
+
+	for _, tc := range tcs {
+		if got := filterAddrs(tc.addrs, tc.ipv6); !reflect.DeepEqual(got, tc.want) {
+			t.Errorf("%v with IPv6 %t: got %v, but wanted %v", tc.addrs, tc.ipv6, got, tc.want)
+		}
+	}
+}
diff --git a/test/iptables/iptables_util.go b/test/iptables/iptables_util.go
index 7146edbb9..d4bc55b24 100644
--- a/test/iptables/iptables_util.go
+++ b/test/iptables/iptables_util.go
@@ -18,6 +18,7 @@ import (
 	"fmt"
 	"net"
 	"os/exec"
+	"strings"
 	"time"
 
 	"gvisor.dev/gvisor/pkg/test/testutil"
@@ -157,8 +158,10 @@ func connectTCP(ip net.IP, port int, timeout time.Duration) error {
 	return nil
 }
 
-// localAddrs returns a list of local network interface addresses.
-func localAddrs() ([]string, error) {
+// localAddrs returns a list of local network interface addresses. When ipv6 is
+// true, only IPv6 addresses are returned. Otherwise only IPv4 addresses are
+// returned.
+func localAddrs(ipv6 bool) ([]string, error) {
 	addrs, err := net.InterfaceAddrs()
 	if err != nil {
 		return nil, err
@@ -167,7 +170,19 @@ func localAddrs() ([]string, error) {
 	for _, addr := range addrs {
 		addrStrs = append(addrStrs, addr.String())
 	}
-	return addrStrs, nil
+	return filterAddrs(addrStrs, ipv6), nil
+}
+
+func filterAddrs(addrs []string, ipv6 bool) []string {
+	addrStrs := make([]string, 0, len(addrs))
+	for _, addr := range addrs {
+		// Add only IPv4 or only IPv6 addresses.
+		parts := strings.Split(addr, "/")
+		if isIPv6 := net.ParseIP(parts[0]).To4() == nil; isIPv6 == ipv6 {
+			addrStrs = append(addrStrs, parts[0])
+		}
+	}
+	return addrStrs
 }
 
 // getInterfaceName returns the name of the interface other than loopback.
diff --git a/test/iptables/nat.go b/test/iptables/nat.go
index 5e54a3963..8562b0820 100644
--- a/test/iptables/nat.go
+++ b/test/iptables/nat.go
@@ -241,7 +241,7 @@ func (NATPreRedirectIP) Name() string {
 
 // ContainerAction implements TestCase.ContainerAction.
 func (NATPreRedirectIP) ContainerAction(ip net.IP) error {
-	addrs, err := localAddrs()
+	addrs, err := localAddrs(false)
 	if err != nil {
 		return err
 	}