Merge pull request #2187 from kevinGC:deflake-connectTCP

PiperOrigin-RevId: 301716568

2 files changed, 17 insertions, 5 deletions

diff --git a/test/iptables/filter_input.go b/test/iptables/filter_input.go
index 05647de33..4ccd4cce7 100644
--- a/test/iptables/filter_input.go
+++ b/test/iptables/filter_input.go
@@ -194,8 +194,14 @@ func (FilterInputDropTCPDestPort) ContainerAction(ip net.IP) error {
 
 // LocalAction implements TestCase.LocalAction.
 func (FilterInputDropTCPDestPort) LocalAction(ip net.IP) error {
-	if err := connectTCP(ip, dropPort, sendloopDuration); err == nil {
-		return fmt.Errorf("connection destined to port %d should not be accepted, but got accepted", dropPort)
+	// After the container sets its DROP rule, we shouldn't be able to connect.
+	// However, we may succeed in connecting if this runs before the container
+	// sets the rule. To avoid this race, we retry connecting until
+	// sendloopDuration has elapsed, ignoring whether the connect succeeds. The
+	// test works becuase the container will error if a connection is
+	// established after the rule is set.
+	for start := time.Now(); time.Since(start) < sendloopDuration; {
+		connectTCP(ip, dropPort, sendloopDuration-time.Since(start))
 	}
 
 	return nil
@@ -226,8 +232,14 @@ func (FilterInputDropTCPSrcPort) ContainerAction(ip net.IP) error {
 
 // LocalAction implements TestCase.LocalAction.
 func (FilterInputDropTCPSrcPort) LocalAction(ip net.IP) error {
-	if err := connectTCP(ip, acceptPort, sendloopDuration); err == nil {
-		return fmt.Errorf("connection should not be accepted, but was")
+	// After the container sets its DROP rule, we shouldn't be able to connect.
+	// However, we may succeed in connecting if this runs before the container
+	// sets the rule. To avoid this race, we retry connecting until
+	// sendloopDuration has elapsed, ignoring whether the connect succeeds. The
+	// test works becuase the container will error if a connection is
+	// established after the rule is set.
+	for start := time.Now(); time.Since(start) < sendloopDuration; {
+		connectTCP(ip, acceptPort, sendloopDuration-time.Since(start))
 	}
 
 	return nil
diff --git a/test/iptables/iptables_util.go b/test/iptables/iptables_util.go
index e8ae65c5a..134391e8d 100644
--- a/test/iptables/iptables_util.go
+++ b/test/iptables/iptables_util.go
@@ -144,7 +144,7 @@ func connectTCP(ip net.IP, port int, timeout time.Duration) error {
 	// The container may not be listening when we first connect, so retry
 	// upon error.
 	callback := func() error {
-		conn, err := net.DialTCP("tcp4", nil, &contAddr)
+		conn, err := net.DialTimeout("tcp", contAddr.String(), timeout)
 		if conn != nil {
 			conn.Close()
 		}