1 files changed, 89 insertions, 29 deletions
diff --git a/test/iptables/iptables_util.go b/test/iptables/iptables_util.go
index 7146edbb9..ca80a4b5f 100644
--- a/test/iptables/iptables_util.go
+++ b/test/iptables/iptables_util.go
@@ -18,27 +18,29 @@ import (
 	"fmt"
 	"net"
 	"os/exec"
+	"strings"
 	"time"
 
 	"gvisor.dev/gvisor/pkg/test/testutil"
 )
 
-const iptablesBinary = "iptables"
-const localIP = "127.0.0.1"
-
-// filterTable calls `iptables -t filter` with the given args.
-func filterTable(args ...string) error {
-	return tableCmd("filter", args)
+// filterTable calls `ip{6}tables -t filter` with the given args.
+func filterTable(ipv6 bool, args ...string) error {
+	return tableCmd(ipv6, "filter", args)
 }
 
-// natTable calls `iptables -t nat` with the given args.
-func natTable(args ...string) error {
-	return tableCmd("nat", args)
+// natTable calls `ip{6}tables -t nat` with the given args.
+func natTable(ipv6 bool, args ...string) error {
+	return tableCmd(ipv6, "nat", args)
 }
 
-func tableCmd(table string, args []string) error {
+func tableCmd(ipv6 bool, table string, args []string) error {
 	args = append([]string{"-t", table}, args...)
-	cmd := exec.Command(iptablesBinary, args...)
+	binary := "iptables"
+	if ipv6 {
+		binary = "ip6tables"
+	}
+	cmd := exec.Command(binary, args...)
 	if out, err := cmd.CombinedOutput(); err != nil {
 		return fmt.Errorf("error running iptables with args %v\nerror: %v\noutput: %s", args, err, string(out))
 	}
@@ -46,18 +48,18 @@ func tableCmd(table string, args []string) error {
 }
 
 // filterTableRules is like filterTable, but runs multiple iptables commands.
-func filterTableRules(argsList [][]string) error {
-	return tableRules("filter", argsList)
+func filterTableRules(ipv6 bool, argsList [][]string) error {
+	return tableRules(ipv6, "filter", argsList)
 }
 
 // natTableRules is like natTable, but runs multiple iptables commands.
-func natTableRules(argsList [][]string) error {
-	return tableRules("nat", argsList)
+func natTableRules(ipv6 bool, argsList [][]string) error {
+	return tableRules(ipv6, "nat", argsList)
 }
 
-func tableRules(table string, argsList [][]string) error {
+func tableRules(ipv6 bool, table string, argsList [][]string) error {
 	for _, args := range argsList {
-		if err := tableCmd(table, args); err != nil {
+		if err := tableCmd(ipv6, table, args); err != nil {
 			return err
 		}
 	}
@@ -70,7 +72,7 @@ func listenUDP(port int, timeout time.Duration) error {
 	localAddr := net.UDPAddr{
 		Port: port,
 	}
-	conn, err := net.ListenUDP(network, &localAddr)
+	conn, err := net.ListenUDP("udp", &localAddr)
 	if err != nil {
 		return err
 	}
@@ -83,17 +85,42 @@ func listenUDP(port int, timeout time.Duration) error {
 // sendUDPLoop sends 1 byte UDP packets repeatedly to the IP and port specified
 // over a duration.
 func sendUDPLoop(ip net.IP, port int, duration time.Duration) error {
-	// Send packets for a few seconds.
+	conn, err := connectUDP(ip, port)
+	if err != nil {
+		return err
+	}
+	defer conn.Close()
+	loopUDP(conn, duration)
+	return nil
+}
+
+// spawnUDPLoop works like sendUDPLoop, but returns immediately and sends
+// packets in another goroutine.
+func spawnUDPLoop(ip net.IP, port int, duration time.Duration) error {
+	conn, err := connectUDP(ip, port)
+	if err != nil {
+		return err
+	}
+	go func() {
+		defer conn.Close()
+		loopUDP(conn, duration)
+	}()
+	return nil
+}
+
+func connectUDP(ip net.IP, port int) (net.Conn, error) {
 	remote := net.UDPAddr{
 		IP:   ip,
 		Port: port,
 	}
-	conn, err := net.DialUDP(network, nil, &remote)
+	conn, err := net.DialUDP("udp", nil, &remote)
 	if err != nil {
-		return err
+		return nil, err
 	}
-	defer conn.Close()
+	return conn, nil
+}
 
+func loopUDP(conn net.Conn, duration time.Duration) {
 	to := time.After(duration)
 	for timedOut := false; !timedOut; {
 		// This may return an error (connection refused) if the remote
@@ -108,8 +135,6 @@ func sendUDPLoop(ip net.IP, port int, duration time.Duration) error {
 			time.Sleep(200 * time.Millisecond)
 		}
 	}
-
-	return nil
 }
 
 // listenTCP listens for connections on a TCP port.
@@ -119,7 +144,7 @@ func listenTCP(port int, timeout time.Duration) error {
 	}
 
 	// Starts listening on port.
-	lConn, err := net.ListenTCP("tcp4", &localAddr)
+	lConn, err := net.ListenTCP("tcp", &localAddr)
 	if err != nil {
 		return err
 	}
@@ -157,17 +182,38 @@ 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
 	}
 	addrStrs := make([]string, 0, len(addrs))
 	for _, addr := range addrs {
-		addrStrs = append(addrStrs, addr.String())
+		// Add only IPv4 or only IPv6 addresses.
+		parts := strings.Split(addr.String(), "/")
+		if len(parts) != 2 {
+			return nil, fmt.Errorf("bad interface address: %q", addr.String())
+		}
+		if isIPv6 := net.ParseIP(parts[0]).To4() == nil; isIPv6 == ipv6 {
+			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.
@@ -184,3 +230,17 @@ func getInterfaceName() (string, bool) {
 
 	return ifname, ifname != ""
 }
+
+func localIP(ipv6 bool) string {
+	if ipv6 {
+		return "::1"
+	}
+	return "127.0.0.1"
+}
+
+func nowhereIP(ipv6 bool) string {
+	if ipv6 {
+		return "2001:db8::1"
+	}
+	return "192.0.2.1"
+}