Merge pull request #3371 from kevinGC:ipt-tests-udp-faster

PiperOrigin-RevId: 323426155

3 files changed, 49 insertions, 26 deletions

diff --git a/test/iptables/filter_input.go b/test/iptables/filter_input.go
index 068f228bd..af4355ba8 100644
--- a/test/iptables/filter_input.go
+++ b/test/iptables/filter_input.go
@@ -81,7 +81,7 @@ func (FilterInputDropUDP) ContainerAction(ip net.IP) error {
 
 // LocalAction implements TestCase.LocalAction.
 func (FilterInputDropUDP) LocalAction(ip net.IP) error {
-	return sendUDPLoop(ip, dropPort, sendloopDuration)
+	return spawnUDPLoop(ip, dropPort, sendloopDuration)
 }
 
 // FilterInputDropOnlyUDP tests that "-p udp -j DROP" only affects UDP traffic.
@@ -141,7 +141,7 @@ func (FilterInputDropUDPPort) ContainerAction(ip net.IP) error {
 
 // LocalAction implements TestCase.LocalAction.
 func (FilterInputDropUDPPort) LocalAction(ip net.IP) error {
-	return sendUDPLoop(ip, dropPort, sendloopDuration)
+	return spawnUDPLoop(ip, dropPort, sendloopDuration)
 }
 
 // FilterInputDropDifferentUDPPort tests that dropping traffic for a single UDP port
@@ -169,7 +169,7 @@ func (FilterInputDropDifferentUDPPort) ContainerAction(ip net.IP) error {
 
 // LocalAction implements TestCase.LocalAction.
 func (FilterInputDropDifferentUDPPort) LocalAction(ip net.IP) error {
-	return sendUDPLoop(ip, acceptPort, sendloopDuration)
+	return spawnUDPLoop(ip, acceptPort, sendloopDuration)
 }
 
 // FilterInputDropTCPDestPort tests that connections are not accepted on specified source ports.
@@ -269,7 +269,7 @@ func (FilterInputDropAll) ContainerAction(ip net.IP) error {
 
 // LocalAction implements TestCase.LocalAction.
 func (FilterInputDropAll) LocalAction(ip net.IP) error {
-	return sendUDPLoop(ip, dropPort, sendloopDuration)
+	return spawnUDPLoop(ip, dropPort, sendloopDuration)
 }
 
 // FilterInputMultiUDPRules verifies that multiple UDP rules are applied
@@ -365,7 +365,7 @@ func (FilterInputDefaultPolicyAccept) ContainerAction(ip net.IP) error {
 
 // LocalAction implements TestCase.LocalAction.
 func (FilterInputDefaultPolicyAccept) LocalAction(ip net.IP) error {
-	return sendUDPLoop(ip, acceptPort, sendloopDuration)
+	return spawnUDPLoop(ip, acceptPort, sendloopDuration)
 }
 
 // FilterInputDefaultPolicyDrop tests the default DROP policy.
@@ -396,7 +396,7 @@ func (FilterInputDefaultPolicyDrop) ContainerAction(ip net.IP) error {
 
 // LocalAction implements TestCase.LocalAction.
 func (FilterInputDefaultPolicyDrop) LocalAction(ip net.IP) error {
-	return sendUDPLoop(ip, acceptPort, sendloopDuration)
+	return spawnUDPLoop(ip, acceptPort, sendloopDuration)
 }
 
 // FilterInputReturnUnderflow tests that -j RETURN in a built-in chain causes
@@ -428,7 +428,7 @@ func (FilterInputReturnUnderflow) ContainerAction(ip net.IP) error {
 
 // LocalAction implements TestCase.LocalAction.
 func (FilterInputReturnUnderflow) LocalAction(ip net.IP) error {
-	return sendUDPLoop(ip, acceptPort, sendloopDuration)
+	return spawnUDPLoop(ip, acceptPort, sendloopDuration)
 }
 
 // FilterInputSerializeJump verifies that we can serialize jumps.
@@ -482,7 +482,7 @@ func (FilterInputJumpBasic) ContainerAction(ip net.IP) error {
 
 // LocalAction implements TestCase.LocalAction.
 func (FilterInputJumpBasic) LocalAction(ip net.IP) error {
-	return sendUDPLoop(ip, acceptPort, sendloopDuration)
+	return spawnUDPLoop(ip, acceptPort, sendloopDuration)
 }
 
 // FilterInputJumpReturn jumps, returns, and executes a rule.
@@ -512,7 +512,7 @@ func (FilterInputJumpReturn) ContainerAction(ip net.IP) error {
 
 // LocalAction implements TestCase.LocalAction.
 func (FilterInputJumpReturn) LocalAction(ip net.IP) error {
-	return sendUDPLoop(ip, acceptPort, sendloopDuration)
+	return spawnUDPLoop(ip, acceptPort, sendloopDuration)
 }
 
 // FilterInputJumpReturnDrop jumps to a chain, returns, and DROPs packets.
@@ -549,7 +549,7 @@ func (FilterInputJumpReturnDrop) ContainerAction(ip net.IP) error {
 
 // LocalAction implements TestCase.LocalAction.
 func (FilterInputJumpReturnDrop) LocalAction(ip net.IP) error {
-	return sendUDPLoop(ip, dropPort, sendloopDuration)
+	return spawnUDPLoop(ip, dropPort, sendloopDuration)
 }
 
 // FilterInputJumpBuiltin verifies that jumping to a top-levl chain is illegal.
@@ -604,7 +604,7 @@ func (FilterInputJumpTwice) ContainerAction(ip net.IP) error {
 
 // LocalAction implements TestCase.LocalAction.
 func (FilterInputJumpTwice) LocalAction(ip net.IP) error {
-	return sendUDPLoop(ip, acceptPort, sendloopDuration)
+	return spawnUDPLoop(ip, acceptPort, sendloopDuration)
 }
 
 // FilterInputDestination verifies that we can filter packets via `-d
@@ -638,7 +638,7 @@ func (FilterInputDestination) ContainerAction(ip net.IP) error {
 
 // LocalAction implements TestCase.LocalAction.
 func (FilterInputDestination) LocalAction(ip net.IP) error {
-	return sendUDPLoop(ip, acceptPort, sendloopDuration)
+	return spawnUDPLoop(ip, acceptPort, sendloopDuration)
 }
 
 // FilterInputInvertDestination verifies that we can filter packets via `! -d
@@ -667,7 +667,7 @@ func (FilterInputInvertDestination) ContainerAction(ip net.IP) error {
 
 // LocalAction implements TestCase.LocalAction.
 func (FilterInputInvertDestination) LocalAction(ip net.IP) error {
-	return sendUDPLoop(ip, acceptPort, sendloopDuration)
+	return spawnUDPLoop(ip, acceptPort, sendloopDuration)
 }
 
 // FilterInputSource verifies that we can filter packets via `-s
@@ -696,7 +696,7 @@ func (FilterInputSource) ContainerAction(ip net.IP) error {
 
 // LocalAction implements TestCase.LocalAction.
 func (FilterInputSource) LocalAction(ip net.IP) error {
-	return sendUDPLoop(ip, acceptPort, sendloopDuration)
+	return spawnUDPLoop(ip, acceptPort, sendloopDuration)
 }
 
 // FilterInputInvertSource verifies that we can filter packets via `! -s
@@ -725,5 +725,5 @@ func (FilterInputInvertSource) ContainerAction(ip net.IP) error {
 
 // LocalAction implements TestCase.LocalAction.
 func (FilterInputInvertSource) LocalAction(ip net.IP) error {
-	return sendUDPLoop(ip, acceptPort, sendloopDuration)
+	return spawnUDPLoop(ip, acceptPort, sendloopDuration)
 }
diff --git a/test/iptables/iptables_util.go b/test/iptables/iptables_util.go
index d4bc55b24..174694002 100644
--- a/test/iptables/iptables_util.go
+++ b/test/iptables/iptables_util.go
@@ -84,17 +84,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)
 	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
@@ -109,8 +134,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.
diff --git a/test/iptables/nat.go b/test/iptables/nat.go
index 149dec2bb..23288577d 100644
--- a/test/iptables/nat.go
+++ b/test/iptables/nat.go
@@ -67,7 +67,7 @@ func (NATPreRedirectUDPPort) ContainerAction(ip net.IP) error {
 
 // LocalAction implements TestCase.LocalAction.
 func (NATPreRedirectUDPPort) LocalAction(ip net.IP) error {
-	return sendUDPLoop(ip, acceptPort, sendloopDuration)
+	return spawnUDPLoop(ip, acceptPort, sendloopDuration)
 }
 
 // NATPreRedirectTCPPort tests that connections are redirected on specified ports.
@@ -187,7 +187,7 @@ func (NATDropUDP) ContainerAction(ip net.IP) error {
 
 // LocalAction implements TestCase.LocalAction.
 func (NATDropUDP) LocalAction(ip net.IP) error {
-	return sendUDPLoop(ip, acceptPort, sendloopDuration)
+	return spawnUDPLoop(ip, acceptPort, sendloopDuration)
 }
 
 // NATAcceptAll tests that all UDP packets are accepted.
@@ -213,7 +213,7 @@ func (NATAcceptAll) ContainerAction(ip net.IP) error {
 
 // LocalAction implements TestCase.LocalAction.
 func (NATAcceptAll) LocalAction(ip net.IP) error {
-	return sendUDPLoop(ip, acceptPort, sendloopDuration)
+	return spawnUDPLoop(ip, acceptPort, sendloopDuration)
 }
 
 // NATOutRedirectIP uses iptables to select packets based on destination IP and
@@ -310,7 +310,7 @@ func (NATPreRedirectIP) ContainerAction(ip net.IP) error {
 
 // LocalAction implements TestCase.LocalAction.
 func (NATPreRedirectIP) LocalAction(ip net.IP) error {
-	return sendUDPLoop(ip, dropPort, sendloopDuration)
+	return spawnUDPLoop(ip, dropPort, sendloopDuration)
 }
 
 // NATPreDontRedirectIP tests that iptables matching with "-d" does not match
@@ -332,7 +332,7 @@ func (NATPreDontRedirectIP) ContainerAction(ip net.IP) error {
 
 // LocalAction implements TestCase.LocalAction.
 func (NATPreDontRedirectIP) LocalAction(ip net.IP) error {
-	return sendUDPLoop(ip, acceptPort, sendloopDuration)
+	return spawnUDPLoop(ip, acceptPort, sendloopDuration)
 }
 
 // NATPreRedirectInvert tests that iptables can match with "! -d".
@@ -353,7 +353,7 @@ func (NATPreRedirectInvert) ContainerAction(ip net.IP) error {
 
 // LocalAction implements TestCase.LocalAction.
 func (NATPreRedirectInvert) LocalAction(ip net.IP) error {
-	return sendUDPLoop(ip, dropPort, sendloopDuration)
+	return spawnUDPLoop(ip, dropPort, sendloopDuration)
 }
 
 // NATRedirectRequiresProtocol tests that use of the --to-ports flag requires a