2 files changed, 11 insertions, 5 deletions
diff --git a/test/iptables/filter_input.go b/test/iptables/filter_input.go
index 05647de33..41e0cfa8d 100644
--- a/test/iptables/filter_input.go
+++ b/test/iptables/filter_input.go
@@ -194,8 +194,11 @@ 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)
+	// Ensure we cannot connect to the container.
+	for start := time.Now(); time.Since(start) < sendloopDuration; {
+		if err := connectTCP(ip, dropPort, sendloopDuration-time.Since(start)); err == nil {
+			return fmt.Errorf("expected not to connect, but was able to connect on port %d", dropPort)
+		}
 	}
 
 	return nil
@@ -226,8 +229,11 @@ 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")
+	// Ensure we cannot connect to the container.
+	for start := time.Now(); time.Since(start) < sendloopDuration; {
+		if err := connectTCP(ip, acceptPort, sendloopDuration-time.Since(start)); err == nil {
+			return fmt.Errorf("expected not to connect, but was able to connect on port %d", acceptPort)
+		}
 	}
 
 	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()
 		}